560 lines
21 KiB
Python
560 lines
21 KiB
Python
"""
|
||
阿里通义千问提供商适配器
|
||
"""
|
||
import logging
|
||
from typing import AsyncGenerator, Dict, List, Any
|
||
|
||
from .base import BaseLLMProvider, LLMConfig, LLMMessage, LLMResponse, LLMStreamChunk, ToolCall
|
||
from .registry import ProviderRegistry
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@ProviderRegistry.register
|
||
class QwenProvider(BaseLLMProvider):
|
||
"""
|
||
阿里通义千问提供商适配器
|
||
|
||
使用 OpenAI 兼容接口
|
||
"""
|
||
|
||
provider_type = 'qwen'
|
||
provider_name = '通义千问'
|
||
supported_model_types = ['chat', 'embedding']
|
||
|
||
DEFAULT_API_BASE = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
|
||
|
||
def __init__(self, api_key: str = '', api_base: str = '', **kwargs):
|
||
super().__init__(api_key, api_base, **kwargs)
|
||
self.api_base = api_base or self.DEFAULT_API_BASE
|
||
self._client = None
|
||
self._async_client = None
|
||
|
||
def _get_client(self):
|
||
"""获取同步客户端"""
|
||
if self._client is None:
|
||
try:
|
||
from openai import OpenAI
|
||
self._client = OpenAI(
|
||
api_key=self.api_key,
|
||
base_url=self.api_base,
|
||
)
|
||
except ImportError:
|
||
raise ImportError('请安装 openai 库: pip install openai')
|
||
return self._client
|
||
|
||
def _get_async_client(self):
|
||
"""获取异步客户端"""
|
||
if self._async_client is None:
|
||
try:
|
||
from openai import AsyncOpenAI
|
||
self._async_client = AsyncOpenAI(
|
||
api_key=self.api_key,
|
||
base_url=self.api_base,
|
||
)
|
||
except ImportError:
|
||
raise ImportError('请安装 openai 库: pip install openai')
|
||
return self._async_client
|
||
|
||
def chat(
|
||
self,
|
||
messages: List[LLMMessage],
|
||
config: LLMConfig,
|
||
) -> LLMResponse:
|
||
"""同步对话"""
|
||
client = self._get_client()
|
||
|
||
# 构建请求参数
|
||
kwargs = {
|
||
'model': config.model,
|
||
'messages': [m.to_dict() for m in messages],
|
||
'temperature': config.temperature,
|
||
'top_p': config.top_p,
|
||
'max_tokens': config.max_tokens,
|
||
'stop': config.stop,
|
||
**config.extra_params,
|
||
}
|
||
|
||
# 添加 Function Calling 参数
|
||
if config.tools:
|
||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||
if config.tool_choice == 'required':
|
||
kwargs['tool_choice'] = 'required'
|
||
elif config.tool_choice == 'none':
|
||
kwargs['tool_choice'] = 'none'
|
||
elif config.tool_choice != 'auto':
|
||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||
else:
|
||
kwargs['tool_choice'] = 'auto'
|
||
|
||
response = client.chat.completions.create(**kwargs)
|
||
|
||
choice = response.choices[0]
|
||
usage = response.usage
|
||
|
||
# 解析工具调用
|
||
tool_calls = None
|
||
if choice.message.tool_calls:
|
||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||
|
||
return LLMResponse(
|
||
content=choice.message.content or '',
|
||
model=response.model,
|
||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||
completion_tokens=usage.completion_tokens if usage else 0,
|
||
total_tokens=usage.total_tokens if usage else 0,
|
||
finish_reason=choice.finish_reason or '',
|
||
raw_response=response.model_dump(),
|
||
tool_calls=tool_calls,
|
||
)
|
||
|
||
def _parse_tool_calls(self, tool_calls) -> List[ToolCall]:
|
||
"""解析工具调用"""
|
||
import json
|
||
result = []
|
||
for tc in tool_calls:
|
||
try:
|
||
arguments = json.loads(tc.function.arguments) if tc.function.arguments else {}
|
||
except json.JSONDecodeError:
|
||
arguments = {'raw': tc.function.arguments}
|
||
|
||
result.append(ToolCall(
|
||
id=tc.id,
|
||
name=tc.function.name,
|
||
arguments=arguments,
|
||
))
|
||
return result
|
||
|
||
async def chat_async(
|
||
self,
|
||
messages: List[LLMMessage],
|
||
config: LLMConfig,
|
||
) -> LLMResponse:
|
||
"""异步对话"""
|
||
client = self._get_async_client()
|
||
|
||
# 构建请求参数
|
||
kwargs = {
|
||
'model': config.model,
|
||
'messages': [m.to_dict() for m in messages],
|
||
'temperature': config.temperature,
|
||
'top_p': config.top_p,
|
||
'max_tokens': config.max_tokens,
|
||
'stop': config.stop,
|
||
**config.extra_params,
|
||
}
|
||
|
||
# 添加 Function Calling 参数
|
||
if config.tools:
|
||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||
if config.tool_choice == 'required':
|
||
kwargs['tool_choice'] = 'required'
|
||
elif config.tool_choice == 'none':
|
||
kwargs['tool_choice'] = 'none'
|
||
elif config.tool_choice != 'auto':
|
||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||
else:
|
||
kwargs['tool_choice'] = 'auto'
|
||
|
||
response = await client.chat.completions.create(**kwargs)
|
||
|
||
choice = response.choices[0]
|
||
usage = response.usage
|
||
|
||
# 解析工具调用
|
||
tool_calls = None
|
||
if choice.message.tool_calls:
|
||
tool_calls = self._parse_tool_calls(choice.message.tool_calls)
|
||
|
||
return LLMResponse(
|
||
content=choice.message.content or '',
|
||
model=response.model,
|
||
prompt_tokens=usage.prompt_tokens if usage else 0,
|
||
completion_tokens=usage.completion_tokens if usage else 0,
|
||
total_tokens=usage.total_tokens if usage else 0,
|
||
finish_reason=choice.finish_reason or '',
|
||
raw_response=response.model_dump(),
|
||
tool_calls=tool_calls,
|
||
)
|
||
|
||
async def chat_stream(
|
||
self,
|
||
messages: List[LLMMessage],
|
||
config: LLMConfig,
|
||
) -> AsyncGenerator[LLMStreamChunk, None]:
|
||
"""流式对话"""
|
||
client = self._get_async_client()
|
||
|
||
# 构建请求参数
|
||
kwargs = {
|
||
'model': config.model,
|
||
'messages': [m.to_dict() for m in messages],
|
||
'temperature': config.temperature,
|
||
'top_p': config.top_p,
|
||
'max_tokens': config.max_tokens,
|
||
'stop': config.stop,
|
||
'stream': True,
|
||
'stream_options': {"include_usage": True},
|
||
**config.extra_params,
|
||
}
|
||
|
||
# 添加 Function Calling 参数
|
||
if config.tools:
|
||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||
if config.tool_choice == 'required':
|
||
kwargs['tool_choice'] = 'required'
|
||
elif config.tool_choice == 'none':
|
||
kwargs['tool_choice'] = 'none'
|
||
elif config.tool_choice != 'auto':
|
||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||
else:
|
||
kwargs['tool_choice'] = 'auto'
|
||
|
||
stream = await client.chat.completions.create(**kwargs)
|
||
|
||
# 用于累积工具调用
|
||
tool_call_accumulator = {}
|
||
|
||
async for chunk in stream:
|
||
if chunk.choices:
|
||
choice = chunk.choices[0]
|
||
delta = choice.delta
|
||
|
||
# 处理工具调用增量
|
||
tool_call_delta = None
|
||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||
for tc_delta in delta.tool_calls:
|
||
tc_index = tc_delta.index
|
||
if tc_index not in tool_call_accumulator:
|
||
tool_call_accumulator[tc_index] = {
|
||
'id': tc_delta.id or '',
|
||
'name': '',
|
||
'arguments': '',
|
||
}
|
||
if tc_delta.id:
|
||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||
if tc_delta.function:
|
||
if tc_delta.function.name:
|
||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||
if tc_delta.function.arguments:
|
||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||
tool_call_delta = tool_call_accumulator[tc_index]
|
||
|
||
# 检查是否完成
|
||
tool_calls = None
|
||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||
|
||
yield LLMStreamChunk(
|
||
content=delta.content or '',
|
||
is_finished=choice.finish_reason is not None,
|
||
finish_reason=choice.finish_reason or '',
|
||
tool_calls=tool_calls,
|
||
tool_call_delta=tool_call_delta,
|
||
)
|
||
|
||
if chunk.usage:
|
||
yield LLMStreamChunk(
|
||
content='',
|
||
is_finished=True,
|
||
prompt_tokens=chunk.usage.prompt_tokens,
|
||
completion_tokens=chunk.usage.completion_tokens,
|
||
total_tokens=chunk.usage.total_tokens,
|
||
)
|
||
|
||
def _parse_accumulated_tool_calls(self, accumulator: Dict) -> List[ToolCall]:
|
||
"""解析累积的工具调用"""
|
||
import json
|
||
result = []
|
||
for idx in sorted(accumulator.keys()):
|
||
tc = accumulator[idx]
|
||
try:
|
||
arguments = json.loads(tc['arguments']) if tc['arguments'] else {}
|
||
except json.JSONDecodeError:
|
||
arguments = {'raw': tc['arguments']}
|
||
|
||
result.append(ToolCall(
|
||
id=tc['id'],
|
||
name=tc['name'],
|
||
arguments=arguments,
|
||
))
|
||
return result
|
||
|
||
def chat_stream_sync(
|
||
self,
|
||
messages: List[LLMMessage],
|
||
config: LLMConfig,
|
||
):
|
||
"""同步流式对话"""
|
||
client = self._get_client()
|
||
|
||
# 构建请求参数
|
||
kwargs = {
|
||
'model': config.model,
|
||
'messages': [m.to_dict() for m in messages],
|
||
'temperature': config.temperature,
|
||
'top_p': config.top_p,
|
||
'max_tokens': config.max_tokens,
|
||
'stop': config.stop,
|
||
'stream': True,
|
||
'stream_options': {"include_usage": True},
|
||
**config.extra_params,
|
||
}
|
||
|
||
# 添加 Function Calling 参数
|
||
if config.tools:
|
||
kwargs['tools'] = [t.to_openai_format() for t in config.tools]
|
||
if config.tool_choice == 'required':
|
||
kwargs['tool_choice'] = 'required'
|
||
elif config.tool_choice == 'none':
|
||
kwargs['tool_choice'] = 'none'
|
||
elif config.tool_choice != 'auto':
|
||
kwargs['tool_choice'] = {'type': 'function', 'function': {'name': config.tool_choice}}
|
||
else:
|
||
kwargs['tool_choice'] = 'auto'
|
||
|
||
stream = client.chat.completions.create(**kwargs)
|
||
|
||
# 用于累积工具调用
|
||
tool_call_accumulator = {}
|
||
|
||
for chunk in stream:
|
||
if chunk.choices:
|
||
choice = chunk.choices[0]
|
||
delta = choice.delta
|
||
|
||
# 处理工具调用增量
|
||
tool_call_delta = None
|
||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||
for tc_delta in delta.tool_calls:
|
||
tc_index = tc_delta.index
|
||
if tc_index not in tool_call_accumulator:
|
||
tool_call_accumulator[tc_index] = {
|
||
'id': tc_delta.id or '',
|
||
'name': '',
|
||
'arguments': '',
|
||
}
|
||
if tc_delta.id:
|
||
tool_call_accumulator[tc_index]['id'] = tc_delta.id
|
||
if tc_delta.function:
|
||
if tc_delta.function.name:
|
||
tool_call_accumulator[tc_index]['name'] = tc_delta.function.name
|
||
if tc_delta.function.arguments:
|
||
tool_call_accumulator[tc_index]['arguments'] += tc_delta.function.arguments
|
||
tool_call_delta = tool_call_accumulator[tc_index]
|
||
|
||
# 检查是否完成
|
||
tool_calls = None
|
||
if choice.finish_reason == 'tool_calls' or (choice.finish_reason and tool_call_accumulator):
|
||
tool_calls = self._parse_accumulated_tool_calls(tool_call_accumulator)
|
||
|
||
yield LLMStreamChunk(
|
||
content=delta.content or '',
|
||
is_finished=choice.finish_reason is not None,
|
||
finish_reason=choice.finish_reason or '',
|
||
tool_calls=tool_calls,
|
||
tool_call_delta=tool_call_delta,
|
||
)
|
||
|
||
if chunk.usage:
|
||
yield LLMStreamChunk(
|
||
content='',
|
||
is_finished=True,
|
||
prompt_tokens=chunk.usage.prompt_tokens,
|
||
completion_tokens=chunk.usage.completion_tokens,
|
||
total_tokens=chunk.usage.total_tokens,
|
||
)
|
||
|
||
_EMBEDDING_KEYWORDS = ('embed', 'bge')
|
||
_RERANK_KEYWORDS = ('rerank',)
|
||
_SKIP_KEYWORDS = (
|
||
'paraformer', 'sambert', 'wordart', 'wanx', 'cosyvoice',
|
||
'qwen-vl-ocr', 'tts', 'asr', 'realtime',
|
||
'tongyi-xiaomi', 'flux', 'stable-diffusion',
|
||
'qwen-image-', 'z-image-', 'gui-plus',
|
||
'qwen-mt-', 'livetranslate', 'deep-search',
|
||
'flash-character',
|
||
)
|
||
|
||
async def fetch_models_from_api(self) -> List[Dict[str, Any]]:
|
||
"""通过 DashScope 兼容的 /v1/models 端点在线拉取模型列表"""
|
||
base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
|
||
try:
|
||
return await self.fetch_models_from_api_strict()
|
||
except Exception as e:
|
||
logger.warning(f'在线拉取模型列表失败 ({base_url}): {e}')
|
||
return []
|
||
|
||
async def fetch_models_from_api_strict(self) -> List[Dict[str, Any]]:
|
||
"""通过 DashScope 兼容端点严格拉取模型列表，失败时保留上游异常"""
|
||
import httpx
|
||
|
||
base_url = (self.api_base or self.DEFAULT_API_BASE).rstrip('/')
|
||
headers = {}
|
||
if self.api_key:
|
||
headers['Authorization'] = f'Bearer {self.api_key}'
|
||
|
||
async with httpx.AsyncClient(timeout=15) as client:
|
||
resp = await client.get(f'{base_url}/models', headers=headers)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
model_list = data.get('data', [])
|
||
if not model_list:
|
||
return []
|
||
|
||
results = []
|
||
for item in model_list:
|
||
model_id = item.get('id', '')
|
||
if not model_id:
|
||
continue
|
||
|
||
model_lower = model_id.lower()
|
||
|
||
if any(kw in model_lower for kw in self._SKIP_KEYWORDS):
|
||
continue
|
||
|
||
if any(kw in model_lower for kw in self._RERANK_KEYWORDS):
|
||
model_type = 'rerank'
|
||
elif any(kw in model_lower for kw in self._EMBEDDING_KEYWORDS):
|
||
model_type = 'embedding'
|
||
else:
|
||
model_type = 'chat'
|
||
|
||
results.append({
|
||
'model_name': model_id,
|
||
'display_name': model_id,
|
||
'model_type': model_type,
|
||
'max_tokens': 4096,
|
||
'context_window': 4096,
|
||
'supports_vision': False,
|
||
'supports_function_call': model_type == 'chat',
|
||
'input_price': 0,
|
||
'output_price': 0,
|
||
})
|
||
|
||
# DashScope /v1/models 可能只返回 chat 模型，
|
||
# 将默认列表中的 embedding/rerank 模型合并进去
|
||
fetched_names = {m['model_name'] for m in results}
|
||
for dm in self.__class__.get_default_models():
|
||
if dm['model_type'] in ('embedding', 'rerank') and dm['model_name'] not in fetched_names:
|
||
results.append(dm)
|
||
|
||
type_order = {'chat': 0, 'embedding': 1, 'rerank': 2}
|
||
results.sort(key=lambda m: (type_order.get(m['model_type'], 9), m['model_name']))
|
||
logger.info(f'从 {base_url} 拉取到 {len(results)} 个模型（含补充的 embedding/rerank）')
|
||
return results
|
||
|
||
@classmethod
|
||
def get_default_models(cls) -> List[Dict[str, Any]]:
|
||
"""获取默认模型列表"""
|
||
return [
|
||
# ---- Chat 模型 ----
|
||
{
|
||
'model_name': 'qwen-max',
|
||
'display_name': '通义千问 Max',
|
||
'model_type': 'chat',
|
||
'max_tokens': 8192,
|
||
'context_window': 32768,
|
||
'supports_vision': False,
|
||
'supports_function_call': True,
|
||
'input_price': 0.02,
|
||
'output_price': 0.06,
|
||
},
|
||
{
|
||
'model_name': 'qwen-plus',
|
||
'display_name': '通义千问 Plus',
|
||
'model_type': 'chat',
|
||
'max_tokens': 8192,
|
||
'context_window': 131072,
|
||
'supports_vision': False,
|
||
'supports_function_call': True,
|
||
'input_price': 0.0008,
|
||
'output_price': 0.002,
|
||
},
|
||
{
|
||
'model_name': 'qwen-turbo',
|
||
'display_name': '通义千问 Turbo',
|
||
'model_type': 'chat',
|
||
'max_tokens': 8192,
|
||
'context_window': 131072,
|
||
'supports_vision': False,
|
||
'supports_function_call': True,
|
||
'input_price': 0.0003,
|
||
'output_price': 0.0006,
|
||
},
|
||
{
|
||
'model_name': 'qwen3-235b-a22b',
|
||
'display_name': 'Qwen3 235B-A22B',
|
||
'model_type': 'chat',
|
||
'max_tokens': 8192,
|
||
'context_window': 131072,
|
||
'supports_vision': False,
|
||
'supports_function_call': True,
|
||
'input_price': 0.004,
|
||
'output_price': 0.012,
|
||
},
|
||
{
|
||
'model_name': 'qwen3-32b',
|
||
'display_name': 'Qwen3 32B',
|
||
'model_type': 'chat',
|
||
'max_tokens': 8192,
|
||
'context_window': 131072,
|
||
'supports_vision': False,
|
||
'supports_function_call': True,
|
||
'input_price': 0.002,
|
||
'output_price': 0.006,
|
||
},
|
||
{
|
||
'model_name': 'qwen-vl-max',
|
||
'display_name': '通义千问 VL Max',
|
||
'model_type': 'chat',
|
||
'max_tokens': 2048,
|
||
'context_window': 32768,
|
||
'supports_vision': True,
|
||
'supports_function_call': False,
|
||
'input_price': 0.02,
|
||
'output_price': 0.06,
|
||
},
|
||
{
|
||
'model_name': 'qwen-vl-plus',
|
||
'display_name': '通义千问 VL Plus',
|
||
'model_type': 'chat',
|
||
'max_tokens': 2048,
|
||
'context_window': 32768,
|
||
'supports_vision': True,
|
||
'supports_function_call': False,
|
||
'input_price': 0.008,
|
||
'output_price': 0.02,
|
||
},
|
||
# ---- Embedding 模型 ----
|
||
{
|
||
'model_name': 'text-embedding-v3',
|
||
'display_name': '通义文本向量 V3',
|
||
'model_type': 'embedding',
|
||
'max_tokens': 8192,
|
||
'context_window': 8192,
|
||
'input_price': 0.0007,
|
||
'output_price': 0,
|
||
},
|
||
{
|
||
'model_name': 'text-embedding-v2',
|
||
'display_name': '通义文本向量 V2',
|
||
'model_type': 'embedding',
|
||
'max_tokens': 2048,
|
||
'context_window': 2048,
|
||
'input_price': 0.0007,
|
||
'output_price': 0,
|
||
},
|
||
# ---- Rerank 模型 ----
|
||
{
|
||
'model_name': 'gte-rerank',
|
||
'display_name': 'GTE Rerank',
|
||
'model_type': 'rerank',
|
||
'max_tokens': 4096,
|
||
'context_window': 4096,
|
||
'input_price': 0.001,
|
||
'output_price': 0,
|
||
},
|
||
]
|