Update app.py

app.py

@@ -1015,27 +1015,37 @@ def retriever_neo4j(question: str):
     structured_data = structured_retriever(question)
     return structured_data
 
-template = """Answer the question based only on the following context:
-{context}
-Question: {question}
-Use natural language and be concise.
-Answer:"""
-
-qa_prompt = ChatPromptTemplate.from_template(template)
-
-chain_neo4j = (
-    RunnableParallel(
-        {
-            "context": _search_query | retriever_neo4j,
-            "question": RunnablePassthrough(),
-        }
-    )
-    | qa_prompt
-    | chat_model
-    | StrOutputParser()
+_template = """Given the following conversation and a follow-up question, rephrase the follow-up question to be a standalone question,
+in its original language.
+Chat History:
+{chat_history}
+Follow Up Input: {question}
+Standalone question:"""
+
+CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
+
+def _format_chat_history(chat_history: list[tuple[str, str]]) -> list:
+    buffer = []
+    for human, ai in chat_history:
+        buffer.append(HumanMessage(content=human))
+        buffer.append(AIMessage(content=ai))
+    return buffer
+
+_search_query = RunnableBranch(
+    (
+        RunnableLambda(lambda x: bool(x.get("chat_history"))).with_config(
+            run_name="HasChatHistoryCheck"
+        ),
+        RunnablePassthrough.assign(
+            chat_history=lambda x: _format_chat_history(x["chat_history"])
+        )
+        | CONDENSE_QUESTION_PROMPT
+        | ChatOpenAI(temperature=0, api_key=os.environ['OPENAI_API_KEY'])
+        | StrOutputParser(),
+    ),
+    RunnableLambda(lambda x : x["question"]),
 )
 
-# Define a function to select between Pinecone and Neo4j
 def generate_answer(message, choice, retrieval_mode):
     logging.debug(f"generate_answer called with choice: {choice} and retrieval_mode: {retrieval_mode}")
 
@@ -1049,12 +1059,21 @@ def generate_answer(message, choice, retrieval_mode):
             chain_type_kwargs={"prompt": prompt_template}
         )
         response = qa_chain({"query": message})
-        return response['result'], extract_addresses(response['result'])
+        return response['output'], extract_addresses(response['output'])
     elif retrieval_mode == "Knowledge-Graph":
-        context = retriever_neo4j(message)
-        qa_chain = ChatPromptTemplate.from_template(prompt_template.template)
-        response = qa_chain.invoke({"context": context, "question": message})
-        return response['result'], extract_addresses(response['result'])
+        chain_neo4j = (
+            RunnableParallel(
+                {
+                    "context": _search_query | retriever_neo4j,
+                    "question": RunnablePassthrough(),
+                }
+            )
+            | prompt_template
+            | chat_model
+            | StrOutputParser()
+        )
+        response = chain_neo4j.invoke({"question": message})
+        return response['output'], extract_addresses(response['output'])
     else:
         return "Invalid retrieval mode selected.", []
 
@@ -1255,7 +1274,7 @@ def show_map_if_details(history, choice):
     if choice in ["Details", "Conversational"]:
         return gr.update(visible=True), update_map_with_response(history)
     else:
-        return gr.update(visible(False)), ""
+        return gr.update(visible=False), ""
 
 def generate_audio_elevenlabs(text):
     XI_API_KEY = os.environ['ELEVENLABS_API']