Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,404 @@
|
||||
"""
|
||||
LLM 服务
|
||||
"""
|
||||
import logging
|
||||
from typing import AsyncGenerator, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.providers import (
|
||||
BaseLLMProvider,
|
||||
LLMConfig,
|
||||
LLMMessage,
|
||||
LLMResponse,
|
||||
ProviderRegistry,
|
||||
)
|
||||
from ai_platform.providers.base import LLMStreamChunk, ToolDefinition
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMService:
|
||||
"""
|
||||
LLM 服务
|
||||
|
||||
统一管理 LLM 调用，支持多种提供商
|
||||
"""
|
||||
|
||||
def __init__(self, db: Optional[AsyncSession] = None):
|
||||
self._provider_cache: Dict[str, BaseLLMProvider] = {}
|
||||
self._db = db
|
||||
|
||||
async def _get_provider_async(self, model_id: str) -> tuple:
|
||||
"""
|
||||
异步获取模型对应的提供商实例
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
|
||||
Returns:
|
||||
(provider, model_name)
|
||||
"""
|
||||
from ai_platform.models import LLMModel, LLMProvider
|
||||
|
||||
if not self._db:
|
||||
raise ValueError("数据库会话未初始化")
|
||||
|
||||
# 查询模型
|
||||
result = await self._db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_active == True,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
model = result.scalar_one_or_none()
|
||||
if not model:
|
||||
raise ValueError(f'模型不存在或已禁用: {model_id}')
|
||||
|
||||
# 查询提供商
|
||||
provider_result = await self._db.execute(
|
||||
select(LLMProvider).where(LLMProvider.id == model.provider_id)
|
||||
)
|
||||
provider = provider_result.scalar_one_or_none()
|
||||
if not provider or not provider.is_active:
|
||||
raise ValueError(f'提供商不存在或已禁用')
|
||||
|
||||
# 缓存提供商实例
|
||||
cache_key = str(provider.id)
|
||||
if cache_key not in self._provider_cache:
|
||||
provider_instance = ProviderRegistry.create_instance(
|
||||
provider_type=provider.provider_type,
|
||||
api_key=provider.api_key,
|
||||
api_base=provider.api_base or "",
|
||||
ollama_host=provider.ollama_host or "",
|
||||
)
|
||||
if not provider_instance:
|
||||
raise ValueError(f'不支持的提供商类型: {provider.provider_type}')
|
||||
self._provider_cache[cache_key] = provider_instance
|
||||
|
||||
return self._provider_cache[cache_key], model.model_name
|
||||
|
||||
def _get_provider_sync(self, model_id: str, model_data: dict) -> tuple:
|
||||
"""
|
||||
同步获取提供商实例（使用预加载的数据）
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
model_data: 预加载的模型和提供商数据
|
||||
|
||||
Returns:
|
||||
(provider, model_name)
|
||||
"""
|
||||
provider_type = model_data.get('provider_type')
|
||||
api_key = model_data.get('api_key', '')
|
||||
api_base = model_data.get('api_base', '')
|
||||
ollama_host = model_data.get('ollama_host', '')
|
||||
model_name = model_data.get('model_name', '')
|
||||
provider_id = model_data.get('provider_id', '')
|
||||
|
||||
# 缓存提供商实例
|
||||
cache_key = str(provider_id)
|
||||
if cache_key not in self._provider_cache:
|
||||
provider_instance = ProviderRegistry.create_instance(
|
||||
provider_type=provider_type,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
ollama_host=ollama_host,
|
||||
)
|
||||
if not provider_instance:
|
||||
raise ValueError(f'不支持的提供商类型: {provider_type}')
|
||||
self._provider_cache[cache_key] = provider_instance
|
||||
|
||||
return self._provider_cache[cache_key], model_name
|
||||
|
||||
def chat_with_provider(
|
||||
self,
|
||||
provider: BaseLLMProvider,
|
||||
model_name: str,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
tools: List[Dict] = None,
|
||||
tool_choice: str = 'auto',
|
||||
**kwargs
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
使用指定提供商进行同步对话
|
||||
|
||||
Args:
|
||||
provider: 提供商实例
|
||||
model_name: 模型名称
|
||||
messages: 消息列表 [{"role": "user", "content": "..."}]
|
||||
temperature: 温度参数
|
||||
max_tokens: 最大 Token
|
||||
tools: 工具定义列表 [{"name": "...", "description": "...", "parameters": {...}}]
|
||||
tool_choice: 工具选择策略 (auto, none, required, 或具体工具名)
|
||||
**kwargs: 其他参数
|
||||
|
||||
Returns:
|
||||
LLMResponse
|
||||
"""
|
||||
llm_messages = self._convert_messages(messages)
|
||||
|
||||
# 转换工具定义
|
||||
tool_definitions = None
|
||||
if tools:
|
||||
tool_definitions = [
|
||||
ToolDefinition(
|
||||
name=t['name'],
|
||||
description=t.get('description', ''),
|
||||
parameters=t.get('parameters', {}),
|
||||
)
|
||||
for t in tools
|
||||
]
|
||||
|
||||
config = LLMConfig(
|
||||
model=model_name,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tool_definitions,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
return provider.chat(llm_messages, config)
|
||||
|
||||
def _convert_messages(self, messages: List[Dict]) -> List[LLMMessage]:
|
||||
"""转换消息格式，支持 tool 消息"""
|
||||
llm_messages = []
|
||||
for m in messages:
|
||||
msg = LLMMessage(
|
||||
role=m['role'],
|
||||
content=m.get('content', ''),
|
||||
name=m.get('name'),
|
||||
tool_calls=m.get('tool_calls'),
|
||||
tool_call_id=m.get('tool_call_id'),
|
||||
)
|
||||
llm_messages.append(msg)
|
||||
return llm_messages
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
model_id: str,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
tools: List[Dict] = None,
|
||||
tool_choice: str = 'auto',
|
||||
**kwargs
|
||||
) -> LLMResponse:
|
||||
"""
|
||||
异步对话
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
messages: 消息列表
|
||||
temperature: 温度参数
|
||||
max_tokens: 最大 Token
|
||||
tools: 工具定义列表
|
||||
tool_choice: 工具选择策略
|
||||
**kwargs: 其他参数
|
||||
|
||||
Returns:
|
||||
LLMResponse
|
||||
"""
|
||||
provider, model_name = await self._get_provider_async(model_id)
|
||||
|
||||
llm_messages = self._convert_messages(messages)
|
||||
|
||||
# 转换工具定义
|
||||
tool_definitions = None
|
||||
if tools:
|
||||
tool_definitions = [
|
||||
ToolDefinition(
|
||||
name=t['name'],
|
||||
description=t.get('description', ''),
|
||||
parameters=t.get('parameters', {}),
|
||||
)
|
||||
for t in tools
|
||||
]
|
||||
|
||||
config = LLMConfig(
|
||||
model=model_name,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tool_definitions,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
return await provider.chat_async(llm_messages, config)
|
||||
|
||||
async def chat_stream(
|
||||
self,
|
||||
model_id: str,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
tools: List[Dict] = None,
|
||||
tool_choice: str = 'auto',
|
||||
**kwargs
|
||||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||||
"""
|
||||
异步流式对话
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
messages: 消息列表
|
||||
temperature: 温度参数
|
||||
max_tokens: 最大 Token
|
||||
tools: 工具定义列表
|
||||
tool_choice: 工具选择策略
|
||||
**kwargs: 其他参数
|
||||
|
||||
Yields:
|
||||
LLMStreamChunk
|
||||
"""
|
||||
provider, model_name = await self._get_provider_async(model_id)
|
||||
|
||||
llm_messages = self._convert_messages(messages)
|
||||
|
||||
# 转换工具定义
|
||||
tool_definitions = None
|
||||
if tools:
|
||||
tool_definitions = [
|
||||
ToolDefinition(
|
||||
name=t['name'],
|
||||
description=t.get('description', ''),
|
||||
parameters=t.get('parameters', {}),
|
||||
)
|
||||
for t in tools
|
||||
]
|
||||
|
||||
config = LLMConfig(
|
||||
model=model_name,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tool_definitions,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
async for chunk in provider.chat_stream(llm_messages, config):
|
||||
yield chunk
|
||||
|
||||
def chat_stream_sync(
|
||||
self,
|
||||
model_id: str,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
tools: List[Dict] = None,
|
||||
tool_choice: str = 'auto',
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
同步流式对话（生成器）
|
||||
|
||||
注意：此方法会在内部运行异步代码来获取 provider，
|
||||
需要确保 LLMService 初始化时传入了 db_session
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
messages: 消息列表
|
||||
temperature: 温度参数
|
||||
max_tokens: 最大 Token
|
||||
tools: 工具定义列表
|
||||
tool_choice: 工具选择策略
|
||||
**kwargs: 其他参数
|
||||
|
||||
Yields:
|
||||
LLMStreamChunk
|
||||
"""
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
|
||||
# 在同步上下文中运行异步代码获取 provider
|
||||
async def get_provider():
|
||||
return await self._get_provider_async(model_id)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
if loop.is_running():
|
||||
# 如果已有事件循环在运行，使用线程池
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future = executor.submit(asyncio.run, get_provider())
|
||||
provider, model_name = future.result()
|
||||
else:
|
||||
provider, model_name = loop.run_until_complete(get_provider())
|
||||
|
||||
# 使用获取到的 provider 进行流式调用
|
||||
yield from self.chat_stream_sync_with_provider(
|
||||
provider=provider,
|
||||
model_name=model_name,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def chat_stream_sync_with_provider(
|
||||
self,
|
||||
provider: BaseLLMProvider,
|
||||
model_name: str,
|
||||
messages: List[Dict[str, str]],
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 2048,
|
||||
tools: List[Dict] = None,
|
||||
tool_choice: str = 'auto',
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
使用指定提供商进行同步流式对话（生成器）
|
||||
|
||||
Args:
|
||||
provider: 提供商实例
|
||||
model_name: 模型名称
|
||||
messages: 消息列表
|
||||
temperature: 温度参数
|
||||
max_tokens: 最大 Token
|
||||
tools: 工具定义列表
|
||||
tool_choice: 工具选择策略
|
||||
**kwargs: 其他参数
|
||||
|
||||
Yields:
|
||||
LLMStreamChunk
|
||||
"""
|
||||
llm_messages = self._convert_messages(messages)
|
||||
|
||||
# 转换工具定义
|
||||
tool_definitions = None
|
||||
if tools:
|
||||
tool_definitions = [
|
||||
ToolDefinition(
|
||||
name=t['name'],
|
||||
description=t.get('description', ''),
|
||||
parameters=t.get('parameters', {}),
|
||||
)
|
||||
for t in tools
|
||||
]
|
||||
|
||||
config = LLMConfig(
|
||||
model=model_name,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
tools=tool_definitions,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
for chunk in provider.chat_stream_sync(llm_messages, config):
|
||||
yield chunk
|
||||
|
||||
@staticmethod
|
||||
def get_available_providers() -> List[Dict]:
|
||||
"""获取可用的提供商列表"""
|
||||
return ProviderRegistry.get_all_types()
|
||||
|
||||
@staticmethod
|
||||
def get_default_models(provider_type: str) -> List[Dict]:
|
||||
"""获取提供商的默认模型列表"""
|
||||
return ProviderRegistry.get_default_models(provider_type)
|
||||
Reference in New Issue
Block a user