新增属性图谱

2024-09-14 09:44:43 +08:00
parent e18e4f22db
commit df7ebf3d4f
7 changed files with 316 additions and 29 deletions
@@ -0,0 +1,87 @@
+from llama_index.core.indices.property_graph import LLMSynonymRetriever,VectorContextRetriever,PGRetriever
+from llama_index.core.indices.property_graph.transformations.schema_llm import *
+from llama_index.core import SimpleDirectoryReader
+from llama_index.core import settings
+from llama_index.core import PropertyGraphIndex
+from typing import List,Tuple,Literal
+from app.settings import init_settings
+import os
+from llama_index.core.storage.storage_context import StorageContext
+from llama_index.core import load_index_from_storage
+from app.observability import init_observability
+from app.engine.vectordb import get_Neo4j_Graph_Store
+from llama_index.core.response_synthesizers import ResponseMode
+from util.register import *
+from llama_index.core.query_engine import RetrieverQueryEngine
+from app.engine.prompt import text_qa_template, refine_template, summary_template, simple_template
+from app.engine.engine import get_node_postprocessors
+
+class PropertyGraph:
+    def __init__(self,prjFlag:str) -> None:
+        self._prjFlag = prjFlag
+
+    def create_query_engine(self,retriever):
+        postprocess = get_node_postprocessors()
+        query_engine = RetrieverQueryEngine.from_args(
+            retriever = retriever,
+            text_qa_template=text_qa_template,
+            refine_template=refine_template,
+            summary_template = summary_template,
+            simple_template = simple_template,
+            node_postprocessors=postprocess,
+            use_async=True,
+            streaming=False,
+            response_mode = ResponseMode.TREE_SUMMARIZE
+        )
+        return query_engine
+    
+    def getPropertyGraphIndex(self):
+        GRAPH_STORE_TYPE = os.getenv("GRAPH_STORE_TYPE", "")
+        if GRAPH_STORE_TYPE == 'neo4j':
+            index = PropertyGraphIndex.from_existing(property_graph_store= get_Neo4j_Graph_Store(self._prjFlag))
+        else:
+            GRAPH_STORAGE_DIR = os.getenv("GRAPH_STORAGE_DIR", "storage_graph")
+            prjCachePath = GRAPH_STORAGE_DIR + f"/{self._prjFlag}"
+            if not os.path.exists(prjCachePath):
+                return None
+            storeContext = StorageContext.from_defaults(persist_dir = prjCachePath)
+            index = load_index_from_storage(storeContext)
+        return index
+    
+    def query(self,query_str:str):
+        index = self.getPropertyGraphIndex()
+        synonym_retriver = LLMSynonymRetriever(index.property_graph_store,
+                                               llm=settings.Settings.llm,
+                                               max_keywords=10,
+                                               include_text=False
+                                               )
+        if index.property_graph_store.supports_vector_queries:
+            vector_store = None
+        else:
+            vector_store = index.vector_store
+        vector_retriver = VectorContextRetriever(index.property_graph_store, 
+                                                vector_store = vector_store,
+                                                embed_model=settings.Settings.embed_model,
+                                                similarity_top_k=5,
+                                                include_text=False
+                                                )
+        
+        retriever = index.as_retriever(sub_retrievers=[synonym_retriver,vector_retriver])
+        query_engine = self.create_query_engine(retriever)
+
+        response = query_engine.query(query_str)
+        print(response)
+        return str(response)
+        
+        
+if __name__ == "__main__":
+    init_settings()
+    init_observability()
+    # graph = PropertyGraph('projects_1b20bbf4-3243-4ac3-bcf0-8a91e9157521')
+    # graph.query('代码为XLBT的金额是')
+            
+
+           
+
+
+