add main.py

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)