split graph_history_manager and title_generator

This commit is contained in:
dimitrievgs 2025-09-24 22:58:38 +03:00
parent 5a1fb4d482
commit cb195fddc6
4 changed files with 129 additions and 103 deletions

View File

@ -8,17 +8,12 @@ from flask import Flask, request, jsonify
from flask_cors import CORS
from workflows import run_agent, graph_history_manager
from llm_client import MODELS # Добавляем импорт списка моделей
from title_generator import TitleGenerator
api = Flask(__name__)
CORS(
api
) # Разрешаем CORS для всех доменов (в production нужно настроить более строго)
# Инициализируем сервис генерации заголовков
title_generator = TitleGenerator(graph_history_manager)
title_generator.start()
@api.route('/api/chat', methods=['POST'])
def chat():
@ -90,9 +85,3 @@ def get_messages_from_root_to_node(graph_id, node_id):
def get_available_models():
"""API endpoint для получения списка доступных моделей."""
return jsonify(list(MODELS.keys()))
@api.teardown_appcontext
def cleanup(error):
"""Очистка ресурсов при завершении приложения."""
if hasattr(api, '_title_generator'):
title_generator.stop()

View File

@ -4,6 +4,7 @@ from typing import Dict, Any, List, Optional
import json
import sqlite3
import uuid # Для генерации UUID
from title_generator import TitleGenerator
# ----------------------------------------------- Exceptions ------------------------------------------------------------
@ -24,6 +25,22 @@ class GraphHistoryManager:
self.db_path = db_path
self._create_tables()
# Инициализируем сервис генерации заголовков
self.title_generator = TitleGenerator(self)
self.title_generator.start()
def __del__(self):
"""
Деструктор класса. Останавливает генератор заголовков при уничтожении объекта.
"""
print("🚩 Завершение работы GraphHistoryManager. Останавливаем генератор заголовков...")
self.stop_title_generator()
def stop_title_generator(self):
"""Останавливает процесс генерации заголовков."""
if self.title_generator:
self.title_generator.stop()
# ----------------------------------------------- Internal Methods ------------------------------------------------------------
def _get_connection(self):
"""
@ -70,17 +87,6 @@ class GraphHistoryManager:
FOREIGN KEY (target_node_id) REFERENCES graph_nodes_data (node_id) ON DELETE CASCADE
)
""")
# Очередь для обработки заголовков
cursor.execute("""
CREATE TABLE IF NOT EXISTS title_generation_queue (
id INTEGER PRIMARY KEY AUTOINCREMENT,
item_type TEXT NOT NULL, -- 'graph' or 'node'
graph_id TEXT NOT NULL,
node_id TEXT DEFAULT NULL, -- только для узлов
priority INTEGER DEFAULT 0, -- для приоритета обработки
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
)
""")
conn.commit()
# _get_max_graph_id - удален, так как используем UUID для graph_id
@ -174,10 +180,7 @@ class GraphHistoryManager:
# Добавляем новые узлы
for node in new_nodes:
self._add_graph_node(conn, graph_id, node)
# Добавляем в очередь заголовков с высоким приоритетом
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
("node", graph_id, node["id"], 10))
self.title_generator.add_node_to_queue(cursor, graph_id, node["id"])
# Добавляем новые ребра
for edge in new_edges:
@ -189,9 +192,7 @@ class GraphHistoryManager:
# Если новый граф, добавляем его в очередь заголовков
if is_new_graph:
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
("graph", graph_id, 10))
self.title_generator.add_graph_to_queue(cursor, graph_id)
conn.commit() # Фиксируем все изменения
print(f"Изменения графа {graph_id} сохранены.")
@ -356,74 +357,6 @@ class GraphHistoryManager:
print(f"Ошибка при удалении графа: {e}")
return False
# ----------------------------------------------- Title Generation Queue Methods ------------------------------------------------------------
def add_to_title_queue(self,
item_type: str,
graph_id: str,
node_id: str = None,
priority: int = 0):
"""Добавляет элемент в очередь генерации заголовков."""
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
(item_type, graph_id, node_id, priority))
conn.commit()
def get_next_from_title_queue(self) -> Optional[Dict[str, Any]]:
"""Получает следующий элемент из очереди для обработки."""
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("""
SELECT id, item_type, graph_id, node_id
FROM title_generation_queue
ORDER BY priority DESC, id ASC
LIMIT 1
""")
result = cursor.fetchone()
if result:
queue_id, item_type, graph_id, node_id = result
# Удаляем из очереди
cursor.execute(
"DELETE FROM title_generation_queue WHERE id = ?",
(queue_id, ))
conn.commit()
return {
"item_type": item_type,
"graph_id": graph_id,
"node_id": node_id
}
return None
def populate_initial_title_queue(self):
"""Заполняет очередь элементами без сгенерированных заголовков при старте сервера."""
with self._get_connection() as conn:
cursor = conn.cursor()
# Очищаем существующую очередь
cursor.execute("DELETE FROM title_generation_queue")
# Добавляем графы без заголовков
cursor.execute(
"SELECT id FROM graphs WHERE title_generated = FALSE OR title_generated IS NULL"
)
for (graph_id, ) in cursor.fetchall():
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
("graph", graph_id, 1))
# Добавляем узлы без заголовков
cursor.execute(
"SELECT graph_id, node_id FROM graph_nodes_data WHERE title_generated = FALSE OR title_generated IS NULL"
)
for graph_id, node_id in cursor.fetchall():
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
("node", graph_id, node_id, 1))
conn.commit()
# ----------------------------------------------- Title Update Methods ------------------------------------------------------------
def update_graph_title(self, graph_id: str, title: str):

View File

@ -3,10 +3,10 @@
Работает в фоновом режиме, обрабатывая очередь элементов без заголовков.
"""
import sqlite3
import threading
import time
from typing import Optional
from graph_history_manager import GraphHistoryManager
from typing import Any, Dict, Optional
from nodes import DEFAULT_SUMMARIZATION_LLM_NAME
from llm_client import get_llm
from langchain_core.messages import SystemMessage, HumanMessage
@ -14,7 +14,9 @@ from langchain_core.messages import SystemMessage, HumanMessage
class TitleGenerator:
"""Сервис для асинхронной генерации заголовков."""
def __init__(self, history_manager: GraphHistoryManager):
def __init__(self, history_manager, db_path="graph_history.db"):
self.db_path = db_path
self._create_table()
self.history_manager = history_manager
self.llm = get_llm(DEFAULT_SUMMARIZATION_LLM_NAME)
self.running = False
@ -27,6 +29,31 @@ class TitleGenerator:
self.node_title_prompt = """Создай краткий заголовок (максимум 40 символов) для этого сообщения/действия.
Заголовок должен кратко описывать суть сообщения или действия. Отвечай только заголовком, без дополнительных объяснений. Ты должен сделать саммери, а не ответить на вопросы, если они есть в сообщении."""
def _get_connection(self):
"""
Получает соединение с базой данных.
Устанавливаем таймаут для ожидания блокировки базы данных при конкурентной записи.
"""
return sqlite3.connect(self.db_path,
timeout=5.0) # Увеличил таймаут до 5 секунд
def _create_table(self):
"""Создает таблицы для хранения графов, если они не существуют."""
with self._get_connection() as conn:
cursor = conn.cursor()
# Очередь для обработки заголовков
cursor.execute("""
CREATE TABLE IF NOT EXISTS title_generation_queue (
id INTEGER PRIMARY KEY AUTOINCREMENT,
item_type TEXT NOT NULL, -- 'graph' or 'node'
graph_id TEXT NOT NULL,
node_id TEXT DEFAULT NULL, -- только для узлов
priority INTEGER DEFAULT 0, -- для приоритета обработки
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
)
""")
conn.commit()
def start(self):
"""Запускает фоновый процесс генерации заголовков."""
if self.running:
@ -35,13 +62,41 @@ class TitleGenerator:
self.running = True
# Заполняем очередь при старте
self.history_manager.populate_initial_title_queue()
self.populate_initial_title_queue()
# Запускаем фоновый поток
self.thread = threading.Thread(target=self._process_queue, daemon=True)
self.thread.start()
print("Сервис генерации заголовков запущен")
def populate_initial_title_queue(self):
"""Заполняет очередь элементами без сгенерированных заголовков при старте сервера."""
with self._get_connection() as conn:
cursor = conn.cursor()
# Очищаем существующую очередь
cursor.execute("DELETE FROM title_generation_queue")
# Добавляем графы без заголовков
cursor.execute(
"SELECT id FROM graphs WHERE title_generated = FALSE OR title_generated IS NULL"
)
for (graph_id, ) in cursor.fetchall():
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
("graph", graph_id, 1))
# Добавляем узлы без заголовков
cursor.execute(
"SELECT graph_id, node_id FROM graph_nodes_data WHERE title_generated = FALSE OR title_generated IS NULL"
)
for graph_id, node_id in cursor.fetchall():
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
("node", graph_id, node_id, 1))
conn.commit()
def stop(self):
"""Останавливает фоновый процесс."""
self.running = False
@ -53,7 +108,7 @@ class TitleGenerator:
"""Основной цикл обработки очереди заголовков."""
while self.running:
try:
item = self.history_manager.get_next_from_title_queue()
item = self.get_next_from_title_queue()
if item:
if item["item_type"] == "graph":
@ -68,6 +123,31 @@ class TitleGenerator:
print(f"Ошибка при обработке очереди заголовков: {e}")
time.sleep(10)
def get_next_from_title_queue(self) -> Optional[Dict[str, Any]]:
"""Получает следующий элемент из очереди для обработки."""
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute("""
SELECT id, item_type, graph_id, node_id
FROM title_generation_queue
ORDER BY priority DESC, id ASC
LIMIT 1
""")
result = cursor.fetchone()
if result:
queue_id, item_type, graph_id, node_id = result
# Удаляем из очереди
cursor.execute(
"DELETE FROM title_generation_queue WHERE id = ?",
(queue_id, ))
conn.commit()
return {
"item_type": item_type,
"graph_id": graph_id,
"node_id": node_id
}
return None
def _generate_graph_title(self, graph_id: str):
"""Генерирует заголовок для графа."""
try:
@ -148,4 +228,28 @@ class TitleGenerator:
print(f"Сгенерирован заголовок узла {node_id}: {title}")
except Exception as e:
print(f"Ошибка генерации заголовка узла {node_id}: {e}")
print(f"Ошибка генерации заголовка узла {node_id}: {e}")
'''def add_to_title_queue(self,
item_type: str,
graph_id: str,
node_id: str = None,
priority: int = 0):
"""Добавляет элемент в очередь генерации заголовков."""
with self._get_connection() as conn:
cursor = conn.cursor()
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
(item_type, graph_id, node_id, priority))
conn.commit()'''
def add_node_to_queue(self, cursor, graph_id, node_id):
# Добавляем в очередь заголовков с высоким приоритетом
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, node_id, priority) VALUES (?, ?, ?, ?)",
("node", graph_id, node_id, 10))
def add_graph_to_queue(self, cursor, graph_id):
cursor.execute(
"INSERT INTO title_generation_queue (item_type, graph_id, priority) VALUES (?, ?, ?)",
("graph", graph_id, 10))

Binary file not shown.