feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'

src/graph.py

@@ -1,28 +1,46 @@
+"""
+Graph definition using LangGraph.
+"""
+
 from typing import Dict, Any
-from langgraph.graph import StateGraph
-from src.nodes import ReflectState, draft_answer, reflect, rewrite
+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:
-    graph = StateGraph(ReflectState)
+    """
+    Builds and returns the LangGraph graph.
+    """
+    graph = StateGraph(State)
 
     # Add nodes
-    graph.add_node("draft_answer", draft_answer)
-    graph.add_node("reflect", reflect)
-    graph.add_node("rewrite", rewrite)
+    graph.add_node("ask", ask_llm)
+    graph.add_node("final", final)
 
-    # Define transitions
-    graph.set_entry_point("draft_answer")
-    graph.add_edge("draft_answer", "reflect")
-
-    # Conditional edge after reflect
-    def decide_next(state: ReflectState) -> str:
-        if state["verdict"] == "ok":
-            return "end"
-        if state["round"] < state["max_rounds"]:
-            return "rewrite"
-        return "end"
-
-    graph.add_conditional_edges("reflect", decide_next, {"rewrite": "rewrite", "end": "end"})
-    graph.add_edge("rewrite", "reflect")
+    # Define edges
+    graph.set_entry_point("ask")
+    graph.add_edge("ask", "final")
+    graph.add_edge("final", END)
 
     return graph