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

backend/app/engine/index.py

@@ -1,15 +1,38 @@
-import logging
+import logging,os
 from llama_index.core.indices import VectorStoreIndex
-from app.engine.vectordb import get_vector_store
-from app.engine.loaders import get_document_Types
+from app.engine.vectordb import get_vector_store,get_Neo4j_Graph_Store
 from typing import Dict,Any
+from llama_index.core import PropertyGraphIndex
+from llama_index.core.storage.storage_context import StorageContext
+from llama_index.core import load_index_from_storage
+
 logger = logging.getLogger("uvicorn")
 
+
 def get_index(prjFlag:str):
     if prjFlag is None or prjFlag == '':
         raise ValueError('无效的工程标识')
     logger.info("Connecting vector store...")
-    store = get_vector_store(prjFlag)
-    index = VectorStoreIndex.from_vector_store(store)
+    index = None
+    llm_query = os.getenv('LLM_QUERY_WAY')
+    if llm_query == 'graph':
+       index = getPropertyGraphIndex(prjFlag)
+    else:
+        store = get_vector_store(prjFlag)
+        index = VectorStoreIndex.from_vector_store(store)
     logger.info("Finished load index from vector store.")
     return index
+
+
+def getPropertyGraphIndex(prjFlag:str):
+    GRAPH_STORE_TYPE = os.getenv("GRAPH_STORE_TYPE", "")
+    if GRAPH_STORE_TYPE == 'neo4j':
+        index = PropertyGraphIndex.from_existing(property_graph_store= get_Neo4j_Graph_Store(prjFlag))
+    else:
+        GRAPH_STORAGE_DIR = os.getenv("GRAPH_STORAGE_PATH", "storage_graph")
+        prjCachePath = GRAPH_STORAGE_DIR + f"/{prjFlag}"
+        if not os.path.exists(prjCachePath):
+            return None
+        storeContext = StorageContext.from_defaults(persist_dir = prjCachePath,vector_store = get_vector_store(prjFlag))
+        index = load_index_from_storage(storeContext)
+    return index