feat: solution for 'Экзамен: Самокорректирующийся агент'

src/graph.py

@@ -1,46 +1,14 @@
-"""
-Graph definition using LangGraph.
-"""
-
+from langgraph.graph import StateGraph
+from src.nodes import generate_response
 from typing import Dict, Any
-from langgraph.graph import StateGraph, END
-from langchain_core.messages import AIMessage, HumanMessage
-from src.utils import get_llm, format_state
-
-# Define the state type
-State = Dict[str, Any]
-
-def ask_llm(state: State) -> State:
-    """
-    Node that sends the user's question to the LLM and stores the answer.
-    """
-    llm = get_llm()
-    question = state.get("question", "")
-    # Create a conversation with the LLM
-    response = llm.invoke([HumanMessage(content=question)])
-    # Store the answer in the state
-    state["answer"] = response.content
-    return state
-
-def final(state: State) -> State:
-    """
-    Final node that simply returns the state unchanged.
-    """
-    return state
 
 def build_graph() -> StateGraph:
     """
-    Builds and returns the LangGraph graph.
+    Builds a simple StateGraph with a single node that echoes user input.
     """
-    graph = StateGraph(State)
-
-    # Add nodes
-    graph.add_node("ask", ask_llm)
-    graph.add_node("final", final)
-
-    # Define edges
-    graph.set_entry_point("ask")
-    graph.add_edge("ask", "final")
-    graph.add_edge("final", END)
-
+    graph = StateGraph()
+    # Add the echo node
+    graph.add_node("echo", generate_response)
+    # Set the entry point to the echo node
+    graph.set_entry_point("echo")
     return graph