diff --git a/README.md b/README.md new file mode 100644 index 0000000..2ebee46 --- /dev/null +++ b/README.md @@ -0,0 +1,79 @@ +# LangGraph Reflection Demo + +This project demonstrates a simple LangGraph that generates an answer to a question, reflects on it, and rewrites it if necessary. The graph loops until the answer is deemed satisfactory or a maximum number of rounds is reached. + +## Features + +- **Draft generation** – 5–10 sentence answer to a user question. +- **Reflection** – LLM critiques the draft and decides if it is acceptable. +- **Rewrite** – If the draft needs improvement, the LLM rewrites it based on the critique. +- **Loop control** – The process repeats until the answer is good enough or the maximum number of rounds is exceeded. +- **CLI** – Run the graph from the command line. + +## Installation + +```bash +# Clone the repository +git clone https://github.com/yourusername/langgraph-reflection-demo.git +cd langgraph-reflection-demo + +# Create a virtual environment (optional but recommended) +python -m venv venv +source venv/bin/activate # On Windows use `venv\Scripts\activate` + +# Install dependencies +pip install -r requirements.txt + +# Set your OpenAI API key +export OPENAI_API_KEY="your-openai-key" +``` + +> **Note**: If you prefer to use Ollama instead of OpenAI, replace `langchain-openai` with `langchain-ollama` in `requirements.txt` and adjust the LLM import in `nodes.py`. + +## Usage + +```bash +python main.py "Explain the theory of relativity in simple terms." +``` + +Optional arguments: + +- `--max-rounds N` – Maximum number of rewrite attempts (default: 2). +- `--model MODEL` – LLM model name (default: `gpt-3.5-turbo`). + +Example: + +```bash +python main.py "What is quantum computing?" --max-rounds 3 --model gpt-4 +``` + +The script will print: + +``` +Initial draft: +... + +Reflection verdict: needs_revision +Critique: +... + +Rewritten draft: +... + +Final answer: +... +``` + +## Testing + +Run the unit tests with: + +```bash +pytest +``` + +The tests require a valid OpenAI API key set in the environment. + +## License + +MIT License diff --git a/graph.py b/graph.py new file mode 100644 index 0000000..91f292d --- /dev/null +++ b/graph.py @@ -0,0 +1,26 @@ +from langgraph.graph import StateGraph, END +from nodes import draft_answer, reflect, rewrite, ReflectState + +def build_graph(): + graph = StateGraph(ReflectState) + + graph.add_node("draft_answer", draft_answer) + graph.add_node("reflect", reflect) + graph.add_node("rewrite", rewrite) + + def condition(state: ReflectState): + if state["verdict"] == "ok": + return "ok" + if state["round"] >= state["max_rounds"]: + return "maxed" + return "needs_revision" + + graph.set_entry_point("draft_answer") + graph.add_conditional_edges("reflect", condition, { + "ok": END, + "needs_revision": "rewrite", + "maxed": END, + }) + graph.add_edge("rewrite", "reflect") + + return graph.compile() diff --git a/main.py b/main.py new file mode 100644 index 0000000..03dfc5e --- /dev/null +++ b/main.py @@ -0,0 +1,41 @@ +import argparse +import os +from dotenv import load_dotenv + +from graph import build_graph +from nodes import llm + +def parse_args(): + parser = argparse.ArgumentParser(description="LangGraph Reflection Demo") + 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 attempts") + parser.add_argument("--model", type=str, default="gpt-3.5-turbo", help="LLM model name") + return parser.parse_args() + +def main(): + load_dotenv() + args = parse_args() + + # Configure LLM + llm.model = args.model + + graph = build_graph() + initial_state = { + "question": args.question, + "draft": "", + "critique": "", + "verdict": "", + "round": 1, + "max_rounds": args.max_rounds, + } + + result = graph.invoke(initial_state) + + print("\n=== Final Result ===") + print(f"Draft:\n{result['draft']}\n") + print(f"Verdict: {result['verdict']}") + print(f"Critique:\n{result['critique']}\n") + print(f"Round: {result['round']}") + +if __name__ == "__main__": + main() diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..db6abec --- /dev/null +++ b/nodes.py @@ -0,0 +1,55 @@ +from typing import TypedDict, Dict, Any +from langchain_community.llms import OpenAI +import re + +# Global LLM instance +llm = OpenAI(temperature=0.7) + +class ReflectState(TypedDict): + question: str + draft: str + critique: str + verdict: str + round: int + max_rounds: int + +def draft_answer(state: ReflectState) -> ReflectState: + prompt = ( + f"Answer the following question in 5–10 sentences:\n\n" + f"Question: {state['question']}\n\n" + f"Answer:" + ) + answer = llm.invoke(prompt).strip() + state["draft"] = answer + return state + +def reflect(state: ReflectState) -> ReflectState: + prompt = ( + f"You are a critical reviewer. Evaluate the following draft answer.\n\n" + f"Draft:\n{state['draft']}\n\n" + f"Provide a verdict ('ok' or 'needs_revision') and 2–3 remarks.\n" + f"Format:\n" + f"Verdict: \n" + f"Remarks:\n" + ) + response = llm.invoke(prompt).strip() + # Parse verdict + verdict_match = re.search(r"Verdict:\s*(\w+)", response, re.IGNORECASE) + remarks_match = re.search(r"Remarks:\s*(.*)", response, re.DOTALL | re.IGNORECASE) + verdict = verdict_match.group(1).lower() if verdict_match else "needs_revision" + remarks = remarks_match.group(1).strip() if remarks_match else "No remarks provided." + state["verdict"] = verdict + state["critique"] = remarks + return state + +def rewrite(state: ReflectState) -> ReflectState: + prompt = ( + f"Rewrite the following draft answer to address the remarks below.\n\n" + f"Draft:\n{state['draft']}\n\n" + f"Remarks:\n{state['critique']}\n\n" + f"Provide the revised answer in 5–10 sentences." + ) + revised = llm.invoke(prompt).strip() + state["draft"] = revised + state["round"] += 1 + return state diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9c0c9c2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +langchain-openai==0.0.3 +langchain==0.1.0 +langgraph==0.0.1 +openai==1.3.0 +python-dotenv==1.0.0