Update agent.py

agent.py

@@ -1,36 +1,22 @@
-from typing import TypedDict, Annotated
-from langchain_core.messages import AnyMessage, HumanMessage
-from langgraph.graph.message import add_messages
-from langgraph.graph import START, StateGraph
-from langgraph.prebuilt import ToolNode, tools_condition
 from transformers import pipeline
-from tools import search_tool, weather_info_tool, hub_stats_tool
-from dotenv import load_dotenv
-
-load_dotenv()
-
-# Initialize conversational pipeline (DialoGPT)
-chat_pipe = pipeline("conversational", model="microsoft/DialoGPT-medium")
-
-tools = [search_tool, weather_info_tool, hub_stats_tool]
+from langchain_core.messages import AnyMessage, HumanMessage
+from typing import TypedDict, List
 
-# Wrapper function to call conversational pipeline
-def chat_llm(messages: list[AnyMessage]) -> str:
-    last_message = messages[-1].content if messages else ""
-    response = chat_pipe(last_message)
-    return response[0]['generated_text']
+# Load text-generation pipeline with DialoGPT-medium (chat-like model)
+chat_pipe = pipeline("text-generation", model="microsoft/DialoGPT-medium")
 
 class AgentState(TypedDict):
-    messages: Annotated[list[AnyMessage], add_messages]
+    messages: List[AnyMessage]
 
 def assistant(state: AgentState):
-    reply = chat_llm(state["messages"])
-    return {"messages": [HumanMessage(content=reply)]}
-
-builder = StateGraph(AgentState)
-builder.add_node("assistant", assistant)
-builder.add_node("tools", ToolNode(tools))
-builder.add_edge(START, "assistant")
-builder.add_conditional_edges("assistant", tools_condition)
-builder.add_edge("tools", "assistant")
-alfred = builder.compile()
+    # Extract last user message content
+    last_message = state["messages"][-1].content if state["messages"] else ""
+    
+    # Generate response from the model
+    outputs = chat_pipe(last_message, max_length=100, do_sample=True, top_p=0.9, temperature=0.8)
+    generated_text = outputs[0]['generated_text']
+    
+    # Remove the prompt from output if present
+    response = generated_text[len(last_message):].strip()
+    
+    return {"messages": [HumanMessage(content=response)]}