武阳
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from pathlib import Path
|
|
from tempfile import TemporaryDirectory
|
|
from unittest.mock import patch
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from backend.api.server import AgentState, _normalize_model_list, create_app
|
|
|
|
|
|
class ModelListTests(unittest.TestCase):
|
|
def test_normalize_model_list_uses_provider_prefixed_llm_ids(self) -> None:
|
|
payload = {
|
|
'data': [
|
|
{
|
|
'id': 'ASR_NonStreaming',
|
|
'object': 'model',
|
|
'owned_by': 'xiaomi',
|
|
'model_type': 'speech2text',
|
|
},
|
|
{
|
|
'id': 'DeepSeek-R1-0528',
|
|
'object': 'model',
|
|
'owned_by': 'xiaomi',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'gpt-5',
|
|
'object': 'model',
|
|
'owned_by': 'azure_openai',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'gpt-4o-audio-preview',
|
|
'object': 'model',
|
|
'owned_by': 'azure_openai',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'embedding-v1',
|
|
'object': 'model',
|
|
'owned_by': 'example',
|
|
'model_type': 'text-embedding',
|
|
},
|
|
{
|
|
'id': 'ernie-4.0-turbo-128k',
|
|
'object': 'model',
|
|
'owned_by': 'baidu_qianfan',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'Pro/deepseek-ai/DeepSeek-V3',
|
|
'object': 'model',
|
|
'owned_by': 'siliconflow',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'pa/claude-opus-4-7',
|
|
'object': 'model',
|
|
'owned_by': 'ppio',
|
|
'model_type': 'llm',
|
|
},
|
|
{
|
|
'id': 'gemini-2.0-flash-20250609',
|
|
'object': 'model',
|
|
'owned_by': 'ppio',
|
|
'model_type': 'llm',
|
|
},
|
|
]
|
|
}
|
|
|
|
models = _normalize_model_list(payload)
|
|
|
|
self.assertEqual(
|
|
[model['id'] for model in models],
|
|
[
|
|
'azure_openai/gpt-5',
|
|
'ppio/gemini-2.0-flash-20250609',
|
|
'siliconflow/Pro/deepseek-ai/DeepSeek-V3',
|
|
'xiaomi/DeepSeek-R1-0528',
|
|
],
|
|
)
|
|
self.assertEqual(models[0]['provider'], 'azure_openai')
|
|
self.assertEqual(models[0]['model_type'], 'llm')
|
|
|
|
def test_models_endpoint_returns_provider_prefixed_llm_models(self) -> None:
|
|
class FakeResponse:
|
|
def __enter__(self) -> 'FakeResponse':
|
|
return self
|
|
|
|
def __exit__(self, *args: object) -> None:
|
|
return None
|
|
|
|
def read(self) -> bytes:
|
|
return (
|
|
b'{"object":"list","data":['
|
|
b'{"id":"mimo-v2-flash","object":"model","owned_by":"xiaomi","model_type":"llm"},'
|
|
b'{"id":"ASR_Streaming","object":"model","owned_by":"xiaomi","model_type":"speech2text"},'
|
|
b'{"id":"gpt-5","object":"model","owned_by":"azure_openai","model_type":"llm"}'
|
|
b']}'
|
|
)
|
|
|
|
with TemporaryDirectory() as tmp_dir:
|
|
state = AgentState(
|
|
cwd=Path(tmp_dir),
|
|
model='xiaomi/mimo-v2-flash',
|
|
base_url='http://model.example/v1',
|
|
api_key='token',
|
|
allow_shell=False,
|
|
allow_write=False,
|
|
session_directory=Path(tmp_dir) / 'sessions',
|
|
)
|
|
client = TestClient(create_app(state))
|
|
with patch('backend.api.server.request.urlopen', return_value=FakeResponse()):
|
|
response = client.get('/api/models')
|
|
|
|
self.assertEqual(response.status_code, 200)
|
|
payload = response.json()
|
|
self.assertEqual(payload['raw_count'], 3)
|
|
self.assertEqual(payload['filtered_count'], 2)
|
|
self.assertEqual(
|
|
[model['id'] for model in payload['models']],
|
|
['azure_openai/gpt-5', 'xiaomi/mimo-v2-flash'],
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|