Update app.py

app.py

@@ -57,37 +57,42 @@ class ZephyrPipelineModel:
 
 class LangGraphAgent:
     def __init__(self):
-        # Load Zephyr with correct config
-        self.model = ZephyrPipelineModel()
+        # ✅ 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,
+        )
 
-        # Define a simple graph with just one node
+        # ✅ Simple LangGraph setup
         builder = StateGraph()
 
         def call_model(state):
             messages = state.get("messages", [])
-            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)]}
+            response = self.model.invoke(messages)
+            return {"messages": messages + [response]}
 
         builder.add_node("chat", call_model)
         builder.set_entry_point("chat")
         builder.add_edge("chat", END)
 
-        # Compile the graph
         self.graph = builder.compile()
 
     def __call__(self, question: str) -> str:
-        # Wrap input in HumanMessage format
         result = self.graph.invoke({
             "messages": [HumanMessage(content=question)]
         })
 
-        # Extract final response
-        messages = result.get("messages", [])
-        for msg in reversed(messages):
+        for msg in reversed(result.get("messages", [])):
             if isinstance(msg, AIMessage):
                 return msg.content