Update app.py

app.py

@@ -5,6 +5,12 @@ import inspect
 import pandas as pd
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain_community.llms import HuggingFacePipeline
+from tools import tools
+from langchain_core.messages import HumanMessage
+from langgraph.prebuilt import ToolExecutor, chat_agent_executor
+
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -23,31 +29,26 @@ class BasicAgent:
 
 class ZephyrAgent:
     def __init__(self):
-        self.api_url = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
-        self.headers = {
-            "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
-        }
-        print("ZephyrAPI initialized using Inference API.")
+        print("Initializing local Zephyr model with tools...")
 
-    def __call__(self, question: str) -> str:
-        prompt = f"<|system|>\nYou are a helpful assistant.\n<|user|>\n{question}\n<|assistant|>\n"
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 256,
-                "temperature": 0.7,
-                "top_p": 0.9,
-            }
-        }
+        tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
+        model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.float16, device_map="auto")
+
+        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
+        llm = HuggingFacePipeline(pipeline=pipe)
 
+        # Wrap tools
+        tool_executor = ToolExecutor(tools)
+        self.agent_executor = chat_agent_executor.create_chat_agent_executor(llm=llm, tools=tools)
+
+    def __call__(self, question: str) -> str:
         try:
-            response = requests.post(self.api_url, headers=self.headers, json=payload, timeout=60)
-            response.raise_for_status()
-            result = response.json()
-            return result[0]["generated_text"].split("<|assistant|>")[-1].strip()
+            message = HumanMessage(content=question)
+            response = self.agent_executor.invoke({"messages": [message]})
+            return response.content
         except Exception as e:
-            print(f"Error: {e}")
-            return "⚠️ Model could not respond. Check API access or token."
+            print(f"Tool-augmented ZephyrAgent error: {e}")
+            return "⚠️ Agent failed to answer due to tool or model error."
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """