work on streaming chat
This commit is contained in:
dimitrievgs
parent
7bd9fc92f7
commit
ee84617a1c
99
app/api.py
99
app/api.py
|
|
@ -4,9 +4,10 @@
|
|||
включая обработку сообщений, управление графами и получение истории.
|
||||
"""
|
||||
|
||||
import json
|
||||
from flask import Flask, request, jsonify
|
||||
from flask_cors import CORS
|
||||
from workflows import run_agent, graph_history_manager
|
||||
from workflows import run_agent, graph_history_manager, run_agent_streaming
|
||||
from llm_client import MODELS # Добавляем импорт списка моделей
|
||||
|
||||
api = Flask(__name__)
|
||||
|
|
@ -101,3 +102,99 @@ def get_messages_from_root_to_node(graph_id, node_id):
|
|||
def get_available_models():
|
||||
"""API endpoint для получения списка доступных моделей."""
|
||||
return jsonify(list(MODELS.keys()))
|
||||
|
||||
@api.route('/api/chat/send', methods=['POST'])
|
||||
def send_user_message():
|
||||
"""API endpoint для создания узла пользователя и получения graph_id."""
|
||||
data = request.get_json()
|
||||
message = data.get("message")
|
||||
graph_id = data.get("graph_id")
|
||||
parent_node_id = data.get("parent_node_id")
|
||||
|
||||
if not message:
|
||||
return jsonify({"error": "Сообщение не может быть пустым."}, 400)
|
||||
|
||||
try:
|
||||
result = graph_history_manager.create_user_node(message, graph_id, parent_node_id)
|
||||
return jsonify(result)
|
||||
except Exception as e:
|
||||
print(f"Ошибка при создании узла пользователя: {e}")
|
||||
return jsonify({"error": str(e)}, 500)
|
||||
|
||||
|
||||
from flask import Response, stream_with_context
|
||||
|
||||
@api.route('/api/chat/stream', methods=['POST'])
|
||||
def chat_stream():
|
||||
"""API endpoint для стриминга ответа LLM."""
|
||||
data = request.get_json()
|
||||
graph_id = data.get("graph_id")
|
||||
user_node_id = data.get("user_node_id")
|
||||
system_prompt = data.get("system_prompt")
|
||||
model = data.get("model")
|
||||
|
||||
if not graph_id or not user_node_id:
|
||||
return jsonify({"error": "Не указаны graph_id или user_node_id."}, 400)
|
||||
|
||||
def generate():
|
||||
try:
|
||||
# Создаём узел-заглушку для ответа LLM
|
||||
assistant_node_id = graph_history_manager.create_assistant_placeholder_node(graph_id, user_node_id)
|
||||
|
||||
# Отправляем ID нового узла клиенту
|
||||
yield f"data: {json.dumps({'type': 'node_created', 'node_id': assistant_node_id})}\n\n"
|
||||
|
||||
# Запускаем стриминг ответа от LLM
|
||||
accumulated_content = ""
|
||||
for chunk_data in run_agent_streaming(graph_id, user_node_id, assistant_node_id, system_prompt, model):
|
||||
if chunk_data.get("type") == "chunk":
|
||||
accumulated_content += chunk_data.get("content", "")
|
||||
yield f"data: {json.dumps(chunk_data)}\n\n"
|
||||
elif chunk_data.get("type") == "error":
|
||||
yield f"data: {json.dumps(chunk_data)}\n\n"
|
||||
return
|
||||
|
||||
# После завершения стриминга обновляем узел полным контентом
|
||||
graph_history_manager.update_assistant_node_content(graph_id, assistant_node_id, accumulated_content)
|
||||
|
||||
# Инициируем генерацию заголовка
|
||||
graph_history_manager.title_generator.add_node_to_queue_direct(graph_id, assistant_node_id)
|
||||
|
||||
yield f"data: {json.dumps({'type': 'done', 'node_id': assistant_node_id})}\n\n"
|
||||
|
||||
except Exception as e:
|
||||
print(f"Ошибка при стриминге ответа: {e}")
|
||||
error_data = {'type': 'error', 'error': str(e)}
|
||||
yield f"data: {json.dumps(error_data)}\n\n"
|
||||
|
||||
return Response(stream_with_context(generate()), content_type='text/event-stream')
|
||||
|
||||
|
||||
@api.route('/api/chat/regenerate', methods=['POST'])
|
||||
def regenerate_message():
|
||||
"""API endpoint для регенерации сообщения."""
|
||||
data = request.get_json()
|
||||
graph_id = data.get("graph_id")
|
||||
node_id = data.get("node_id")
|
||||
system_prompt = data.get("system_prompt")
|
||||
model = data.get("model")
|
||||
|
||||
if not graph_id or not node_id:
|
||||
return jsonify({"error": "Не указаны graph_id или node_id."}, 400)
|
||||
|
||||
try:
|
||||
# Определяем родительский узел
|
||||
parent_node_id = graph_history_manager.get_parent_node_id(graph_id, node_id)
|
||||
if not parent_node_id:
|
||||
return jsonify({"error": "Не удалось найти родительский узел."}, 404)
|
||||
|
||||
# Создаём новый узел-заглушку для регенерации
|
||||
new_assistant_node_id = graph_history_manager.create_assistant_placeholder_node(graph_id, parent_node_id)
|
||||
|
||||
return jsonify({
|
||||
"new_node_id": new_assistant_node_id,
|
||||
"parent_node_id": parent_node_id
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"Ошибка при регенерации сообщения: {e}")
|
||||
return jsonify({"error": str(e)}, 500)
|
||||
|
|
|
|||
|
|
@ -425,10 +425,11 @@ class GraphHistoryManager:
|
|||
new_current_node_id = None
|
||||
if parents:
|
||||
# Если есть родители, делаем текущим узлом одного из родителей
|
||||
# Для простоты, берем первого родителя. Можно добавить более сложную логику,
|
||||
# например, последнего добавленного родителя, если timestamp доступен.
|
||||
new_current_node_id = parents[0]
|
||||
# Если детей нет и родителей тоже, current_node_id останется None.
|
||||
elif children:
|
||||
# Если нет родителей, но есть дети, делаем текущим узлом одного из детей
|
||||
new_current_node_id = children[0]
|
||||
# Если нет ни родителей, ни детей, current_node_id останется None.
|
||||
|
||||
cursor.execute("UPDATE graphs SET current_node_id = ? WHERE id = ?",
|
||||
(new_current_node_id, graph_id))
|
||||
|
|
@ -469,3 +470,145 @@ class GraphHistoryManager:
|
|||
updated_rows = cursor.rowcount
|
||||
print(f"Обновлено строк: {updated_rows}")
|
||||
conn.commit()
|
||||
|
||||
def create_user_node(self, message: str, graph_id: Optional[str] = None, parent_node_id: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Создаёт узел пользователя и возвращает graph_id и node_id.
|
||||
"""
|
||||
if not graph_id:
|
||||
graph_id = str(uuid.uuid4())
|
||||
|
||||
user_node_id = str(uuid.uuid4())
|
||||
|
||||
with self._get_connection() as conn:
|
||||
try:
|
||||
conn.execute("INSERT OR IGNORE INTO graphs (id) VALUES (?)", (graph_id,))
|
||||
|
||||
node_data = {
|
||||
"label": message[:30] + "..." if len(message) > 30 else message,
|
||||
"message": {"role": "user", "content": message, "node_id": user_node_id}
|
||||
}
|
||||
|
||||
self._add_graph_node(conn, graph_id, {
|
||||
"id": user_node_id,
|
||||
"type": "user",
|
||||
"data": node_data
|
||||
})
|
||||
|
||||
if parent_node_id:
|
||||
edge_id = str(uuid.uuid4())
|
||||
self._add_graph_edge(conn, graph_id, {
|
||||
"id": edge_id,
|
||||
"source": parent_node_id,
|
||||
"target": user_node_id
|
||||
})
|
||||
|
||||
self._update_graph_current_node_id(conn, graph_id, user_node_id)
|
||||
conn.commit()
|
||||
|
||||
return {
|
||||
"graph_id": graph_id,
|
||||
"node_id": user_node_id,
|
||||
"parent_node_id": parent_node_id
|
||||
}
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"Ошибка при создании узла пользователя: {e}")
|
||||
raise
|
||||
|
||||
def create_assistant_placeholder_node(self, graph_id: str, parent_node_id: str) -> str:
|
||||
"""
|
||||
Создаёт узел-заглушку для ответа ассистента.
|
||||
"""
|
||||
assistant_node_id = str(uuid.uuid4())
|
||||
|
||||
with self._get_connection() as conn:
|
||||
try:
|
||||
node_data = {
|
||||
"label": "Получение ответа...",
|
||||
"message": {"role": "assistant", "content": "", "node_id": assistant_node_id},
|
||||
"is_placeholder": True
|
||||
}
|
||||
|
||||
self._add_graph_node(conn, graph_id, {
|
||||
"id": assistant_node_id,
|
||||
"type": "llm",
|
||||
"data": node_data
|
||||
})
|
||||
|
||||
edge_id = str(uuid.uuid4())
|
||||
self._add_graph_edge(conn, graph_id, {
|
||||
"id": edge_id,
|
||||
"source": parent_node_id,
|
||||
"target": assistant_node_id
|
||||
})
|
||||
|
||||
self._update_graph_current_node_id(conn, graph_id, assistant_node_id)
|
||||
conn.commit()
|
||||
|
||||
return assistant_node_id
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"Ошибка при создании узла-заглушки: {e}")
|
||||
raise
|
||||
|
||||
def update_assistant_node_content(self, graph_id: str, node_id: str, content: str):
|
||||
"""
|
||||
Обновляет содержимое узла ассистента после завершения стриминга.
|
||||
"""
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute("SELECT node_data FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?",
|
||||
(graph_id, node_id))
|
||||
result = cursor.fetchone()
|
||||
|
||||
if result:
|
||||
node_data = json.loads(result[0])
|
||||
node_data["message"]["content"] = content
|
||||
node_data["label"] = content[:30] + "..." if len(content) > 30 else content
|
||||
node_data["is_placeholder"] = False
|
||||
|
||||
cursor.execute("UPDATE graph_nodes_data SET node_data = ? WHERE graph_id = ? AND node_id = ?",
|
||||
(json.dumps(node_data), graph_id, node_id))
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"Ошибка при обновлении узла: {e}")
|
||||
raise
|
||||
|
||||
def mark_node_as_error(self, graph_id: str, node_id: str, error_message: str):
|
||||
"""
|
||||
Помечает узел как ошибочный.
|
||||
"""
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
cursor.execute("SELECT node_data FROM graph_nodes_data WHERE graph_id = ? AND node_id = ?",
|
||||
(graph_id, node_id))
|
||||
result = cursor.fetchone()
|
||||
|
||||
if result:
|
||||
node_data = json.loads(result[0])
|
||||
node_data["message"]["content"] = error_message
|
||||
node_data["label"] = "Ошибка"
|
||||
node_data["is_placeholder"] = False
|
||||
|
||||
cursor.execute("UPDATE graph_nodes_data SET node_type = ?, node_data = ? WHERE graph_id = ? AND node_id = ?",
|
||||
("error", json.dumps(node_data), graph_id, node_id))
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
conn.rollback()
|
||||
print(f"Ошибка при маркировке узла как ошибочного: {e}")
|
||||
raise
|
||||
|
||||
def get_parent_node_id(self, graph_id: str, node_id: str) -> Optional[str]:
|
||||
"""
|
||||
Возвращает ID родительского узла.
|
||||
"""
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT source_node_id FROM graph_edges_data WHERE graph_id = ? AND target_node_id = ?",
|
||||
(graph_id, node_id))
|
||||
result = cursor.fetchone()
|
||||
return result[0] if result else None
|
||||
|
|
@ -97,7 +97,7 @@ class CustomLLM:
|
|||
self.config = config
|
||||
self.name = config["name"]
|
||||
self.provider = config["provider"]
|
||||
|
||||
|
||||
try:
|
||||
|
||||
if self.provider == "openai":
|
||||
|
|
@ -109,7 +109,7 @@ class CustomLLM:
|
|||
|
||||
elif self.provider == "mistralai":
|
||||
# Для Mistral используем прямые HTTP-запросы
|
||||
self.client = Mistral(api_key=config["apiKey"],
|
||||
self.client = Mistral(api_key=config["apiKey"],
|
||||
server_url=config["apiBase"])
|
||||
self.model_name = config["model_name"]
|
||||
self.stream = config["stream"]
|
||||
|
|
@ -239,6 +239,93 @@ class CustomLLM:
|
|||
raise
|
||||
return ""
|
||||
|
||||
def stream2(self, messages: List[Any]):
|
||||
"""
|
||||
Генератор для стриминга ответа от LLM.
|
||||
Возвращает чанки контента по мере их получения.
|
||||
"""
|
||||
if self.provider == "openai":
|
||||
openai_messages = self._prepare_openai_messages(messages)
|
||||
|
||||
try:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model_name,
|
||||
messages=openai_messages,
|
||||
stream=True, # Всегда True для стриминга
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
chunk_message = chunk.choices[0].delta.content
|
||||
if chunk_message is not None:
|
||||
yield chunk_message
|
||||
|
||||
except Exception as e:
|
||||
print(f"Ошибка при стриминге OpenAI API: {e}")
|
||||
raise
|
||||
|
||||
elif self.provider == "mistralai":
|
||||
mistralai_messages = self._prepare_mistralai_messages(messages)
|
||||
|
||||
try:
|
||||
response = self.client.chat.stream(
|
||||
model=self.model_name,
|
||||
messages=mistralai_messages
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
delta = chunk.data.choices[0].delta
|
||||
chunk_message = getattr(delta, "content", None)
|
||||
|
||||
if chunk_message is not None:
|
||||
yield chunk_message
|
||||
|
||||
if getattr(delta, "finish_reason", None) is not None:
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Ошибка при стриминге Mistral API: {e}")
|
||||
raise
|
||||
|
||||
def _prepare_openai_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
|
||||
"""Подготавливает сообщения в формате OpenAI."""
|
||||
openai_messages = []
|
||||
if isinstance(messages, str):
|
||||
openai_messages.append({"role": "user", "content": messages})
|
||||
elif isinstance(messages, list):
|
||||
for msg in messages:
|
||||
if isinstance(msg, HumanMessage):
|
||||
openai_messages.append({"role": "user", "content": msg.content})
|
||||
elif isinstance(msg, SystemMessage):
|
||||
openai_messages.append({"role": "system", "content": msg.content})
|
||||
elif isinstance(msg, AIMessage):
|
||||
openai_messages.append({"role": "assistant", "content": msg.content})
|
||||
else:
|
||||
openai_messages.append({
|
||||
"role": msg.type if hasattr(msg, 'type') else 'user',
|
||||
"content": msg.content
|
||||
})
|
||||
return openai_messages
|
||||
|
||||
def _prepare_mistralai_messages(self, messages: List[Any]) -> List[Dict[str, str]]:
|
||||
"""Подготавливает сообщения в формате Mistral AI."""
|
||||
mistralai_messages = []
|
||||
if isinstance(messages, str):
|
||||
mistralai_messages.append({"role": "user", "content": messages})
|
||||
elif isinstance(messages, list):
|
||||
for msg in messages:
|
||||
if isinstance(msg, HumanMessage):
|
||||
mistralai_messages.append({"role": "user", "content": msg.content})
|
||||
elif isinstance(msg, SystemMessage):
|
||||
mistralai_messages.append({"role": "system", "content": msg.content})
|
||||
elif isinstance(msg, AIMessage):
|
||||
mistralai_messages.append({"role": "assistant", "content": msg.content})
|
||||
else:
|
||||
mistralai_messages.append({
|
||||
"role": msg.type if hasattr(msg, 'type') else 'user',
|
||||
"content": msg.content
|
||||
})
|
||||
return mistralai_messages
|
||||
|
||||
|
||||
def get_llm(name: str) -> CustomLLM:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -256,3 +256,21 @@ class TitleGenerator:
|
|||
cursor.execute(
|
||||
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
|
||||
("graph", graph_id, 10))
|
||||
|
||||
def add_node_to_queue_direct(self, graph_id: str, node_id: str, priority: int = 10):
|
||||
"""Добавляет узел в очередь генерации заголовков напрямую (с собственным соединением)."""
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
|
||||
("node", graph_id, node_id, priority))
|
||||
conn.commit()
|
||||
|
||||
def add_graph_to_queue_direct(self, graph_id: str, priority: int = 10):
|
||||
"""Добавляет граф в очередь генерации заголовков напрямую (с собственным соединением)."""
|
||||
with self._get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
|
||||
("graph", graph_id, priority))
|
||||
conn.commit()
|
||||
|
|
@ -10,6 +10,7 @@ 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()
|
||||
|
|
@ -182,4 +183,42 @@ def run_agent(user_input: str,
|
|||
"current_node_id": final_state.current_node_id
|
||||
}
|
||||
}
|
||||
return response_data
|
||||
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)}
|
||||
Loading…
Reference in New Issue
Block a user