qzl
70 lines
1.7 KiB
Python
70 lines
1.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_uses_output_schema_title() -> None:
|
|
instruction = build_json_finalize_instruction(
|
|
schema_json="{}",
|
|
attempt=1,
|
|
)
|
|
assert "[输出结构Output Schema]" in instruction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_json_response_returns_alias_keys() -> None:
|
|
model = _Model(payload={"ganzhi": {"yearGanZhi": "丙午"}})
|
|
_, payload = await finalize_json_response(
|
|
model=model,
|
|
formatter=_Formatter(),
|
|
base_messages=[],
|
|
output_model=_Output,
|
|
retries=0,
|
|
)
|
|
|
|
assert payload == {"ganzhi": {"yearGanZhi": "丙午"}}
|