Files
agent/src/app/modules/agent/engine/graphs/project_edits_graph.py
T
2026-03-12 16:55:23 +03:00

103 lines
4.1 KiB
Python

import logging
from langgraph.graph import END, START, StateGraph
from app.modules.agent.engine.graphs.progress import emit_progress_sync
from app.modules.agent.engine.graphs.project_edits_logic import ProjectEditsLogic
from app.modules.agent.engine.graphs.state import AgentGraphState
from app.modules.agent.llm import AgentLlmService
LOGGER = logging.getLogger(__name__)
class ProjectEditsGraphFactory:
_max_validation_attempts = 2
def __init__(self, llm: AgentLlmService) -> None:
self._logic = ProjectEditsLogic(llm)
def build(self, checkpointer=None):
graph = StateGraph(AgentGraphState)
graph.add_node("collect_context", self._collect_context)
graph.add_node("plan_changes", self._plan_changes)
graph.add_node("generate_changeset", self._generate_changeset)
graph.add_node("self_check", self._self_check)
graph.add_node("build_result", self._build_result)
graph.add_edge(START, "collect_context")
graph.add_edge("collect_context", "plan_changes")
graph.add_edge("plan_changes", "generate_changeset")
graph.add_edge("generate_changeset", "self_check")
graph.add_conditional_edges(
"self_check",
self._route_after_self_check,
{"retry": "generate_changeset", "ready": "build_result"},
)
graph.add_edge("build_result", END)
return graph.compile(checkpointer=checkpointer)
def _collect_context(self, state: AgentGraphState) -> dict:
emit_progress_sync(
state,
stage="graph.project_edits.collect_context",
message="Собираю контекст и релевантные файлы для правок.",
)
result = self._logic.collect_context(state)
self._log_step_result("collect_context", result)
return result
def _plan_changes(self, state: AgentGraphState) -> dict:
emit_progress_sync(
state,
stage="graph.project_edits.plan_changes",
message="Определяю, что именно нужно изменить и в каких файлах.",
)
result = self._logic.plan_changes(state)
self._log_step_result("plan_changes", result)
return result
def _generate_changeset(self, state: AgentGraphState) -> dict:
emit_progress_sync(
state,
stage="graph.project_edits.generate_changeset",
message="Формирую предлагаемые правки по выбранным файлам.",
)
result = self._logic.generate_changeset(state)
self._log_step_result("generate_changeset", result)
return result
def _self_check(self, state: AgentGraphState) -> dict:
emit_progress_sync(
state,
stage="graph.project_edits.self_check",
message="Проверяю, что правки соответствуют запросу и не трогают лишнее.",
)
result = self._logic.self_check(state)
self._log_step_result("self_check", result)
return result
def _build_result(self, state: AgentGraphState) -> dict:
emit_progress_sync(
state,
stage="graph.project_edits.build_result",
message="Формирую итоговый changeset и краткий обзор.",
)
result = self._logic.build_result(state)
self._log_step_result("build_result", result)
return result
def _route_after_self_check(self, state: AgentGraphState) -> str:
if state.get("validation_passed"):
return "ready"
attempts = int(state.get("validation_attempts", 0) or 0)
return "ready" if attempts >= self._max_validation_attempts else "retry"
def _log_step_result(self, step: str, result: dict) -> None:
LOGGER.warning(
"graph step result: graph=project_edits step=%s keys=%s changeset_items=%s answer_len=%s",
step,
sorted(result.keys()),
len(result.get("changeset", []) or []),
len(str(result.get("answer", "") or "")),
)