156 lines
4.9 KiB
Python
156 lines
4.9 KiB
Python
"""
|
|
OpenRouter Provider 구현
|
|
OpenRouter는 다양한 AI 모델에 대한 통합 API를 제공합니다.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
from ..base import BaseAIProvider, AIProviderType, AIMessage, AIResponse
|
|
|
|
|
|
class OpenRouterProvider(BaseAIProvider):
|
|
"""OpenRouter API 프로바이더"""
|
|
|
|
BASE_URL = "https://openrouter.ai/api/v1"
|
|
|
|
AVAILABLE_MODELS = [
|
|
"openai/gpt-4o",
|
|
"openai/gpt-4o-mini",
|
|
"anthropic/claude-3.5-sonnet",
|
|
"anthropic/claude-3-opus",
|
|
"google/gemini-pro-1.5",
|
|
"google/gemini-flash-1.5",
|
|
"meta-llama/llama-3.1-70b-instruct",
|
|
"meta-llama/llama-3.1-8b-instruct",
|
|
"mistralai/mixtral-8x7b-instruct",
|
|
"deepseek/deepseek-chat",
|
|
]
|
|
|
|
def __init__(self, api_key: str, model: Optional[str] = None):
|
|
super().__init__(api_key, model)
|
|
self._client = None
|
|
self._async_client = None
|
|
|
|
@property
|
|
def provider_type(self) -> AIProviderType:
|
|
return AIProviderType.OPENROUTER
|
|
|
|
@property
|
|
def provider_name(self) -> str:
|
|
return "OpenRouter"
|
|
|
|
@property
|
|
def default_model(self) -> str:
|
|
return "openai/gpt-4o-mini"
|
|
|
|
@property
|
|
def available_models(self) -> List[str]:
|
|
return self.AVAILABLE_MODELS.copy()
|
|
|
|
def initialize(self) -> bool:
|
|
"""OpenRouter 클라이언트 초기화 (OpenAI 호환 API 사용)"""
|
|
if not self.validate_api_key():
|
|
return False
|
|
|
|
try:
|
|
from openai import OpenAI, AsyncOpenAI
|
|
|
|
self._client = OpenAI(
|
|
api_key=self._api_key,
|
|
base_url=self.BASE_URL
|
|
)
|
|
self._async_client = AsyncOpenAI(
|
|
api_key=self._api_key,
|
|
base_url=self.BASE_URL
|
|
)
|
|
self._is_initialized = True
|
|
return True
|
|
except ImportError:
|
|
print("OpenAI 라이브러리가 설치되지 않았습니다. pip install openai 를 실행하세요.")
|
|
return False
|
|
except Exception as e:
|
|
print(f"OpenRouter 초기화 실패: {e}")
|
|
return False
|
|
|
|
def _convert_messages(self, messages: List[AIMessage]) -> List[Dict[str, str]]:
|
|
"""AIMessage를 OpenAI 형식으로 변환"""
|
|
return [{"role": msg.role, "content": msg.content} for msg in messages]
|
|
|
|
async def chat(
|
|
self,
|
|
messages: List[AIMessage],
|
|
temperature: float = 0.7,
|
|
max_tokens: Optional[int] = None,
|
|
**kwargs
|
|
) -> AIResponse:
|
|
"""비동기 채팅 요청"""
|
|
if not self._is_initialized:
|
|
if not self.initialize():
|
|
raise RuntimeError("OpenRouter 초기화 실패")
|
|
|
|
converted_messages = self._convert_messages(messages)
|
|
|
|
params = {
|
|
"model": self._model,
|
|
"messages": converted_messages,
|
|
"temperature": temperature,
|
|
}
|
|
if max_tokens:
|
|
params["max_tokens"] = max_tokens
|
|
|
|
# OpenRouter 특화 헤더는 extra_headers로 전달 가능
|
|
params.update(kwargs)
|
|
|
|
response = await self._async_client.chat.completions.create(**params)
|
|
|
|
return AIResponse(
|
|
content=response.choices[0].message.content,
|
|
model=response.model,
|
|
provider=self.provider_name,
|
|
usage={
|
|
"prompt_tokens": response.usage.prompt_tokens,
|
|
"completion_tokens": response.usage.completion_tokens,
|
|
"total_tokens": response.usage.total_tokens,
|
|
} if response.usage else None,
|
|
raw_response=response
|
|
)
|
|
|
|
def chat_sync(
|
|
self,
|
|
messages: List[AIMessage],
|
|
temperature: float = 0.7,
|
|
max_tokens: Optional[int] = None,
|
|
**kwargs
|
|
) -> AIResponse:
|
|
"""동기 채팅 요청"""
|
|
if not self._is_initialized:
|
|
if not self.initialize():
|
|
raise RuntimeError("OpenRouter 초기화 실패")
|
|
|
|
converted_messages = self._convert_messages(messages)
|
|
|
|
params = {
|
|
"model": self._model,
|
|
"messages": converted_messages,
|
|
"temperature": temperature,
|
|
}
|
|
if max_tokens:
|
|
params["max_tokens"] = max_tokens
|
|
params.update(kwargs)
|
|
|
|
response = self._client.chat.completions.create(**params)
|
|
|
|
return AIResponse(
|
|
content=response.choices[0].message.content,
|
|
model=response.model,
|
|
provider=self.provider_name,
|
|
usage={
|
|
"prompt_tokens": response.usage.prompt_tokens,
|
|
"completion_tokens": response.usage.completion_tokens,
|
|
"total_tokens": response.usage.total_tokens,
|
|
} if response.usage else None,
|
|
raw_response=response
|
|
)
|
|
|