Update app.py

app.py

@@ -33,23 +33,32 @@ class BasicAgent:
         print(f"Agent returning fixed answer: {fixed_answer}")
         return fixed_answer
 
+class ZephyrPipelineModel:
+    def __init__(self):
+        self.pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta")
+
+    def __call__(self, prompt: str) -> str:
+        # Zephyr wants a list of messages
+        messages = [{"role": "user", "content": prompt}]
+        result = self.pipe(messages, max_new_tokens=512)
+        return result[0]["generated_text"]
+
 class LangGraphAgent:
     def __init__(self):
         # Load Zephyr with correct config
-        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 = ZephyrPipelineModel()
 
         # Define a simple graph with just one node
         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")