llm-agent-backend/app/llm_client.py

"""
Этот модуль предоставляет инструменты для взаимодействия с различными LLM,
включая OpenAI (Gemini) и MistralAI, обеспечивая унифицированный интерфейс.
"""

import requests
from typing import Dict, Any, List
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage

# --- Конфигурация LLM ---
# Пустой словарь для локальных моделей (Ollama)
LOCAL_MODELS: Dict[str, Any] = {}

# Конфигурация внешних моделей
MODELS: Dict[str, Dict[str, Any]] = {
    "gemini-2.0-flash": {
        "name": "gemini-2.0-flash",
        "provider": "openai",  # Изменено на "openai"
        "model_name": "gemini-2.0-flash",  # Добавлено имя модели для LangChain
        "apiBase": "https://render-service-gsu7.onrender.com/g/v1beta",
        "apiKey": "AIzaSyDpueKFWVqknVKlQn6TdasLmJ2lvAUiBik",
        "stream": True,
        "capabilities": ["vision"],
    },
    "gemini-2.5-flash": {
        "name": "gemini-2.5-flash",
        "provider": "openai",  # Изменено на "openai"
        "model_name": "gemini-2.5-flash",  # Добавлено имя модели для LangChain
        "apiBase": "https://render-service-gsu7.onrender.com/g/v1beta",
        "apiKey": "AIzaSyDpueKFWVqknVKlQn6TdasLmJ2lvAUiBik",
        "stream": True,
        "capabilities": ["vision"],
    },
    "gemini-2.5-flash-preview-05-20": {
        "name": "gemini-2.5-flash-preview-05-20",
        "provider": "openai",  # Изменено на "openai"
        "model_name":
        "gemini-2.5-flash-preview-05-20",  # Добавлено имя модели для LangChain
        "apiBase": "https://render-service-gsu7.onrender.com/g/v1beta",
        "apiKey": "AIzaSyDpueKFWVqknVKlQn6TdasLmJ2lvAUiBik",
        "stream": True,
        "capabilities": ["vision"],
    },
    "mistral-small-latest": {
        "name": "mistral-small-latest",
        "provider": "mistralai",
        "baseUrl": "https://render-service-gsu7.onrender.com/m",
        "apiKey": "Q0m29fvxBY0Cfdj4sjHaKqccy1NjonLW",
        "stream": True,
        "capabilities": ["vision"],
    },
}


class CustomLLM:
    """
    Класс-обертка для взаимодействия с различными LLM.
    """

    def __init__(self, config: Dict[str, Any]):
        self.name = config["name"]
        self.provider = config["provider"]
        self.config = config

        if self.provider == "openai":
            # Инициализация Gemini через OpenAI API
            from openai import OpenAI
            self.client = OpenAI(api_key=config["apiKey"],
                                 base_url=config["apiBase"])
            self.model_name = config["model_name"]
            self.stream = config["stream"]

        elif self.provider == "mistralai":
            # Для Mistral используем прямые HTTP-запросы
            self.base_url = config["baseUrl"]
            self.api_key = config["apiKey"]
            self.stream = config.get("stream", False)
            self.capabilities = config.get("capabilities", [])
        else:
            raise ValueError(f"Неизвестный провайдер LLM: {self.provider}")

    def invoke(self, messages: List[Any]) -> str:
        """
        Отправляет запрос к LLM и возвращает ответ.
        Messages могут быть строкой или списком объектов Langchain Message.
        """
        if self.provider == "openai":
            # OpenAI ожидает список сообщений в определенном формате
            openai_messages = []
            if isinstance(messages, str):
                openai_messages.append({"role": "user", "content": messages})
            elif isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, HumanMessage):
                        openai_messages.append({
                            "role": "user",
                            "content": msg.content
                        })
                    elif isinstance(msg, SystemMessage):
                        openai_messages.append({
                            "role": "system",
                            "content": msg.content
                        })
                    elif isinstance(msg, AIMessage): # Добавляем обработку AIMessage
                        openai_messages.append({
                            "role": "assistant",
                            "content": msg.content
                        })
                    else:  # Предполагаем, что это другие типы сообщений Langchain или словари
                        openai_messages.append({
                            "role":
                            msg.type if hasattr(msg, 'type') else 'user',
                            "content":
                            msg.content
                        })

            try:
                response = self.client.chat.completions.create(
                    model=self.model_name,
                    messages=openai_messages,
                    stream=self.stream,
                )

                if self.stream:
                    # Handle streaming responses
                    collected_chunks = []
                    collected_messages = []
                    for chunk in response:
                        collected_chunks.append(chunk)
                        chunk_message = chunk.choices[0].delta.content
                        if chunk_message is not None:
                            collected_messages.append(chunk_message)
                    full_reply_content = ''.join(collected_messages)
                    return full_reply_content
                else:
                    return response.choices[0].message.content

            except Exception as e:
                print(f"Ошибка при вызове OpenAI API: {e}")
                raise
        elif self.provider == "mistralai":
            # Отступы исправлены
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            }
            # Формируем сообщения для Mistral API
            mistral_messages = []
            if isinstance(messages, str):
                mistral_messages.append({"role": "user", "content": messages})
            elif isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, HumanMessage):
                        mistral_messages.append({
                            "role": "user",
                            "content": msg.content
                        })
                    elif isinstance(msg, SystemMessage):
                        mistral_messages.append({
                            "role": "system",
                            "content": msg.content
                        })
                    elif isinstance(msg, AIMessage): # Добавляем обработку AIMessage
                        mistral_messages.append({
                            "role": "assistant",
                            "content": msg.content
                        })
                    else:  # Предполагаем, что это другие типы сообщений Langchain или словари
                        mistral_messages.append({
                            "role":
                            msg.type if hasattr(msg, 'type') else 'user',
                            "content":
                            msg.content
                        })

            payload = {
                "model": self.name,
                "messages": mistral_messages,
                "stream": self.stream,
            }
            try:
                response = requests.post(self.base_url,
                                         json=payload,
                                         headers=headers)
                response.raise_for_status(
                )  # Вызывает исключение для HTTP ошибок
                data = response.json()
                # Извлекаем контент ответа
                return data["choices"][0]["message"]["content"]
            except requests.exceptions.RequestException as e:
                print(f"Ошибка при вызове Mistral API: {e}")
                raise
            except KeyError:
                print(f"Неверный формат ответа от Mistral API: {data}")
                raise ValueError("Неверный формат ответа от Mistral API")
        return ""


def get_llm(name: str) -> CustomLLM:
    """
    Возвращает экземпляр CustomLLM для заданной модели.
    """
    config = MODELS.get(name)
    if not config:
        raise ValueError(f"Модель '{name}' не сконфигурирована.")
    return CustomLLM(config)