Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import os
 import numpy as np
 from groq import Groq
-import spaces
+import spaces  # Import spaces
 from transformers import AutoModel, AutoTokenizer
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from parler_tts import ParlerTTSForConditionalGeneration
@@ -11,9 +11,9 @@ import soundfile as sf
 from langchain_community.embeddings import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
-from langchain.agents import initialize_agent, Tool
-from langchain.llms import OpenAI
+from langchain.chains import RetrievalQA, LLMChain
+from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
+from langchain.llms import Groq as GroqLlm  # Import GroqLlm
 from PIL import Image
 from decord import VideoReader, cpu
 from tavily import TavilyClient
@@ -24,6 +24,7 @@ from safetensors.torch import load_file
 # Initialize models and clients
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
+llm = GroqLlm(client=client, model=MODEL)  # Initialize GroqLlm
 
 vqa_model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2', trust_remote_code=True,
                                        device_map="auto", torch_dtype=torch.bfloat16)
@@ -43,7 +44,7 @@ image_pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dt
 image_pipe.scheduler = EulerDiscreteScheduler.from_config(image_pipe.scheduler.config, timestep_spacing="trailing")
 
 # Tavily Client for web search
-tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API"))
+tavily_client = TavilyClient(api_key=os.environ.get("TAVILY_API_KEY"))  # Corrected API key
 
 # Function to play voice output
 def play_voice_output(response):
@@ -59,8 +60,6 @@ def play_voice_output(response):
 def numpy_code_calculator(query):
     """Generates and executes NumPy code for mathematical operations."""
     try:
-        # You might need to use a more sophisticated approach to generate NumPy code
-        # based on the user's query. This is a simple example.
         llm_response = client.chat.completions.create(
             model=MODEL,
             messages=[
@@ -68,7 +67,7 @@ def numpy_code_calculator(query):
             ]
         )
         code = llm_response.choices[0].message.content
-        print(f"Generated NumPy code:\n{code}")  # Print the generated code
+        print(f"Generated NumPy code:\n{code}")
 
         # Execute the code in a safe environment
         local_dict = {"np": np}
@@ -104,9 +103,12 @@ def doc_question_answering(query, file_path):
     return qa.run(query)
 
 # Function to handle different input types and choose the right tool
-def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, websearch=False):
+def handle_input(user_prompt, image=None, audio=None, doc=None, websearch=False):
     # Voice input handling
     if audio:
+        # Make sure 'audio' is a file object
+        if isinstance(audio, str):
+            audio = open(audio, "rb") 
         transcription = client.audio.transcriptions.create(
             file=(audio.name, audio.read()),
             model="whisper-large-v3"
@@ -143,12 +145,8 @@ def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, webs
         )
 
     # Initialize agent
-    agent = initialize_agent(
-        tools,
-        client,
-        agent="zero-shot-react-description",
-        verbose=True,
-    )
+    agent = ZeroShotAgent(llm_chain=LLMChain(llm=llm, prompt=None), tools=tools, verbose=True)
+    agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
 
     # If user uploaded an image and text, use MiniCPM model
     if image:
@@ -159,9 +157,9 @@ def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, webs
 
     # Use the agent to determine the best tool and get the response
     if websearch:
-        response = agent.run(f"{user_prompt} Use the Web Search tool if necessary.")
+        response = agent_executor.run(f"{user_prompt} Use the Web Search tool if necessary.")
     else:
-        response = agent.run(user_prompt)
+        response = agent_executor.run(user_prompt)
 
     return response
 
@@ -209,7 +207,7 @@ def create_ui():
 def main_interface(user_prompt, image=None, audio=None, doc=None, voice_only=False, websearch=False):
     vqa_model.to(device='cuda', dtype=torch.bfloat16)
     tts_model.to("cuda")
-    unet.to("cuda", torch.float16)
+    unet.to("cuda")
     image_pipe.to("cuda")
 
     response = handle_input(user_prompt, image=image, audio=audio, doc=doc, websearch=websearch)