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.
- 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
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
| 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 |
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
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
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
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
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)
The ReflectionAgent is designed to be highly customizable:
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
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
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
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
- 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
git clone https://github.com/codebywiam/reflection-agent-langgraph.git
cd reflection-agent-langgraph
pip install -r requirements.txtThe 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 ===================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")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
- 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
- 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
Contributions are welcome! Please open an issue or submit a pull request.
This project is licensed under the MIT License - see the LICENSE file for details.