127 lines
4.9 KiB
Python
127 lines
4.9 KiB
Python
from typing import Any, Dict, Optional
|
||
|
||
import httpx
|
||
from config.settings import settings
|
||
|
||
from services.model_config import model_config
|
||
|
||
|
||
class LLMRouter:
|
||
def __init__(self):
|
||
self.ollama_base_url = settings.OLLAMA_BASE_URL
|
||
self.openrouter_base_url = settings.OPENROUTER_BASE_URL
|
||
|
||
async def generate_ollama(
|
||
self,
|
||
prompt: str,
|
||
model: Optional[str] = None,
|
||
preset: str = "creative_draft",
|
||
options: Optional[Dict[str, Any]] = None,
|
||
) -> str:
|
||
"""Генерация через локальную Ollama с пресетами"""
|
||
from services.presets import get_preset
|
||
|
||
# Используем модель из конфига, если не указана
|
||
model = model or model_config.get_default_model("ollama")
|
||
|
||
# Валидация модели
|
||
if not model_config.validate_model("ollama", model):
|
||
raise ValueError(f"Model {model} not available for Ollama")
|
||
|
||
# Получаем базовые параметры из пресета
|
||
preset_options = get_preset(preset, "ollama")
|
||
# Объединяем с пользовательскими опциями (пользовательские переопределяют пресет)
|
||
merged_options = {**preset_options, **(options or {})}
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
payload = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"stream": False,
|
||
"options": merged_options,
|
||
}
|
||
response = await client.post(
|
||
f"{self.ollama_base_url}/api/generate", json=payload, timeout=60.0
|
||
)
|
||
response.raise_for_status()
|
||
return response.json()["response"]
|
||
|
||
async def generate_openrouter(
|
||
self,
|
||
prompt: str,
|
||
model: Optional[str] = None,
|
||
api_key: Optional[str] = None,
|
||
preset: str = "creative_draft",
|
||
options: Optional[Dict[str, Any]] = None,
|
||
) -> str:
|
||
"""Генерация через OpenRouter API с пресетами"""
|
||
from services.presets import get_preset
|
||
|
||
# Используем ключ из settings, если не передан явно
|
||
actual_api_key = api_key or settings.OPENROUTER_API_KEY
|
||
|
||
# Валидация API ключа
|
||
if not actual_api_key or actual_api_key == "your_api_key_here":
|
||
raise ValueError(
|
||
"OpenRouter API key not configured. "
|
||
"Please set OPENROUTER_API_KEY in backend/.env file"
|
||
)
|
||
|
||
# Используем модель из конфига, если не указана
|
||
model = model or model_config.get_default_model("openrouter")
|
||
|
||
# Валидация модели
|
||
if not model_config.validate_model("openrouter", model):
|
||
raise ValueError(f"Model {model} not available for OpenRouter")
|
||
|
||
# Получаем базовые параметры из пресета
|
||
preset_options = get_preset(preset, "openrouter")
|
||
# Объединяем с пользовательскими опциями
|
||
merged_options = {**preset_options, **(options or {})}
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
payload = {
|
||
"model": model,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
}
|
||
# Добавляем объединённые опции (включая frequency_penalty, presence_penalty и temperature)
|
||
if merged_options:
|
||
payload.update(merged_options)
|
||
headers = {
|
||
"Authorization": f"Bearer {actual_api_key}",
|
||
"Content-Type": "application/json",
|
||
"HTTP-Referer": "http://localhost:3000", # Для OpenRouter
|
||
"X-Title": "Writer IDE", # Для OpenRouter
|
||
}
|
||
response = await client.post(
|
||
f"{self.openrouter_base_url}/chat/completions",
|
||
json=payload,
|
||
headers=headers,
|
||
timeout=60.0,
|
||
)
|
||
response.raise_for_status()
|
||
return response.json()["choices"][0]["message"]["content"]
|
||
|
||
async def generate(
|
||
self,
|
||
prompt: str,
|
||
provider: str = "ollama",
|
||
model: Optional[str] = None,
|
||
api_key: Optional[str] = None,
|
||
preset: str = "creative_draft",
|
||
options: Optional[Dict[str, Any]] = None,
|
||
) -> str:
|
||
"""Универсальный метод генерации"""
|
||
if provider == "ollama":
|
||
return await self.generate_ollama(prompt, model, preset, options)
|
||
elif provider == "openrouter":
|
||
return await self.generate_openrouter(
|
||
prompt, model, api_key, preset, options
|
||
)
|
||
else:
|
||
raise ValueError(f"Unknown provider: {provider}")
|
||
|
||
|
||
# Singleton instance
|
||
llm_router = LLMRouter()
|