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
@@ -13,7 +13,6 @@ from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 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,7 +23,6 @@ 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)
@@ -103,7 +101,7 @@ 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, audio=None, doc=None, websearch=False):
+def handle_input(user_prompt, image=None, video=None, audio=None, doc=None, websearch=False):
     # Voice input handling
     if audio:
         # Make sure 'audio' is a file object
@@ -144,8 +142,16 @@ def handle_input(user_prompt, image=None, audio=None, doc=None, websearch=False)
             )
         )
 
+    # Function for the agent's LLM
+    def llm_function(query):
+        response = client.chat.completions.create(
+            model=MODEL,
+            messages=[{"role": "user", "content": query}]
+        )
+        return response.choices[0].message.content
+
     # Initialize agent
-    agent = ZeroShotAgent(llm_chain=LLMChain(llm=llm, prompt=None), tools=tools, verbose=True)
+    agent = ZeroShotAgent(llm_chain=LLMChain(llm=llm_function, 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