From 707c35d49de681113625fd51687d3d4674c2cb1c Mon Sep 17 00:00:00 2001 From: kuzakhmetovartur Date: Mon, 29 Jun 2026 12:16:00 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20solution=20for=20'=D0=9F=D0=BE=D0=B2?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BD=D1=8B=D0=B9=20=D1=8D=D0=BA=D0=B7=D0=B0?= =?UTF-8?q?=D0=BC=D0=B5=D0=BD:=20=D0=93=D1=80=D0=B0=D1=84=20=D1=81=20?= =?UTF-8?q?=D1=80=D0=B5=D1=84=D0=BB=D0=B5=D0=BA=D1=81=D0=B8=D0=B5=D0=B9=20?= =?UTF-8?q?=D0=B8=20=D0=B4=D0=BE=D1=80=D0=B0=D0=B1=D0=BE=D1=82=D0=BA=D0=BE?= =?UTF-8?q?=D0=B9'?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 94 +++++++++++++++++------------------- requirements.txt | 4 +- src/graph.py | 123 ++++++++++++++++++++++++++++++++++++++--------- src/main.py | 63 ++++++++++++++---------- src/state.py | 15 +++--- 5 files changed, 189 insertions(+), 110 deletions(-) diff --git a/README.md b/README.md index 0a0da4d..7c37b01 100644 --- a/README.md +++ b/README.md @@ -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 === + -=== 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 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 2d94960..c341710 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,3 @@ langgraph langchain-openai -langchain-tavily -tavily-python -python-dotenv \ No newline at end of file +openai \ No newline at end of file diff --git a/src/graph.py b/src/graph.py index 1dc7269..69ba2e6 100644 --- a/src/graph.py +++ b/src/graph.py @@ -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() \ No newline at end of file + return graph \ No newline at end of file diff --git a/src/main.py b/src/main.py index 7beaf86..f605088 100644 --- a/src/main.py +++ b/src/main.py @@ -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() \ No newline at end of file diff --git a/src/state.py b/src/state.py index 0d92a0e..afc15db 100644 --- a/src/state.py +++ b/src/state.py @@ -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 \ No newline at end of file +class ReflectState(TypedDict): + question: str + draft: str + critique: List[str] + verdict: str # "ok" | "needs_revision" + round: int + max_rounds: int \ No newline at end of file