add main.py

This commit is contained in:
2026-05-28 17:42:58 +00:00
parent 9963000a1b
commit 50b90593df
+152
View File
@@ -0,0 +1,152 @@
"""
Main entry point for the Planning Agent assignment.
The program demonstrates 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.
Three example tasks are executed when run as a script.
"""
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
from rich.console import Console
# ---------------------------------------------------------------------------
# 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.
# In a real scenario this could invoke tools or perform computation.
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
# ---------------------------------------------------------------------------
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},
)
# Start from planning
graph.set_entry_point("planning")
agent = graph.compile()
# ---------------------------------------------------------------------------
# 7. Helper to run a task and print results
# ---------------------------------------------------------------------------
def run_task(task: str) -> None:
console = Console()
console.print(f"\n[bold cyan]Running task:[/bold cyan] {task}")
result_state = agent.invoke({"task": task})
plan = result_state.get("plan", [])
results = result_state.get("results", [])
console.print("\n[green]Plan:\n[/green]")
for i, step in enumerate(plan, 1):
console.print(f"{i}. {step}")
console.print("\n[blue]Execution results:\n[/blue]")
for r in 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(results) + "\nProvide a concise final answer.")
summary_response = llm.invoke([HumanMessage(content=summary_prompt)])
console.print(summary_response.content)
# ---------------------------------------------------------------------------
# 8. Main three example tasks
# ---------------------------------------------------------------------------
if __name__ == "__main__":
examples = [
"Compare Python and JavaScript in terms of performance, syntax simplicity, and ecosystem support.",
"Explain how to set up a basic Flask application with a single route.",
"Outline the steps required to deploy a Dockerized Node.js app to AWS Elastic Beanstalk.",
]
for ex in examples:
run_task(ex)
"""