hotels/llm_client.py

"""
Универсальный клиент для работы с разными 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)