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

src/nodes.py

@@ -1,80 +1,21 @@
-from typing import TypedDict, Dict, Any
-from langchain_openai import ChatOpenAI
-from langchain.prompts import PromptTemplate
+from langchain_core.messages import HumanMessage, AIMessage
+from typing import Dict, Any
 
-# Define the state structure
-class ReflectState(TypedDict):
-    question: str
-    draft: str
-    critique: str
-    verdict: str  # "ok" or "needs_revision"
-    round: int
-    max_rounds: int
+def generate_response(state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Simple node that echoes the user's message as an AI response.
+    """
+    messages = state.get("messages", [])
+    if not messages:
+        return state
 
-# Initialize the LLM (requires OPENAI_API_KEY environment variable)
-llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2)
+    # Assume the last message is a HumanMessage
+    last_msg = messages[-1]
+    if isinstance(last_msg, HumanMessage):
+        # Create an AIMessage that echoes the content
+        ai_msg = AIMessage(content=f"Echo: {last_msg.content}")
+        messages.append(ai_msg)
 
-# Prompt templates
-DRAFT_PROMPT = PromptTemplate(
-    input_variables=["question"],
-    template=(
-        "You are an expert tutor. Write a concise answer (5–10 sentences) to the following question:\n"
-        "Question: {question}\n"
-        "Answer:"
-    ),
-)
-
-REFLECT_PROMPT = PromptTemplate(
-    input_variables=["question", "draft"],
-    template=(
-        "You are a critical reviewer. Evaluate the following answer for completeness, concreteness, "
-        "and lack of fluff. Provide a verdict ('ok' or 'needs_revision') and 2–3 critique points.\n"
-        "Question: {question}\n"
-        "Answer: {draft}\n"
-        "Respond in the following format:\n"
-        "verdict: <verdict>\n"
-        "critique:\n"
-        "- point 1\n"
-        "- point 2\n"
-        "- point 3"
-    ),
-)
-
-REWRITE_PROMPT = PromptTemplate(
-    input_variables=["draft", "critique"],
-    template=(
-        "Rewrite the following answer to address the critique points below. "
-        "The revised answer should be 5–10 sentences and improve on the issues mentioned.\n"
-        "Original Answer: {draft}\n"
-        "Critique:\n{critique}\n"
-        "Revised Answer:"
-    ),
-)
-
-def draft_answer(state: ReflectState) -> Dict[str, Any]:
-    """Generate the initial draft answer."""
-    question = state["question"]
-    response = llm.invoke(DRAFT_PROMPT.format(question=question))
-    draft = response.content.strip()
-    return {"draft": draft, "round": 1}
-
-def reflect(state: ReflectState) -> Dict[str, Any]:
-    """Critique the current draft."""
-    question = state["question"]
-    draft = state["draft"]
-    response = llm.invoke(REFLECT_PROMPT.format(question=question, draft=draft))
-    text = response.content.strip()
-    # Parse verdict and critique
-    verdict_line, critique_section = text.split("critique:", 1)
-    verdict = verdict_line.replace("verdict:", "").strip().lower()
-    critique = critique_section.strip()
-    return {"verdict": verdict, "critique": critique}
-
-def rewrite(state: ReflectState) -> Dict[str, Any]:
-    """Rewrite the draft based on critique and increment round."""
-    draft = state["draft"]
-    critique = state["critique"]
-    response = llm.invoke(REWRITE_PROMPT.format(draft=draft, critique=critique))
-    new_draft = response.content.strip()
-    new_round = state["round"] + 1
-    return {"draft": new_draft, "round": new_round}
+    # Update the state with the new messages list
+    state["messages"] = messages
+    return state