224 lines
11 KiB
Python
224 lines
11 KiB
Python
"""
|
||
Определяет логику рабочих процессов агента,
|
||
используя LangGraph для управления состояниями
|
||
и переходами между узлами обработки.
|
||
"""
|
||
|
||
from typing import Dict, Any, Optional, Set
|
||
from langgraph.graph import StateGraph, END
|
||
from graph_history_manager import GraphHistoryManager
|
||
from models import AgentState
|
||
from nodes import parse_command_node, execute_command_node, call_llm_node, generate_images_node, analyze_image_node, get_meet_subtitles_node, get_teams_subtitles_node, summarize_history_node, handle_error_node, help_node
|
||
import uuid # Добавлено для генерации UUID
|
||
from llm_client import get_llm
|
||
# --- LangGraph: Построение графа ---
|
||
|
||
graph_history_manager = GraphHistoryManager()
|
||
|
||
workflow = StateGraph(AgentState)
|
||
|
||
# Добавляем узлы
|
||
workflow.add_node("parse_command", parse_command_node)
|
||
workflow.add_node("execute_command",
|
||
execute_command_node) # Это по сути роутер
|
||
workflow.add_node("call_llm", call_llm_node)
|
||
workflow.add_node("generate_images", generate_images_node)
|
||
workflow.add_node("analyze_image", analyze_image_node)
|
||
workflow.add_node("get_meet_subtitles", get_meet_subtitles_node)
|
||
workflow.add_node("get_teams_subtitles", get_teams_subtitles_node)
|
||
workflow.add_node("summarize_history", summarize_history_node)
|
||
workflow.add_node("handle_error", handle_error_node)
|
||
workflow.add_node("help", help_node)
|
||
|
||
# Устанавливаем точку входа
|
||
workflow.set_entry_point("parse_command")
|
||
|
||
|
||
# Определяем маршрутизацию после parse_command
|
||
def route_command(state: AgentState) -> str:
|
||
"""Маршрутизирует выполнение на основе распознанной команды."""
|
||
if state.error:
|
||
return "handle_error" # Если была ошибка парсинга команды
|
||
elif state.command == "imagine":
|
||
return "generate_images"
|
||
elif state.command == "analyze":
|
||
return "analyze_image"
|
||
elif state.command == "subtitles_meet":
|
||
return "get_meet_subtitles"
|
||
elif state.command == "subtitles_teams":
|
||
return "get_teams_subtitles"
|
||
elif state.command == "summarize":
|
||
return "summarize_history"
|
||
elif state.command == "help":
|
||
return "help"
|
||
elif state.command == "chat":
|
||
return "call_llm"
|
||
elif state.command == "error": # Команда была, но не известная
|
||
return "handle_error"
|
||
else:
|
||
# Fallback на LLM, если команда не распознана или пуста (хотя parse_command должен был бы ее поймать)
|
||
return "call_llm"
|
||
|
||
|
||
workflow.add_conditional_edges(
|
||
"parse_command", # Откуда
|
||
route_command, # Функция, определяющая куда идти
|
||
{
|
||
"generate_images": "generate_images",
|
||
"analyze_image": "analyze_image",
|
||
"get_meet_subtitles": "get_meet_subtitles",
|
||
"get_teams_subtitles": "get_teams_subtitles",
|
||
"summarize_history": "summarize_history",
|
||
"help": "help",
|
||
"call_llm": "call_llm",
|
||
"handle_error":
|
||
"handle_error", # Для ошибок, найденных в parse_command
|
||
},
|
||
)
|
||
|
||
# Все узлы операций (кроме ошибок) ведут к END после завершения
|
||
workflow.add_edge("call_llm", END)
|
||
workflow.add_edge("generate_images", END)
|
||
workflow.add_edge("analyze_image", END)
|
||
workflow.add_edge("get_meet_subtitles", END)
|
||
workflow.add_edge("get_teams_subtitles", END)
|
||
workflow.add_edge("summarize_history", END)
|
||
workflow.add_edge("help", END)
|
||
workflow.add_edge("handle_error",
|
||
END) # Ошибка - это тоже конечный пункт для текущего запроса
|
||
|
||
# Компилируем граф
|
||
app = workflow.compile()
|
||
|
||
|
||
def run_agent(user_input: str,
|
||
existing_graph_id: Optional[str] = None,
|
||
parent_node_id: Optional[str] = None,
|
||
system_prompt: Optional[str] = None,
|
||
model: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
Запускает агента с заданным пользовательским вводом.
|
||
Может продолжить существующий граф по graph_id.
|
||
Возвращает обновленное состояние и ID графа.
|
||
"""
|
||
|
||
# ---------------------------------------------------- Initial State Setup ----------------------------------------------------------------
|
||
initial_state = AgentState(
|
||
input=user_input,
|
||
parent_node_id=parent_node_id,
|
||
system_prompt=system_prompt,
|
||
selected_model=model or "gemini-2.5-flash"
|
||
)
|
||
|
||
# Сохраняем исходные ID узлов и ребер для определения новых после выполнения графа
|
||
original_node_ids: Set[str] = set()
|
||
original_edge_ids: Set[str] = set()
|
||
|
||
# Если есть существующий ID графа, загружаем его историю и структуру
|
||
loaded_graph_data = None
|
||
if existing_graph_id:
|
||
# Получаем данные графа, включая сгенерированные сообщения до parent_node_id (либо current_node_id из БД)
|
||
# Если parent_node_id передан, используем его для построения истории.
|
||
loaded_graph_data = graph_history_manager.get_graph(
|
||
existing_graph_id, target_node_id=parent_node_id)
|
||
|
||
if loaded_graph_data:
|
||
print(f"Продолжаю существующий граф {existing_graph_id}")
|
||
# Загружаем узлы и ребра для продолжения графа
|
||
initial_state.graph_nodes = loaded_graph_data.get("graph_nodes", [])
|
||
initial_state.graph_edges = loaded_graph_data.get("graph_edges", [])
|
||
|
||
# Сохраняем ID загруженных узлов и ребер
|
||
original_node_ids = {node['id'] for node in initial_state.graph_nodes}
|
||
original_edge_ids = {edge['id'] for edge in initial_state.graph_edges}
|
||
|
||
# Устанавливаем parent_node_id для нового входного узла
|
||
# Если parent_node_id был передан в запросе, используем его,
|
||
# иначе берем current_node_id из загруженного графа.
|
||
initial_state.parent_node_id = parent_node_id or loaded_graph_data.get("current_node_id")
|
||
|
||
# История чата для текущего раунда работы LLM формируется из загруженных сообщений.
|
||
initial_state.temporary_chat_history = loaded_graph_data.get("messages", [])
|
||
|
||
# Определяем graph_id, который будет использоваться для сохранения.
|
||
# Если существующий ID не передан, GraphHistoryManager сгенерирует новый.
|
||
graph_id_to_save = existing_graph_id
|
||
|
||
|
||
# ---------------------------------------------------- run_agent ----------------------------------------------------------------
|
||
# Запускаем граф
|
||
result = app.invoke(initial_state)
|
||
final_state = AgentState(**result)
|
||
|
||
# ---------------------------------------------------- Save Graph ----------------------------------------------------------------
|
||
# Определяем новые узлы и ребра для сохранения
|
||
newly_added_nodes = [node for node in final_state.graph_nodes if node['id'] not in original_node_ids]
|
||
newly_added_edges = [edge for edge in final_state.graph_edges if edge['id'] not in original_edge_ids]
|
||
|
||
# Сохраняем обновленное состояние графа.
|
||
# Обратите внимание: chat_history больше не сохраняется напрямую,
|
||
# она является частью данных узлов в graph_nodes.
|
||
final_graph_id = graph_history_manager.save_graph_changes(
|
||
graph_id_to_save,
|
||
newly_added_nodes,
|
||
newly_added_edges,
|
||
final_state.current_node_id
|
||
)
|
||
|
||
# ---------------------------------------------------- prepare response ----------------------------------------------------------------
|
||
# Для `messages` в `response_data` используем `final_state.temporary_chat_history`,
|
||
# так как она отражает только сообщения текущей ветки, добавленные в ходе этого выполнения.
|
||
response_data = {
|
||
"graph_id": final_graph_id,
|
||
"messages": final_state.temporary_chat_history,
|
||
"llm_response": final_state.llm_response,
|
||
"image_urls": final_state.image_urls,
|
||
"analysis_result": final_state.analysis_result,
|
||
"subtitles": final_state.subtitles,
|
||
"summarized_history_text": final_state.summarized_history_text,
|
||
"error": final_state.error,
|
||
"graph_visualization_data": {
|
||
"nodes": final_state.graph_nodes,
|
||
"edges": final_state.graph_edges,
|
||
"current_node_id": final_state.current_node_id
|
||
}
|
||
}
|
||
return response_data
|
||
|
||
def run_agent_streaming(graph_id: str, user_node_id: str, assistant_node_id: str, system_prompt: Optional[str] = None, model: Optional[str] = None):
|
||
"""
|
||
Генератор для стримингового ответа LLM.
|
||
Возвращает чанки контента по мере их получения.
|
||
"""
|
||
try:
|
||
# Загружаем граф и историю до user_node_id
|
||
loaded_graph_data = graph_history_manager.get_graph(graph_id, target_node_id=user_node_id)
|
||
if not loaded_graph_data:
|
||
yield {"type": "error", "error": "Граф не найден."}
|
||
return
|
||
|
||
messages = loaded_graph_data.get("messages", [])
|
||
|
||
# Формируем сообщения для LLM
|
||
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
|
||
|
||
messages_for_llm = []
|
||
if system_prompt:
|
||
messages_for_llm.append(SystemMessage(content=system_prompt))
|
||
|
||
for msg in messages:
|
||
if msg["role"] == "user":
|
||
messages_for_llm.append(HumanMessage(content=msg["content"]))
|
||
elif msg["role"] == "assistant":
|
||
messages_for_llm.append(AIMessage(content=msg["content"]))
|
||
|
||
# Получаем LLM и стримим ответ
|
||
llm = get_llm(model or "gemini-2.5-flash")
|
||
|
||
for chunk in llm.stream2(messages_for_llm):
|
||
content = chunk.content if hasattr(chunk, 'content') else str(chunk)
|
||
yield {"type": "chunk", "content": content}
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка при стриминге: {e}")
|
||
yield {"type": "error", "error": str(e)} |