Files
PlantAI-dev/tests/unit/test_deepseek_json_protocol.py

89 lines
2.8 KiB
Python

import asyncio
from types import SimpleNamespace
from gfa_core.mas.agents.conversation_tools import ConversationToolRegistry
from gfa_core.mas.agents.router_runtime import RouterRuntime
from gfa_core.mas.llm import LLMMessage, LLMResponse
from gfa_core.mas.llm.providers.deepseek import DeepSeekProvider
def test_deepseek_provider_forwards_requested_json_mode():
captured = {}
class _Completions:
def create(self, **kwargs):
captured.update(kwargs)
return SimpleNamespace(
choices=[
SimpleNamespace(
message=SimpleNamespace(content='{"ok":true}', tool_calls=[]),
finish_reason="stop",
)
],
usage=SimpleNamespace(
prompt_tokens=1,
completion_tokens=1,
total_tokens=2,
),
)
provider = DeepSeekProvider(api_key="test-key")
provider.client = SimpleNamespace(
chat=SimpleNamespace(completions=_Completions())
)
provider.call(
[LLMMessage(role="user", content="Return JSON")],
response_format={"type": "json_object"},
)
assert captured["response_format"] == {"type": "json_object"}
def test_router_route_but_not_freeform_reply_requests_native_json_mode():
class _RecordingLLM:
def __init__(self):
self.calls = []
self.responses = [
LLMResponse(
content=(
'{"goal":"general","needs":[],"subject":"编写脚本",'
'"referent":null,"clarification_question":null,'
'"interaction_kind":"supplement"}'
),
model="test-model",
provider="test-provider",
usage={},
),
LLMResponse(
content=(
'{"status":"ok","action":"reply",'
'"message":"已生成脚本。"}'
),
model="test-model",
provider="test-provider",
usage={},
),
]
async def acall(self, _messages, **kwargs):
self.calls.append(kwargs)
return self.responses.pop(0)
llm = _RecordingLLM()
outcome = asyncio.run(
RouterRuntime(llm, ConversationToolRegistry(), max_iterations=5).run(
query="帮我写一个 FASTA 脚本",
history=[],
session_id="deepseek-json-protocol",
job_id="",
goal_ledger=None,
)
)
assert outcome.action == "reply"
assert [call.get("response_format") for call in llm.calls] == [
{"type": "json_object"},
None,
]