146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
"""
|
||
Agent implementation for the Planning Agent assignment.
|
||
|
||
This module contains all logic for building a LangGraph agent that:
|
||
1. Plans a task into discrete steps using an LLM.
|
||
2. Executes each step sequentially, collecting results.
|
||
3. Returns a final summary of all results.
|
||
|
||
The graph is built in :func:`build_agent` and returned as a compiled object.
|
||
"""
|
||
|
||
import os
|
||
from typing import TypedDict, List
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import HumanMessage, SystemMessage
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. State definition
|
||
# ---------------------------------------------------------------------------
|
||
class PlanningState(TypedDict):
|
||
task: str
|
||
plan: List[str] | None
|
||
current_step: int
|
||
results: List[str]
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. LLM configuration – BroJS provider
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOpenAI(
|
||
model="openai/gpt-oss-20b:free",
|
||
base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
|
||
api_key=os.getenv("JOURNAL_MCP_PAT"),
|
||
temperature=0.2,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. Planning node – split the task into steps
|
||
# ---------------------------------------------------------------------------
|
||
def planning(state: PlanningState) -> PlanningState:
|
||
"""Ask LLM to produce a numbered list of steps.
|
||
|
||
The prompt asks for JSON output with a single key ``steps`` containing an array
|
||
of strings. This guarantees deterministic parsing.
|
||
"""
|
||
task = state["task"]
|
||
system_prompt = (
|
||
"You are a helpful assistant that breaks down a user request into a list of concrete, actionable steps."
|
||
)
|
||
user_prompt = (
|
||
f"Task: {task}\n\nReturn a JSON object with a single key 'steps' containing an array of strings. Each string should be a concise step. Provide between 3 and 6 steps.")
|
||
|
||
response = llm.invoke([SystemMessage(content=system_prompt), HumanMessage(content=user_prompt)])
|
||
# Parse JSON safely
|
||
import json, re
|
||
try:
|
||
data = json.loads(response.content)
|
||
steps = data.get("steps", [])
|
||
except Exception:
|
||
# Fallback: extract numbered list via regex
|
||
pattern = r"\d+\.\s*(.+)"
|
||
steps = [m.group(1).strip() for m in re.finditer(pattern, response.content)]
|
||
return {
|
||
"task": task,
|
||
"plan": steps,
|
||
"current_step": 0,
|
||
"results": [],
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Execution node – run one step and record result
|
||
# ---------------------------------------------------------------------------
|
||
def execution(state: PlanningState) -> PlanningState:
|
||
idx = state["current_step"]
|
||
plan = state["plan"] or []
|
||
if idx >= len(plan):
|
||
return state
|
||
step_text = plan[idx]
|
||
# For demonstration, we simply echo the step as result.
|
||
result = f"Result of step {idx+1}: {step_text}"
|
||
new_results = state["results"] + [result]
|
||
return {
|
||
"task": state["task"],
|
||
"plan": plan,
|
||
"current_step": idx + 1,
|
||
"results": new_results,
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. Decision node – continue or finish
|
||
# ---------------------------------------------------------------------------
|
||
def should_continue(state: PlanningState) -> str:
|
||
if state["current_step"] >= len(state.get("plan", [])):
|
||
return "finish"
|
||
return "execute"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6. Build graph and compile
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def build_agent() -> StateGraph:
|
||
graph = StateGraph(PlanningState)
|
||
graph.add_node("planning", planning)
|
||
graph.add_node("execution", execution)
|
||
graph.add_conditional_edges(
|
||
"planning",
|
||
lambda _: "execute" if _["plan"] else "finish",
|
||
)
|
||
graph.add_edge("execution", "should_continue")
|
||
graph.add_conditional_edges(
|
||
"should_continue",
|
||
should_continue,
|
||
{"execute": "execution", "finish": END},
|
||
)
|
||
graph.set_entry_point("planning")
|
||
return graph
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 7. Public helper to get compiled agent
|
||
# ---------------------------------------------------------------------------
|
||
def get_agent():
|
||
"""Return a compiled LangGraph agent ready for invocation."""
|
||
return build_agent().compile()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# If run as script, demonstrate usage with one example task.
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
from rich.console import Console
|
||
console = Console()
|
||
agent = get_agent()
|
||
task_text = "Compare Python and JavaScript in terms of performance, syntax simplicity, and ecosystem support."
|
||
result_state = agent.invoke({"task": task_text})
|
||
console.print("\n[bold cyan]Plan:\n[/bold cyan]")
|
||
for i, step in enumerate(result_state.get("plan", []), 1):
|
||
console.print(f"{i}. {step}")
|
||
console.print("\n[blue]Execution results:\n[/blue]")
|
||
for r in result_state.get("results", []):
|
||
console.print(r)
|
||
console.print("\n[bold magenta]Final summary:[/bold magenta]\n")
|
||
# Final LLM summarization
|
||
summary_prompt = (
|
||
"You have executed the following steps: \n" + "\n".join(result_state.get("results", [])) + "\nProvide a concise final answer.")
|
||
summary_response = llm.invoke([HumanMessage(content=summary_prompt)])
|
||
console.print(summary_response.content)
|