adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'hotfix/v0.2.10' into develop

Browse Source
This commit is contained in:
Ke Sun
2026-04-10 10:16:39 +08:00
parent 3a62d50048 3e0f59adc6
commit 807dee8460
26 changed files with 421 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
api/app/controllers/chunk_controller.py
View File

@@ -23,6 +23,7 @@ from app.models.user_model import User
from app.schemas import chunk_schema
from app.schemas.response_schema import ApiResponse
from app.services import knowledge_service, document_service, file_service, knowledgeshare_service
from app.services.model_service import ModelApiKeyService
# Obtain a dedicated API logger
api_logger = get_api_logger()
@@ -460,18 +461,20 @@ async def retrieve_chunks(
if retrieve_data.retrieve_type == chunk_schema.RetrieveType.Graph:
kb_ids = [str(kb_id) for kb_id in private_kb_ids]
workspace_ids = [str(workspace_id) for workspace_id in private_workspace_ids]
llm_key = ModelApiKeyService.get_available_api_key(db, db_knowledge.llm_id)
emb_key = ModelApiKeyService.get_available_api_key(db, db_knowledge.embedding_id)
# Prepare to configure chat_mdl、embedding_model、vision_model information
chat_model = Base(
key=db_knowledge.llm.api_keys[0].api_key,
model_name=db_knowledge.llm.api_keys[0].model_name,
base_url=db_knowledge.llm.api_keys[0].api_base
key=llm_key.api_key,
model_name=llm_key.model_name,
base_url=llm_key.api_base
)
embedding_model = OpenAIEmbed(
key=db_knowledge.embedding.api_keys[0].api_key,
model_name=db_knowledge.embedding.api_keys[0].model_name,
base_url=db_knowledge.embedding.api_keys[0].api_base
key=emb_key.api_key,
model_name=emb_key.model_name,
base_url=emb_key.api_base
)
doc = kg_retriever.retrieval(question=retrieve_data.query, workspace_ids=workspace_ids, kb_ids= kb_ids, emb_mdl=embedding_model, llm=chat_model)
doc = kg_retriever.retrieval(question=retrieve_data.query, workspace_ids=workspace_ids, kb_ids=kb_ids, emb_mdl=embedding_model, llm=chat_model)
if doc:
rs.insert(0, doc)
return success(data=jsonable_encoder(rs), msg="retrieval successful")

7
api/app/core/rag/deepdoc/parser/mineru_parser.py
View File

@@ -292,9 +292,10 @@ class MinerUParser(RAGPdfParser):
self.page_from = page_from
self.page_to = page_to
try:
with pdfplumber.open(fnm) if isinstance(fnm, (str, PathLike)) else pdfplumber.open(BytesIO(fnm)) as pdf:
self.pdf = pdf
self.page_images = [p.to_image(resolution=72 * zoomin, antialias=True).original for _, p in enumerate(self.pdf.pages[page_from:page_to])]
with sys.modules[LOCK_KEY_pdfplumber]: # ← 加这一行，获取全局锁
with pdfplumber.open(fnm) if isinstance(fnm, (str, PathLike)) else pdfplumber.open(BytesIO(fnm)) as pdf:
self.pdf = pdf
self.page_images = [p.to_image(resolution=72 * zoomin, antialias=True).original for _, p in enumerate(self.pdf.pages[page_from:page_to])]
except Exception as e:
self.page_images = None
self.total_page = 0

46
api/app/core/rag/nlp/search.py
View File

@@ -28,6 +28,7 @@ from app.core.rag.common.float_utils import get_float
from app.core.rag.common.constants import PAGERANK_FLD, TAG_FLD
from app.core.rag.llm.chat_model import Base
from app.core.rag.llm.embedding_model import OpenAIEmbed
from app.services.model_service import ModelApiKeyService
import logging
logger = logging.getLogger(__name__)
@@ -114,9 +115,8 @@ def knowledge_retrieval(
# Use the specified reranker for re-ranking
if reranker_id:
try:
return rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
all_results = rerank(db=db, reranker_id=reranker_id, query=query, docs=all_results, top_k=reranker_top_k)
except Exception as rerank_error:
# If reranker fails, log warning and continue with original results
logger.warning(
"Reranker failed, falling back to original results",
extra={
@@ -132,7 +132,10 @@ def knowledge_retrieval(
from app.core.rag.common.settings import kg_retriever
doc = kg_retriever.retrieval(question=query, workspace_ids=workspace_ids, kb_ids=kb_ids, emb_mdl=embedding_model, llm=chat_model)
if doc:
all_results.insert(0, doc)
all_results.insert(0, DocumentChunk(
page_content=doc.get("page_content", ""),
metadata=doc.get("metadata", {})
))
except Exception as graph_error:
print(f"Failed to retrieve from knowledge graph: {str(graph_error)}")
@@ -198,16 +201,18 @@ def _retrieve_for_knowledge(
workspace_ids.append(str(db_knowledge.workspace_id))
if not chat_model:
llm_key = ModelApiKeyService.get_available_api_key(db, db_knowledge.llm_id)
chat_model = Base(
key=db_knowledge.llm.api_keys[0].api_key,
model_name=db_knowledge.llm.api_keys[0].model_name,
base_url=db_knowledge.llm.api_keys[0].api_base,
key=llm_key.api_key,
model_name=llm_key.model_name,
base_url=llm_key.api_base,
)
if not embedding_model:
emb_key = ModelApiKeyService.get_available_api_key(db, db_knowledge.embedding_id)
embedding_model = OpenAIEmbed(
key=db_knowledge.embedding.api_keys[0].api_key,
model_name=db_knowledge.embedding.api_keys[0].model_name,
base_url=db_knowledge.embedding.api_keys[0].api_base,
key=emb_key.api_key,
model_name=emb_key.model_name,
base_url=emb_key.api_base,
)
vector_service = ElasticSearchVectorFactory().init_vector(knowledge=db_knowledge)
@@ -248,6 +253,29 @@ def _retrieve_for_knowledge(
seen_ids.add(doc.metadata["doc_id"])
unique_rs.append(doc)
rs = unique_rs
if unique_rs:
rs = vector_service.rerank(
query=kb_config["query"],
docs=unique_rs,
top_k=kb_config["top_k"]
)
if kb_config["retrieve_type"] == "graph":
try:
from app.core.rag.common.settings import kg_retriever
graph_doc = kg_retriever.retrieval(
question=kb_config["query"],
workspace_ids=[str(db_knowledge.workspace_id)],
kb_ids=[str(db_knowledge.id)],
emb_mdl=embedding_model,
llm=chat_model,
)
if graph_doc:
rs.insert(0, DocumentChunk(
page_content=graph_doc.get("page_content", ""),
metadata=graph_doc.get("metadata", {})
))
except Exception as graph_error:
logger.warning(f"Graph retrieval failed for kb {db_knowledge.id}: {graph_error}")
results.extend(rs)
return results, chat_model, embedding_model

6
api/app/core/tools/builtin/datetime_tool.py
View File

@@ -230,7 +230,7 @@ class DateTimeTool(BuiltinTool):
@staticmethod
def _datetime_to_timestamp(kwargs) -> dict:
"""日期时间转时间戳"""
input_value = kwargs.get("input_value")
input_value = kwargs.get("input_value").strip()
input_format = kwargs.get("input_format", "%Y-%m-%d %H:%M:%S")
timezone_str = kwargs.get("from_timezone", "Asia/Shanghai")
@@ -253,9 +253,9 @@ class DateTimeTool(BuiltinTool):
return {
"datetime": input_value,
"timezone": timezone_str,
"timestamp": int(dt.timestamp()),
"timestamp": int(dt.timestamp()) * 1000,
"iso_format": dt.isoformat(),
"result_data": int(dt.timestamp())
"result_data": int(dt.timestamp()) * 1000
}
def _calculate_datetime(self, kwargs) -> dict:

23
api/app/core/tools/builtin/operation_tool.py
View File

@@ -138,6 +138,29 @@ class OperationTool(BaseTool):
default="Asia/Shanghai"
)
]
elif self.operation == "datetime_to_timestamp":
return [
ToolParameter(
name="input_value",
type=ParameterType.STRING,
description="输入值（时间字符串，如：2026-04-07 10:30:25）",
required=True
),
ToolParameter(
name="input_format",
type=ParameterType.STRING,
description="输入时间格式（如：%Y-%m-%d %H:%M:%S）",
required=False,
default="%Y-%m-%d %H:%M:%S"
),
ToolParameter(
name="from_timezone",
type=ParameterType.STRING,
description="源时区（如：UTC, Asia/Shanghai）",
required=False,
default="Asia/Shanghai"
)
]
else:
return []

64
api/app/core/workflow/nodes/knowledge/node.py
View File

@@ -8,6 +8,8 @@ from langchain_core.documents import Document
from app.core.error_codes import BizCode
from app.core.exceptions import BusinessException
from app.core.models import RedBearRerank, RedBearModelConfig
from app.core.rag.llm.chat_model import Base
from app.core.rag.llm.embedding_model import OpenAIEmbed
from app.core.rag.models.chunk import DocumentChunk
from app.core.rag.vdb.elasticsearch.elasticsearch_vector import ElasticSearchVectorFactory
from app.core.workflow.engine.state_manager import WorkflowState
@@ -39,8 +41,9 @@ class KnowledgeRetrievalNode(BaseNode):
if isinstance(business_result, dict) and "chunks" in business_result:
return business_result["chunks"]
return business_result
def _extract_citations(self, business_result: Any) -> list:
@staticmethod
def _extract_citations(business_result: Any) -> list:
if isinstance(business_result, dict):
return business_result.get("citations", [])
return []
@@ -230,23 +233,23 @@ class KnowledgeRetrievalNode(BaseNode):
}
)
)
case RetrieveType.HYBRID:
case retrieve_type if retrieve_type in (RetrieveType.HYBRID, RetrieveType.Graph):
rs1_task = asyncio.to_thread(
vector_service.search_by_vector, **{
"query": query,
"top_k": kb_config.top_k,
"indices": indices,
"score_threshold": kb_config.vector_similarity_weight
}
)
vector_service.search_by_vector, **{
"query": query,
"top_k": kb_config.top_k,
"indices": indices,
"score_threshold": kb_config.vector_similarity_weight
}
)
rs2_task = asyncio.to_thread(
vector_service.search_by_full_text, **{
"query": query,
"top_k": kb_config.top_k,
"indices": indices,
"score_threshold": kb_config.similarity_threshold
}
)
vector_service.search_by_full_text, **{
"query": query,
"top_k": kb_config.top_k,
"indices": indices,
"score_threshold": kb_config.similarity_threshold
}
)
rs1, rs2 = await asyncio.gather(rs1_task, rs2_task)
# Deduplicate hybrid retrieval results
@@ -266,6 +269,33 @@ class KnowledgeRetrievalNode(BaseNode):
key=lambda d: d.metadata.get("score", 0),
reverse=True
)[:kb_config.top_k])
if kb_config.retrieve_type == RetrieveType.Graph:
from app.core.rag.common.settings import kg_retriever
llm_key = self.model_balance(db_knowledge.llm)
emb_key = self.model_balance(db_knowledge.embedding)
chat_model = Base(
key=llm_key.api_key,
model_name=llm_key.model_name,
base_url=llm_key.api_base
)
embedding_model = OpenAIEmbed(
key=emb_key.api_key,
model_name=emb_key.model_name,
base_url=emb_key.api_base
)
doc = await asyncio.to_thread(
kg_retriever.retrieval,
question=query,
workspace_ids=[str(db_knowledge.workspace_id)],
kb_ids=[str(kb_config.kb_id)],
emb_mdl=embedding_model,
llm=chat_model
)
if doc:
rs.insert(0, DocumentChunk(
page_content=doc.get("page_content", ""),
metadata=doc.get("metadata", {})
))
case _:
raise RuntimeError("Unknown retrieval type")
return rs

23
api/app/services/tool_service.py
View File

@@ -574,6 +574,29 @@ class ToolService:
"default": "Asia/Shanghai"
}
]
elif operation == "datetime_to_timestamp":
return [
{
"name": "input_value",
"type": "string",
"description": "输入值（时间字符串，如：2026-04-07 10:30:25）",
"required": True
},
{
"name": "input_format",
"type": "string",
"description": "输入时间格式（如：%Y-%m-%d %H:%M:%S）",
"required": False,
"default": "%Y-%m-%d %H:%M:%S"
},
{
"name": "from_timezone",
"type": "string",
"description": "源时区（如：UTC, Asia/Shanghai）",
"required": False,
"default": "Asia/Shanghai"
}
]
else:
# 默认返回所有参数（除了operation）
return [