Update app.py

app.py

@@ -31,55 +31,49 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
-class ZephyrPipelineModel:
+class ZephyrAPI:
     def __init__(self):
-        hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            raise ValueError("HF_TOKEN environment variable not set.")
-            
-        self.client = InferenceClient(
-            model="HuggingFaceH4/zephyr-7b-beta",
-            token=hf_token,
-        )
+        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.")
+
+    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,
+            }
+        }
 
-    def __call__(self, prompt: str) -> str:
-        messages = [{"role": "user", "content": prompt}]
         try:
-            completion = self.client.chat.completions.create(
-                model="HuggingFaceH4/zephyr-7b-beta",
-                messages=messages,
-                max_tokens=512,
-                temperature=0.7,
-            )
-            return completion.choices[0].message.content
+            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()
         except Exception as e:
-            return f"❌ Inference failed: {str(e)}"
+            print(f"Error: {e}")
+            return "⚠️ Model could not respond. Check API access or token."
+
 
 class LangGraphAgent:
     def __init__(self):
-        # ✅ Add token from environment
-        hf_token = os.environ.get("HF_TOKEN")
-        if not hf_token:
-            raise ValueError("HF_TOKEN is not set.")
-
-        # ✅ Restore this structure with token
-        self.model = ChatHuggingFace.from_model_id(
-            model_id="HuggingFaceH4/zephyr-7b-beta",
-            task="text-generation",
-            model_kwargs={
-                "temperature": 0.7,
-                "max_new_tokens": 512
-            },
-            huggingfacehub_api_token=hf_token,
-        )
+        self.model = ZephyrAPI()
 
-        # ✅ Simple LangGraph setup
         builder = StateGraph()
 
         def call_model(state):
             messages = state.get("messages", [])
-            response = self.model.invoke(messages)
-            return {"messages": messages + [response]}
+            user_msg = next((m for m in messages if isinstance(m, HumanMessage)), None)
+            if not user_msg:
+                return {"messages": messages + [AIMessage(content="❌ No user input found.")]}
+            
+            response = self.model(user_msg.content)
+            return {"messages": messages + [AIMessage(content=response)]}
 
         builder.add_node("chat", call_model)
         builder.set_entry_point("chat")
@@ -92,7 +86,8 @@ class LangGraphAgent:
             "messages": [HumanMessage(content=question)]
         })
 
-        for msg in reversed(result.get("messages", [])):
+        messages = result.get("messages", [])
+        for msg in reversed(messages):
             if isinstance(msg, AIMessage):
                 return msg.content