diff --git a/main.py b/main.py new file mode 100644 index 0000000..a1ae8a5 --- /dev/null +++ b/main.py @@ -0,0 +1,173 @@ +""" +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 TypedDict, List, Optional + +from langgraph.graph import StateGraph, START, END +from langchain_core.messages import HumanMessage, SystemMessage +from langchain_openai import ChatOpenAI +from rich.console import Console +from rich.table import Table + +# --------------------------------------------------------------------------- +# 1. State definition +# --------------------------------------------------------------------------- +class PlanningState(TypedDict): + task: str + plan: List[str] | None + current_step: int + results: List[str] + +# --------------------------------------------------------------------------- +# 2. 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, +) + +# --------------------------------------------------------------------------- +# 3. Planning node – split task into steps +# --------------------------------------------------------------------------- +def planning(state: PlanningState) -> PlanningState: + """Ask the LLM to produce a numbered list of steps. + + The prompt forces JSON output for reliable parsing. + """ + 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 to parse JSON; if fails, fall back to simple split by lines. + try: + import json + data = json.loads(text) + plan: List[str] = data.get("plan", []) + except Exception: + # Fallback – split on newlines and strip numbering + plan = [] + for line in text.splitlines(): + line = line.strip() + if not line: + continue + # Remove leading numbers like "1. " or "- " + 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": [], + } + +# --------------------------------------------------------------------------- +# 4. Execution node – perform one step (here we just echo the step) +# --------------------------------------------------------------------------- +def execution(state: PlanningState) -> PlanningState: + """Execute a single step. + + In a real assignment you would replace this with calls to tools or other logic. + For demonstration, we simply record the step text as the result. + """ + idx = state["current_step"] + if idx >= len(state["plan"]): + return state + step_text = state["plan"][idx] + # Simulate execution – in practice you might call a tool here. + result = f"Executed: {step_text}" + new_results = state["results"].copy() + new_results.append(result) + return { + **state, + "current_step": idx + 1, + "results": new_results, + } + +# --------------------------------------------------------------------------- +# 5. Decision node – should we continue? +# --------------------------------------------------------------------------- +def should_continue(state: PlanningState) -> str: + if state["current_step"] >= len(state.get("plan", [])): + return "finish" + return "execute" + +# --------------------------------------------------------------------------- +# 6. 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() + +# --------------------------------------------------------------------------- +# 7. 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]}"}} + + # Run the agent – we capture intermediate states via a callback. + results: List[PlanningState] = [] + for event in agent.stream(state, config): + if isinstance(event, dict) and "messages" in event: + continue # ignore final message + results.append(event) + + # Print plan + 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) + + # Print execution results + console.print("\n[bold underline]Execution Results:[/]") + for r in state.get("results", []): + console.print(r) + + console.print("\n[green]Finished.[/]\n") + +# --------------------------------------------------------------------------- +# 8. 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)