65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from collections.abc import Iterable as ABCIterable
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from pydantic import BaseModel
|
||
|
|
|
||
|
|
from instructor.dsl import ListResponse
|
||
|
|
from instructor.dsl.iterable import IterableBase
|
||
|
|
from instructor.mode import Mode
|
||
|
|
from instructor.processing.response import process_response
|
||
|
|
from instructor.utils.core import prepare_response_model
|
||
|
|
|
||
|
|
|
||
|
|
class User(BaseModel):
|
||
|
|
name: str
|
||
|
|
|
||
|
|
|
||
|
|
def test_listresponse_preserves_raw_response_on_slice() -> None:
|
||
|
|
raw: Any = {"provider": "test"}
|
||
|
|
resp = ListResponse([User(name="a"), User(name="b")], _raw_response=raw)
|
||
|
|
|
||
|
|
assert resp.get_raw_response() is raw
|
||
|
|
assert resp[0].name == "a"
|
||
|
|
|
||
|
|
sliced = resp[1:]
|
||
|
|
assert isinstance(sliced, ListResponse)
|
||
|
|
assert sliced.get_raw_response() is raw
|
||
|
|
assert sliced[0].name == "b"
|
||
|
|
|
||
|
|
|
||
|
|
def test_process_response_wraps_iterablebase_tasks_with_raw_response() -> None:
|
||
|
|
class FakeIterableResponse(BaseModel, IterableBase):
|
||
|
|
tasks: list[User]
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_response( # type: ignore[override]
|
||
|
|
cls, _response: Any, **_kwargs: Any
|
||
|
|
) -> FakeIterableResponse:
|
||
|
|
return cls(tasks=[User(name="x"), User(name="y")])
|
||
|
|
|
||
|
|
# `process_response()` is typed with a BaseModel-bounded type variable for `response`,
|
||
|
|
# so use a BaseModel instance here to keep `ty` happy.
|
||
|
|
raw_response: Any = User(name="raw")
|
||
|
|
out = process_response(
|
||
|
|
raw_response,
|
||
|
|
response_model=FakeIterableResponse,
|
||
|
|
stream=False,
|
||
|
|
mode=Mode.TOOLS,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert isinstance(out, ListResponse)
|
||
|
|
assert [u.name for u in out] == ["x", "y"]
|
||
|
|
assert out.get_raw_response() is raw_response
|
||
|
|
|
||
|
|
|
||
|
|
def test_prepare_response_model_supports_list_and_iterable() -> None:
|
||
|
|
prepared_list = prepare_response_model(list[User])
|
||
|
|
assert prepared_list is not None
|
||
|
|
assert issubclass(prepared_list, IterableBase)
|
||
|
|
|
||
|
|
prepared_iterable = prepare_response_model(ABCIterable[User]) # type: ignore[index]
|
||
|
|
assert prepared_iterable is not None
|
||
|
|
assert issubclass(prepared_iterable, IterableBase)
|