149 lines
5.1 KiB
Python
149 lines
5.1 KiB
Python
"""
|
||
Main entry point for the "Planning Agent" assignment.
|
||
|
||
The project demonstrates a LangGraph agent that first plans a task into discrete steps and then executes those steps one by one.
|
||
|
||
Examples are provided in the ``__main__`` section – run the script with different tasks to see how the planner works.
|
||
|
||
Requirements:
|
||
- langgraph>=0.2.0
|
||
- langchain-openai>=0.3.0
|
||
- python-dotenv>=1.0.0
|
||
- rich>=13.0.0
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import List
|
||
|
||
from langgraph.graph import StateGraph, END
|
||
from langchain_core.messages import HumanMessage, SystemMessage
|
||
from langchain_openai import ChatOpenAI
|
||
from rich.console import Console
|
||
from rich.table import Table
|
||
|
||
# Import shared state definition
|
||
from models import PlanningState
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. LLM configuration – BroJS endpoint
|
||
# ---------------------------------------------------------------------------
|
||
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,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. Planning node – split task into steps
|
||
# ---------------------------------------------------------------------------
|
||
def planning(state: PlanningState) -> PlanningState:
|
||
system = SystemMessage(
|
||
content="You are a helpful assistant that splits a task into clear, actionable steps. Return a JSON array of strings under the key `plan`."
|
||
)
|
||
user = HumanMessage(content=f"Plan the following task: {state['task']}")
|
||
|
||
response = llm.invoke([system, user])
|
||
text = response.content.strip()
|
||
|
||
try:
|
||
import json
|
||
data = json.loads(text)
|
||
plan: List[str] = data.get("plan", [])
|
||
except Exception:
|
||
plan = []
|
||
for line in text.splitlines():
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
if line[0].isdigit() and (len(line) > 2 and line[1] in ".-"):
|
||
line = line.split("", 1)[1]
|
||
plan.append(line)
|
||
return {
|
||
**state,
|
||
"plan": plan,
|
||
"current_step": 0,
|
||
"results": [],
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. Execution node – perform one step (here we just echo the step)
|
||
# ---------------------------------------------------------------------------
|
||
def execution(state: PlanningState) -> PlanningState:
|
||
idx = state["current_step"]
|
||
if idx >= len(state.get("plan", [])):
|
||
return state
|
||
step_text = state["plan"][idx]
|
||
result = f"Executed: {step_text}"
|
||
new_results = state["results"].copy()
|
||
new_results.append(result)
|
||
return {
|
||
**state,
|
||
"current_step": idx + 1,
|
||
"results": new_results,
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 4. Decision node – should we continue?
|
||
# ---------------------------------------------------------------------------
|
||
def should_continue(state: PlanningState) -> str:
|
||
if state["current_step"] >= len(state.get("plan", [])):
|
||
return "finish"
|
||
return "execute"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 5. Build the graph
|
||
# ---------------------------------------------------------------------------
|
||
graph = StateGraph(PlanningState)
|
||
graph.add_node("planning", planning)
|
||
graph.add_node("execution", execution)
|
||
graph.add_conditional_edges(
|
||
"planning",
|
||
lambda _: "execute" if _.get("plan") else "finish",
|
||
)
|
||
graph.add_edge("execution", "should_continue")
|
||
graph.add_conditional_edges(
|
||
"should_continue",
|
||
should_continue,
|
||
{"execute": "execution", "finish": END},
|
||
)
|
||
agent = graph.compile()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 6. Helper to run the agent and pretty‑print results
|
||
# ---------------------------------------------------------------------------
|
||
def run_agent(task: str) -> None:
|
||
console = Console()
|
||
state: PlanningState = {"task": task, "plan": None, "current_step": 0, "results": []}
|
||
config = {"configurable": {"thread_id": f"{task[:8]}"}}
|
||
|
||
final_state = agent.invoke(state, config)
|
||
state = final_state
|
||
|
||
console.print("\n[bold underline]Task:[/]", task)
|
||
if state.get("plan"):
|
||
table = Table(title="Plan", show_header=False, box=None)
|
||
for i, step in enumerate(state["plan"], 1):
|
||
table.add_row(f"{i}. {step}")
|
||
console.print(table)
|
||
|
||
console.print("\n[bold underline]Execution Results:[/]")
|
||
for r in state.get("results", []):
|
||
console.print(r)
|
||
|
||
console.print("\n[green]Finished.[/]\n")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 7. Demo examples
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
examples = [
|
||
"Compare Python and JavaScript in web development.",
|
||
"Plan a weekend trip to the mountains.",
|
||
"Explain how a neural network learns.",
|
||
]
|
||
for ex in examples:
|
||
run_agent(ex)
|