Update app.py

app.py

@@ -14,13 +14,11 @@ from langchain_experimental.graph_transformers import LLMGraphTransformer
 from langchain.chains.graph_qa.cypher import GraphCypherQAChain
 from neo4j import GraphDatabase
 
-# Add Llama-Index imports
+# Llama-Index imports
 from llama_index.core import SimpleDirectoryReader, KnowledgeGraphIndex, Settings
 from llama_index.core.graph_stores import SimpleGraphStore
 from llama_index.core import StorageContext
-from llama_index.llms.huggingface import HuggingFaceInferenceAPI
-from langchain.embeddings import HuggingFaceEmbeddings
-from llama_index.embeddings.langchain import LangchainEmbedding
+from llama_index.embeddings import OpenAIEmbedding
 
 def main():
     st.set_page_config(
@@ -117,7 +115,7 @@ def main():
                     st.write("PDF file uploaded and saved to temporary file.")
 
                 # Process document using Llama-Index
-                index = process_document(tmp_file_path, graph)
+                index = process_document(tmp_file_path, graph, st.session_state['OPENAI_API_KEY'])
                 
                 # Store the index in session state
                 st.session_state['index'] = index
@@ -188,41 +186,31 @@ def main():
                 res = st.session_state['qa'].invoke({"query": question})
                 st.write("\n**Answer:**\n" + res['result'])
 
-def process_document(file_path, graph):
-    # Initialize Llama-Index components
+def process_document(file_path, graph, openai_api_key):
+    # Configure OpenAI
+    os.environ["OPENAI_API_KEY"] = openai_api_key
     Settings.chunk_size = 512
+    Settings.llm = ChatOpenAI(temperature=0, model="gpt-4")
+    
+    # Setup embeddings
+    embed_model = OpenAIEmbedding()
     
     # Create graph store
     graph_store = SimpleGraphStore()
     storage_context = StorageContext.from_defaults(graph_store=graph_store)
     
-    # Load document
-    documents = SimpleDirectoryReader(file_path).load_data()
+    # Load and process document
+    documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
     
     # Create Knowledge Graph Index
     index = KnowledgeGraphIndex.from_documents(
         documents=documents,
         max_triplets_per_chunk=3,
         storage_context=storage_context,
+        embed_model=embed_model,
         include_embeddings=True
     )
     
-    # Convert to Neo4j
-    g = index.get_networkx_graph()
-    for node in g.nodes():
-        cypher = f"""
-        CREATE (n:{node['type']} {{id: '{node['id']}', text: '{node['text']}'}})
-        """
-        graph.query(cypher)
-    
-    for edge in g.edges():
-        cypher = f"""
-        MATCH (a), (b)
-        WHERE a.id = '{edge[0]}' AND b.id = '{edge[1]}'
-        CREATE (a)-[r:{edge['relationship']}]->(b)
-        """
-        graph.query(cypher)
-    
     return index
 
 if __name__ == "__main__":