Update app.py

app.py

@@ -1,39 +1,15 @@
 import gradio as gr
 import os
 import logging
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from langchain_openai import ChatOpenAI
-from langchain_community.graphs import Neo4jGraph
-from typing import List, Tuple
-from pydantic import BaseModel, Field
-from langchain_core.messages import AIMessage, HumanMessage
-from langchain_core.runnables import (
-    RunnableBranch,
-    RunnableLambda,
-    RunnablePassthrough,
-    RunnableParallel,
-)
-from langchain_core.prompts.prompt import PromptTemplate
 import requests
 import tempfile
-from langchain.memory import ConversationBufferWindowMemory
-import time
-import logging
-from langchain.chains import ConversationChain
+from langchain_openai import ChatOpenAI
+from langchain_community.graphs import Neo4jGraph
 import torch
-import torchaudio
-from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 import numpy as np
+from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor
 import threading
 
-# Setup conversational memory
-conversational_memory = ConversationBufferWindowMemory(
-    memory_key='chat_history',
-    k=10,
-    return_messages=True
-)
-
 # Setup Neo4j connection
 graph = Neo4jGraph(
     url="neo4j+s://6457770f.databases.neo4j.io",
@@ -41,20 +17,7 @@ graph = Neo4jGraph(
     password="Z10duoPkKCtENuOukw3eIlvl0xJWKtrVSr-_hGX1LQ4"
 )
 
-# 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)
-
+# Function to clean input for Neo4j full-text query
 def remove_lucene_chars(input: str) -> str:
     return input.translate(str.maketrans({
         "\\": r"\\", "+": r"\+", "-": r"\-", "&": r"\&", "|": r"\|", "!": r"\!",
@@ -63,6 +26,7 @@ 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]
@@ -71,60 +35,29 @@ def generate_full_text_query(input: str) -> str:
     full_text_query += f" {words[-1]}~2"
     return full_text_query.strip()
 
-# 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:
+# 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
         response = graph.query(
-            """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
+            """
+            CALL db.index.fulltext.queryNodes('entity', $query)
+            YIELD node, score
+            RETURN node.content AS content, score
+            ORDER BY score DESC LIMIT 1
             """,
-            {"query": generate_full_text_query(entity)},
+            {"query": query}
         )
-        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
-
-# 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 get the response
-def get_response(question):
-    try:
-        return chain_neo4j.invoke({"question": question})
+        # 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."
     except Exception as e:
-        return f"Error: {str(e)}"
-
-# Define the function to clear input and output
-def clear_fields():
-    return [], "", None
+        logging.error(f"Error querying Neo4j: {e}")
+        return "An error occurred while fetching data from the database."
 
 # Function to generate audio with Eleven Labs TTS
 def generate_audio_elevenlabs(text):
@@ -152,75 +85,56 @@ def generate_audio_elevenlabs(text):
                 if chunk:
                     f.write(chunk)
             audio_path = f.name
-        logging.debug(f"Audio saved to {audio_path}")
-        return audio_path  # Return audio path for automatic playback
+        return audio_path
     else:
-        logging.error(f"Error generating audio: {response.text}")
         return None
 
-# Function to handle voice to voice conversation
-def handle_voice_to_voice(chat_history, question):
-    response = get_response(question)
-    audio_path = generate_audio_elevenlabs(response)
-    chat_history.append(("[Voice Input]", "[Voice Response]"))
-    return chat_history, "", audio_path
-
-# Function to transcribe audio input
-def transcribe_function(stream, new_chunk):
-    try:
-        sr, y = new_chunk[0], new_chunk[1]
-    except TypeError:
-        print(f"Error chunk structure: {type(new_chunk)}, content: {new_chunk}")
-        return stream, "", None
-
-    if y is None or len(y) == 0:
-        return stream, "", None
-
-    y = y.astype(np.float32)
-    max_abs_y = np.max(np.abs(y))
-    if max_abs_y > 0:
-        y = y / max_abs_y
-
-    if stream is not None and len(stream) > 0:
-        stream = np.concatenate([stream, y])
-    else:
-        stream = y
+# Define ASR model for speech-to-text
+model_id = 'openai/whisper-large-v3'
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype).to(device)
+processor = AutoProcessor.from_pretrained(model_id)
+
+pipe_asr = pipeline(
+    "automatic-speech-recognition",
+    model=model,
+    tokenizer=processor.tokenizer,
+    feature_extractor=processor.feature_extractor,
+    max_new_tokens=128,
+    chunk_length_s=15,
+    batch_size=16,
+    torch_dtype=torch_dtype,
+    device=device,
+    return_timestamps=True
+)
 
-    result = pipe_asr({"array": stream, "sampling_rate": sr}, return_timestamps=False)
-    full_text = result.get("text", "")
+# Function to handle voice input, generate response from Neo4j, and return audio output
+def handle_voice_to_voice(audio):
+    # Transcribe audio input to text
+    sr, y = audio
+    result = pipe_asr({"array": y, "sampling_rate": sr}, return_timestamps=False)
+    question = result.get("text", "")
 
-    threading.Thread(target=auto_reset_state).start()
+    # Get response using the transcribed question
+    response = get_response(question)
 
-    return stream, full_text, full_text
+    # Generate audio from the response
+    audio_path = generate_audio_elevenlabs(response)
+    return audio_path
 
 # Define the Gradio interface
-with gr.Blocks(theme="rawrsor1/Everforest") as demo:
-    chatbot = gr.Chatbot([], elem_id="RADAR", bubble_full_width=False)
-    mode_selection = gr.Radio(
-        choices=["Normal Chatbot", "Voice to Voice Conversation"],
-        label="Mode Selection",
-        value="Normal Chatbot"
-    )
-    question_input = gr.Textbox(label="Ask a Question", placeholder="Type your question here...")
-    audio_input = gr.Audio(sources=["microphone"], streaming=True, type='numpy', every=0.1, label="Speak to Ask")
+with gr.Blocks() as demo:
+    audio_input = gr.Audio(sources=["microphone"], type='numpy', streaming=False, label="Speak to Ask")
     submit_voice_btn = gr.Button("Submit Voice")
-    audio_output = gr.Audio(label="Audio", type="filepath", autoplay=True, interactive=False)
+    audio_output = gr.Audio(label="Response Audio", type="filepath", autoplay=True, interactive=False)
 
     # Interactions for Submit Voice Button
     submit_voice_btn.click(
         fn=handle_voice_to_voice,
-        inputs=[chatbot, question_input],
-        outputs=[chatbot, question_input, audio_output],
-        api_name="api_voice_to_voice_translation"
-    )
-
-    # Speech-to-Text functionality
-    state = gr.State()
-    audio_input.stream(
-        transcribe_function,
-        inputs=[state, audio_input],
-        outputs=[state, question_input],
-        api_name="api_voice_to_text"
+        inputs=audio_input,
+        outputs=audio_output
     )
 
-    demo.launch(show_error=True, share=True)
+# Launch the Gradio interface
+demo.launch(show_error=True, share=True)