Files
hotels/user_settings_manager.py
Фёдор 0cf3297290 Проект аудита отелей: основные скрипты и документация
- Краулеры: smart_crawler.py, regional_crawler.py
- Аудит: audit_orel_to_excel.py, audit_chukotka_to_excel.py
- РКН проверка: check_rkn_registry.py, recheck_unclear_rkn.py
- Отчёты: create_orel_horizontal_report.py
- Обработка: process_all_hotels_embeddings.py
- Документация: README.md, DB_SCHEMA_REFERENCE.md
2025-10-16 10:52:09 +03:00

217 lines
7.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
Класс для работы с настройками пользователей в БД
"""
import psycopg2
from psycopg2.extras import RealDictCursor
from typing import Dict, List, Optional
import json
class UserSettingsManager:
"""Менеджер настроек пользователей"""
def __init__(self):
self.db_config = {
'host': '147.45.189.234',
'port': 5432,
'database': 'default_db',
'user': 'gen_user',
'password': '2~~9_^kVsU?2\\S'
}
def get_connection(self):
"""Получить подключение к БД"""
return psycopg2.connect(**self.db_config, cursor_factory=RealDictCursor)
def get_user_setting(self, user_id: str, setting_key: str, default_value: str = None) -> str:
"""Получить настройку пользователя"""
try:
conn = self.get_connection()
cur = conn.cursor()
cur.execute("""
SELECT setting_value
FROM user_settings
WHERE user_id = %s AND setting_key = %s
""", (user_id, setting_key))
result = cur.fetchone()
cur.close()
conn.close()
return result['setting_value'] if result else default_value
except Exception as e:
print(f"Ошибка получения настройки: {e}")
return default_value
def set_user_setting(self, user_id: str, setting_key: str, setting_value: str) -> bool:
"""Установить настройку пользователя"""
try:
conn = self.get_connection()
cur = conn.cursor()
cur.execute("""
INSERT INTO user_settings (user_id, setting_key, setting_value, updated_at)
VALUES (%s, %s, %s, CURRENT_TIMESTAMP)
ON CONFLICT (user_id, setting_key)
DO UPDATE SET
setting_value = EXCLUDED.setting_value,
updated_at = CURRENT_TIMESTAMP
""", (user_id, setting_key, setting_value))
conn.commit()
cur.close()
conn.close()
return True
except Exception as e:
print(f"Ошибка сохранения настройки: {e}")
return False
def get_user_llm_settings(self, user_id: str) -> Dict[str, str]:
"""Получить все LLM настройки пользователя"""
settings = {}
# Получаем основные настройки
settings['provider'] = self.get_user_setting(user_id, 'llm_provider', 'openai')
settings['model'] = self.get_user_setting(user_id, 'llm_model', 'gpt-4o-mini')
settings['temperature'] = self.get_user_setting(user_id, 'llm_temperature', '0.3')
settings['max_tokens'] = self.get_user_setting(user_id, 'llm_max_tokens', '800')
return settings
def set_user_llm_settings(self, user_id: str, provider: str, model: str,
temperature: str = '0.3', max_tokens: str = '800') -> bool:
"""Сохранить LLM настройки пользователя"""
try:
conn = self.get_connection()
cur = conn.cursor()
settings = {
'llm_provider': provider,
'llm_model': model,
'llm_temperature': temperature,
'llm_max_tokens': max_tokens
}
for key, value in settings.items():
cur.execute("""
INSERT INTO user_settings (user_id, setting_key, setting_value, updated_at)
VALUES (%s, %s, %s, CURRENT_TIMESTAMP)
ON CONFLICT (user_id, setting_key)
DO UPDATE SET
setting_value = EXCLUDED.setting_value,
updated_at = CURRENT_TIMESTAMP
""", (user_id, key, value))
conn.commit()
cur.close()
conn.close()
return True
except Exception as e:
print(f"Ошибка сохранения LLM настроек: {e}")
return False
def get_available_models(self, provider: str = None) -> List[Dict]:
"""Получить доступные модели"""
try:
conn = self.get_connection()
cur = conn.cursor()
if provider:
cur.execute("""
SELECT provider, model_id, model_name, description,
context_length, pricing_input, pricing_output
FROM llm_models
WHERE provider = %s AND is_active = true
ORDER BY model_name
""", (provider,))
else:
cur.execute("""
SELECT provider, model_id, model_name, description,
context_length, pricing_input, pricing_output
FROM llm_models
WHERE is_active = true
ORDER BY provider, model_name
""")
models = cur.fetchall()
cur.close()
conn.close()
return [dict(model) for model in models]
except Exception as e:
print(f"Ошибка получения моделей: {e}")
return []
def get_providers(self) -> List[Dict]:
"""Получить список провайдеров с количеством моделей"""
try:
conn = self.get_connection()
cur = conn.cursor()
cur.execute("""
SELECT provider, COUNT(*) as model_count
FROM llm_models
WHERE is_active = true
GROUP BY provider
ORDER BY provider
""")
providers = cur.fetchall()
cur.close()
conn.close()
return [dict(provider) for provider in providers]
except Exception as e:
print(f"Ошибка получения провайдеров: {e}")
return []
# Глобальный экземпляр
user_settings_manager = UserSettingsManager()
if __name__ == "__main__":
print("=" * 70)
print("🧪 ТЕСТ МЕНЕДЖЕРА НАСТРОЕК ПОЛЬЗОВАТЕЛЕЙ")
print("=" * 70)
# Тест получения провайдеров
print("📊 Провайдеры:")
providers = user_settings_manager.get_providers()
for provider in providers:
print(f" {provider['provider']}: {provider['model_count']} моделей")
# Тест получения моделей OpenAI
print(f"\n🤖 Модели OpenAI:")
models = user_settings_manager.get_available_models('openai')
for model in models[:5]: # Показываем первые 5
print(f" {model['model_id']}: {model['model_name']}")
# Тест настроек пользователя
test_user = "user_test_123"
print(f"\n👤 Тест настроек пользователя {test_user}:")
# Сохраняем настройки
success = user_settings_manager.set_user_llm_settings(
test_user, 'openrouter', 'anthropic/claude-3-haiku', '0.5', '1000'
)
print(f"✅ Сохранение настроек: {'Успешно' if success else 'Ошибка'}")
# Получаем настройки
settings = user_settings_manager.get_user_llm_settings(test_user)
print(f"📋 Настройки: {settings}")
print("=" * 70)