Update app.py

app.py

@@ -17,7 +17,20 @@ graph = Neo4jGraph(
     password="Z10duoPkKCtENuOukw3eIlvl0xJWKtrVSr-_hGX1LQ4"
 )
 
-# Function to clean input for Neo4j full-text query
+# Define entity extraction and retrieval functions
+class Entities(BaseModel):
+    names: List[str] = Field(
+        ..., description="All the person, organization, or business entities that appear in the text"
+    )
+
+entity_prompt = ChatPromptTemplate.from_messages([
+    ("system", "You are extracting organization and person entities from the text."),
+    ("human", "Use the given format to extract information from the following input: {question}"),
+])
+
+chat_model = ChatOpenAI(temperature=0, model_name="gpt-4o", api_key=os.environ['OPENAI_API_KEY'])
+entity_chain = entity_prompt | chat_model.with_structured_output(Entities)
+
 def remove_lucene_chars(input: str) -> str:
     return input.translate(str.maketrans({
         "\\": r"\\", "+": r"\+", "-": r"\-", "&": r"\&", "|": r"\|", "!": r"\!",
@@ -26,7 +39,6 @@ def remove_lucene_chars(input: str) -> str:
         ";": r"\;", " ": r"\ "
     }))
 
-# Function to generate a full-text query
 def generate_full_text_query(input: str) -> str:
     full_text_query = ""
     words = [el for el in remove_lucene_chars(input).split() if el]
@@ -35,29 +47,100 @@ def generate_full_text_query(input: str) -> str:
     full_text_query += f" {words[-1]}~2"
     return full_text_query.strip()
 
-# Define the function to query Neo4j and get a response
-def get_response(question):
-    query = generate_full_text_query(question)
-    try:
-        # Query the Neo4j database using a full-text search
+# Setup logging to a file to capture debug information
+logging.basicConfig(filename='neo4j_retrieval.log', level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+
+def structured_retriever(question: str) -> str:
+    result = ""
+    entities = entity_chain.invoke({"question": question})
+    for entity in entities.names:
         response = graph.query(
-            """
-            CALL db.index.fulltext.queryNodes('entity', $query)
-            YIELD node, score
-            RETURN node.content AS content, score
-            ORDER BY score DESC LIMIT 1
+            """CALL db.index.fulltext.queryNodes('entity', $query, {limit:2})
+            YIELD node,score
+            CALL {
+              WITH node
+              MATCH (node)-[r:!MENTIONS]->(neighbor)
+              RETURN node.id + ' - ' + type(r) + ' -> ' + neighbor.id AS output
+              UNION ALL
+              WITH node
+              MATCH (node)<-[r:!MENTIONS]-(neighbor)
+              RETURN neighbor.id + ' - ' + type(r) + ' -> ' +  node.id AS output
+            }
+            RETURN output LIMIT 50
             """,
-            {"query": query}
+            {"query": generate_full_text_query(entity)},
         )
-        # Extract the content from the top response
-        if response:
-            result = response[0]['content']
-            return result
-        else:
-            return "Sorry, I couldn't find any relevant information in the database."
+        result += "\n".join([el['output'] for el in response])
+    return result
+
+def retriever_neo4j(question: str):
+    structured_data = structured_retriever(question)
+    logging.debug(f"Structured data: {structured_data}")
+    return structured_data
+
+# Setup for condensing the follow-up questions
+_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"]),
+)
+
+
+template = """I am a guide for Birmingham, Alabama. I can provide recommendations and insights about the city, including events and activities.
+Ask your question directly, and I'll provide a precise and quick,short and crisp response in a conversational way without any Greet.
+{context}
+Question: {question}
+Answer:"""
+
+
+qa_prompt = ChatPromptTemplate.from_template(template)
+
+# Define the chain for Neo4j-based retrieval and response generation
+chain_neo4j = (
+    RunnableParallel(
+        {
+            "context": RunnableLambda(lambda x: retriever_neo4j(x["question"])),
+            "question": RunnablePassthrough(),
+        }
+    )
+    | ChatPromptTemplate.from_template("Answer: {context} Question: {question}")
+    | chat_model
+    | StrOutputParser()
+)
+
+# Define the function to query Neo4j and get a response
+def get_response(question):
+    try:
+        return chain_neo4j.invoke({"question": question})
     except Exception as e:
-        logging.error(f"Error querying Neo4j: {e}")
-        return "An error occurred while fetching data from the database."
+        return f"Error: {str(e)}"
+
+
 
 # Function to generate audio with Eleven Labs TTS
 def generate_audio_elevenlabs(text):
@@ -131,8 +214,8 @@ def handle_voice_to_voice(audio):
 
 
 # Define the Gradio interface
-with gr.Blocks() as demo:
-    audio_input = gr.Audio(sources=["microphone"], type='numpy', streaming=False, label="Speak to Ask")
+with gr.Blocks(theme="rawrsor1/Everforest") as demo:
+    audio_input = gr.Audio(sources=["microphone"], type='numpy', streaming=True, label="Speak to Ask")
     submit_voice_btn = gr.Button("Submit Voice")
     audio_output = gr.Audio(label="Response Audio", type="filepath", autoplay=True, interactive=False)