修改chunker接口，增加openapi

OpenSPG · Dec 25, 2024 · fd272d7 · fd272d7
1 parent 03bb962
commit fd272d7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/kag/solver/retriever/chunk_retriever.py b/kag/solver/retriever/chunk_retriever.py
@@ -12,19 +12,22 @@
 
 class ChunkRetriever(KagBaseModule, ABC):
     def __init__(self, recall_num: int = 10,
-                 rerank_topk: int = 10, **kwargs):
+                 rerank_topk: int = 10,
+                 graph_api: GraphApiABC = None,
+                 search_api: SearchApiABC = None,
+                 **kwargs):
         super().__init__(**kwargs)
         self.recall_num = recall_num
         self.rerank_topk = rerank_topk
         self.schema: SchemaUtils = SchemaUtils(LogicFormConfiguration({
             "KAG_PROJECT_ID": KAG_PROJECT_CONF.project_id,
             "KAG_PROJECT_HOST_ADDR": KAG_PROJECT_CONF.host_addr
         }))
-        self.graph_api = GraphApiABC.from_config({
+        self.graph_api = graph_api or GraphApiABC.from_config({
             "type": "openspg"}
         )
 
-        self.search_api = SearchApiABC.from_config({
+        self.search_api = search_api or SearchApiABC.from_config({
             "type": "openspg"
         })
 

diff --git a/kag/solver/retriever/impl/default_chunk_retrieval.py b/kag/solver/retriever/impl/default_chunk_retrieval.py
@@ -18,6 +18,9 @@
 
 import numpy as np
 import logging
+
+from kag.solver.tools.graph_api.graph_api_abc import GraphApiABC
+from kag.solver.tools.search_api.search_api_abc import SearchApiABC
 from knext.schema.client import CHUNK_TYPE, OTHER_TYPE
 from kag.common.utils import processing_phrases
 from kag.common.conf import KAG_CONFIG
@@ -53,9 +56,11 @@ def __init__(
             rerank_topk: int = 10,
             reranker_model_path: str = None,
             vectorize_model: Vectorizer = None,
+            graph_api: GraphApiABC = None,
+            search_api: SearchApiABC = None,
             **kwargs,
     ):
-        super().__init__(recall_num, rerank_topk, **kwargs)
+        super().__init__(recall_num, rerank_topk, graph_api, search_api, **kwargs)
         if vectorize_model is None:
             vectorize_model = Vectorizer.from_config(
                 KAG_CONFIG.all_config["vectorize_model"]
@@ -527,6 +532,8 @@ def __init__(
             rerank_topk: int = 10,
             reranker_model_path: str = None,
             vectorize_model: VectorizeModelABC = None,
+            graph_api: GraphApiABC = None,
+            search_api: SearchApiABC = None,
             **kwargs,
     ):
         super().__init__(
@@ -539,6 +546,8 @@ def __init__(
             rerank_topk,
             reranker_model_path,
             vectorize_model,
+            graph_api,
+            search_api,
             **kwargs
         )
         self.text_sim = TextSimilarity(vectorizer=self.vectorize_model)