File size: 2,363 Bytes
bba179e 03b5791 bba179e 03b5791 bba179e 54d4a54 bba179e c62892c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from typing import List
from abc import ABC
from langchain_openai import ChatOpenAI
class BaseLLM(ABC):
_api_key: str
_support_models: List[str]
_default_model: str
_base_url: str
_default_temperature: float = 0.5
_default_max_tokens: int = 4096
def __init__(self, *, api_key: str):
self._api_key = api_key
@property
def support_models(self) -> List[str]:
return self._support_models
@property
def default_model(self) -> str:
return self._default_model
@property
def base_url(self) -> str:
return self._base_url
@property
def api_key(self) -> str:
return self._api_key
@property
def default_temperature(self) -> float:
return self._default_temperature
@property
def default_max_tokens(self) -> int:
return self._default_max_tokens
def get_chat_engine(self, *, model: str = None, temperature: float = None, max_tokens: int = None):
model = model or self.default_model
temperature = temperature or self.default_temperature
max_tokens = max_tokens or self.default_max_tokens
return ChatOpenAI(
model=model,
api_key=self.api_key,
base_url=self.base_url,
temperature=temperature,
max_tokens=max_tokens,
)
class DeepSeekLLM(BaseLLM):
_support_models = ['deepseek-chat', 'deepseek-coder']
_base_url = 'https://api.deepseek.com/v1'
_default_model = 'deepseek-chat'
_default_max_tokens = 4096
class OpenRouterLLM(BaseLLM):
_support_models = [
'openai/gpt-4o-mini', 'anthropic/claude-3.5-sonnet', 'openai/gpt-4o',
'nvidia/nemotron-4-340b-instruct', 'deepseek/deepseek-coder',
'google/gemini-flash-1.5', 'deepseek/deepseek-chat',
'liuhaotian/llava-yi-34b', 'qwen/qwen-110b-chat',
'qwen/qwen-72b-chat', 'google/gemini-pro-1.5',
'cohere/command-r-plus', 'anthropic/claude-3-haiku',
]
_base_url = 'https://openrouter.ai/api/v1'
_default_model = 'anthropic/claude-3.5-sonnet'
_default_max_tokens = 4096
class TongYiLLM(BaseLLM):
_support_models = ['qwen-turbo', 'qwen-plus', 'qwen-max', 'qwen-long']
_default_model = 'qwen-turbo'
_base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
_default_max_tokens: int = 2000
|