93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""
|
||
LLM 服务模块：支持 DeepSeek 和其他兼容 OpenAI 的 LLM
|
||
"""
|
||
import os
|
||
from typing import Optional
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
|
||
class LLMService:
|
||
"""LLM 服务，支持 DeepSeek 和其他兼容 OpenAI 的 LLM"""
|
||
|
||
def __init__(self):
|
||
"""初始化 LLM 服务"""
|
||
# 优先使用 DEEPSEEK_API_KEY，如果没有则使用 LLM_API_KEY
|
||
api_key = (
|
||
os.getenv("DEEPSEEK_API_KEY") or
|
||
os.getenv("LLM_API_KEY", "")
|
||
)
|
||
|
||
# 获取 base_url，优先使用 DEEPSEEK_BASE_URL
|
||
base_url = (
|
||
os.getenv("DEEPSEEK_BASE_URL") or
|
||
os.getenv("LLM_BASE_URL", "")
|
||
)
|
||
|
||
# 获取模型名称，优先使用 DEEPSEEK_MODEL
|
||
model_name = (
|
||
os.getenv("DEEPSEEK_MODEL") or
|
||
os.getenv("LLM_MODEL") or
|
||
"deepseek-chat" # DeepSeek 默认模型
|
||
)
|
||
|
||
# 如果没有提供 base_url 但有 DEEPSEEK_API_KEY，使用 DeepSeek 默认地址
|
||
if not base_url and os.getenv("DEEPSEEK_API_KEY"):
|
||
base_url = "https://api.deepseek.com"
|
||
|
||
if not api_key:
|
||
self.llm = None
|
||
return
|
||
|
||
# 配置 LLM 参数
|
||
# 获取 temperature，处理字符串默认值
|
||
temp_str = os.getenv("LLM_TEMPERATURE", "0.7")
|
||
try:
|
||
temperature = float(temp_str)
|
||
except (ValueError, TypeError):
|
||
temperature = 0.7
|
||
|
||
llm_kwargs = {
|
||
"temperature": temperature,
|
||
"model": model_name,
|
||
"api_key": api_key,
|
||
}
|
||
|
||
if base_url:
|
||
# 移除可能的 /v1/chat/completions 路径，langchain 会自动添加
|
||
if base_url.endswith("/v1/chat/completions"):
|
||
base_url = base_url[:-20] # 移除 "/v1/chat/completions"
|
||
elif base_url.endswith("/v1"):
|
||
base_url = base_url[:-3] # 移除 "/v1"
|
||
# 确保 base_url 不以 / 结尾（langchain 会自动添加）
|
||
if base_url.endswith("/"):
|
||
base_url = base_url[:-1]
|
||
llm_kwargs["base_url"] = base_url
|
||
|
||
try:
|
||
self.llm = ChatOpenAI(**llm_kwargs)
|
||
except Exception as e:
|
||
print(f"初始化 LLM 失败: {e}")
|
||
self.llm = None
|
||
|
||
def get_llm(self) -> Optional[ChatOpenAI]:
|
||
"""
|
||
获取 LLM 实例
|
||
|
||
Returns:
|
||
ChatOpenAI 实例，如果未配置则返回 None
|
||
"""
|
||
return self.llm
|
||
|
||
def is_available(self) -> bool:
|
||
"""
|
||
检查 LLM 服务是否可用
|
||
|
||
Returns:
|
||
True 如果 LLM 已配置且可用，否则 False
|
||
"""
|
||
return self.llm is not None
|
||
|
||
|
||
# 创建全局实例
|
||
llm_service = LLMService()
|