Update app.py

app.py

@@ -8,16 +8,20 @@ from transformers import AutoModel, AutoTokenizer
 from diffusers import StableDiffusion3Pipeline
 from parler_tts import ParlerTTSForConditionalGeneration
 import soundfile as sf
-from llama_index.core.agent import ReActAgent
-from llama_index.core.tools import FunctionTool
-from llama_index.llms.groq import Groq
+from langchain.agents import AgentExecutor, create_react_agent
+from langchain.tools import BaseTool
+from langchain_groq import ChatGroq
 from PIL import Image
 from tavily import TavilyClient
 import requests
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from llama_index.core.chat_engine.types import AgentChatResponse
-from llama_index.core import VectorStoreIndex
+from langchain.schema import AIMessage
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.document_loaders import TextLoader
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.chains import RetrievalQA
 
 # Initialize models and clients
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
@@ -48,38 +52,71 @@ def play_voice_output(response):
     return "output.wav"
 
 # NumPy Code Calculator Tool
-def numpy_code_calculator(query):
-    try:
-        # Assume query is a request for a numpy computation
-        local_dict = {"np": np}
-        exec(query, local_dict)
-        result = local_dict.get("result", "No result found")
-        return str(result)
-    except Exception as e:
-        return f"Error: {e}"
+class NumpyCodeCalculator(BaseTool):
+    name = "Numpy"
+    description = "Useful for performing numpy computations"
+
+    def _run(self, query: str) -> str:
+        try:
+            local_dict = {"np": np}
+            exec(query, local_dict)
+            result = local_dict.get("result", "No result found")
+            return str(result)
+        except Exception as e:
+            return f"Error: {e}"
 
 # Web Search Tool
-def web_search(query):
-    answer = tavily_client.qna_search(query=query)
-    return answer
+class WebSearch(BaseTool):
+    name = "Web"
+    description = "Useful for searching the web for information"
+
+    def _run(self, query: str) -> str:
+        answer = tavily_client.qna_search(query=query)
+        return answer
 
 # Image Generation Tool
-def image_generation(query):
-    image = pipe(
-        query,
-        negative_prompt="",
-        num_inference_steps=15,
-        guidance_scale=7.0,
-    ).images[0]
-    image.save("output.jpg")
-    return "output.jpg"
+class ImageGeneration(BaseTool):
+    name = "Image"
+    description = "Useful for generating images based on text descriptions"
+
+    def _run(self, query: str) -> str:
+        image = pipe(
+            query,
+            negative_prompt="",
+            num_inference_steps=15,
+            guidance_scale=7.0,
+        ).images[0]
+        image.save("output.jpg")
+        return "output.jpg"
 
 # Document Question Answering Tool
-def document_question_answering(query, docs):
-    index = VectorStoreIndex.from_documents(docs)
-    query_engine = index.as_query_engine(similarity_top_k=3)
-    response = query_engine.query(query)
-    return str(response)
+class DocumentQuestionAnswering(BaseTool):
+    name = "Document"
+    description = "Useful for answering questions about a specific document"
+
+    def __init__(self, document):
+        super().__init__()
+        self.document = document
+        self.qa_chain = self._setup_qa_chain()
+
+    def _setup_qa_chain(self):
+        loader = TextLoader(self.document)
+        documents = loader.load()
+        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+        texts = text_splitter.split_documents(documents)
+        embeddings = HuggingFaceEmbeddings()
+        db = FAISS.from_documents(texts, embeddings)
+        retriever = db.as_retriever()
+        qa_chain = RetrievalQA.from_chain_type(
+            llm=ChatGroq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY")),
+            chain_type="stuff",
+            retriever=retriever,
+        )
+        return qa_chain
+
+    def _run(self, query: str) -> str:
+        response = self.qa_chain.run(query)
+        return str(response)
 
 # Function to handle different input types and choose the right tool
 def handle_input(user_prompt, image=None, audio=None, websearch=False, document=None):
@@ -93,43 +130,38 @@ def handle_input(user_prompt, image=None, audio=None, websearch=False, document=
         user_prompt = transcription.text
 
     tools = [
-        FunctionTool.from_defaults(fn=numpy_code_calculator, name="Numpy"),
-        FunctionTool.from_defaults(fn=image_generation, name="Image"),
+        NumpyCodeCalculator(),
+        ImageGeneration(),
     ]
 
     # Add the web search tool only if websearch mode is enabled
     if websearch:
-        tools.append(FunctionTool.from_defaults(fn=web_search, name="Web"))
+        tools.append(WebSearch())
 
     # Add the document question answering tool only if a document is provided
     if document:
-        docs = LlamaParse(result_type="text").load_data(document)
-        tools.append(FunctionTool.from_defaults(fn=document_question_answering, name="Document", docs=docs))
+        tools.append(DocumentQuestionAnswering(document))
 
-    llm = Groq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY"))
-    agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)
+    llm = ChatGroq(model=MODEL, api_key=os.environ.get("GROQ_API_KEY"))
+    agent = create_react_agent(llm, tools, verbose=True)
+    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
     if image:
         image = Image.open(image).convert('RGB')
         messages = [{"role": "user", "content": [image, user_prompt]}]
         response = vqa_model.chat(image=None, msgs=messages, tokenizer=tokenizer)
     else:
-        response = agent.chat(user_prompt)
-
-    # Extract the content from AgentChatResponse to return as a string
-    if isinstance(response, AgentChatResponse):
-        response = response.response
+        response = agent_executor.run(user_prompt)
 
     return response
 
-
-# Gradio UI Setup
+    
 def create_ui():
     with gr.Blocks(css="""
         /* Overall Styling */
         body {
-            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-            background-color: #f4f4f4;
+            font-family: 'Poppins', sans-serif;
+            background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
             margin: 0;
             padding: 0;
             color: #333;
@@ -139,8 +171,14 @@ def create_ui():
         .gradio-container h1 {
             text-align: center;
             padding: 20px 0;
-            background-color: #007bff; /* Example color */
+            background: linear-gradient(45deg, #007bff, #00c6ff);
             color: white;
+            font-size: 2.5em;
+            font-weight: bold;
+            letter-spacing: 1px;
+            text-transform: uppercase;
+            margin: 0;
+            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.2);
         }
 
         /* Input Area Styling */
@@ -149,6 +187,10 @@ def create_ui():
             justify-content: space-around;
             align-items: center;
             padding: 20px;
+            background-color: white;
+            border-radius: 10px;
+            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.1);
+            margin-bottom: 20px;
         }
 
         .gradio-container .gr-column {
@@ -159,40 +201,135 @@ def create_ui():
         /* Textbox Styling */
         .gradio-container textarea {
             width: calc(100% - 20px);
-            padding: 10px;
-            border: 2px solid #ccc;
-            border-radius: 5px;
-            font-size: 16px;
+            padding: 15px;
+            border: 2px solid #007bff;
+            border-radius: 8px;
+            font-size: 1.1em;
+            transition: border-color 0.3s, box-shadow 0.3s;
+        }
+
+        .gradio-container textarea:focus {
+            border-color: #00c6ff;
+            box-shadow: 0px 0px 8px rgba(0, 198, 255, 0.5);
+            outline: none;
         }
 
         /* Button Styling */
         .gradio-container button {
-            background-color: #007bff; /* Example color */
+            background: linear-gradient(45deg, #007bff, #00c6ff);
             color: white;
-            padding: 12px 20px;
+            padding: 15px 25px;
             border: none;
-            border-radius: 5px;
+            border-radius: 8px;
             cursor: pointer;
-            font-size: 16px;
-            transition: background-color 0.3s;
+            font-size: 1.2em;
+            font-weight: bold;
+            transition: background 0.3s, transform 0.3s;
+            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
         }
 
         .gradio-container button:hover {
-            background-color: #0056b3; /* Example darker color */
+            background: linear-gradient(45deg, #0056b3, #009bff);
+            transform: translateY(-3px);
+        }
+
+        .gradio-container button:active {
+            transform: translateY(0);
         }
 
         /* Output Area Styling */
         .gradio-container .output-area {
             padding: 20px;
             text-align: center;
+            background-color: #f7f9fc;
+            border-radius: 10px;
+            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.1);
+            margin-top: 20px;
         }
 
         /* Image Styling */
         .gradio-container img {
             max-width: 100%;
             height: auto;
-            border-radius: 5px;
+            border-radius: 10px;
+            box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
+            transition: transform 0.3s, box-shadow 0.3s;
+        }
+
+        .gradio-container img:hover {
+            transform: scale(1.05);
+            box-shadow: 0px 6px 12px rgba(0, 0, 0, 0.2);
+        }
+
+        /* Checkbox Styling */
+        .gradio-container input[type="checkbox"] {
+            width: 20px;
+            height: 20px;
+            cursor: pointer;
+            accent-color: #007bff;
+            transition: transform 0.3s;
+        }
+
+        .gradio-container input[type="checkbox"]:checked {
+            transform: scale(1.2);
+        }
+
+        /* Audio and Document Upload Styling */
+        .gradio-container .gr-file-upload input[type="file"] {
+            width: 100%;
+            padding: 10px;
+            border: 2px solid #007bff;
+            border-radius: 8px;
+            cursor: pointer;
+            background-color: white;
+            transition: border-color 0.3s, background-color 0.3s;
+        }
+
+        .gradio-container .gr-file-upload input[type="file"]:hover {
+            border-color: #00c6ff;
+            background-color: #f0f8ff;
+        }
+
+        /* Advanced Tooltip Styling */
+        .gradio-container .gr-tooltip {
+            position: relative;
+            display: inline-block;
+            cursor: pointer;
+        }
+
+        .gradio-container .gr-tooltip .tooltiptext {
+            visibility: hidden;
+            width: 200px;
+            background-color: black;
+            color: #fff;
+            text-align: center;
+            border-radius: 6px;
+            padding: 5px;
+            position: absolute;
+            z-index: 1;
+            bottom: 125%;
+            left: 50%;
+            margin-left: -100px;
+            opacity: 0;
+            transition: opacity 0.3s;
+        }
+
+        .gradio-container .gr-tooltip:hover .tooltiptext {
+            visibility: visible;
+            opacity: 1;
         }
+
+        /* Footer Styling */
+        .gradio-container footer {
+            text-align: center;
+            padding: 10px;
+            background: #007bff;
+            color: white;
+            font-size: 0.9em;
+            border-radius: 0 0 10px 10px;
+            box-shadow: 0px -2px 8px rgba(0, 0, 0, 0.1);
+        }
+
     """) as demo:
         gr.Markdown("# AI Assistant")
         with gr.Row():
@@ -257,7 +394,6 @@ def main_interface(user_prompt, image=None, audio=None, voice_only=False, websea
     else:
         return response, None
 
-
 # Launch the UI
 demo = create_ui()
 demo.launch()