mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-09 01:36:57 +08:00
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
from pydantic import BaseModel
|
|
from typing import Optional, Dict, Any, Type
|
|
from abc import ABC, abstractmethod
|
|
from langchain_core.language_models import BaseLLM as LangchainBaseLLM
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
from apps.system.models.system_model import AiModelDetail
|
|
|
|
|
|
# from langchain_community.llms import Tongyi, VLLM
|
|
|
|
class LLMConfig(BaseModel):
|
|
"""Base configuration class for large language models"""
|
|
model_type: str # Model type: openai/tongyi/vllm etc.
|
|
model_name: str # Specific model name
|
|
api_key: Optional[str] = None
|
|
api_base_url: Optional[str] = None
|
|
additional_params: Dict[str, Any] = {}
|
|
|
|
|
|
class BaseLLM(ABC):
|
|
"""Abstract base class for large language models"""
|
|
|
|
def __init__(self, config: LLMConfig):
|
|
self.config = config
|
|
self._llm = self._init_llm()
|
|
|
|
@abstractmethod
|
|
def _init_llm(self) -> LangchainBaseLLM:
|
|
"""Initialize specific large language model instance"""
|
|
pass
|
|
|
|
@property
|
|
def llm(self) -> LangchainBaseLLM:
|
|
"""Return the langchain LLM instance"""
|
|
return self._llm
|
|
|
|
class OpenAILLM(BaseLLM):
|
|
def _init_llm(self) -> LangchainBaseLLM:
|
|
return ChatOpenAI(
|
|
model=self.config.model_name,
|
|
api_key=self.config.api_key,
|
|
base_url=self.config.api_base_url,
|
|
stream_usage=True,
|
|
**self.config.additional_params
|
|
)
|
|
|
|
def generate(self, prompt: str) -> str:
|
|
return self.llm.invoke(prompt)
|
|
|
|
""" class TongyiLLM(BaseLLM):
|
|
def _init_llm(self) -> LangchainBaseLLM:
|
|
return Tongyi(
|
|
model_name=self.config.model_name,
|
|
dashscope_api_key=self.config.api_key,
|
|
**self.config.additional_params
|
|
)
|
|
|
|
def generate(self, prompt: str) -> str:
|
|
return self.llm.invoke(prompt)
|
|
|
|
class VLLMLLM(BaseLLM):
|
|
def _init_llm(self) -> LangchainBaseLLM:
|
|
return VLLM(
|
|
model=self.config.model_name,
|
|
**self.config.additional_params
|
|
)
|
|
|
|
def generate(self, prompt: str) -> str:
|
|
return self.llm.invoke(prompt) """
|
|
|
|
|
|
class LLMFactory:
|
|
"""Large Language Model Factory Class"""
|
|
|
|
_llm_types: Dict[str, Type[BaseLLM]] = {
|
|
"openai": OpenAILLM,
|
|
"tongyi": OpenAILLM,
|
|
"vllm": OpenAILLM
|
|
}
|
|
|
|
@classmethod
|
|
def create_llm(cls, config: LLMConfig) -> BaseLLM:
|
|
llm_class = cls._llm_types.get(config.model_type)
|
|
if not llm_class:
|
|
raise ValueError(f"Unsupported LLM type: {config.model_type}")
|
|
return llm_class(config)
|
|
|
|
@classmethod
|
|
def register_llm(cls, model_type: str, llm_class: Type[BaseLLM]):
|
|
"""Register new model type"""
|
|
cls._llm_types[model_type] = llm_class
|
|
|
|
|
|
# todo
|
|
def get_llm_config(aimodel: AiModelDetail) -> LLMConfig:
|
|
config = LLMConfig(
|
|
model_type="openai",
|
|
model_name=aimodel.name,
|
|
api_key=aimodel.api_key,
|
|
api_base_url=aimodel.endpoint,
|
|
additional_params={"temperature": aimodel.temperature}
|
|
)
|
|
return config |