Fix/agent external knowledge retrieval (langgenius#9241)

api/configs/middleware/__init__.py

@@ -191,6 +191,22 @@ def BROKER_USE_SSL(self) -> bool:
  return self.CELERY_BROKER_URL.startswith("rediss://") if self.CELERY_BROKER_URL else False
 
 
+class InternalTestConfig(BaseSettings):
+ """
+ Configuration settings for Internal Test
+ """
+
+ AWS_SECRET_ACCESS_KEY: Optional[str] = Field(
+ description="Internal test AWS secret access key",
+ default=None,
+ )
+
+ AWS_ACCESS_KEY_ID: Optional[str] = Field(
+ description="Internal test AWS access key ID",
+ default=None,
+ )
+
+
 class MiddlewareConfig(
  # place the configs in alphabet order
  CeleryConfig,
@@ -224,5 +240,6 @@ class MiddlewareConfig(
  TiDBVectorConfig,
  WeaviateConfig,
  ElasticsearchConfig,
+ InternalTestConfig,
 ):
  pass

api/controllers/console/datasets/external.py

@@ -13,6 +13,7 @@
 from services.dataset_service import DatasetService
 from services.external_knowledge_service import ExternalDatasetService
 from services.hit_testing_service import HitTestingService
+from services.knowledge_service import ExternalDatasetTestService
 
 
 def _validate_name(name):
@@ -232,8 +233,31 @@ def post(self, dataset_id):
  raise InternalServerError(str(e))
 
 
+class BedrockRetrievalApi(Resource):
+ # this api is only for internal testing
+ def post(self):
+ parser = reqparse.RequestParser()
+ parser.add_argument("retrieval_setting", nullable=False, required=True, type=dict, location="json")
+ parser.add_argument(
+ "query",
+ nullable=False,
+ required=True,
+ type=str,
+ )
+ parser.add_argument("knowledge_id", nullable=False, required=True, type=str)
+ args = parser.parse_args()
+
+ # Call the knowledge retrieval service
+ result = ExternalDatasetTestService.knowledge_retrieval(
+ args["retrieval_setting"], args["query"], args["knowledge_id"]
+ )
+ return result, 200
+
+
 api.add_resource(ExternalKnowledgeHitTestingApi, "/datasets/<uuid:dataset_id>/external-hit-testing")
 api.add_resource(ExternalDatasetCreateApi, "/datasets/external")
 api.add_resource(ExternalApiTemplateListApi, "/datasets/external-knowledge-api")
 api.add_resource(ExternalApiTemplateApi, "/datasets/external-knowledge-api/<uuid:external_knowledge_api_id>")
 api.add_resource(ExternalApiUseCheckApi, "/datasets/external-knowledge-api/<uuid:external_knowledge_api_id>/use-check")
+# this api is only for internal test
+api.add_resource(BedrockRetrievalApi, "/test/retrieval")

api/core/rag/retrieval/dataset_retrieval.py

@@ -539,7 +539,7 @@ def to_dataset_retriever_tool(
  continue
 
  # pass if dataset is not available
- if dataset and dataset.available_document_count == 0:
+ if dataset and dataset.provider != "external" and dataset.available_document_count == 0:
  continue
 
  available_datasets.append(dataset)

api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py

@@ -1,10 +1,12 @@
 from pydantic import BaseModel, Field
 
 from core.rag.datasource.retrieval_service import RetrievalService
+from core.rag.models.document import Document as RetrievalDocument
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.tools.tool.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
 from extensions.ext_database import db
 from models.dataset import Dataset, Document, DocumentSegment
+from services.external_knowledge_service import ExternalDatasetService
 
 default_retrieval_model = {
  "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
@@ -53,97 +55,137 @@ def _run(self, query: str) -> str:
 
  for hit_callback in self.hit_callbacks:
  hit_callback.on_query(query, dataset.id)
-
- # get retrieval model , if the model is not setting , using default
- retrieval_model = dataset.retrieval_model or default_retrieval_model
- if dataset.indexing_technique == "economy":
- # use keyword table query
- documents = RetrievalService.retrieve(
- retrieval_method="keyword_search", dataset_id=dataset.id, query=query, top_k=self.top_k
+ if dataset.provider == "external":
+  results = []
+  external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval(
+  tenant_id=dataset.tenant_id,
+  dataset_id=dataset.id,
+  query=query,
+ external_retrieval_parameters=dataset.retrieval_model,
  )
- return str("\n".join([document.page_content for document in documents]))
+ for external_document in external_documents:
+ document = RetrievalDocument(
+ page_content=external_document.get("content"),
+ metadata=external_document.get("metadata"),
+ provider="external",
+ )
+ document.metadata["score"] = external_document.get("score")
+ document.metadata["title"] = external_document.get("title")
+ document.metadata["dataset_id"] = dataset.id
+ document.metadata["dataset_name"] = dataset.name
+ results.append(document)
+ # deal with external documents
+ context_list = []
+ for position, item in enumerate(results, start=1):
+ source = {
+ "position": position,
+ "dataset_id": item.metadata.get("dataset_id"),
+ "dataset_name": item.metadata.get("dataset_name"),
+ "document_name": item.metadata.get("title"),
+ "data_source_type": "external",
+ "retriever_from": self.retriever_from,
+ "score": item.metadata.get("score"),
+ "title": item.metadata.get("title"),
+ "content": item.page_content,
+ }
+ context_list.append(source)
+ for hit_callback in self.hit_callbacks:
+ hit_callback.return_retriever_resource_info(context_list)
+
+ return str("\n".join([item.page_content for item in results]))
  else:
- if self.top_k > 0:
- # retrieval source
+ # get retrieval model , if the model is not setting , using default
+ retrieval_model = dataset.retrieval_model or default_retrieval_model
+ if dataset.indexing_technique == "economy":
+ # use keyword table query
  documents = RetrievalService.retrieve(
- retrieval_method=retrieval_model.get("search_method", "semantic_search"),
- dataset_id=dataset.id,
- query=query,
- top_k=self.top_k,
- score_threshold=retrieval_model.get("score_threshold", 0.0)
- if retrieval_model["score_threshold_enabled"]
- else 0.0,
- reranking_model=retrieval_model.get("reranking_model", None)
- if retrieval_model["reranking_enable"]
- else None,
- reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model",
- weights=retrieval_model.get("weights", None),
+ retrieval_method="keyword_search", dataset_id=dataset.id, query=query, top_k=self.top_k
  )
+ return str("\n".join([document.page_content for document in documents]))
  else:
- documents = []
-
- for hit_callback in self.hit_callbacks:
- hit_callback.on_tool_end(documents)
- document_score_list = {}
- if dataset.indexing_technique != "economy":
- for item in documents:
- if item.metadata.get("score"):
- document_score_list[item.metadata["doc_id"]] = item.metadata["score"]
- document_context_list = []
- index_node_ids = [document.metadata["doc_id"] for document in documents]
- segments = DocumentSegment.query.filter(
- DocumentSegment.dataset_id == self.dataset_id,
- DocumentSegment.completed_at.isnot(None),
- DocumentSegment.status == "completed",
- DocumentSegment.enabled == True,
- DocumentSegment.index_node_id.in_(index_node_ids),
- ).all()
-
- if segments:
- index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}
- sorted_segments = sorted(
- segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf"))
- )
- for segment in sorted_segments:
- if segment.answer:
- document_context_list.append(f"question:{segment.get_sign_content()} answer:{segment.answer}")
- else:
- document_context_list.append(segment.get_sign_content())
- if self.return_resource:
- context_list = []
- resource_number = 1
+ if self.top_k > 0:
+ # retrieval source
+ documents = RetrievalService.retrieve(
+ retrieval_method=retrieval_model.get("search_method", "semantic_search"),
+ dataset_id=dataset.id,
+ query=query,
+ top_k=self.top_k,
+ score_threshold=retrieval_model.get("score_threshold", 0.0)
+ if retrieval_model["score_threshold_enabled"]
+ else 0.0,
+ reranking_model=retrieval_model.get("reranking_model", None)
+ if retrieval_model["reranking_enable"]
+ else None,
+ reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model",
+ weights=retrieval_model.get("weights", None),
+ )
+ else:
+ documents = []
+
+ for hit_callback in self.hit_callbacks:
+ hit_callback.on_tool_end(documents)
+ document_score_list = {}
+ if dataset.indexing_technique != "economy":
+ for item in documents:
+ if item.metadata.get("score"):
+ document_score_list[item.metadata["doc_id"]] = item.metadata["score"]
+ document_context_list = []
+ index_node_ids = [document.metadata["doc_id"] for document in documents]
+ segments = DocumentSegment.query.filter(
+ DocumentSegment.dataset_id == self.dataset_id,
+ DocumentSegment.completed_at.isnot(None),
+ DocumentSegment.status == "completed",
+ DocumentSegment.enabled == True,
+ DocumentSegment.index_node_id.in_(index_node_ids),
+ ).all()
+
+ if segments:
+ index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}
+ sorted_segments = sorted(
+ segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf"))
+ )
  for segment in sorted_segments:
- context = {}
- document = Document.query.filter(
- Document.id == segment.document_id,
- Document.enabled == True,
- Document.archived == False,
- ).first()
- if dataset and document:
- source = {
- "position": resource_number,
- "dataset_id": dataset.id,
- "dataset_name": dataset.name,
- "document_id": document.id,
- "document_name": document.name,
- "data_source_type": document.data_source_type,
- "segment_id": segment.id,
- "retriever_from": self.retriever_from,
- "score": document_score_list.get(segment.index_node_id, None),
- }
- if self.retriever_from == "dev":
- source["hit_count"] = segment.hit_count
- source["word_count"] = segment.word_count
- source["segment_position"] = segment.position
- source["index_node_hash"] = segment.index_node_hash
- if segment.answer:
- source["content"] = f"question:{segment.content} \nanswer:{segment.answer}"
- else:
- source["content"] = segment.content
- context_list.append(source)
- resource_number += 1
-
- for hit_callback in self.hit_callbacks:
- hit_callback.return_retriever_resource_info(context_list)
-
- return str("\n".join(document_context_list))
+ if segment.answer:
+ document_context_list.append(
+ f"question:{segment.get_sign_content()} answer:{segment.answer}"
+ )
+ else:
+ document_context_list.append(segment.get_sign_content())
+ if self.return_resource:
+ context_list = []
+ resource_number = 1
+ for segment in sorted_segments:
+ context = {}
+ document = Document.query.filter(
+ Document.id == segment.document_id,
+ Document.enabled == True,
+ Document.archived == False,
+ ).first()
+ if dataset and document:
+ source = {
+ "position": resource_number,
+ "dataset_id": dataset.id,
+ "dataset_name": dataset.name,
+ "document_id": document.id,
+ "document_name": document.name,
+ "data_source_type": document.data_source_type,
+ "segment_id": segment.id,
+ "retriever_from": self.retriever_from,
+ "score": document_score_list.get(segment.index_node_id, None),
+ }
+ if self.retriever_from == "dev":
+ source["hit_count"] = segment.hit_count
+ source["word_count"] = segment.word_count
+ source["segment_position"] = segment.position
+ source["index_node_hash"] = segment.index_node_hash
+ if segment.answer:
+ source["content"] = f"question:{segment.content} \nanswer:{segment.answer}"
+ else:
+ source["content"] = segment.content
+ context_list.append(source)
+ resource_number += 1
+
+ for hit_callback in self.hit_callbacks:
+ hit_callback.return_retriever_resource_info(context_list)
+
+ return str("\n".join(document_context_list))

api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

@@ -79,8 +79,9 @@ def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query:
 
  results = (
  db.session.query(Dataset)
- .join(subquery, Dataset.id == subquery.c.dataset_id)
+ .outerjoin(subquery, Dataset.id == subquery.c.dataset_id)
  .filter(Dataset.tenant_id == self.tenant_id, Dataset.id.in_(dataset_ids))
+ .filter((subquery.c.available_document_count > 0) | (Dataset.provider == "external"))
  .all()
  )
 
@@ -121,10 +122,13 @@ def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query:
  )
  elif node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE.value:
  if node_data.multiple_retrieval_config.reranking_mode == "reranking_model":
- reranking_model = {
- "reranking_provider_name": node_data.multiple_retrieval_config.reranking_model.provider,
- "reranking_model_name": node_data.multiple_retrieval_config.reranking_model.model,
- }
+ if node_data.multiple_retrieval_config.reranking_model:
+ reranking_model = {
+ "reranking_provider_name": node_data.multiple_retrieval_config.reranking_model.provider,
+ "reranking_model_name": node_data.multiple_retrieval_config.reranking_model.model,
+ }
+ else:
+ reranking_model = None
  weights = None
  elif node_data.multiple_retrieval_config.reranking_mode == "weighted_score":
  reranking_model = None

api/services/dataset_service.py

@@ -234,6 +234,7 @@ def update_dataset(dataset_id, data, user):
  dataset.name = data.get("name", dataset.name)
  dataset.description = data.get("description", "")
  external_knowledge_id = data.get("external_knowledge_id", None)
+ dataset.permission = data.get("permission")
  db.session.add(dataset)
  if not external_knowledge_id:
  raise ValueError("External knowledge id is required.")