mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
import json
|
|
from abc import ABC, abstractmethod
|
|
from typing import Optional, Dict, Any, Type
|
|
|
|
from langchain.chat_models.base import BaseChatModel
|
|
from pydantic import BaseModel
|
|
from sqlmodel import Session, select
|
|
|
|
from apps.ai_model.openai.llm import BaseChatOpenAI
|
|
from apps.system.models.system_model import AiModelDetail
|
|
from common.core.db import engine
|
|
|
|
|
|
# from langchain_community.llms import Tongyi, VLLM
|
|
|
|
class LLMConfig(BaseModel):
|
|
"""Base configuration class for large language models"""
|
|
model_id: Optional[int] = None
|
|
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) -> BaseChatModel:
|
|
"""Initialize specific large language model instance"""
|
|
pass
|
|
|
|
@property
|
|
def llm(self) -> BaseChatModel:
|
|
"""Return the langchain LLM instance"""
|
|
return self._llm
|
|
|
|
|
|
class OpenAILLM(BaseLLM):
|
|
def _init_llm(self) -> BaseChatModel:
|
|
return BaseChatOpenAI(
|
|
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,
|
|
extra_body={"enable_thinking": True},
|
|
)
|
|
|
|
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 """
|
|
|
|
|
|
def get_default_config() -> LLMConfig:
|
|
with Session(engine) as session:
|
|
db_model = session.exec(
|
|
select(AiModelDetail).where(AiModelDetail.default_model == True)
|
|
).first()
|
|
if not db_model:
|
|
raise ValueError("The system default model has not been set")
|
|
|
|
additional_params = {}
|
|
if db_model.config:
|
|
try:
|
|
config_raw = json.loads(db_model.config)
|
|
additional_params = {item["key"]: item["val"] for item in config_raw if "key" in item and "val" in item}
|
|
except Exception:
|
|
pass
|
|
|
|
# 构造 LLMConfig
|
|
return LLMConfig(
|
|
model_id=db_model.id,
|
|
model_type="openai",
|
|
model_name=db_model.base_model,
|
|
api_key=db_model.api_key,
|
|
api_base_url=db_model.api_domain,
|
|
additional_params=additional_params,
|
|
)
|