247 lines
9.7 KiB
Python
247 lines
9.7 KiB
Python
# orchestrator.py
|
||
from typing import Dict, Any, Literal
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.checkpoint.memory import MemorySaver
|
||
import uuid
|
||
from datetime import datetime
|
||
from models import Task, TaskStatus, CodeSolution, ReviewResult
|
||
from agents import DeveloperAgent, ReviewerAgent
|
||
from config import MAX_REVIEW_RETRIES
|
||
|
||
class Orchestrator:
|
||
"""Оркестратор - управляет процессом разработки"""
|
||
|
||
def __init__(self):
|
||
self.developer = DeveloperAgent()
|
||
self.reviewer = ReviewerAgent()
|
||
self.max_retries = MAX_REVIEW_RETRIES
|
||
|
||
# Создаем граф состояний
|
||
self.workflow = self._build_workflow()
|
||
self.checkpointer = MemorySaver()
|
||
self.app = self.workflow.compile()
|
||
|
||
def _build_workflow(self) -> StateGraph:
|
||
"""Строит граф процесса разработки"""
|
||
|
||
# Определяем состояния
|
||
workflow = StateGraph(Dict)
|
||
|
||
# Добавляем узлы
|
||
workflow.add_node("developer", self._developer_node)
|
||
workflow.add_node("reviewer", self._reviewer_node)
|
||
workflow.add_node("finalize", self._finalize_node)
|
||
|
||
# Определяем переходы
|
||
workflow.set_entry_point("developer")
|
||
workflow.add_edge("developer", "reviewer")
|
||
|
||
# Условные переходы после ревью
|
||
workflow.add_conditional_edges(
|
||
"reviewer",
|
||
self._decide_next_step,
|
||
{
|
||
"approved": "finalize",
|
||
"rework": "developer",
|
||
"rejected": "finalize",
|
||
"failed": "finalize"
|
||
}
|
||
)
|
||
|
||
workflow.add_edge("finalize", END)
|
||
|
||
return workflow
|
||
|
||
def _developer_node(self, state: Dict) -> Dict:
|
||
"""Узел разработчика"""
|
||
task = state["task"]
|
||
print(f"\n🔧 [Iteration {task.iteration}] Разработчик начал работу...")
|
||
|
||
# Получаем замечания с предыдущего ревью
|
||
previous_comments = []
|
||
if task.last_review and task.last_review.comments:
|
||
previous_comments = [
|
||
{"file": c.file, "line": c.line, "text": c.text}
|
||
for c in task.last_review.comments
|
||
]
|
||
|
||
try:
|
||
# Генерируем код
|
||
solution = self.developer.develop(task, previous_comments)
|
||
task.solution = solution
|
||
task.status = TaskStatus.REVIEW
|
||
|
||
print(f"✅ Разработчик создал решение: {len(solution.files)} файлов")
|
||
for filename in solution.files:
|
||
print(f" - {filename}")
|
||
|
||
except Exception as e:
|
||
task.status = TaskStatus.FAILED
|
||
task.error_log.append(f"Developer error: {str(e)}")
|
||
print(f"❌ Ошибка разработчика: {e}")
|
||
|
||
return {"task": task}
|
||
|
||
def _reviewer_node(self, state: Dict) -> Dict:
|
||
"""Узел ревьювера"""
|
||
task = state["task"]
|
||
print(f"\n🔍 [Iteration {task.iteration}] Ревьювер проверяет код...")
|
||
|
||
if not task.solution:
|
||
task.status = TaskStatus.FAILED
|
||
task.error_log.append("No solution to review")
|
||
return {"task": task}
|
||
|
||
try:
|
||
# Проводим ревью
|
||
review = self.reviewer.review(task, task.solution)
|
||
task.last_review = review
|
||
task.iteration += 1
|
||
|
||
print(f"📊 Вердикт: {review.status}")
|
||
if review.comments:
|
||
print(f" Найдено замечаний: {len(review.comments)}")
|
||
for comment in review.comments[:3]: # Показываем первые 3
|
||
print(f" - [{comment.severity}] {comment.text[:100]}")
|
||
|
||
# Обновляем статус
|
||
if review.status == "approved":
|
||
task.status = TaskStatus.APPROVED
|
||
print("✅ Код одобрен!")
|
||
elif review.status == "rejected":
|
||
task.status = TaskStatus.REJECTED
|
||
print("❌ Код отклонен архитектурно")
|
||
else:
|
||
task.status = TaskStatus.REWORK
|
||
print("🔄 Требуются доработки")
|
||
|
||
except Exception as e:
|
||
task.status = TaskStatus.FAILED
|
||
task.error_log.append(f"Reviewer error: {str(e)}")
|
||
print(f"❌ Ошибка ревьювера: {e}")
|
||
|
||
return {"task": task}
|
||
|
||
def _decide_next_step(self, state: Dict) -> Literal["approved", "rework", "rejected", "failed"]:
|
||
"""Решает, что делать дальше"""
|
||
task = state["task"]
|
||
|
||
if task.status == TaskStatus.APPROVED:
|
||
return "approved"
|
||
elif task.status == TaskStatus.REJECTED:
|
||
return "rejected"
|
||
elif task.status == TaskStatus.FAILED:
|
||
return "failed"
|
||
elif task.iteration >= self.max_retries:
|
||
print(f"\n⚠️ Достигнут лимит итераций ({self.max_retries}). Останавливаемся.")
|
||
task.status = TaskStatus.FAILED
|
||
return "failed"
|
||
else:
|
||
return "rework"
|
||
|
||
def _finalize_node(self, state: Dict) -> Dict:
|
||
"""Финальный узел - подведение итогов"""
|
||
task = state["task"]
|
||
|
||
print("\n" + "="*50)
|
||
print("📋 ИТОГИ РАБОТЫ")
|
||
print("="*50)
|
||
print(f"Задача: {task.requirement[:100]}...")
|
||
print(f"Статус: {task.status}")
|
||
print(f"Итераций: {task.iteration}")
|
||
|
||
if task.status == TaskStatus.APPROVED and task.solution:
|
||
print(f"\n✅ Решение принято!")
|
||
print(f" Файлов: {len(task.solution.files)}")
|
||
print(f" Тестов: {len(task.solution.tests)}")
|
||
|
||
# Сохраняем решение в файлы
|
||
self._save_solution(task.solution)
|
||
|
||
elif task.last_review:
|
||
print(f"\n📝 Финальный вердикт: {task.last_review.summary}")
|
||
|
||
if task.error_log:
|
||
print(f"\n❌ Ошибки: {len(task.error_log)}")
|
||
for error in task.error_log:
|
||
print(f" - {error}")
|
||
|
||
return {"task": task, "finalized": True}
|
||
|
||
def _save_solution(self, solution: CodeSolution):
|
||
"""Сохраняет решение в файлы"""
|
||
import os
|
||
output_dir = "generated_code"
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
for filename, content in solution.files.items():
|
||
filepath = os.path.join(output_dir, filename)
|
||
with open(filepath, "w", encoding="utf-8") as f:
|
||
f.write(content)
|
||
print(f" 💾 Сохранен файл: {filepath}")
|
||
|
||
for testname, test_content in solution.tests.items():
|
||
testpath = os.path.join(output_dir, testname)
|
||
with open(testpath, "w", encoding="utf-8") as f:
|
||
f.write(test_content)
|
||
print(f" 💾 Сохранен тест: {testpath}")
|
||
|
||
def run(self, requirement: str, task_id: str = None) -> Dict:
|
||
"""Запускает процесс разработки"""
|
||
|
||
if not task_id:
|
||
task_id = str(uuid.uuid4())
|
||
|
||
initial_task = Task(
|
||
id=task_id,
|
||
requirement=requirement,
|
||
status=TaskStatus.PENDING,
|
||
iteration=0
|
||
)
|
||
|
||
print(f"\n🚀 Запуск мультиагентной системы")
|
||
print(f"📝 Задача: {requirement}")
|
||
print(f"🆔 ID: {task_id}")
|
||
|
||
# Запускаем граф
|
||
initial_state = {"task": initial_task}
|
||
final_state = self.app.invoke(initial_state)
|
||
|
||
return final_state
|
||
|
||
# main.py
|
||
def main():
|
||
"""Пример использования"""
|
||
|
||
# Создаем оркестратора
|
||
orchestrator = Orchestrator()
|
||
|
||
# Пример задачи
|
||
requirement = """
|
||
Реализуй класс BankAccount с методами:
|
||
- deposit(amount): пополнение счета
|
||
- withdraw(amount): снятие средств (нельзя снять больше, чем есть)
|
||
- get_balance(): получение баланса
|
||
- add_interest(rate): добавление процентов (rate в процентах)
|
||
|
||
Требования:
|
||
- Нельзя создать счет с отрицательным балансом
|
||
- Все операции должны быть потокобезопасными (используй threading.Lock)
|
||
- Напиши юнит-тесты (pytest) для всех методов, включая граничные случаи
|
||
- Добавь docstring для всех методов
|
||
"""
|
||
|
||
# Запускаем процесс
|
||
result = orchestrator.run(requirement)
|
||
|
||
# Анализируем результат
|
||
task = result["task"]
|
||
if task.status == TaskStatus.APPROVED:
|
||
print("\n🎉 Успех! Код принят и сохранен в папке 'generated_code'")
|
||
else:
|
||
print(f"\n💔 Неудача. Статус: {task.status}")
|
||
if task.last_review:
|
||
print(f"Причина: {task.last_review.summary}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |