141 lines
4.0 KiB
Python
141 lines
4.0 KiB
Python
"""
|
|
Self-Correcting Agent implementation using LangGraph.
|
|
|
|
This module defines a simple LangGraph that:
|
|
1. Generates an answer to a user question.
|
|
2. Checks the quality of the answer.
|
|
3. Corrects the answer if needed.
|
|
4. Returns the final answer.
|
|
|
|
The graph is intentionally simple to satisfy the assignment specification
|
|
and to remain fully importable without external API keys.
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict
|
|
|
|
# Import LangGraph components
|
|
try:
|
|
from langgraph.graph import StateGraph, State, END
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"langgraph is required. Install it via 'pip install langgraph==0.0.1'"
|
|
) from exc
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# State definition
|
|
# --------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class AgentState(State):
|
|
"""
|
|
Holds the state of the agent during execution.
|
|
"""
|
|
question: str = ""
|
|
answer: str = ""
|
|
feedback: str = ""
|
|
final_answer: str = ""
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Node implementations
|
|
# --------------------------------------------------------------------------- #
|
|
def ask(state: AgentState) -> AgentState:
|
|
"""
|
|
Generates an answer to the provided question.
|
|
"""
|
|
# In a real implementation, this would call an LLM.
|
|
# Here we use a deterministic placeholder.
|
|
state.answer = f"Answer to: {state.question}"
|
|
return state
|
|
|
|
def check(state: AgentState) -> AgentState:
|
|
"""
|
|
Checks the quality of the generated answer.
|
|
"""
|
|
# Simple heuristic: if the answer contains the word 'bad', flag it.
|
|
if "bad" in state.answer.lower():
|
|
state.feedback = "Needs correction"
|
|
else:
|
|
state.feedback = "Good"
|
|
return state
|
|
|
|
def correct(state: AgentState) -> AgentState:
|
|
"""
|
|
Corrects the answer if the feedback indicates a problem.
|
|
"""
|
|
if state.feedback == "Needs correction":
|
|
# In a real scenario, this would call an LLM to rewrite the answer.
|
|
state.final_answer = f"Corrected: {state.answer}"
|
|
else:
|
|
state.final_answer = state.answer
|
|
return state
|
|
|
|
def final(state: AgentState) -> str:
|
|
"""
|
|
Returns the final answer to the user.
|
|
"""
|
|
return state.final_answer
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Graph construction
|
|
# --------------------------------------------------------------------------- #
|
|
def build_agent_graph() -> StateGraph:
|
|
"""
|
|
Builds and returns the LangGraph for the self-correcting agent.
|
|
"""
|
|
graph = StateGraph(AgentState)
|
|
|
|
# Add nodes
|
|
graph.add_node("ask", ask)
|
|
graph.add_node("check", check)
|
|
graph.add_node("correct", correct)
|
|
graph.add_node("final", final)
|
|
|
|
# Define edges
|
|
graph.set_entry_point("ask")
|
|
graph.add_edge("ask", "check")
|
|
|
|
# Conditional transition from check to either correct or final
|
|
def check_transition(state: AgentState) -> str:
|
|
return "correct" if state.feedback != "Good" else "final"
|
|
|
|
graph.add_conditional_edges("check", check_transition)
|
|
|
|
graph.add_edge("correct", "final")
|
|
graph.add_edge("final", END)
|
|
|
|
return graph
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Public API
|
|
# --------------------------------------------------------------------------- #
|
|
def run_agent(question: str) -> str:
|
|
"""
|
|
Runs the self-correcting agent on the given question.
|
|
|
|
Parameters
|
|
----------
|
|
question : str
|
|
The user question to answer.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
The final answer produced by the agent.
|
|
"""
|
|
graph = build_agent_graph()
|
|
# Initialize state
|
|
init_state = AgentState(question=question)
|
|
# Run the graph
|
|
result = graph.invoke(init_state)
|
|
# The result is the final answer string
|
|
return result
|
|
|
|
__all__ = [
|
|
"AgentState",
|
|
"ask",
|
|
"check",
|
|
"correct",
|
|
"final",
|
|
"build_agent_graph",
|
|
"run_agent",
|
|
] |