wuyang
208 lines
7.7 KiB
Python
208 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import httpx
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agent_platform.runtime.app import create_app
|
|
from agent_platform.store import RuntimeStore
|
|
|
|
|
|
class FakeModelProvider:
|
|
def __init__(self) -> None:
|
|
self.completed = []
|
|
self.forwarded = []
|
|
self.forward_options = []
|
|
|
|
async def complete(self, **kwargs):
|
|
self.completed.append(kwargs)
|
|
return {
|
|
"choices": [{"message": {"role": "assistant", "content": "work complete"}, "finish_reason": "stop"}],
|
|
"usage": {"total_tokens": 10},
|
|
}
|
|
|
|
async def forward(self, payload, **kwargs):
|
|
self.forwarded.append(payload)
|
|
self.forward_options.append(kwargs)
|
|
request = httpx.Request("POST", "https://model.example.test/v1/chat/completions")
|
|
if payload.get("stream"):
|
|
return httpx.Response(
|
|
200,
|
|
request=request,
|
|
headers={"content-type": "text/event-stream"},
|
|
content=(
|
|
'data: {"model":"ChatGPT-5.6:Luna","choices":[{"index":0,"delta":'
|
|
'{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":'
|
|
'{"name":"read_file","arguments":"{}"}}]},"finish_reason":null}]}\n\n'
|
|
"data: [DONE]\n\n"
|
|
),
|
|
)
|
|
return httpx.Response(
|
|
200,
|
|
request=request,
|
|
headers={"content-type": "application/json"},
|
|
json={
|
|
"model": payload["model"],
|
|
"choices": [{"message": {"role": "assistant", "content": "chat reply"}, "finish_reason": "stop"}],
|
|
},
|
|
)
|
|
|
|
async def close(self):
|
|
return None
|
|
|
|
|
|
class FakeTools:
|
|
async def close(self):
|
|
return None
|
|
|
|
def specs(self, *, read_only=False, allow_delegate=True):
|
|
return []
|
|
|
|
async def execute(self, name, arguments, context):
|
|
raise AssertionError("No tool should be called")
|
|
|
|
|
|
def headers(settings, identity_jwt, chat_id="chat-1"):
|
|
return {
|
|
"Authorization": f"Bearer {settings.internal_provider_key}",
|
|
"X-OpenWebUI-User-Jwt": identity_jwt,
|
|
"X-OpenWebUI-Chat-Id": chat_id,
|
|
}
|
|
|
|
|
|
def test_model_list_requires_internal_key(settings) -> None:
|
|
app = create_app(
|
|
settings,
|
|
store=RuntimeStore(settings.database_url),
|
|
provider=FakeModelProvider(),
|
|
tools=FakeTools(),
|
|
)
|
|
with TestClient(app) as client:
|
|
assert client.get("/v1/models").status_code == 401
|
|
response = client.get(
|
|
"/v1/models",
|
|
headers={"Authorization": f"Bearer {settings.internal_provider_key}"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert len(response.json()["data"]) == 8
|
|
|
|
|
|
def test_chat_maps_provider_model_and_work_upgrade_is_one_way(settings, identity_jwt) -> None:
|
|
provider = FakeModelProvider()
|
|
app = create_app(settings, store=RuntimeStore(settings.database_url), provider=provider, tools=FakeTools())
|
|
with TestClient(app) as client:
|
|
chat = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt),
|
|
json={"model": "chat-light", "messages": [{"role": "user", "content": "hello"}]},
|
|
)
|
|
assert chat.status_code == 200
|
|
assert provider.forwarded[0]["model"] == "ChatGPT-5.6:Luna"
|
|
assert chat.json()["model"] == "chat-light"
|
|
|
|
extreme_chat = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt, "extreme-chat"),
|
|
json={"model": "chat-extreme", "messages": [{"role": "user", "content": "think"}]},
|
|
)
|
|
assert extreme_chat.status_code == 200
|
|
assert provider.forwarded[-1]["model"] == "deepseek-v4-pro"
|
|
assert provider.forward_options[-1] == {
|
|
"provider": "deepseek",
|
|
"thinking_enabled": True,
|
|
"reasoning_effort": "max",
|
|
"max_output_tokens": 16_384,
|
|
}
|
|
|
|
work = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt),
|
|
json={"model": "work-high", "messages": [{"role": "user", "content": "do it"}]},
|
|
)
|
|
assert work.status_code == 200
|
|
assert provider.completed[-1]["model"] == "ChatGPT-5.6:Sol"
|
|
|
|
extreme_work = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt, "extreme-work"),
|
|
json={"model": "work-extreme", "messages": [{"role": "user", "content": "do it carefully"}]},
|
|
)
|
|
assert extreme_work.status_code == 200
|
|
assert provider.completed[-1]["model"] == "deepseek-v4-pro"
|
|
assert provider.completed[-1]["provider"] == "deepseek"
|
|
assert provider.completed[-1]["reasoning_effort"] == "max"
|
|
|
|
downgrade = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt),
|
|
json={"model": "chat-medium", "messages": [{"role": "user", "content": "back"}]},
|
|
)
|
|
assert downgrade.status_code == 409
|
|
|
|
|
|
def test_chat_stream_hides_provider_model_and_preserves_native_tool_calls(settings, identity_jwt) -> None:
|
|
provider = FakeModelProvider()
|
|
app = create_app(settings, store=RuntimeStore(settings.database_url), provider=provider, tools=FakeTools())
|
|
with TestClient(app) as client:
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt, "chat-stream"),
|
|
json={
|
|
"model": "chat-light",
|
|
"stream": True,
|
|
"messages": [{"role": "user", "content": "inspect"}],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert '"model": "chat-light"' in response.text
|
|
assert "ChatGPT-5.6:Luna" not in response.text
|
|
assert '"tool_calls"' in response.text
|
|
assert '"name": "read_file"' in response.text
|
|
assert "data: [DONE]" in response.text
|
|
|
|
|
|
def test_openwebui_background_task_bypasses_work_loop(settings, identity_jwt) -> None:
|
|
provider = FakeModelProvider()
|
|
app = create_app(settings, store=RuntimeStore(settings.database_url), provider=provider, tools=FakeTools())
|
|
with TestClient(app) as client:
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt, "background-task"),
|
|
json={
|
|
"model": "work-medium",
|
|
"messages": [{"role": "user", "content": "Generate tags"}],
|
|
"metadata": {"task": "tags_generation"},
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["model"] == "work-medium"
|
|
assert provider.completed == []
|
|
assert provider.forwarded[-1]["model"] == "ChatGPT-5.6:Terra"
|
|
|
|
|
|
def test_work_stream_never_emits_delta_tool_calls(settings, identity_jwt) -> None:
|
|
app = create_app(
|
|
settings,
|
|
store=RuntimeStore(settings.database_url),
|
|
provider=FakeModelProvider(),
|
|
tools=FakeTools(),
|
|
)
|
|
with TestClient(app) as client:
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
headers=headers(settings, identity_jwt, "stream-chat"),
|
|
json={
|
|
"model": "work-medium",
|
|
"stream": True,
|
|
"messages": [{"role": "user", "content": "do it"}],
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert "work complete" in response.text
|
|
for line in response.text.splitlines():
|
|
if line.startswith("data: {"):
|
|
chunk = json.loads(line.removeprefix("data: "))
|
|
assert "tool_calls" not in chunk["choices"][0]["delta"]
|
|
assert '"finish_reason": "stop"' in response.text
|