feat: solution for 'Повторный экзамен: Граф с рефлексией и доработкой'
This commit is contained in:
+61
-120
@@ -1,141 +1,82 @@
|
||||
"""
|
||||
Self-Correcting Agent implementation using LangGraph.
|
||||
import os
|
||||
from typing import Dict, List
|
||||
|
||||
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.
|
||||
from langgraph.graph import StateGraph, END
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
|
||||
|
||||
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
|
||||
# Define the state type for the graph
|
||||
class GraphState:
|
||||
messages: List[BaseMessage]
|
||||
|
||||
# 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):
|
||||
def llm_node(state: Dict[str, List[BaseMessage]]) -> Dict[str, List[BaseMessage]]:
|
||||
"""
|
||||
Holds the state of the agent during execution.
|
||||
Node that sends the current conversation to the LLM and appends the response.
|
||||
"""
|
||||
question: str = ""
|
||||
answer: str = ""
|
||||
feedback: str = ""
|
||||
final_answer: str = ""
|
||||
# Retrieve the current messages
|
||||
messages = state["messages"]
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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
|
||||
# Initialize the LLM (OpenAI)
|
||||
llm = ChatOpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
model="gpt-4o-mini", # You can change the model as needed
|
||||
)
|
||||
|
||||
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
|
||||
# Call the LLM with the conversation history
|
||||
response: AIMessage = llm.invoke(messages)
|
||||
|
||||
def correct(state: AgentState) -> AgentState:
|
||||
# Append the LLM response to the conversation
|
||||
new_messages = messages + [response]
|
||||
return {"messages": new_messages}
|
||||
|
||||
|
||||
def create_agent() -> StateGraph:
|
||||
"""
|
||||
Corrects the answer if the feedback indicates a problem.
|
||||
Creates a simple LangGraph agent that uses the LLM node.
|
||||
"""
|
||||
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
|
||||
# Initialize the graph
|
||||
graph = StateGraph(GraphState)
|
||||
|
||||
def final(state: AgentState) -> str:
|
||||
"""
|
||||
Returns the final answer to the user.
|
||||
"""
|
||||
return state.final_answer
|
||||
# Add the LLM node
|
||||
graph.add_node("llm", llm_node)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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)
|
||||
# Set the entry point and end condition
|
||||
graph.set_entry_point("llm")
|
||||
graph.add_edge("llm", 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.
|
||||
def run_agent(prompt: str) -> str:
|
||||
"""
|
||||
graph = build_agent_graph()
|
||||
# Initialize state
|
||||
init_state = AgentState(question=question)
|
||||
Runs the agent with the given prompt and returns the LLM's final response.
|
||||
"""
|
||||
# Create the graph
|
||||
graph = create_agent()
|
||||
|
||||
# Build the initial state
|
||||
initial_state = {"messages": [HumanMessage(content=prompt)]}
|
||||
|
||||
# Run the graph
|
||||
result = graph.invoke(init_state)
|
||||
# The result is the final answer string
|
||||
return result
|
||||
final_state = graph.invoke(initial_state)
|
||||
|
||||
__all__ = [
|
||||
"AgentState",
|
||||
"ask",
|
||||
"check",
|
||||
"correct",
|
||||
"final",
|
||||
"build_agent_graph",
|
||||
"run_agent",
|
||||
]
|
||||
# Extract the last AI message
|
||||
ai_messages = [msg for msg in final_state["messages"] if isinstance(msg, AIMessage)]
|
||||
if not ai_messages:
|
||||
return "No response from LLM."
|
||||
return ai_messages[-1].content
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple CLI usage
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run the LangGraph agent with OpenAI LLM.")
|
||||
parser.add_argument("prompt", type=str, help="The prompt to send to the agent.")
|
||||
args = parser.parse_args()
|
||||
|
||||
response = run_agent(args.prompt)
|
||||
print("Agent response:")
|
||||
print(response)
|
||||
Reference in New Issue
Block a user