259 lines
9.3 KiB
Python
259 lines
9.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Тест системы памяти Telegram бота.
|
||
|
||
Проверяет:
|
||
1. Сохранение сообщений в SQLite
|
||
2. Сохранение сообщений в ChromaDB
|
||
3. Загрузку истории из БД
|
||
4. RAG-поиск по векторной базе
|
||
5. Извлечение фактов
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import asyncio
|
||
from pathlib import Path
|
||
|
||
# Добавляем путь к боту
|
||
BOT_DIR = Path(__file__).parent
|
||
sys.path.insert(0, str(BOT_DIR))
|
||
|
||
# Импортируем компоненты памяти
|
||
from vector_memory import (
|
||
hybrid_memory_manager,
|
||
save_message,
|
||
get_context,
|
||
load_history_to_state,
|
||
search_memory,
|
||
get_profile
|
||
)
|
||
|
||
from bot.models.user_state import UserState, StateManager
|
||
|
||
|
||
# Тестовый пользователь
|
||
TEST_USER_ID = 999999
|
||
|
||
|
||
def test_sqlite_save():
|
||
"""Тест 1: Сохранение сообщений в SQLite."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 1: Сохранение сообщений в SQLite")
|
||
print("="*60)
|
||
|
||
# Сохраняем тестовые сообщения
|
||
save_message(TEST_USER_ID, "user", "Привет! Меня зовут Владимир.")
|
||
save_message(TEST_USER_ID, "assistant", "Привет, Владимир! Чем могу помочь?")
|
||
save_message(TEST_USER_ID, "user", "Я работаю системным администратором.")
|
||
save_message(TEST_USER_ID, "assistant", "Отлично! Какие технологии вы используете?")
|
||
|
||
# Проверяем сохранение — используем гибридный менеджер напрямую
|
||
context = hybrid_memory_manager.get_context(TEST_USER_ID, max_messages=10)
|
||
print(f"✅ Сохранено сообщений в SQLite: {len(context)}")
|
||
|
||
for msg in context:
|
||
print(f" - [{msg.role}]: {msg.content[:50]}...")
|
||
|
||
assert len(context) >= 4, "Сообщения не сохранились в SQLite!"
|
||
print("\n✅ ТЕСТ 1: УСПЕШНО\n")
|
||
return True
|
||
|
||
|
||
def test_vector_save():
|
||
"""Тест 2: Сохранение сообщений в векторную базу."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 2: Сохранение сообщений в ChromaDB")
|
||
print("="*60)
|
||
|
||
# Проверяем векторную базу
|
||
if hybrid_memory_manager.vector:
|
||
stats = hybrid_memory_manager.vector.get_stats()
|
||
print(f"✅ Векторная база активна")
|
||
print(f" - Документы: {stats.get('total_documents', 0)}")
|
||
print(f" - Модель: {stats.get('model', 'unknown')}")
|
||
|
||
assert stats.get('total_documents', 0) >= 4, "Сообщения не сохранились в ChromaDB!"
|
||
print("\n✅ ТЕСТ 2: УСПЕШНО\n")
|
||
return True
|
||
else:
|
||
print("⚠️ Векторная база не активна")
|
||
return False
|
||
|
||
|
||
def test_history_loading():
|
||
"""Тест 3: Загрузка истории из БД в состояние."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 3: Загрузка истории из БД")
|
||
print("="*60)
|
||
|
||
# Создаём тестовое состояние
|
||
state = UserState()
|
||
state_manager = StateManager()
|
||
|
||
# Загружаем историю
|
||
history = load_history_to_state(TEST_USER_ID, state, state_manager)
|
||
|
||
print(f"✅ Загружено сообщений из БД: {len(history)}")
|
||
print(f" - История в state.ai_chat_history: {len(state.ai_chat_history)}")
|
||
print(f" - Флаг загрузки: {state_manager.is_history_loaded(TEST_USER_ID)}")
|
||
|
||
if history:
|
||
print("\n Последние сообщения:")
|
||
for msg in history[-3:]:
|
||
print(f" {msg[:80]}...")
|
||
|
||
assert len(state.ai_chat_history) >= 4, "История не загрузилась из БД!"
|
||
assert state_manager.is_history_loaded(TEST_USER_ID), "Флаг загрузки не установлен!"
|
||
print("\n✅ ТЕСТ 3: УСПЕШНО\n")
|
||
return True
|
||
|
||
|
||
async def test_rag_search():
|
||
"""Тест 4: RAG-поиск по векторной базе."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 4: RAG-поиск по векторной базе")
|
||
print("="*60)
|
||
|
||
# Ищем по запросу
|
||
query = "Владимир работа"
|
||
results = search_memory(TEST_USER_ID, query, limit=5)
|
||
|
||
print(f"✅ Найдено результатов по запросу '{query}': {len(results)}")
|
||
|
||
for msg, score in results:
|
||
print(f" - [{score:.2f}] [{msg.role}]: {msg.content[:60]}...")
|
||
|
||
assert len(results) >= 1, "RAG-поиск не нашёл результатов!"
|
||
print("\n✅ ТЕСТ 4: УСПЕШНО\n")
|
||
return True
|
||
|
||
|
||
def test_fact_extraction():
|
||
"""Тест 5: Извлечение фактов."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 5: Извлечение фактов")
|
||
print("="*60)
|
||
|
||
# Сохраняем сообщение с фактом
|
||
save_message(TEST_USER_ID, "user", "Меня зовут Владимир, я живу в городе Ангарск.")
|
||
|
||
# Получаем профиль
|
||
profile = get_profile(TEST_USER_ID)
|
||
|
||
print(f"✅ Факты в профиле:")
|
||
total_facts = 0
|
||
for fact_type, facts in profile.items():
|
||
print(f" [{fact_type}]:")
|
||
for fact in facts:
|
||
print(f" - {fact}")
|
||
total_facts += 1
|
||
|
||
print(f"\n Всего фактов: {total_facts}")
|
||
|
||
# Факты могут не извлечься эвристиками, это нормально
|
||
print("\n✅ ТЕСТ 5: ЗАВЕРШЁН\n")
|
||
return True
|
||
|
||
|
||
async def test_memory_gradient():
|
||
"""Тест 6: Градиентная память (STM → LTM → RAG)."""
|
||
print("\n" + "="*60)
|
||
print("ТЕСТ 6: Градиентная память (STM → LTM → RAG)")
|
||
print("="*60)
|
||
|
||
# Получаем контекст с градиентной памятью
|
||
from vector_memory import get_context as get_formatted_context
|
||
|
||
query = "Владимир"
|
||
context = get_formatted_context(TEST_USER_ID, query=query, stm_size=3, ltm_size=5)
|
||
|
||
print(f"✅ Сформирован контекст для ИИ:")
|
||
print(f" - Длина: {len(context)} символов")
|
||
print(f" - STM размер: 3 сообщения")
|
||
print(f" - LTM размер: 5 сообщений")
|
||
print(f" - RAG поиск по запросу: {query}")
|
||
|
||
# Проверяем наличие секций
|
||
has_stm = "💬 STM" in context
|
||
has_ltm = "🕰️ LTM" in context
|
||
has_rag = "🔍 RAG" in context
|
||
has_profile = "📋 ПРОФИЛЬ" in context
|
||
|
||
print(f"\n Секции:")
|
||
print(f" - Профиль: {'✅' if has_profile else '❌'}")
|
||
print(f" - STM: {'✅' if has_stm else '❌'}")
|
||
print(f" - LTM: {'✅' if has_ltm else '❌'}")
|
||
print(f" - RAG: {'✅' if has_rag else '❌'}")
|
||
|
||
assert has_stm or has_ltm, "Градиентная память не работает!"
|
||
print("\n✅ ТЕСТ 6: УСПЕШНО\n")
|
||
return True
|
||
|
||
|
||
async def main():
|
||
"""Запуск всех тестов."""
|
||
print("\n" + "="*60)
|
||
print("🧠 ТЕСТИРОВАНИЕ СИСТЕМЫ ПАМЯТИ TELEGRAM БОТА")
|
||
print("="*60)
|
||
print(f"Тестовый пользователь ID: {TEST_USER_ID}")
|
||
print(f"Дата: {__import__('datetime').datetime.now()}")
|
||
|
||
results = {
|
||
"SQLite Save": False,
|
||
"Vector Save": False,
|
||
"History Loading": False,
|
||
"RAG Search": False,
|
||
"Fact Extraction": False,
|
||
"Memory Gradient": False
|
||
}
|
||
|
||
try:
|
||
# Тест 1: SQLite
|
||
results["SQLite Save"] = test_sqlite_save()
|
||
|
||
# Тест 2: Vector
|
||
results["Vector Save"] = test_vector_save()
|
||
|
||
# Тест 3: Загрузка истории
|
||
results["History Loading"] = test_history_loading()
|
||
|
||
# Тест 4: RAG-поиск
|
||
results["RAG Search"] = await test_rag_search()
|
||
|
||
# Тест 5: Извлечение фактов
|
||
results["Fact Extraction"] = test_fact_extraction()
|
||
|
||
# Тест 6: Градиентная память
|
||
results["Memory Gradient"] = await test_memory_gradient()
|
||
|
||
except Exception as e:
|
||
print(f"\n❌ ОШИБКА ПРИ ТЕСТИРОВАНИИ: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
# Итоговый отчёт
|
||
print("\n" + "="*60)
|
||
print("📊 ИТОГОВЫЙ ОТЧЁТ")
|
||
print("="*60)
|
||
|
||
for test_name, result in results.items():
|
||
icon = "✅" if result else "❌"
|
||
print(f" {icon} {test_name}: {'УСПЕШНО' if result else 'ПРОВАЛ'}")
|
||
|
||
total_passed = sum(results.values())
|
||
total_tests = len(results)
|
||
|
||
print(f"\n Всего пройдено: {total_passed}/{total_tests}")
|
||
|
||
if total_passed == total_tests:
|
||
print("\n🎉 ВСЕ ТЕСТЫ ПРОЙДЕНЫ!")
|
||
else:
|
||
print(f"\n⚠️ {total_tests - total_passed} тест(а) не пройдено")
|
||
|
||
print("\n" + "="*60)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|