Update app.py

app.py

@@ -179,24 +179,53 @@ class VideoTranscriptionTool(Tool):
         except Exception as e:
             return f"Error in extracting YouTube transcript: {str(e)}"
 
+
+class HFModelWrapper:
+    def __init__(self, model_id: str, token: str, temperature: float = 0.1):
+        self.api_url = f"https://api-inference.huggingface.co/models/{model_id}"
+        self.headers = {
+            "Authorization": f"Bearer {token}",
+            "Content-Type": "application/json"
+        }
+        self.temperature = temperature
+
+    def __call__(self, prompt: str) -> str:
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "temperature": self.temperature,
+                "max_new_tokens": 512,
+                "return_full_text": False
+            }
+        }
+        response = requests.post(self.api_url, headers=self.headers, json=payload)
+        if response.status_code == 200:
+            return response.json()[0]["generated_text"]
+        else:
+            raise RuntimeError(f"Hugging Face API error: {response.status_code} - {response.text}")
+
+
 class BasicAgent:
     def __init__(self):
         token = os.environ.get("HF_API_TOKEN")
-        self.api_token = os.environ.get("HF_API_TOKEN")
-        self.api_url = "https://api-inference.huggingface.co/models/"
-        self.model_id = "mistralai/Mistral-7B-Instruct-v0.3"
-        
-
-        search_tool = DuckDuckGoSearchTool()
-        wiki_search_tool = WikiSearchTool()
-        str_reverse_tool = StringReverseTool()
-        keywords_extract_tool = KeywordsExtractorTool()
-        speech_to_text_tool = SpeechToTextTool()
-        visit_webpage_tool = VisitWebpageTool()
-        final_answer_tool = FinalAnswerTool()
-        video_transcription_tool = VideoTranscriptionTool()
-
-        system_prompt = f"""
+        model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+        model = HFModelWrapper(model_id=model_id, token=token)
+
+        # Tool initialization
+        tools = [
+            DuckDuckGoSearchTool(),
+            WikiSearchTool(),
+            StringReverseTool(),
+            KeywordsExtractorTool(),
+            SpeechToTextTool(),
+            VisitWebpageTool(),
+            FinalAnswerTool(),
+            parse_excel_to_json,
+            VideoTranscriptionTool()
+        ]
+
+        # Custom system prompt
+        system_prompt = """
 You are my general AI assistant. Your task is to answer the question I asked.
 First, provide an explanation of your reasoning, step by step, to arrive at the answer.
 Then, return your final answer in a single line, formatted as follows: "FINAL ANSWER: [YOUR FINAL ANSWER]".
@@ -205,12 +234,20 @@ If the answer is a number, do not use commas or units (e.g., $, %) unless specif
 If the answer is a string, do not use articles or abbreviations (e.g., for cities), and write digits in plain text unless specified.
 If the answer is a comma-separated list, apply the above rules for each element based on whether it is a number or a string.
 """
+
+        # Instantiate CodeAgent with wrapped model
         self.agent = CodeAgent(
-            model=model_id,
-            tools=[search_tool, wiki_search_tool, str_reverse_tool, keywords_extract_tool, speech_to_text_tool, visit_webpage_tool, final_answer_tool, parse_excel_to_json, video_transcription_tool],
+            model=model,
+            tools=tools,
             add_base_tools=True
         )
-        self.agent.prompt_templates["system_prompt"] = self.agent.prompt_templates["system_prompt"] + system_prompt
+
+        # Add system prompt safely
+        if "system_prompt" in self.agent.prompt_templates:
+            self.agent.prompt_templates["system_prompt"] += system_prompt
+        else:
+            self.agent.prompt_templates["system_prompt"] = system_prompt
+
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")