Skip to content

An implementation of a self-improving AI agent using LangGraph that can generate content, evaluate its quality, and refine outputs through reflection.

License

Notifications You must be signed in to change notification settings

codebywiam/reflection-agent-langgraph

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Reflection Agent with LangGraph 🤖

A sophisticated AI agent that leverages LangGraph's graph-based workflow structure to create self-improving content through iterative reflection and refinement. This project demonstrates advanced AI capabilities by building an agent that can evaluate its own outputs, identify weaknesses, and continuously improve through feedback loops.

Features

  • Self-Improving Workflow: Generates, evaluates, and refines content iteratively
  • Modular Design: Easily adaptable to different tasks (writing, coding, etc.)
  • Multiple LLM Support: Works with OpenAI, Anthropic, and other LangChain-compatible models
  • Visualization: Built-in workflow visualization
  • Customizable: Adjust reflection depth and criteria

Architecture

The Reflection Agent implements a sophisticated workflow with four main components:

graph TD
    A[Initial Input] --> B[Generation Node]
    B --> C[Evaluation Node]
    C --> D[Final Output]
    C --> E[Reflection Node]
    E --> B
Loading

Workflow Components

Node Type Purpose Implementation Details
1. Generation Node Creates initial content based on input prompt • Uses LLM (OpenAI, Anthropic, etc.) with task-specific prompt
• Generates first draft quickly (System 1 thinking)
2. Evaluation Node Assesses output quality • Checks against quality criteria
• Determines if reflection is needed
• Can use simple rules or LLM-based evaluation
3. Reflection Node Critiques and suggests improvements • LLM analyzes strengths/weaknesses
• Provides specific improvement suggestions
• Returns feedback as HumanMessage
4. State Management Tracks conversation state MessageGraph: Tracks conversation state
Message Types:
  - HumanMessage: User input/feedback
  - AIMessage: Generated content
  - SystemMessage: Reflection criteria

Workflow Logic

1. Initialization

def _build_workflow(self) -> MessageGraph:
    graph = MessageGraph()
    graph.add_node("generate", self._generation_node)
    graph.add_node("reflect", self._reflection_node)

-> Implementation:

  • Creates a MessageGraph structure to manage conversation flow
  • Adds two main processing nodes: "generate" and "reflect"
  • Each node is linked to its corresponding processing function

2. Conditional Routing

graph.add_conditional_edges(
    "generate",
    self._should_continue,
    {True: "reflect", False: END}
)

-> Implementation:

  • After generation, _should_continue() evaluates whether to continue
  • Routes to "reflect" if improvement is needed
  • Routes to END if content meets quality standards or max iterations reached
  • Prevents unnecessary processing and infinite loops

3. Iteration Control

def _should_continue(self, state: List[BaseMessage]) -> bool:
    iterations = len([m for m in state if isinstance(m, AIMessage)])
    return iterations < self.max_iterations

-> Implementation:

  • Max iterations: Default of 3 prevents endless improvement cycles
  • Early termination: Stops when iteration limit is reached
  • Iteration tracking: Counts AIMessage instances to track progress

Node Implementation Details

Generation Node

def _generation_node(self, state: Sequence[BaseMessage]) -> List[BaseMessage]:
    generated = self._build_generation_chain().invoke({"messages": state})
    return [AIMessage(content=generated.content)]

-> How it works:

  • Takes current conversation state as input
  • Uses the generation prompt template and LLM to create content
  • Returns generated content as AIMessage for state tracking

Reflection Node

def _reflection_node(self, messages: Sequence[BaseMessage]) -> List[BaseMessage]:
    critique = self._build_reflection_chain().invoke({"messages": messages})
    return [HumanMessage(content=critique.content)]

-> How it works:

  • Analyzes the current conversation history
  • Uses reflection prompt template to generate critique
  • Returns feedback as HumanMessage (simulating human feedback)

Extension Points

The ReflectionAgent is designed to be highly customizable:

1. Custom Prompts

ReflectionAgent(
    llm=your_llm,
    generation_prompt=custom_generation_prompt,
    reflection_prompt=custom_reflection_prompt
)

-> Use cases:

  • Domain-specific content generation (technical docs, marketing copy, etc.)
  • Different writing styles or tones
  • Task-specific instructions and constraints

2. Evaluation Criteria

def _should_continue(self, state: List[BaseMessage]) -> bool:
    # Override this method for custom evaluation logic
    # Examples:
    # - Grammar and spelling validation
    # - Fact-checking against knowledge base
    # - Sentiment analysis
    # - Technical accuracy checks
    # - LLM-based quality scoring

-> Customization options:

  • Quality threshold scoring
  • Domain-specific validation rules
  • External API integration for fact-checking
  • Multi-criteria evaluation

3. Workflow Integration

def _build_generation_chain(self) -> Runnable:
    # Add pre-processing steps here
    return self.generation_prompt | self.llm

def _build_reflection_chain(self) -> Runnable:
    # Add post-processing steps here
    return self.reflection_prompt | self.llm

-> Extension possibilities:

  • Pre-processing: Research, data gathering, context enrichment
  • Post-processing: Formatting, translation, publishing
  • External APIs: Style guides, fact-checking services, databases
  • Human-in-the-loop: Manual review and approval steps

Real-World Applications

Content Creation:

  • Blog posts with SEO optimization
  • Technical documentation with accuracy checks
  • Marketing copy with engagement analysis

Code Generation:

  • Initial code with subsequent review and optimization
  • Documentation generation with clarity improvements
  • Test case creation with coverage analysis

Educational Content:

  • Lesson plans with pedagogical review
  • Exercise creation with difficulty balancing
  • Explanation refinement for different skill levels

Benefits

  • Higher Quality Output: Multiple iterations improve content quality
  • Automated Improvement: Reduces need for manual editing
  • Flexible Framework: Easily adaptable to different domains
  • Transparent Process: Full conversation history shows improvement process
  • Cost Efficient: Stops early when quality threshold is met

Installation

git clone https://github.com/codebywiam/reflection-agent-langgraph.git
cd reflection-agent-langgraph
pip install -r requirements.txt

Testing

The project includes comprehensive unit tests for the ReflectionAgent functionality. Tests are located in the tests/ directory.

Mock Testing

Tests use a MockChatModel class that simulates LLM responses without requiring actual API calls. This allows for:

  • Fast, deterministic test execution
  • Testing of edge cases and error conditions
  • Isolation of agent logic from external dependencies
pytest tests/test_agent.py

=================== test session starts ===================
tests/test_agent.py::TestReflectionAgent::test_initialization PASSED
tests/test_agent.py::TestReflectionAgent::test_run_agent PASSED
tests/test_agent.py::TestReflectionAgent::test_max_iterations_respected PASSED
tests/test_agent.py::TestReflectionAgent::test_custom_prompts PASSED
tests/test_agent.py::TestReflectionAgent::test_should_continue_logic PASSED
tests/test_agent.py::TestReflectionAgent::test_message_flow PASSED
=================== 6 passed in 0.12s ===================

Usage Examples

Basic Usage

from reflection_agent import ReflectionAgent
from langchain_openai import ChatOpenAI

# Initialize with your preferred LLM
llm = ChatOpenAI(model="gpt-4-turbo")
agent = ReflectionAgent(llm)

# Run with input
result = agent.run("Write a LinkedIn post about AI ethics")

Advanced Configuration

Custom Prompts & Reflection Depth

code_agent = ReflectionAgent(
    llm=llm,
    max_iterations=3,
    generation_prompt=CODE_GENERATION_PROMPT,
    reflection_prompt=CODE_REFLECTION_PROMPT
)

See the examples directory for Jupyter notebooks demonstrating:

  • LinkedIn post generation

  • Code refinement

Future Enhancements

  • Multi-modal content support (images, videos)
  • Integration with more LLM providers
  • Real-time collaboration features
  • Advanced analytics and insights
  • API service with rate limiting
  • Web interface for non-technical users
  • Integration with social media platforms
  • A/B testing capabilities for content variants

Acknowledgments

  • LangGraph Team for the excellent graph-based workflow framework
  • LangChain Community for comprehensive LLM tooling
  • OpenAI for providing powerful language models
  • The open-source community for inspiration and contributions

Contributing

Contributions are welcome! Please open an issue or submit a pull request.

License

This project is licensed under the MIT License - see the LICENSE file for details.

About

An implementation of a self-improving AI agent using LangGraph that can generate content, evaluate its quality, and refine outputs through reflection.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages