diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..5751bb7 --- /dev/null +++ b/agent.py @@ -0,0 +1,96 @@ +from typing import TypedDict, List, Optional +import json +import asyncio +from langchain_openai import ChatOpenAI +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_conditional_node + +class PlanningState(TypedDict): + task: str + plan: List[str] | None + current_step: int + results: List[str] + +async def planning(state: PlanningState, llm: ChatOpenAI) -> PlanningState: + """LLM generates a JSON plan with a list of steps.""" + prompt = ( + "Разбей задачу на 3–6 конкретных шагов. Возвращай JSON с ключом \"plan\".\n\n" + f"Задача: {state['task']}" + ) + response = await llm.ainvoke(prompt) + # response may be a string or a ChatResult + text = response if isinstance(response, str) else response.content + plan: List[str] = [] + try: + data = json.loads(text) + plan = data.get("plan", []) + except Exception: + # Fallback: parse numbered list + for line in text.splitlines(): + line = line.strip() + if not line: + continue + # Remove leading number or bullet + if line[0].isdigit(): + parts = line.split('.', 1) + if len(parts) == 2: + step = parts[1].strip() + else: + step = line + plan.append(step) + elif line[0] in ('-','*'): + plan.append(line[1:].strip()) + return { + **state, + "plan": plan, + "current_step": 0, + "results": [] + } + +async def execution(state: PlanningState) -> PlanningState: + """Execute a single step and record the result.""" + if state["plan"] is None or state["current_step"] >= len(state["plan"]): + return state + step = state["plan"][state["current_step"]] + result = f"Шаг {state['current_step']+1}: {step}" + new_results = state["results"] + [result] + return { + **state, + "results": new_results, + "current_step": state["current_step"] + 1 + } + +def should_continue(state: PlanningState) -> str: + """Decide whether to finish or execute another step.""" + if state["current_step"] >= len(state["plan"] or []): + return "finish" + return "execute" + +async def planning_node(state: PlanningState, llm: ChatOpenAI) -> PlanningState: + return await planning(state, llm) + +def create_agent(llm: ChatOpenAI): + workflow = StateGraph(PlanningState) + # Add nodes + workflow.add_node("planning", lambda state: planning_node(state, llm)) + workflow.add_node("execution", execution) + workflow.add_conditional_node("should_continue", should_continue, { + "execute": "execution", + "finish": END + }) + # Set entry point and edges + workflow.set_entry_point("planning") + workflow.add_edge("planning", "execution") + workflow.add_edge("execution", "should_continue") + return workflow.compile() + +async def run_agent(task: str, llm: ChatOpenAI) -> PlanningState: + agent = create_agent(llm) + initial_state: PlanningState = { + "task": task, + "plan": None, + "current_step": 0, + "results": [] + } + result = await agent.ainvoke(initial_state) + return result \ No newline at end of file