56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
|
|
# app/llm_factory.py
|
||
|
|
import os
|
||
|
|
from langchain_community.chat_models import ChatZhipuAI
|
||
|
|
from langchain_openai import ChatOpenAI
|
||
|
|
from pydantic import SecretStr
|
||
|
|
|
||
|
|
class LLMFactory:
|
||
|
|
@staticmethod
|
||
|
|
def create_zhipu():
|
||
|
|
api_key = os.getenv("ZHIPUAI_API_KEY")
|
||
|
|
if not api_key:
|
||
|
|
raise ValueError("ZHIPUAI_API_KEY not set")
|
||
|
|
return ChatZhipuAI(
|
||
|
|
model="glm-4.7-flash",
|
||
|
|
api_key=api_key,
|
||
|
|
temperature=0.1,
|
||
|
|
max_tokens=4096,
|
||
|
|
timeout=120.0,
|
||
|
|
max_retries=3,
|
||
|
|
streaming=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def create_deepseek():
|
||
|
|
api_key = os.getenv("DEEPSEEK_API_KEY")
|
||
|
|
if not api_key:
|
||
|
|
raise ValueError("DEEPSEEK_API_KEY not set")
|
||
|
|
return ChatOpenAI(
|
||
|
|
base_url="https://api.deepseek.com",
|
||
|
|
api_key=SecretStr(api_key),
|
||
|
|
model="deepseek-reasoner",
|
||
|
|
temperature=0.1,
|
||
|
|
max_tokens=4096,
|
||
|
|
timeout=60.0,
|
||
|
|
max_retries=2,
|
||
|
|
streaming=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def create_local():
|
||
|
|
base_url = os.getenv("VLLM_BASE_URL", "http://127.0.0.1:8081/v1")
|
||
|
|
return ChatOpenAI(
|
||
|
|
base_url=base_url,
|
||
|
|
api_key=SecretStr(os.getenv("LLAMACPP_API_KEY", "token-abc123")),
|
||
|
|
model="gemma-4-E4B-it",
|
||
|
|
timeout=60.0,
|
||
|
|
max_retries=2,
|
||
|
|
streaming=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
# 模型创建器映射
|
||
|
|
CREATORS = {
|
||
|
|
"local": create_local,
|
||
|
|
"deepseek": create_deepseek,
|
||
|
|
"zhipu": create_zhipu,
|
||
|
|
}
|