split graph_history_manager and title_generator
This commit is contained in:
parent
5a1fb4d482
commit
cb195fddc6
11
app/api.py
11
app/api.py
|
|
@ -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()
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
@ -149,3 +229,27 @@ class TitleGenerator:
|
|||
|
||||
except Exception as 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))
|
||||
|
|
|
|||
BIN
graph_history.db
BIN
graph_history.db
Binary file not shown.
Loading…
Reference in New Issue
Block a user