将属性图流程同步至主流程中

backend/app/engine/engine.py

@@ -12,6 +12,8 @@ from app.engine.prompt import text_qa_template, refine_template, summary_templat
 from app.engine.retriever.HybridRetriever import HybridRetriever
 from app.engine.response.treeSummResponse import CustomTreeResponse
 from llama_index.core.settings import Settings
+from llama_index.core.indices.property_graph import LLMSynonymRetriever,VectorContextRetriever
+from llama_index.core import PropertyGraphIndex
 
 ModelPlateCategory = '模型平台'
 
@@ -111,16 +113,38 @@ def create_summary_query_engine(index, top_k=3, use_reranker=False, filters=None
     return summary_query_engine
 
 # Create a query engine
-def create_query_engine(index, top_k=3, use_reranker=False, filters=None, response_mode=None):
+def create_query_engine(index,top_k=3, use_reranker=False, filters=None, response_mode=None):
     # 创建向量检索查询工具
     postprocess = None
     if use_reranker:
         postprocess = get_node_postprocessors()
 
+    llm_query = os.getenv('LLM_QUERY_WAY','rag')
+    if llm_query == 'graph':
+        graphIndex:PropertyGraphIndex = index
+        synonym_retriver = LLMSynonymRetriever(graphIndex.property_graph_store,
+                                               llm=Settings.llm,
+                                               include_text=False
+                                               )
+        if graphIndex.property_graph_store.supports_vector_queries:
+            vector_store = None
+        else:
+            vector_store = graphIndex.vector_store
+        vector_retriver = VectorContextRetriever(graphIndex.property_graph_store, 
+                                                vector_store = vector_store,
+                                                embed_model=Settings.embed_model,
+                                                similarity_top_k=top_k,
+                                                include_text=False
+                                                )
+        
+        retriever = graphIndex.as_retriever(sub_retrievers=[synonym_retriver,vector_retriver])
+
+    else:
+        retriever = get_Retriever(index,
+                    similarity_top_k=top_k,
+                    filters=filters),
     query_engine = RetrieverQueryEngine.from_args(
-        get_Retriever(index,
-                      similarity_top_k=top_k,
-                      filters=filters),
+        retriever = retriever,
         text_qa_template=text_qa_template,
         refine_template=refine_template,
         summary_template = summary_template,
@@ -131,4 +155,4 @@ def create_query_engine(index, top_k=3, use_reranker=False, filters=None, respon
         response_mode = response_mode
     )
 
-    return query_engine
+    return query_engine