Enhance agent functionality by integrating message handling and graph structure. Added read_message and answer_message functions to process user input and generate responses using the chat model. Updated AgentState to support both HumanMessage and AIMessage types.

agent.py

@@ -2,7 +2,7 @@ import os
 from typing import TypedDict, List, Dict, Any, Optional
 from langgraph.graph import StateGraph, START, END
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFacePipeline
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, AIMessage
 
 # Base Hugging Face LLM used by the chat wrapper
 base_llm = HuggingFaceEndpoint(
@@ -12,10 +12,43 @@ base_llm = HuggingFaceEndpoint(
     huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
 )
 
+# Chat model that works with LangGraph
 model = ChatHuggingFace(llm=base_llm)
 
+
 class AgentState(TypedDict):
-    messages: Dict[str, Any]
-    answer: Dict[str, Any]
-    
+    messages: List[HumanMessage | AIMessage]
+
+
+def read_message(state: AgentState) -> AgentState:
+    messages = state["messages"]
+    print(f"Processing question: {messages[-1].content if messages else ''}")
+    # Just pass the messages through to the next node
+    return {"messages": messages}
+
+
+def answer_message(state: AgentState) -> AgentState:
+    messages = state["messages"]
+    # Invoke the chat model with the conversation so far
+    response = model.invoke(messages)
+    # Append the model's answer to the messages list
+    return {"messages": messages + [response]}
+
+
+def build_graph():
+    agent_graph = StateGraph(AgentState)
+
+    # Add nodes
+    agent_graph.add_node("read_message", read_message)
+    agent_graph.add_node("answer_message", answer_message)
+
+    # Add edges
+    agent_graph.add_edge(START, "read_message")
+    agent_graph.add_edge("read_message", "answer_message")
+
+    # Final edge
+    agent_graph.add_edge("answer_message", END)
 
+    # Compile and return the executable graph for use in app.py
+    compiled_graph = agent_graph.compile()
+    return compiled_graph

app.py

@@ -3,6 +3,8 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from langchain_core.messages import HumanMessage
+from agent import build_graph 
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -13,11 +15,16 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
+        self.graph = build_graph()
+
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = "This is a default answer."
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer
+        messages = [HumanMessage(content=question)]
+        result = self.graph.invoke({"messages": messages})
+        answer = result['messages'][-1].content
+        print(f"Agent returning answer: {answer}")
+        return answer
+
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """