task-6a1864fd8a94f887e50d4706/main.py

"""Планирующий LangGraph-агент: planning → execution (цикл) → итог."""
from __future__ import annotations

import json
import os
import re
import sys
from typing import Literal, TypedDict

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph

load_dotenv()

BROJS_INFERENCE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
DEFAULT_MODEL = "openai/gpt-oss-20b:free"
DEFAULT_TASK = "Сравни Python и JavaScript"


class PlanningState(TypedDict, total=False):
    task: str
    plan: list[str] | None
    current_step: int
    results: list[str]
    summary: str


def _api_key() -> str:
    return (
        os.getenv("JOURNAL_MCP_PAT")
        or os.getenv("JOURNAL_TOKEN")
        or os.getenv("OPENAI_API_KEY")
        or ""
    )


def _base_url() -> str:
    if os.getenv("OPENAI_BASE_URL"):
        return os.environ["OPENAI_BASE_URL"]
    if os.getenv("OPENAI_API_KEY") and not os.getenv("JOURNAL_MCP_PAT"):
        return os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
    return BROJS_INFERENCE_URL


def build_llm() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", DEFAULT_MODEL),
        base_url=_base_url(),
        api_key=_api_key(),
        temperature=0.2,
    )


def _parse_plan(text: str) -> list[str]:
    text = text.strip()
    if not text:
        return []
    try:
        data = json.loads(text)
        if isinstance(data, dict) and isinstance(data.get("steps"), list):
            return [str(s).strip() for s in data["steps"] if str(s).strip()]
        if isinstance(data, list):
            return [str(s).strip() for s in data if str(s).strip()]
    except json.JSONDecodeError:
        pass
    steps: list[str] = []
    for line in text.splitlines():
        line = line.strip()
        if not line:
            continue
        m = re.match(r"^[\d\-\*\.)]+\s*(.+)$", line)
        steps.append((m.group(1) if m else line).strip())
    return [s for s in steps if s][:6]


def planning(state: PlanningState, llm: ChatOpenAI) -> dict:
    task = state.get("task", DEFAULT_TASK)
    prompt = (
        "Разбей задачу на 3–6 конкретных шагов.\n"
        "Верни ТОЛЬКО JSON: {\"steps\": [\"шаг 1\", \"шаг 2\", ...]}\n\n"
        f"Задача: {task}"
    )
    raw = llm.invoke(
        [
            SystemMessage(content="Ты планировщик. Отвечай только валидным JSON."),
            HumanMessage(content=prompt),
        ]
    ).content
    plan = _parse_plan(str(raw))
    if len(plan) < 3:
        plan = [
            "Собрать ключевые характеристики первого объекта",
            "Собрать ключевые характеристики второго объекта",
            "Сравнить сходства и различия",
            "Сформулировать вывод",
        ]
    return {"plan": plan, "current_step": 0, "results": []}


def execution(state: PlanningState, llm: ChatOpenAI) -> dict:
    plan = state.get("plan") or []
    idx = int(state.get("current_step", 0))
    results = list(state.get("results") or [])
    task = state.get("task", DEFAULT_TASK)
    step_text = plan[idx]

    prompt = (
        f"Общая задача: {task}\n"
        f"Текущий шаг ({idx + 1}/{len(plan)}): {step_text}\n\n"
        "Выполни только этот шаг. Ответ — 2–5 предложений."
    )
    raw = llm.invoke([HumanMessage(content=prompt)]).content
    step_result = str(raw).strip()
    results.append(step_result)
    return {"results": results, "current_step": idx + 1}


def finish(state: PlanningState, llm: ChatOpenAI) -> dict:
    task = state.get("task", DEFAULT_TASK)
    plan = state.get("plan") or []
    results = state.get("results") or []
    parts = [f"Шаг {i + 1}: {plan[i]}\n{r}" for i, r in enumerate(results)]
    body = "\n\n".join(parts)
    prompt = (
        f"Задача: {task}\n\nРезультаты по шагам:\n{body}\n\n"
        "Сформируй краткую итоговую сводку (4–6 предложений)."
    )
    raw = llm.invoke([HumanMessage(content=prompt)]).content
    return {"summary": str(raw).strip()}


def should_continue(state: PlanningState) -> Literal["execute", "finish"]:
    plan = state.get("plan") or []
    if int(state.get("current_step", 0)) >= len(plan):
        return "finish"
    return "execute"


def build_graph(llm: ChatOpenAI | None = None):
    llm = llm or build_llm()

    def _planning(state: PlanningState) -> dict:
        return planning(state, llm)

    def _execution(state: PlanningState) -> dict:
        return execution(state, llm)

    def _finish(state: PlanningState) -> dict:
        return finish(state, llm)

    graph = StateGraph(PlanningState)
    graph.add_node("planning", _planning)
    graph.add_node("execution", _execution)
    graph.add_node("finish", _finish)

    graph.add_edge(START, "planning")
    graph.add_edge("planning", "execution")
    graph.add_conditional_edges(
        "execution",
        should_continue,
        {"execute": "execution", "finish": "finish"},
    )
    graph.add_edge("finish", END)

    return graph.compile()


def run_demo(task: str = DEFAULT_TASK) -> PlanningState:
    app = build_graph()
    initial: PlanningState = {
        "task": task,
        "plan": None,
        "current_step": 0,
        "results": [],
        "summary": "",
    }

    print(f"Задача: {task}\n")
    final: PlanningState = dict(initial)

    for event in app.stream(initial, stream_mode="updates"):
        for node, update in event.items():
            if not isinstance(update, dict):
                continue
            final = {**final, **update}
            if node == "planning" and final.get("plan"):
                print("План:")
                for i, step in enumerate(final["plan"], 1):
                    print(f"{i}. {step}")
                print()
            elif node == "execution" and final.get("results"):
                n = len(final["results"])
                print(f"[Шаг {n}] {final['results'][-1]}\n")
            elif node == "finish":
                print(f"Итог: {final.get('summary', '')}")

    return final


def main() -> int:
    task = " ".join(sys.argv[1:]).strip() or DEFAULT_TASK
    run_demo(task)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())