112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
from typing import TypedDict, List, Optional
|
|
from langchain_openai import ChatOpenAI
|
|
from pydantic import SecretStr
|
|
from langgraph.graph import StateGraph, START, END
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
# ---------- LLM ----------
|
|
llm = ChatOpenAI(
|
|
model="openai/gpt-oss-20b",
|
|
base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
|
|
api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
|
|
temperature=0.2,
|
|
)
|
|
|
|
# ---------- State ----------
|
|
class PlanningState(TypedDict):
|
|
task: str
|
|
plan: List[str] | None
|
|
current_step: int
|
|
results: List[str]
|
|
|
|
# ---------- Nodes ----------
|
|
def planning(state: PlanningState) -> PlanningState:
|
|
prompt = (
|
|
f"Разбей задачу '{state['task']}' на 3–6 конкретных шагов. "
|
|
"Ответ в виде JSON массива строк, например:\n"
|
|
'["Шаг 1", "Шаг 2", ...]'
|
|
)
|
|
response = llm.invoke(prompt).content
|
|
try:
|
|
import json
|
|
plan = json.loads(response)
|
|
if not isinstance(plan, list):
|
|
raise ValueError
|
|
except Exception:
|
|
# fallback: simple split by newlines or numbers
|
|
plan = [line.strip() for line in response.splitlines() if line.strip()]
|
|
return {
|
|
**state,
|
|
"plan": plan,
|
|
"current_step": 0,
|
|
"results": [],
|
|
}
|
|
|
|
def execution(state: PlanningState) -> PlanningState:
|
|
step_idx = state["current_step"]
|
|
if state["plan"] is None or step_idx >= len(state["plan"]):
|
|
return state
|
|
step_text = state["plan"][step_idx]
|
|
# Execute the step (here we just echo it; replace with real logic)
|
|
result = f"[Шаг {step_idx + 1}] Выполнено: {step_text}"
|
|
new_results = state["results"] + [result]
|
|
return {
|
|
**state,
|
|
"current_step": step_idx + 1,
|
|
"results": new_results,
|
|
}
|
|
|
|
def should_continue(state: PlanningState) -> str:
|
|
if state["plan"] is None or state["current_step"] >= len(state["plan"]):
|
|
return "finish"
|
|
return "execute"
|
|
|
|
# ---------- Graph ----------
|
|
graph = StateGraph(PlanningState)
|
|
graph.add_node("planning", planning)
|
|
graph.add_node("execution", execution)
|
|
|
|
# Connect planning to the first execution step
|
|
graph.add_edge("planning", "execution")
|
|
|
|
# Conditional loop: after each execution, decide whether to continue or finish
|
|
graph.add_conditional_edges(
|
|
"execution",
|
|
should_continue,
|
|
{
|
|
"execute": "execution",
|
|
"finish": END,
|
|
},
|
|
)
|
|
|
|
graph.set_entry_point("planning")
|
|
workflow = graph.compile(checkpointer=InMemorySaver())
|
|
|
|
# ---------- Demo ----------
|
|
if __name__ == "__main__":
|
|
task_text = input("Задача: ").strip()
|
|
if not task_text:
|
|
print("Нет задачи.")
|
|
exit(0)
|
|
|
|
initial_state: PlanningState = {
|
|
"task": task_text,
|
|
"plan": None,
|
|
"current_step": 0,
|
|
"results": [],
|
|
}
|
|
|
|
result = workflow.invoke(initial_state)
|
|
plan = result["plan"]
|
|
results = result["results"]
|
|
|
|
print("\nПлан:")
|
|
for i, step in enumerate(plan, start=1):
|
|
print(f"{i}. {step}")
|
|
|
|
print("\nРезультаты выполнения:")
|
|
for r in results:
|
|
print(r)
|
|
|
|
final_summary = "\n".join(results)
|
|
print("\nИтог:\n", final_summary) |