Проект аудита отелей: основные скрипты и документация
- Краулеры: 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
This commit is contained in:
Фёдор
256
llm_client.py
Normal file
256
llm_client.py
Normal file
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
Универсальный клиент для работы с разными LLM провайдерами
|
||||
"""
|
||||
|
||||
import requests
|
||||
import os
|
||||
from typing import List, Dict, Optional
|
||||
from llm_config import (
|
||||
ACTIVE_PROVIDER, CURRENT_CONFIG, CHAT_MODEL,
|
||||
TEMPERATURE, MAX_TOKENS, OPENAI_CONFIG, OPENROUTER_CONFIG, OLLAMA_CONFIG
|
||||
)
|
||||
|
||||
|
||||
class LLMClient:
|
||||
"""Универсальный клиент для OpenAI/OpenRouter/Ollama"""
|
||||
|
||||
def __init__(self):
|
||||
# Читаем настройки из переменных окружения или используем дефолтные
|
||||
self.provider = os.environ.get('ACTIVE_PROVIDER', ACTIVE_PROVIDER)
|
||||
self.model = os.environ.get('LLM_MODEL', CHAT_MODEL)
|
||||
self.temperature = TEMPERATURE
|
||||
self.max_tokens = MAX_TOKENS
|
||||
|
||||
# Определяем конфиг на основе провайдера
|
||||
if self.provider == 'openai':
|
||||
self.config = OPENAI_CONFIG
|
||||
elif self.provider == 'openrouter':
|
||||
self.config = OPENROUTER_CONFIG
|
||||
elif self.provider == 'ollama':
|
||||
self.config = OLLAMA_CONFIG
|
||||
else:
|
||||
self.config = CURRENT_CONFIG
|
||||
|
||||
@property
|
||||
def provider_config(self):
|
||||
"""Получить конфиг текущего провайдера"""
|
||||
return {
|
||||
"provider": self.provider,
|
||||
"model": self.model,
|
||||
"api_base": self.config.get('api_base', ''),
|
||||
"has_key": bool(self.config.get('api_key'))
|
||||
}
|
||||
|
||||
def chat_completion(
|
||||
self,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: Optional[float] = None,
|
||||
max_tokens: Optional[int] = None,
|
||||
model: Optional[str] = None
|
||||
) -> Dict:
|
||||
"""
|
||||
Универсальный метод для chat completion
|
||||
|
||||
Args:
|
||||
messages: Список сообщений [{'role': 'user', 'content': '...'}]
|
||||
temperature: Креативность (0.0-1.0)
|
||||
max_tokens: Макс токенов в ответе
|
||||
model: Модель (опционально)
|
||||
|
||||
Returns:
|
||||
{'text': str, 'usage': dict}
|
||||
"""
|
||||
|
||||
temperature = temperature or self.temperature
|
||||
max_tokens = max_tokens or self.max_tokens
|
||||
model = model or self.model
|
||||
|
||||
if self.provider == 'openai' or self.provider == 'openrouter':
|
||||
return self._openai_style_request(messages, temperature, max_tokens, model)
|
||||
elif self.provider == 'ollama':
|
||||
return self._ollama_request(messages, temperature, max_tokens, model)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {self.provider}")
|
||||
|
||||
def _openai_style_request(
|
||||
self,
|
||||
messages: List[Dict],
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
model: str
|
||||
) -> Dict:
|
||||
"""Запрос к OpenAI или OpenRouter"""
|
||||
|
||||
url = f"{self.config['api_base']}/chat/completions"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.config['api_key']}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
# OpenRouter требует дополнительные заголовки
|
||||
if self.provider == 'openrouter':
|
||||
headers["HTTP-Referer"] = "https://hotel-audit.ru"
|
||||
headers["X-Title"] = "Hotel Audit System"
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens
|
||||
}
|
||||
|
||||
# Настройка прокси
|
||||
proxies = None
|
||||
if self.config.get('proxy'):
|
||||
proxies = {
|
||||
'http': self.config['proxy'],
|
||||
'https': self.config['proxy']
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
proxies=proxies,
|
||||
timeout=60
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
return {
|
||||
'text': data['choices'][0]['message']['content'],
|
||||
'usage': data.get('usage', {}),
|
||||
'model': data.get('model', model)
|
||||
}
|
||||
else:
|
||||
error_msg = f"API Error {response.status_code}: {response.text}"
|
||||
return {
|
||||
'text': f"Ошибка API: {response.status_code}",
|
||||
'error': error_msg
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
'text': f"Ошибка соединения: {str(e)}",
|
||||
'error': str(e)
|
||||
}
|
||||
|
||||
def _ollama_request(
|
||||
self,
|
||||
messages: List[Dict],
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
model: str
|
||||
) -> Dict:
|
||||
"""Запрос к локальной Ollama"""
|
||||
|
||||
url = f"{self.config['api_base']}/api/chat"
|
||||
|
||||
# Конвертируем формат сообщений
|
||||
payload = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"temperature": temperature,
|
||||
"num_predict": max_tokens
|
||||
}
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=payload, timeout=120)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
return {
|
||||
'text': data['message']['content'],
|
||||
'usage': {
|
||||
'prompt_tokens': data.get('prompt_eval_count', 0),
|
||||
'completion_tokens': data.get('eval_count', 0)
|
||||
},
|
||||
'model': model
|
||||
}
|
||||
else:
|
||||
return {
|
||||
'text': f"Ошибка Ollama: {response.status_code}",
|
||||
'error': response.text
|
||||
}
|
||||
|
||||
except requests.exceptions.ConnectionError:
|
||||
return {
|
||||
'text': "Ошибка: Ollama не запущена. Запустите: ollama serve",
|
||||
'error': 'Connection refused'
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'text': f"Ошибка Ollama: {str(e)}",
|
||||
'error': str(e)
|
||||
}
|
||||
|
||||
def simple_chat(self, prompt: str, system: Optional[str] = None) -> str:
|
||||
"""
|
||||
Простой метод для быстрого чата
|
||||
|
||||
Args:
|
||||
prompt: Вопрос пользователя
|
||||
system: Системный промпт (опционально)
|
||||
|
||||
Returns:
|
||||
Ответ модели (только текст)
|
||||
"""
|
||||
|
||||
messages = []
|
||||
|
||||
if system:
|
||||
messages.append({'role': 'system', 'content': system})
|
||||
|
||||
messages.append({'role': 'user', 'content': prompt})
|
||||
|
||||
result = self.chat_completion(messages)
|
||||
return result['text']
|
||||
|
||||
def get_info(self) -> Dict:
|
||||
"""Информация о текущей конфигурации"""
|
||||
return {
|
||||
'provider': self.provider,
|
||||
'model': self.model,
|
||||
'temperature': self.temperature,
|
||||
'max_tokens': self.max_tokens,
|
||||
'api_base': self.config['api_base']
|
||||
}
|
||||
|
||||
|
||||
# Глобальный экземпляр (singleton pattern)
|
||||
llm = LLMClient()
|
||||
|
||||
|
||||
# ==================== ТЕСТЫ ====================
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 70)
|
||||
print("🤖 ТЕСТ LLM КЛИЕНТА")
|
||||
print("=" * 70)
|
||||
|
||||
info = llm.get_info()
|
||||
print(f"\n📊 Конфигурация:")
|
||||
print(f" Провайдер: {info['provider']}")
|
||||
print(f" Модель: {info['model']}")
|
||||
print(f" Temperature: {info['temperature']}")
|
||||
print(f" Max tokens: {info['max_tokens']}")
|
||||
print(f" API: {info['api_base']}")
|
||||
|
||||
print(f"\n💬 Тестовый запрос...")
|
||||
|
||||
response = llm.simple_chat(
|
||||
prompt="Сколько будет 2+2? Ответь одним числом.",
|
||||
system="Ты математический ассистент."
|
||||
)
|
||||
|
||||
print(f" Ответ: {response}")
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print("✅ Клиент работает!")
|
||||
print("=" * 70)
|
||||
|
||||
Reference in New Issue
Block a user