Update app.py

app.py

@@ -121,20 +121,19 @@ graph = Neo4jGraph(
     password="XCSXe1Jl_gjyJqoBGXDqY1UrfgDc4Z_RT5YGrxPAy-g"
 )
 
-dataset_name = "Pijush2023/birmindata07312024"
-page_content_column = 'events_description'
-loader = HuggingFaceDatasetLoader(dataset_name, page_content_column)
-data = loader.load()
+# Avoid pushing the graph documents to Neo4j every time
+# Only push the documents once and comment the code below after the initial push
+# dataset_name = "Pijush2023/birmindata07312024"
+# page_content_column = 'events_description'
+# loader = HuggingFaceDatasetLoader(dataset_name, page_content_column)
+# data = loader.load()
 
-text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=50)
-documents = text_splitter.split_documents(data)
+# text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=50)
+# documents = text_splitter.split_documents(data)
 
-embeddings = OpenAIEmbeddings(api_key=os.environ['OPENAI_API_KEY'])
-llm = ChatOpenAI(temperature=0, model='gpt-4o', api_key=os.environ['OPENAI_API_KEY'])
-
-llm_transformer = LLMGraphTransformer(llm=llm)
-graph_documents = llm_transformer.convert_to_graph_documents(documents)
-graph.add_graph_documents(graph_documents, baseEntityLabel=True, include_source=True)
+# llm_transformer = LLMGraphTransformer(llm=llm)
+# graph_documents = llm_transformer.convert_to_graph_documents(documents)
+# graph.add_graph_documents(graph_documents, baseEntityLabel=True, include_source=True)
 
 class Entities(BaseModel):
     names: list[str] = Field(..., description="All the person, organization, or business entities that appear in the text")
@@ -243,19 +242,23 @@ chain_neo4j = (
 def generate_answer(message, choice, retrieval_mode):
     logging.debug(f"generate_answer called with choice: {choice} and retrieval_mode: {retrieval_mode}")
 
+    prompt_template = QA_CHAIN_PROMPT_1 if choice == "Details" else QA_CHAIN_PROMPT_2
+
     if retrieval_mode == "Vector":
-        qa_chain = build_qa_chain(QA_CHAIN_PROMPT_1 if choice == "Details" else QA_CHAIN_PROMPT_2)
-        agent = initialize_agent_with_prompt(QA_CHAIN_PROMPT_1 if choice == "Details" else QA_CHAIN_PROMPT_2)
-        response = agent(message)
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=chat_model,
+            chain_type="stuff",
+            retriever=retriever,
+            chain_type_kwargs={"prompt": prompt_template}
+        )
+        response = qa_chain({"query": message})
     elif retrieval_mode == "Knowledge-Graph":
         response = chain_neo4j.invoke({"question": message})
     else:
         response = "Invalid retrieval mode selected."
 
-    addresses = extract_addresses(response['output'])
-    return response['output'], addresses
+    return response['output'], extract_addresses(response['output'])
 
-# The rest of your Gradio code...
 def bot(history, choice, tts_choice, retrieval_mode):
     if not history:
         return history