add main.py
This commit is contained in:
@@ -0,0 +1,209 @@
|
||||
"""Планирующий 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())
|
||||
Reference in New Issue
Block a user