qzl
99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from pydantic import BaseModel, ConfigDict, Field
|
|
|
|
from core.agentscope.utils.json_finalize import (
|
|
build_json_finalize_instruction,
|
|
finalize_json_response,
|
|
)
|
|
|
|
|
|
class _Inner(BaseModel):
|
|
model_config = ConfigDict(extra="forbid")
|
|
|
|
year_gan_zhi: str = Field(alias="yearGanZhi")
|
|
|
|
|
|
class _Output(BaseModel):
|
|
model_config = ConfigDict(extra="forbid", populate_by_name=True)
|
|
|
|
ganzhi: _Inner
|
|
|
|
|
|
class _Formatter:
|
|
async def format(self, *args: Any, **kwargs: Any) -> Any:
|
|
del args, kwargs
|
|
return [{"role": "user", "content": "prompt"}]
|
|
|
|
|
|
class _Response:
|
|
def __init__(self, payload: dict[str, Any]) -> None:
|
|
self.content = [
|
|
{"type": "text", "text": json.dumps(payload, ensure_ascii=False)}
|
|
]
|
|
|
|
|
|
class _Model:
|
|
def __init__(self, payload: dict[str, Any]) -> None:
|
|
self._payload = payload
|
|
self.stream = False
|
|
|
|
async def __call__(self, *args: Any, **kwargs: Any) -> _Response:
|
|
del args, kwargs
|
|
return _Response(self._payload)
|
|
|
|
|
|
def test_build_instruction_has_schema() -> None:
|
|
instruction = build_json_finalize_instruction(
|
|
schema_json="{}",
|
|
attempt=1,
|
|
)
|
|
assert "[Output Schema]" in instruction
|
|
|
|
|
|
def test_build_instruction_has_language_constraint_en() -> None:
|
|
instruction = build_json_finalize_instruction(
|
|
schema_json="{}",
|
|
attempt=1,
|
|
language="en-US",
|
|
)
|
|
assert "ENGLISH OUTPUT REQUIRED" in instruction
|
|
assert "═══════════════════════════════════════════════════════════════════════════════" in instruction
|
|
|
|
|
|
def test_build_instruction_has_language_constraint_zh() -> None:
|
|
instruction = build_json_finalize_instruction(
|
|
schema_json="{}",
|
|
attempt=1,
|
|
language="zh-CN",
|
|
)
|
|
assert "返回 JSON。使用简体中文" in instruction
|
|
|
|
|
|
def test_build_instruction_no_language_constraint_when_none() -> None:
|
|
instruction = build_json_finalize_instruction(
|
|
schema_json="{}",
|
|
attempt=1,
|
|
language=None,
|
|
)
|
|
assert "ENGLISH OUTPUT REQUIRED" not in instruction
|
|
assert "返回 JSON" not in instruction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_json_response_returns_alias_keys() -> None:
|
|
model = _Model(payload={"ganzhi": {"yearGanZhi": "丙午"}})
|
|
_, result = await finalize_json_response(
|
|
model=model,
|
|
formatter=_Formatter(),
|
|
base_messages=[],
|
|
output_model=_Output,
|
|
retries=0,
|
|
)
|
|
|
|
assert result.model_dump(mode="json", by_alias=True) == {"ganzhi": {"yearGanZhi": "丙午"}}
|