Add agent.py

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