268 lines
13 KiB
Python
268 lines
13 KiB
Python
import os
|
||
import base64
|
||
import json
|
||
from typing import Dict, Any, Optional
|
||
|
||
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, AIMessageChunk
|
||
from react_agent import build_react_agent
|
||
# Убедитесь, что graph_history_manager импортирован.
|
||
# Если код в workflows.py, то graph_history_manager создается/импортируется там.
|
||
|
||
from graph_history_manager import GraphHistoryManager
|
||
|
||
graph_history_manager = GraphHistoryManager()
|
||
|
||
|
||
def run_agent_streaming(graph_id: str,
|
||
user_node_id: str,
|
||
assistant_node_id: str,
|
||
system_prompt: Optional[str] = None,
|
||
model: Optional[str] = None,
|
||
cache_folder: Optional[str] = None,
|
||
temperature: float = 0.7,
|
||
max_tokens: Optional[int] = None,
|
||
agency_mode: bool = False,
|
||
obsidian_settings: dict = None):
|
||
"""
|
||
Основной генератор: Запускает LangGraph / ReAct агента.
|
||
Поддерживает мультимодальность, чтение файлов из Vault и исполнение Python-тулов.
|
||
Отдает словари, которые api.py запакует в SSE-события.
|
||
"""
|
||
try:
|
||
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", [])
|
||
|
||
# 2. Формируем контекст LangChain с учетом файлов и изображений (вложений)
|
||
messages_for_llm = []
|
||
if system_prompt:
|
||
messages_for_llm.append(SystemMessage(content=system_prompt))
|
||
|
||
# Обрабатываем вложения
|
||
for i, msg in enumerate(messages):
|
||
is_last = (i == len(messages) - 1)
|
||
role = msg.get("role")
|
||
|
||
if role == "user":
|
||
text_content = msg.get("content", "")
|
||
attachments = msg.get("attachments", [])
|
||
|
||
# Проверяем, есть ли мультимодальный контент
|
||
has_multimodal = False
|
||
content_parts = [{"type": "text", "text": text_content}]
|
||
|
||
# Обрабатываем вложения
|
||
if attachments:
|
||
for att in attachments:
|
||
# Прикрепляем файл только если он перманентный (@@/##) или это последнее сообщение
|
||
if att.get("permanent") or is_last:
|
||
attachment_content = prepare_attachment_for_llm(
|
||
att, cache_folder)
|
||
|
||
# Если это изображение, то это мультимодальный контент
|
||
if attachment_content.get("type") == "image_url":
|
||
has_multimodal = True
|
||
|
||
content_parts.append(attachment_content)
|
||
|
||
# Если есть картинки, шлем list of dicts. Иначе склеиваем в один текст.
|
||
if has_multimodal:
|
||
# Мультимодальное сообщение
|
||
messages_for_llm.append(
|
||
HumanMessage(content=content_parts))
|
||
else:
|
||
# Текстовое сообщение (объединяем все текстовые части)
|
||
full_text = text_content
|
||
for part in content_parts[
|
||
1:]: # Пропускаем первую часть (основной текст)
|
||
if part.get("type") == "text":
|
||
full_text += part.get("text", "")
|
||
messages_for_llm.append(HumanMessage(content=full_text))
|
||
|
||
elif role == "assistant":
|
||
messages_for_llm.append(AIMessage(content=msg.get("content", "")))
|
||
|
||
# Локальное хранилище для отладки текущего запуска
|
||
current_run_debug = []
|
||
|
||
def capture_debug(req, res):
|
||
current_run_debug.append({"req": req, "res": res})
|
||
|
||
# 3. Создаем агента. Если agency_mode == False, он соберет 0 тулов и отработает как обычная LLM.
|
||
agent_executor = build_react_agent(
|
||
model_name=model,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
agency_mode=agency_mode,
|
||
obsidian_settings=obsidian_settings or {},
|
||
debug_callback=capture_debug
|
||
)
|
||
|
||
# 4. Выполняем цикл с потоковой передачей событий
|
||
step_counter = 0
|
||
tool_mapping = {} # Хранит связку: ID вызова -> Номер шага
|
||
|
||
# LangGraph Stream Mode 'messages': отдает чанки токенов и вызовы функций
|
||
for chunk, metadata in agent_executor.stream({"messages": messages_for_llm}, stream_mode="messages"):
|
||
|
||
if isinstance(chunk, AIMessageChunk):
|
||
# Агент запросил вызов инструмента (Формирование JSON аргументов)
|
||
if hasattr(chunk, "tool_call_chunks") and chunk.tool_call_chunks:
|
||
for tc in chunk.tool_call_chunks:
|
||
if "name" in tc and tc["name"]:
|
||
step_counter += 1
|
||
tc_id = tc.get("id", f"unknown_{step_counter}")
|
||
# Сохраняем шаг для этого конкретного вызова
|
||
tool_mapping[tc_id] = step_counter
|
||
|
||
yield {
|
||
"type": "tool_start",
|
||
"name": tc["name"],
|
||
"step": step_counter,
|
||
"node_id": assistant_node_id
|
||
}
|
||
|
||
# Обычный стриминг текста финального ответа
|
||
if chunk.content:
|
||
yield {
|
||
"type": "chunk",
|
||
"content": chunk.content,
|
||
"node_id": assistant_node_id
|
||
}
|
||
|
||
# Тул отработал и принес ответ
|
||
elif hasattr(chunk, "type") and chunk.type == "tool":
|
||
req = ""
|
||
res = str(chunk.content)
|
||
|
||
# Достаем JSON отправленных параметров из последнего debug_callback
|
||
if current_run_debug:
|
||
dbg = current_run_debug.pop(0) # Берем первый из очереди
|
||
req = dbg.get("req", "")
|
||
|
||
# Ловушка для определения ошибок
|
||
is_error = False
|
||
if getattr(chunk, "status", "") == "error" or "ОШИБКА ИСПОЛНЕНИЯ" in res or "SCRIPT RUNTIME ERROR" in res:
|
||
is_error = True
|
||
|
||
# Достаем ИМЕННО ТОТ ШАГ, на котором начинался вызов этого инструмента
|
||
tc_id = getattr(chunk, "tool_call_id", "")
|
||
actual_step = tool_mapping.get(tc_id, step_counter) # Fallback, если id нет
|
||
|
||
yield {
|
||
"type": "tool_end",
|
||
"name": getattr(chunk, "name", "tool"),
|
||
"step": actual_step,
|
||
"node_id": assistant_node_id,
|
||
"request": req,
|
||
"response": res,
|
||
"is_error": is_error
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"Ошибка при стриминге: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
yield {"type": "error", "error": str(e)}
|
||
|
||
|
||
def prepare_attachment_for_llm(attachment: Dict[str, Any],
|
||
cache_folder: str) -> Dict[str, Any]:
|
||
"""Подготавливает вложение для отправки в LLM."""
|
||
mime_type = attachment.get("mimeType", "application/octet-stream")
|
||
cached_path = attachment.get("cachedPath")
|
||
att_cache_folder = attachment.get("cacheFolder", cache_folder) # ИЗМЕНЕНО
|
||
source_path = attachment.get("sourcePath")
|
||
att_type = attachment.get("type")
|
||
name = attachment.get("name")
|
||
|
||
# Получаем полный путь к кешированному файлу
|
||
full_cached_path = os.path.join(att_cache_folder, cached_path)
|
||
|
||
# Проверяем существование кеша
|
||
if not os.path.exists(full_cached_path):
|
||
print(f"❌ Кеш не найден: {full_cached_path}")
|
||
return {"type": "text", "text": f"\n[Файл {name} недоступен]\n"}
|
||
|
||
# Проверяем тип файла
|
||
if mime_type.startswith('image/'):
|
||
# Для изображений отправляем как base64
|
||
with open(full_cached_path, 'rb') as f:
|
||
image_data = base64.b64encode(f.read()).decode('utf-8')
|
||
|
||
return {
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:{mime_type};base64,{image_data}"
|
||
}
|
||
}
|
||
elif mime_type.startswith('text/') or \
|
||
mime_type in ['application/json', 'application/xml'] or \
|
||
mime_type.startswith('application/x-') or \
|
||
mime_type == 'application/octet-stream': # TODO: Спорно, но пока вводим - фронтенд не смог сопоставить расширению mime_type
|
||
# неразрывно связано с методом getMimeType() на фронтенде
|
||
# Для текстовых файлов (или потенциально текстовых с неизвестным типом)
|
||
try:
|
||
with open(full_cached_path, 'r', encoding='utf-8') as f:
|
||
content = f.read()
|
||
except UnicodeDecodeError:
|
||
# Если не получилось прочитать как текст
|
||
return {
|
||
"type": "text",
|
||
"text": f"\n[Файл {name} не может быть прочитан как текст]\n"
|
||
}
|
||
|
||
att_type_ru = "промпта" if att_type == "prompt" else "файла"
|
||
header = f"\n\n--- Содержимое {att_type_ru}: {name} ---\n"
|
||
footer = f"\n--- Конец {att_type_ru}: {name} ---\n\n"
|
||
|
||
return {"type": "text", "text": header + content.strip() + footer}
|
||
else:
|
||
# Для других типов файлов возвращаем информацию
|
||
return {
|
||
"type":
|
||
"text",
|
||
"text":
|
||
f"\n[Файл {name} ({mime_type}) не может быть отображен в чате]\n"
|
||
}
|
||
|
||
|
||
def restore_cache_from_source(source_path: str, att_type: str,
|
||
cache_folder: str, cached_path: str) -> str:
|
||
"""Восстанавливает кеш из исходного файла."""
|
||
full_cached_path = os.path.join(cache_folder, cached_path)
|
||
|
||
# Создаем папку кеша, если не существует
|
||
os.makedirs(cache_folder, exist_ok=True)
|
||
|
||
if att_type == 'external':
|
||
# Для внешних файлов source_path - абсолютный
|
||
if not os.path.exists(source_path):
|
||
raise FileNotFoundError(f"Исходный файл не найден: {source_path}")
|
||
|
||
import shutil
|
||
shutil.copy2(source_path, full_cached_path)
|
||
else:
|
||
# Для внутренних файлов и промптов нужно получить полный путь
|
||
# Это требует доступа к vault, что на бэкенде недоступно
|
||
# Поэтому это нужно делать на фронтенде
|
||
raise NotImplementedError(
|
||
"Восстановление кеша для внутренних файлов должно происходить на фронтенде"
|
||
)
|
||
|
||
return full_cached_path
|
||
|
||
|
||
def format_attachment_for_llm(attachment: Dict[str, Any]) -> str:
|
||
"""Форматирует вложение для вставки в сообщение LLM."""
|
||
att_type = "промпта" if attachment["type"] == "prompt" else "файла"
|
||
header = f"\n\n--- Содержимое {att_type}: {attachment['name']} ---\n"
|
||
footer = f"\n--- Конец {att_type}: {attachment['name']} ---\n\n"
|
||
return header + attachment.get("content", "").strip() + footer
|
||
|