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

2026-06-29 12:16:00 +03:00
parent 366ec8dee7
commit 707c35d49d
5 changed files with 189 additions and 110 deletions
@@ -1,90 +1,82 @@
-# LangGraph Research Brief Agent
+# LangGraph Reflection Agent

-This project demonstrates how to build a LangGraph agent that generates a short research brief for a given topic.  
-The agent:
+This project demonstrates a simple LangGraph agent that:

-1. Creates an outline of 4‑5 research steps.  
-2. For each step, performs a web search (via Tavily) and writes a concise note.  
-3. Synthesizes all notes into a coherent brief.
+1. Generates a concise answer to a user‑supplied question.
+2. Critiques the answer using an LLM.
+3. Rewrites the answer if the critic says *needs_revision*, up to a maximum number of rounds.

-## Prerequisites
+The agent is implemented in Python 3.10+ and uses the `langgraph` framework together with `langchain-openai`.

- Python 3.10+
- A **Tavily** API key (free tier available).  
- An **OpenAI** API key (or any compatible LLM provider).
+## Features

-## Setup
+- **Draft generation** – 5–10 sentence answer.
+- **LLM critic** – returns a verdict (`ok` or `needs_revision`) and 2–3 critique points.
+- **Rewrite loop** – rewrites the draft until the verdict is `ok` or the maximum number of rounds is reached.
+- **CLI** – run the agent from the command line.
+
+## Installation

 ```bash
-# Clone the repository
-git clone https://github.com/your-username/langgraph-research-brief.git
-cd langgraph-research-brief
-
 # Create a virtual environment (optional but recommended)
-python -m venv venv
-source venv/bin/activate   # On Windows: venv\Scripts\activate
+python -m venv .venv
+source .venv/bin/activate   # On Windows use `.venv\Scripts\activate`

 # Install dependencies
 pip install -r requirements.txt
 ```

-Create a `.env` file in the project root based on the example:
+### OpenAI API Key
+
+The agent uses OpenAI’s GPT‑3.5‑Turbo by default.  
+Set your API key in the environment:

 ```bash
-cp .env.example .env
+export OPENAI_API_KEY="sk-..."
 ```

-Edit `.env` and replace the placeholders with your actual keys:
+If you prefer to use a local LLM via Ollama, replace the `langchain-openai` dependency with `langchain-ollama` and adjust the LLM initialization in `src/graph.py`.

-```
-OPENAI_API_KEY=sk-...
-TAVILY_API_KEY=your_tavily_key
-```
-
-## Running the Agent
+## Usage

 ```bash
-python src/main.py
+python -m src.main "Explain the difference between a tool and a resource in MCP to a student."
 ```

-You should see output similar to:
+Optional arguments:
+
+- `--max_rounds N` – maximum number of rewrite rounds (default: 2).
+
+Example:
+
+```bash
+python -m src.main "Explain the difference between a tool and a resource in MCP." --max_rounds 3
+```
+
+The script will print:

 ```
-=== Outline ===
-1. Identify the security requirements for MCP integration
-2. Review LangChain's authentication mechanisms
-3. Evaluate secure communication protocols
-4. Test the integration in a sandbox environment
-5. Document best practices and compliance checks
+=== Final Answer ===
+<rewritten answer>

-=== Notes ===
-[Step 1] ... (5‑8 sentence note)
-[Step 2] ... (5‑8 sentence note)
-...
+=== Verdict ===
+ok

-=== Final Brief ===
-...
 ```

+If the final verdict is `needs_revision`, the critique points will also be shown.
+
 ## Project Structure

 ```
 src/
-├── main.py          # Entry point
-├── graph.py         # LangGraph definition
-├── nodes.py         # Node implementations
-├── state.py         # TypedDict for state
-├── .env.example     # Environment variable template
+├── main.py      # CLI entry point
+├── graph.py     # LangGraph graph definition
+└── state.py     # TypedDict for the agent state
 requirements.txt
 README.md
 ```

-## Customization
-
- **Topic**: Change the `default_topic` variable in `src/main.py` to generate a brief on a different subject.  
- **LLM**: Swap `ChatOpenAI` for another provider (e.g., Ollama) by adjusting the imports and initialization in `src/nodes.py`.  
- **Search**: Replace `TavilySearchResults` with another search tool if desired.
-
 ## License

 MIT License
@@ -1,5 +1,3 @@
 langgraph
 langchain-openai
-langchain-tavily
-tavily-python
-python-dotenv
+openai
@@ -1,28 +1,105 @@
-from langgraph import StateGraph
-from src.state import BriefState
-from src.nodes import outline_node, research_step_node, synthesize_node
+import json
+from typing import Dict, Any
+
+from langgraph.graph import StateGraph, END
+from langchain_openai import ChatOpenAI
+from langchain.prompts import PromptTemplate
+
+from .state import ReflectState
+
+# LLM configuration
+llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
+
+# Prompt templates
+DRAFT_PROMPT = PromptTemplate(
+    input_variables=["question"],
+    template=(
+        "You are an assistant that writes a concise answer (5–10 sentences) to the following question:\n\n"
+        "Question: {question}\n\n"
+        "Answer:"
+    ),
+)
+
+REFLECT_PROMPT = PromptTemplate(
+    input_variables=["question", "draft"],
+    template=(
+        "You are a critical reviewer. Evaluate the following draft answer to the question.\n\n"
+        "Question: {question}\n\n"
+        "Draft answer:\n{draft}\n\n"
+        "Provide a JSON object with the following keys:\n"
+        "  verdict: \"ok\" if the answer is complete, concrete, and free of fluff; otherwise \"needs_revision\"\n"
+        "  critique: a list of 2–3 specific points for improvement.\n"
+        "Example output:\n"
+        "{{\"verdict\": \"ok\", \"critique\": []}}\n"
+        "Your output should be valid JSON."
+    ),
+)
+
+REWRITE_PROMPT = PromptTemplate(
+    input_variables=["draft", "critique"],
+    template=(
+        "Rewrite the following draft answer to address the critique points below. "
+        "Make the answer clearer, more concrete, and remove any unnecessary fluff.\n\n"
+        "Critique points:\n{critique}\n\n"
+        "Original draft:\n{draft}\n\n"
+        "Rewritten answer:"
+    ),
+)
+
+def draft_answer(state: ReflectState) -> ReflectState:
+    """Generate the initial draft answer."""
+    prompt = DRAFT_PROMPT.format(question=state["question"])
+    answer = llm.invoke(prompt).content.strip()
+    state["draft"] = answer
+    return state
+
+def reflect(state: ReflectState) -> ReflectState:
+    """Critique the draft and produce verdict and critique list."""
+    prompt = REFLECT_PROMPT.format(question=state["question"], draft=state["draft"])
+    response = llm.invoke(prompt).content.strip()
+    try:
+        data = json.loads(response)
+        verdict = data.get("verdict", "needs_revision")
+        critique = data.get("critique", [])
+    except json.JSONDecodeError:
+        # Fallback if parsing fails
+        verdict = "needs_revision"
+        critique = ["Unable to parse critique."]
+
+    state["verdict"] = verdict
+    state["critique"] = critique
+    return state
+
+def rewrite(state: ReflectState) -> ReflectState:
+    """Rewrite the draft based on critique and increment round."""
+    critique_text = "\n".join(f"- {c}" for c in state["critique"])
+    prompt = REWRITE_PROMPT.format(draft=state["draft"], critique=critique_text)
+    new_draft = llm.invoke(prompt).content.strip()
+    state["draft"] = new_draft
+    state["round"] += 1
+    return state
+
+def decide(state: ReflectState) -> str:
+    """Decide whether to end or rewrite."""
+    if state["verdict"] == "ok":
+        return "end"
+    if state["round"] >= state["max_rounds"]:
+        return "end"
+    return "rewrite"

 def build_graph() -> StateGraph:
-    graph = StateGraph(BriefState)
+    graph = StateGraph(ReflectState)

-    # Add nodes
-    graph.add_node("outline", outline_node)
-    graph.add_node("research_step", research_step_node)
-    graph.add_node("synthesize", synthesize_node)
+    graph.add_node("draft_answer", draft_answer)
+    graph.add_node("reflect", reflect)
+    graph.add_node("rewrite", rewrite)
+    graph.add_node("decide", decide)

-    # Define the condition for looping research steps
-    def condition(state: BriefState):
-        if state["step_index"] < len(state["outline"]):
-            return "research_step"
-        else:
-            return "synthesize"
+    graph.set_entry_point("draft_answer")
+    graph.add_edge("draft_answer", "reflect")
+    graph.add_edge("reflect", "decide")
+    graph.add_conditional_edges("decide", lambda state: state["verdict"] if state["verdict"] == "ok" else ("rewrite" if state["round"] < state["max_rounds"] else "end"))
+    graph.add_edge("rewrite", "reflect")
+    graph.add_edge("end", END)

-    # Build edges
-    graph.add_edge("outline", "research_step")
-    graph.add_conditional_edges("research_step", condition, {
-        "research_step": "research_step",
-        "synthesize": "synthesize"
-    })
-    graph.add_edge("synthesize", "__end__")
-
-    return graph.compile()
+    return graph
@@ -1,38 +1,49 @@
+import argparse
 import os
-from dotenv import load_dotenv
-from src.graph import build_graph
-from src.state import BriefState
+import sys
+
+from .graph import build_graph
+from .state import ReflectState
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="LangGraph reflection agent")
+    parser.add_argument("question", type=str, help="The question to answer")
+    parser.add_argument(
+        "--max_rounds",
+        type=int,
+        default=2,
+        help="Maximum number of rewrite rounds (default: 2)",
+    )
+    return parser.parse_args()

 def main():
-    # Load environment variables
-    load_dotenv()
-    # Default topic
-    default_topic = "Как студенту безопасно подключать MCP к LangChain"
+    args = parse_args()

-    # Initial state
-    initial_state: BriefState = {
-        "topic": default_topic,
-        "outline": None,
-        "step_index": 0,
-        "notes": [],
-        "final_brief": None
+    # Ensure OpenAI key is set
+    if not os.getenv("OPENAI_API_KEY"):
+        print("Error: OPENAI_API_KEY environment variable not set.", file=sys.stderr)
+        sys.exit(1)
+
+    initial_state: ReflectState = {
+        "question": args.question,
+        "draft": "",
+        "critique": [],
+        "verdict": "",
+        "round": 0,
+        "max_rounds": args.max_rounds,
    }

-    # Build and run the graph
    graph = build_graph()
    final_state = graph.invoke(initial_state)

-    # Print results
-    print("\n=== Outline ===")
-    for i, step in enumerate(final_state["outline"], 1):
-        print(f"{i}. {step}")
-
-    print("\n=== Notes ===")
-    for i, note in enumerate(final_state["notes"], 1):
-        print(f"[Step {i}] {note}\n")
-
-    print("\n=== Final Brief ===")
-    print(final_state["final_brief"])
+    print("\n=== Final Answer ===")
+    print(final_state["draft"])
+    print("\n=== Verdict ===")
+    print(final_state["verdict"])
+    if final_state["verdict"] != "ok":
+        print("\n=== Critique ===")
+        for i, point in enumerate(final_state["critique"], 1):
+            print(f"{i}. {point}")

 if __name__ == "__main__":
    main()
@@ -1,8 +1,9 @@
-from typing import TypedDict, List, Optional
+from typing import TypedDict, List

-class BriefState(TypedDict):
-    topic: str
-    outline: List[str] | None
-    step_index: int
-    notes: List[str]
-    final_brief: str | None
+class ReflectState(TypedDict):
+    question: str
+    draft: str
+    critique: List[str]
+    verdict: str  # "ok" | "needs_revision"
+    round: int
+    max_rounds: int