1126 lines
47 KiB
Python
1126 lines
47 KiB
Python
|
# --------------------------------------------------------------------------------
|
||
|
|
# The following code is adapted from a comment on GitHub in the pydantic/pydantic repository by silviumarcu.
|
||
|
|
# Source: https://github.com/pydantic/pydantic/issues/6381#issuecomment-1831607091
|
||
|
|
#
|
||
|
|
# This code is used in accordance with the repository's license, and this reference
|
||
|
|
# serves as an acknowledgment of the original author's contribution to this project.
|
||
|
|
# --------------------------------------------------------------------------------
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import sys
|
||
|
|
import types
|
||
|
|
import warnings
|
||
|
|
from collections.abc import AsyncGenerator, Callable, Generator, Iterable
|
||
|
|
from copy import deepcopy
|
||
|
|
from functools import cache
|
||
|
|
from typing import ( # noqa: UP035
|
||
|
|
Any,
|
||
|
|
Generic,
|
||
|
|
List, # needed for runtime check against typing.List annotations from user code
|
||
|
|
NoReturn,
|
||
|
|
Optional,
|
||
|
|
TypeVar,
|
||
|
|
Union,
|
||
|
|
get_args,
|
||
|
|
get_origin,
|
||
|
|
)
|
||
|
|
|
||
|
|
from jiter import from_json
|
||
|
|
from pydantic import BaseModel, create_model
|
||
|
|
from pydantic.fields import FieldInfo, PydanticUndefined
|
||
|
|
|
||
|
|
from instructor.mode import Mode
|
||
|
|
from instructor.utils import extract_json_from_stream, extract_json_from_stream_async
|
||
|
|
from instructor.dsl.json_tracker import JsonCompleteness, is_json_complete
|
||
|
|
|
||
|
|
T_Model = TypeVar("T_Model", bound=BaseModel)
|
||
|
|
|
||
|
|
if sys.version_info >= (3, 10):
|
||
|
|
# types.UnionType is only available in Python 3.10 and above
|
||
|
|
UNION_ORIGINS = (Union, types.UnionType)
|
||
|
|
else:
|
||
|
|
UNION_ORIGINS = (Union,)
|
||
|
|
|
||
|
|
# Track models currently being processed to prevent infinite recursion
|
||
|
|
# with self-referential models (e.g., TreeNode with children: List["TreeNode"])
|
||
|
|
_processing_models: set[type] = set()
|
||
|
|
|
||
|
|
|
||
|
|
class MakeFieldsOptional:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class PartialLiteralMixin:
|
||
|
|
"""DEPRECATED: This mixin is no longer necessary.
|
||
|
|
|
||
|
|
With completeness-based validation, Literal and Enum types are handled
|
||
|
|
automatically during streaming:
|
||
|
|
- Incomplete JSON: no validation runs, partial values are stored as-is
|
||
|
|
- Complete JSON: full validation against original model
|
||
|
|
|
||
|
|
You can safely remove this mixin from your models.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init_subclass__(cls, **kwargs: Any) -> None:
|
||
|
|
super().__init_subclass__(**kwargs)
|
||
|
|
warnings.warn(
|
||
|
|
"PartialLiteralMixin is deprecated and no longer necessary. "
|
||
|
|
"Completeness-based validation now handles Literal and Enum types "
|
||
|
|
"automatically during streaming. You can safely remove this mixin.",
|
||
|
|
DeprecationWarning,
|
||
|
|
stacklevel=2,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def remove_control_chars(s):
|
||
|
|
return re.sub(r"[\x00-\x1F\x7F-\x9F]", "", s)
|
||
|
|
|
||
|
|
|
||
|
|
def process_potential_object(potential_object, partial_mode, partial_model, **kwargs):
|
||
|
|
"""Process a potential JSON object using completeness-based validation.
|
||
|
|
|
||
|
|
- If JSON is complete (closed braces/brackets): validate against original model
|
||
|
|
- If JSON is incomplete: build partial object using model_construct (no validation)
|
||
|
|
|
||
|
|
Note: Pydantic v2.10+ has `experimental_allow_partial` but it doesn't support
|
||
|
|
BaseModel constraints during partial validation (only TypedDict). If Pydantic
|
||
|
|
adds BaseModel support in the future, this could potentially be simplified.
|
||
|
|
See: https://docs.pydantic.dev/latest/concepts/partial_validation/
|
||
|
|
"""
|
||
|
|
json_str = potential_object.strip() or "{}"
|
||
|
|
parsed = from_json(json_str.encode(), partial_mode=partial_mode)
|
||
|
|
|
||
|
|
tracker = JsonCompleteness()
|
||
|
|
tracker.analyze(json_str)
|
||
|
|
|
||
|
|
# Get original model for validation
|
||
|
|
original_model = getattr(partial_model, "_original_model", None)
|
||
|
|
|
||
|
|
# Check if root is complete AND has actual data (not just empty {})
|
||
|
|
root_complete = tracker.is_root_complete()
|
||
|
|
has_data = bool(parsed) if isinstance(parsed, dict) else True
|
||
|
|
|
||
|
|
if root_complete and has_data and original_model is not None:
|
||
|
|
# Root object is complete with data - validate against original model
|
||
|
|
return original_model.model_validate(parsed, **kwargs)
|
||
|
|
else:
|
||
|
|
# Object is incomplete or empty - build instance using model_construct (no validation)
|
||
|
|
model_for_construct = (
|
||
|
|
original_model if original_model is not None else partial_model
|
||
|
|
)
|
||
|
|
return _build_partial_object(parsed, model_for_construct, tracker, "", **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
def _build_partial_object(
|
||
|
|
data: Any,
|
||
|
|
model: type[BaseModel],
|
||
|
|
tracker: JsonCompleteness,
|
||
|
|
path: str,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> Any:
|
||
|
|
"""Build a partial object using model_construct() to skip validation.
|
||
|
|
|
||
|
|
For each field:
|
||
|
|
- If the field's JSON is complete AND it's a nested BaseModel: validate it
|
||
|
|
- Otherwise: store without validation
|
||
|
|
"""
|
||
|
|
if data is None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
if not isinstance(data, dict):
|
||
|
|
return data
|
||
|
|
|
||
|
|
result = {}
|
||
|
|
|
||
|
|
for field_name in data:
|
||
|
|
field_value = data[field_name]
|
||
|
|
field_path = f"{path}.{field_name}" if path else field_name
|
||
|
|
|
||
|
|
if field_value is None:
|
||
|
|
result[field_name] = None
|
||
|
|
continue
|
||
|
|
|
||
|
|
field_complete = tracker.is_path_complete(field_path)
|
||
|
|
field_info = model.model_fields.get(field_name)
|
||
|
|
field_type = field_info.annotation if field_info else None
|
||
|
|
|
||
|
|
if field_complete and field_type is not None:
|
||
|
|
if isinstance(field_type, type) and issubclass(field_type, BaseModel):
|
||
|
|
result[field_name] = field_type.model_validate(field_value, **kwargs)
|
||
|
|
continue
|
||
|
|
|
||
|
|
if isinstance(field_value, dict):
|
||
|
|
nested_model = None
|
||
|
|
if field_type is not None and isinstance(field_type, type):
|
||
|
|
if issubclass(field_type, BaseModel):
|
||
|
|
nested_model = field_type
|
||
|
|
|
||
|
|
if nested_model:
|
||
|
|
result[field_name] = _build_partial_object(
|
||
|
|
field_value, nested_model, tracker, field_path, **kwargs
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
result[field_name] = field_value
|
||
|
|
elif isinstance(field_value, list):
|
||
|
|
result[field_name] = _build_partial_list(
|
||
|
|
field_value, model, field_name, tracker, field_path, **kwargs
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
result[field_name] = field_value
|
||
|
|
|
||
|
|
# Set missing fields to defaults, None, or empty nested models
|
||
|
|
for field_name, field_info in model.model_fields.items():
|
||
|
|
if field_name not in result:
|
||
|
|
field_type = field_info.annotation
|
||
|
|
if field_info.default is not PydanticUndefined:
|
||
|
|
result[field_name] = deepcopy(field_info.default)
|
||
|
|
elif isinstance(field_type, type) and issubclass(field_type, BaseModel):
|
||
|
|
result[field_name] = _build_partial_object(
|
||
|
|
{}, field_type, tracker, "", **kwargs
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
result[field_name] = None
|
||
|
|
|
||
|
|
return model.model_construct(**result)
|
||
|
|
|
||
|
|
|
||
|
|
def _build_partial_list(
|
||
|
|
items: list,
|
||
|
|
original_model: type[BaseModel] | None,
|
||
|
|
field_name: str,
|
||
|
|
tracker: JsonCompleteness,
|
||
|
|
path: str,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> list:
|
||
|
|
"""Build a partial list, validating complete items."""
|
||
|
|
result = []
|
||
|
|
|
||
|
|
item_type = None
|
||
|
|
if original_model:
|
||
|
|
field_info = original_model.model_fields.get(field_name)
|
||
|
|
if field_info:
|
||
|
|
field_type = field_info.annotation
|
||
|
|
if get_origin(field_type) in (list, List): # noqa: UP006
|
||
|
|
args = get_args(field_type)
|
||
|
|
if args:
|
||
|
|
item_type = args[0]
|
||
|
|
|
||
|
|
for i, item in enumerate(items):
|
||
|
|
item_path = f"{path}[{i}]"
|
||
|
|
item_complete = tracker.is_path_complete(item_path)
|
||
|
|
|
||
|
|
if item_complete and item_type and isinstance(item_type, type):
|
||
|
|
if issubclass(item_type, BaseModel) and isinstance(item, dict):
|
||
|
|
result.append(item_type.model_validate(item, **kwargs))
|
||
|
|
continue
|
||
|
|
|
||
|
|
result.append(item)
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def _process_generic_arg(
|
||
|
|
arg: Any,
|
||
|
|
make_fields_optional: bool = False,
|
||
|
|
) -> Any:
|
||
|
|
arg_origin = get_origin(arg)
|
||
|
|
|
||
|
|
if arg_origin is not None:
|
||
|
|
# Handle any nested generic type (Union, List, Dict, etc.)
|
||
|
|
nested_args = get_args(arg)
|
||
|
|
modified_nested_args = tuple(
|
||
|
|
_process_generic_arg(
|
||
|
|
t,
|
||
|
|
make_fields_optional=make_fields_optional,
|
||
|
|
)
|
||
|
|
for t in nested_args
|
||
|
|
)
|
||
|
|
# Special handling for Union types (types.UnionType isn't subscriptable)
|
||
|
|
if arg_origin in UNION_ORIGINS:
|
||
|
|
return Union[modified_nested_args] # type: ignore
|
||
|
|
|
||
|
|
return arg_origin[modified_nested_args]
|
||
|
|
else:
|
||
|
|
if isinstance(arg, type) and issubclass(arg, BaseModel):
|
||
|
|
# Prevent infinite recursion for self-referential models
|
||
|
|
if arg in _processing_models:
|
||
|
|
return arg # Already processing this model, return unwrapped
|
||
|
|
_processing_models.add(arg)
|
||
|
|
try:
|
||
|
|
return (
|
||
|
|
Partial[arg, MakeFieldsOptional] # type: ignore[valid-type]
|
||
|
|
if make_fields_optional
|
||
|
|
else Partial[arg]
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
_processing_models.discard(arg)
|
||
|
|
else:
|
||
|
|
return arg
|
||
|
|
|
||
|
|
|
||
|
|
def _rebuild_generic_annotation(
|
||
|
|
generic_base: Any,
|
||
|
|
generic_args: tuple[Any, ...],
|
||
|
|
) -> Any:
|
||
|
|
# `types.UnionType` (the runtime origin of `str | int`) cannot be
|
||
|
|
# subscripted directly, so reconstruct unions through `typing.Union`.
|
||
|
|
if generic_base in UNION_ORIGINS:
|
||
|
|
return Union[generic_args] # type: ignore[arg-type]
|
||
|
|
|
||
|
|
return generic_base[generic_args]
|
||
|
|
|
||
|
|
|
||
|
|
def _make_field_optional(
|
||
|
|
field: FieldInfo,
|
||
|
|
) -> tuple[Any, FieldInfo]:
|
||
|
|
tmp_field = deepcopy(field)
|
||
|
|
|
||
|
|
annotation = field.annotation
|
||
|
|
|
||
|
|
# Handle generics (like List, Dict, Union, Literal, etc.)
|
||
|
|
if get_origin(annotation) is not None:
|
||
|
|
# Get the generic base (like List, Dict) and its arguments (like User in List[User])
|
||
|
|
generic_base = get_origin(annotation)
|
||
|
|
generic_args = get_args(annotation)
|
||
|
|
|
||
|
|
modified_args = tuple(
|
||
|
|
_process_generic_arg(arg, make_fields_optional=True) for arg in generic_args
|
||
|
|
)
|
||
|
|
|
||
|
|
# Reconstruct the generic type with modified arguments
|
||
|
|
tmp_field.annotation = (
|
||
|
|
Optional[_rebuild_generic_annotation(generic_base, modified_args)]
|
||
|
|
if generic_base
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
tmp_field.default = None
|
||
|
|
tmp_field.default_factory = None
|
||
|
|
# If the field is a BaseModel, then recursively convert it's
|
||
|
|
# attributes to optionals.
|
||
|
|
elif isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||
|
|
tmp_field.annotation = Optional[Partial[annotation, MakeFieldsOptional]] # type: ignore[assignment, valid-type]
|
||
|
|
tmp_field.default = {}
|
||
|
|
tmp_field.default_factory = None
|
||
|
|
else:
|
||
|
|
tmp_field.annotation = Optional[field.annotation] # type:ignore
|
||
|
|
tmp_field.default = None
|
||
|
|
tmp_field.default_factory = None
|
||
|
|
|
||
|
|
return tmp_field.annotation, tmp_field # type: ignore
|
||
|
|
|
||
|
|
|
||
|
|
class PartialBase(Generic[T_Model]):
|
||
|
|
@classmethod
|
||
|
|
@cache
|
||
|
|
def get_partial_model(cls) -> type[T_Model]:
|
||
|
|
"""Return a partial model for holding incomplete streaming data.
|
||
|
|
|
||
|
|
With completeness-based validation, we use model_construct() to build
|
||
|
|
partial objects without validation. This method creates a model with
|
||
|
|
all fields optional and stores a reference to the original model
|
||
|
|
for validation when JSON is complete.
|
||
|
|
"""
|
||
|
|
assert issubclass(cls, BaseModel), (
|
||
|
|
f"{cls.__name__} must be a subclass of BaseModel"
|
||
|
|
)
|
||
|
|
|
||
|
|
model_name = (
|
||
|
|
cls.__name__
|
||
|
|
if cls.__name__.startswith("Partial")
|
||
|
|
else f"Partial{cls.__name__}"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create partial model with optional fields
|
||
|
|
partial_model = create_model(
|
||
|
|
model_name,
|
||
|
|
__base__=cls,
|
||
|
|
__module__=cls.__module__,
|
||
|
|
**{
|
||
|
|
field_name: _make_field_optional(field_info)
|
||
|
|
for field_name, field_info in cls.model_fields.items()
|
||
|
|
}, # type: ignore[all]
|
||
|
|
)
|
||
|
|
|
||
|
|
# Store reference to original model for validation of complete objects
|
||
|
|
original = getattr(cls, "_original_model", cls)
|
||
|
|
partial_model._original_model = original # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
return partial_model
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_streaming_response(
|
||
|
|
cls,
|
||
|
|
completion: Iterable[Any],
|
||
|
|
mode: Mode,
|
||
|
|
on_event: Callable[..., Any] | None = None,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> Generator[T_Model, None, None]:
|
||
|
|
json_chunks = cls.extract_json(completion, mode, on_event=on_event)
|
||
|
|
|
||
|
|
if mode in {Mode.MD_JSON, Mode.GEMINI_TOOLS}:
|
||
|
|
json_chunks = extract_json_from_stream(json_chunks)
|
||
|
|
|
||
|
|
if mode == Mode.WRITER_TOOLS:
|
||
|
|
yield from cls.writer_model_from_chunks(json_chunks, **kwargs)
|
||
|
|
else:
|
||
|
|
yield from cls.model_from_chunks(json_chunks, **kwargs)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def from_streaming_response_async(
|
||
|
|
cls,
|
||
|
|
completion: AsyncGenerator[Any, None],
|
||
|
|
mode: Mode,
|
||
|
|
on_event: Callable[..., Any] | None = None,
|
||
|
|
**kwargs: Any,
|
||
|
|
) -> AsyncGenerator[T_Model, None]:
|
||
|
|
json_chunks = cls.extract_json_async(completion, mode, on_event=on_event)
|
||
|
|
|
||
|
|
if mode in {Mode.MD_JSON, Mode.GEMINI_TOOLS}:
|
||
|
|
json_chunks = extract_json_from_stream_async(json_chunks)
|
||
|
|
|
||
|
|
if mode == Mode.WRITER_TOOLS:
|
||
|
|
async for item in cls.writer_model_from_chunks_async(json_chunks, **kwargs):
|
||
|
|
yield item
|
||
|
|
else:
|
||
|
|
async for item in cls.model_from_chunks_async(json_chunks, **kwargs):
|
||
|
|
yield item
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def writer_model_from_chunks(
|
||
|
|
cls, json_chunks: Iterable[Any], **kwargs: Any
|
||
|
|
) -> Generator[T_Model, None, None]:
|
||
|
|
potential_object = ""
|
||
|
|
partial_model = cls.get_partial_model()
|
||
|
|
# Always use trailing-strings mode to preserve incomplete data during streaming
|
||
|
|
# PartialLiteralMixin is deprecated - completeness-based validation handles Literals
|
||
|
|
partial_mode = "trailing-strings"
|
||
|
|
final_obj = None
|
||
|
|
for chunk in json_chunks:
|
||
|
|
# Writer mode special handling: chunk might be complete JSON replacing accumulated
|
||
|
|
if (
|
||
|
|
len(chunk) > len(potential_object)
|
||
|
|
and chunk.startswith("{")
|
||
|
|
and chunk.endswith("}")
|
||
|
|
):
|
||
|
|
potential_object = chunk
|
||
|
|
else:
|
||
|
|
potential_object += chunk
|
||
|
|
obj = process_potential_object(
|
||
|
|
potential_object, partial_mode, partial_model, **kwargs
|
||
|
|
)
|
||
|
|
final_obj = obj
|
||
|
|
yield obj
|
||
|
|
|
||
|
|
# Final validation: only validate if the JSON is structurally complete
|
||
|
|
# If JSON is incomplete (stream ended mid-object), skip validation
|
||
|
|
if final_obj is not None:
|
||
|
|
original_model = getattr(cls, "_original_model", None)
|
||
|
|
if original_model is not None:
|
||
|
|
if is_json_complete(potential_object.strip() or "{}"):
|
||
|
|
original_model.model_validate(
|
||
|
|
final_obj.model_dump(exclude_none=True), **kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def writer_model_from_chunks_async(
|
||
|
|
cls, json_chunks: AsyncGenerator[str, None], **kwargs: Any
|
||
|
|
) -> AsyncGenerator[T_Model, None]:
|
||
|
|
potential_object = ""
|
||
|
|
partial_model = cls.get_partial_model()
|
||
|
|
# Always use trailing-strings mode to preserve incomplete data during streaming
|
||
|
|
# PartialLiteralMixin is deprecated - completeness-based validation handles Literals
|
||
|
|
partial_mode = "trailing-strings"
|
||
|
|
final_obj = None
|
||
|
|
async for chunk in json_chunks:
|
||
|
|
# Writer mode special handling: chunk might be complete JSON replacing accumulated
|
||
|
|
if (
|
||
|
|
len(chunk) > len(potential_object)
|
||
|
|
and chunk.startswith("{")
|
||
|
|
and chunk.endswith("}")
|
||
|
|
):
|
||
|
|
potential_object = chunk
|
||
|
|
else:
|
||
|
|
potential_object += chunk
|
||
|
|
obj = process_potential_object(
|
||
|
|
potential_object, partial_mode, partial_model, **kwargs
|
||
|
|
)
|
||
|
|
final_obj = obj
|
||
|
|
yield obj
|
||
|
|
|
||
|
|
# Final validation: only validate if the JSON is structurally complete
|
||
|
|
# If JSON is incomplete (stream ended mid-object), skip validation
|
||
|
|
if final_obj is not None:
|
||
|
|
original_model = getattr(cls, "_original_model", None)
|
||
|
|
if original_model is not None:
|
||
|
|
if is_json_complete(potential_object.strip() or "{}"):
|
||
|
|
original_model.model_validate(
|
||
|
|
final_obj.model_dump(exclude_none=True), **kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def model_from_chunks(
|
||
|
|
cls, json_chunks: Iterable[Any], **kwargs: Any
|
||
|
|
) -> Generator[T_Model, None, None]:
|
||
|
|
potential_object = ""
|
||
|
|
partial_model = cls.get_partial_model()
|
||
|
|
# Always use trailing-strings mode to preserve incomplete data during streaming
|
||
|
|
# PartialLiteralMixin is deprecated - completeness-based validation handles Literals
|
||
|
|
partial_mode = "trailing-strings"
|
||
|
|
final_obj = None
|
||
|
|
for chunk in json_chunks:
|
||
|
|
if chunk is None:
|
||
|
|
continue
|
||
|
|
if not isinstance(chunk, str):
|
||
|
|
try:
|
||
|
|
chunk = str(chunk)
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
potential_object += remove_control_chars(chunk)
|
||
|
|
obj = process_potential_object(
|
||
|
|
potential_object, partial_mode, partial_model, **kwargs
|
||
|
|
)
|
||
|
|
final_obj = obj
|
||
|
|
yield obj
|
||
|
|
|
||
|
|
# Final validation: only validate if the JSON is structurally complete
|
||
|
|
# If JSON is incomplete (stream ended mid-object), skip validation
|
||
|
|
if final_obj is not None:
|
||
|
|
original_model = getattr(cls, "_original_model", None)
|
||
|
|
if original_model is not None:
|
||
|
|
if is_json_complete(potential_object.strip() or "{}"):
|
||
|
|
original_model.model_validate(
|
||
|
|
final_obj.model_dump(exclude_none=True), **kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
async def model_from_chunks_async(
|
||
|
|
cls, json_chunks: AsyncGenerator[str, None], **kwargs: Any
|
||
|
|
) -> AsyncGenerator[T_Model, None]:
|
||
|
|
potential_object = ""
|
||
|
|
partial_model = cls.get_partial_model()
|
||
|
|
# Always use trailing-strings mode to preserve incomplete data during streaming
|
||
|
|
# PartialLiteralMixin is deprecated - completeness-based validation handles Literals
|
||
|
|
partial_mode = "trailing-strings"
|
||
|
|
final_obj = None
|
||
|
|
async for chunk in json_chunks:
|
||
|
|
if chunk is None:
|
||
|
|
continue
|
||
|
|
if not isinstance(chunk, str):
|
||
|
|
try:
|
||
|
|
chunk = str(chunk)
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
potential_object += remove_control_chars(chunk)
|
||
|
|
obj = process_potential_object(
|
||
|
|
potential_object, partial_mode, partial_model, **kwargs
|
||
|
|
)
|
||
|
|
final_obj = obj
|
||
|
|
yield obj
|
||
|
|
|
||
|
|
# Final validation: only validate if the JSON is structurally complete
|
||
|
|
# If JSON is incomplete (stream ended mid-object), skip validation
|
||
|
|
if final_obj is not None:
|
||
|
|
original_model = getattr(cls, "_original_model", None)
|
||
|
|
if original_model is not None:
|
||
|
|
if is_json_complete(potential_object.strip() or "{}"):
|
||
|
|
original_model.model_validate(
|
||
|
|
final_obj.model_dump(exclude_none=True), **kwargs
|
||
|
|
)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def extract_json(
|
||
|
|
completion: Iterable[Any],
|
||
|
|
mode: Mode,
|
||
|
|
on_event: Callable[..., Any] | None = None,
|
||
|
|
) -> Generator[str, None, None]:
|
||
|
|
"""Extract JSON chunks from various LLM provider streaming responses.
|
||
|
|
|
||
|
|
Each provider has a different structure for streaming responses that needs
|
||
|
|
specific handling to extract the relevant JSON data."""
|
||
|
|
json_started = False
|
||
|
|
for chunk in completion:
|
||
|
|
try:
|
||
|
|
if mode in {Mode.COHERE_TOOLS, Mode.COHERE_JSON_SCHEMA}:
|
||
|
|
event_type = getattr(chunk, "event_type", None)
|
||
|
|
if event_type == "text-generation":
|
||
|
|
if text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif event_type == "tool-calls-chunk":
|
||
|
|
delta = getattr(chunk, "tool_call_delta", None)
|
||
|
|
args = getattr(delta, "parameters", None) or getattr(
|
||
|
|
delta, "text", None
|
||
|
|
)
|
||
|
|
if args:
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (args.find("{"), args.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
elif text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif event_type == "tool-calls-generation":
|
||
|
|
tool_calls = getattr(chunk, "tool_calls", None)
|
||
|
|
if tool_calls:
|
||
|
|
args = json.dumps(tool_calls[0].parameters)
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (args.find("{"), args.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
elif text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
else:
|
||
|
|
chunk_type = getattr(chunk, "type", None)
|
||
|
|
if chunk_type == "content-delta":
|
||
|
|
delta = getattr(chunk, "delta", None)
|
||
|
|
message = getattr(delta, "message", None)
|
||
|
|
content = getattr(message, "content", None)
|
||
|
|
if text := getattr(content, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (
|
||
|
|
text.find("{"),
|
||
|
|
text.find("["),
|
||
|
|
)
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif chunk_type == "tool-call-delta":
|
||
|
|
delta = getattr(chunk, "delta", None)
|
||
|
|
message = getattr(delta, "message", None)
|
||
|
|
tool_calls = getattr(message, "tool_calls", None)
|
||
|
|
function = getattr(tool_calls, "function", None)
|
||
|
|
if args := getattr(function, "arguments", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (
|
||
|
|
args.find("{"),
|
||
|
|
args.find("["),
|
||
|
|
)
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
if mode == Mode.MISTRAL_STRUCTURED_OUTPUTS:
|
||
|
|
yield chunk.data.choices[0].delta.content
|
||
|
|
if mode == Mode.MISTRAL_TOOLS:
|
||
|
|
if not chunk.data.choices[0].delta.tool_calls:
|
||
|
|
continue
|
||
|
|
yield chunk.data.choices[0].delta.tool_calls[0].function.arguments
|
||
|
|
if mode == Mode.ANTHROPIC_JSON:
|
||
|
|
if json_chunk := chunk.delta.text:
|
||
|
|
yield json_chunk
|
||
|
|
if mode == Mode.ANTHROPIC_TOOLS:
|
||
|
|
yield chunk.delta.partial_json
|
||
|
|
if mode == Mode.VERTEXAI_JSON:
|
||
|
|
yield chunk.candidates[0].content.parts[0].text
|
||
|
|
if mode == Mode.VERTEXAI_TOOLS:
|
||
|
|
yield json.dumps(
|
||
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
||
|
|
)
|
||
|
|
|
||
|
|
if mode == Mode.GENAI_STRUCTURED_OUTPUTS:
|
||
|
|
try:
|
||
|
|
yield chunk.text
|
||
|
|
except ValueError as e:
|
||
|
|
if "valid `Part`" in str(e):
|
||
|
|
# Skip chunk with invalid Part (e.g., due to finish_reason=1 token limit)
|
||
|
|
continue
|
||
|
|
raise
|
||
|
|
if mode == Mode.GENAI_TOOLS:
|
||
|
|
fc = chunk.candidates[0].content.parts[0].function_call.args
|
||
|
|
yield json.dumps(fc)
|
||
|
|
if mode == Mode.GEMINI_JSON:
|
||
|
|
try:
|
||
|
|
yield chunk.text
|
||
|
|
except ValueError as e:
|
||
|
|
if "valid `Part`" in str(e):
|
||
|
|
# Skip chunk with invalid Part (e.g., due to finish_reason=1 token limit)
|
||
|
|
continue
|
||
|
|
raise
|
||
|
|
if mode == Mode.GEMINI_TOOLS:
|
||
|
|
resp = chunk.candidates[0].content.parts[0].function_call
|
||
|
|
resp_dict = type(resp).to_dict(resp) # type:ignore
|
||
|
|
if "args" in resp_dict:
|
||
|
|
yield json.dumps(resp_dict["args"])
|
||
|
|
elif mode in {
|
||
|
|
Mode.RESPONSES_TOOLS,
|
||
|
|
Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS,
|
||
|
|
}:
|
||
|
|
from openai.types.responses import (
|
||
|
|
ResponseFunctionCallArgumentsDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDoneEvent,
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(chunk, ResponseFunctionCallArgumentsDeltaEvent):
|
||
|
|
yield chunk.delta
|
||
|
|
elif on_event is not None and isinstance(
|
||
|
|
chunk,
|
||
|
|
(
|
||
|
|
ResponseReasoningSummaryTextDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDoneEvent,
|
||
|
|
),
|
||
|
|
):
|
||
|
|
on_event(chunk)
|
||
|
|
|
||
|
|
elif chunk.choices:
|
||
|
|
if mode == Mode.FUNCTIONS:
|
||
|
|
Mode.warn_mode_functions_deprecation()
|
||
|
|
if json_chunk := chunk.choices[0].delta.function_call.arguments:
|
||
|
|
yield json_chunk
|
||
|
|
elif mode in {
|
||
|
|
Mode.JSON,
|
||
|
|
Mode.MD_JSON,
|
||
|
|
Mode.JSON_SCHEMA,
|
||
|
|
Mode.CEREBRAS_JSON,
|
||
|
|
Mode.FIREWORKS_JSON,
|
||
|
|
Mode.PERPLEXITY_JSON,
|
||
|
|
Mode.WRITER_JSON,
|
||
|
|
}:
|
||
|
|
if json_chunk := chunk.choices[0].delta.content:
|
||
|
|
yield json_chunk
|
||
|
|
elif mode in {
|
||
|
|
Mode.TOOLS,
|
||
|
|
Mode.TOOLS_STRICT,
|
||
|
|
Mode.FIREWORKS_TOOLS,
|
||
|
|
Mode.WRITER_TOOLS,
|
||
|
|
}:
|
||
|
|
if json_chunk := chunk.choices[0].delta.tool_calls:
|
||
|
|
if json_chunk[0].function.arguments:
|
||
|
|
yield json_chunk[0].function.arguments
|
||
|
|
else:
|
||
|
|
raise NotImplementedError(
|
||
|
|
f"Mode {mode} is not supported for MultiTask streaming"
|
||
|
|
)
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
async def extract_json_async(
|
||
|
|
completion: AsyncGenerator[Any, None],
|
||
|
|
mode: Mode,
|
||
|
|
on_event: Callable[..., Any] | None = None,
|
||
|
|
) -> AsyncGenerator[str, None]:
|
||
|
|
json_started = False
|
||
|
|
async for chunk in completion:
|
||
|
|
try:
|
||
|
|
if mode in {Mode.COHERE_TOOLS, Mode.COHERE_JSON_SCHEMA}:
|
||
|
|
event_type = getattr(chunk, "event_type", None)
|
||
|
|
if event_type == "text-generation":
|
||
|
|
if text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif event_type == "tool-calls-chunk":
|
||
|
|
delta = getattr(chunk, "tool_call_delta", None)
|
||
|
|
args = getattr(delta, "parameters", None) or getattr(
|
||
|
|
delta, "text", None
|
||
|
|
)
|
||
|
|
if args:
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (args.find("{"), args.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
elif text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif event_type == "tool-calls-generation":
|
||
|
|
tool_calls = getattr(chunk, "tool_calls", None)
|
||
|
|
if tool_calls:
|
||
|
|
args = json.dumps(tool_calls[0].parameters)
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (args.find("{"), args.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
elif text := getattr(chunk, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (text.find("{"), text.find("["))
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
else:
|
||
|
|
chunk_type = getattr(chunk, "type", None)
|
||
|
|
if chunk_type == "content-delta":
|
||
|
|
delta = getattr(chunk, "delta", None)
|
||
|
|
message = getattr(delta, "message", None)
|
||
|
|
content = getattr(message, "content", None)
|
||
|
|
if text := getattr(content, "text", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (
|
||
|
|
text.find("{"),
|
||
|
|
text.find("["),
|
||
|
|
)
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
text = text[json_start:]
|
||
|
|
yield text
|
||
|
|
elif chunk_type == "tool-call-delta":
|
||
|
|
delta = getattr(chunk, "delta", None)
|
||
|
|
message = getattr(delta, "message", None)
|
||
|
|
tool_calls = getattr(message, "tool_calls", None)
|
||
|
|
function = getattr(tool_calls, "function", None)
|
||
|
|
if args := getattr(function, "arguments", None):
|
||
|
|
if not json_started:
|
||
|
|
json_start = min(
|
||
|
|
(
|
||
|
|
pos
|
||
|
|
for pos in (
|
||
|
|
args.find("{"),
|
||
|
|
args.find("["),
|
||
|
|
)
|
||
|
|
if pos != -1
|
||
|
|
),
|
||
|
|
default=-1,
|
||
|
|
)
|
||
|
|
if json_start == -1:
|
||
|
|
continue
|
||
|
|
json_started = True
|
||
|
|
args = args[json_start:]
|
||
|
|
yield args
|
||
|
|
if mode == Mode.ANTHROPIC_JSON:
|
||
|
|
if json_chunk := chunk.delta.text:
|
||
|
|
yield json_chunk
|
||
|
|
if mode == Mode.ANTHROPIC_TOOLS:
|
||
|
|
yield chunk.delta.partial_json
|
||
|
|
if mode == Mode.MISTRAL_STRUCTURED_OUTPUTS:
|
||
|
|
yield chunk.data.choices[0].delta.content
|
||
|
|
if mode == Mode.MISTRAL_TOOLS:
|
||
|
|
if not chunk.data.choices[0].delta.tool_calls:
|
||
|
|
continue
|
||
|
|
yield chunk.data.choices[0].delta.tool_calls[0].function.arguments
|
||
|
|
if mode == Mode.VERTEXAI_JSON:
|
||
|
|
yield chunk.candidates[0].content.parts[0].text
|
||
|
|
if mode == Mode.VERTEXAI_TOOLS:
|
||
|
|
yield json.dumps(
|
||
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
||
|
|
)
|
||
|
|
if mode == Mode.GENAI_STRUCTURED_OUTPUTS:
|
||
|
|
try:
|
||
|
|
yield chunk.text
|
||
|
|
except ValueError as e:
|
||
|
|
if "valid `Part`" in str(e):
|
||
|
|
# Skip chunk with invalid Part (e.g., due to finish_reason=1 token limit)
|
||
|
|
continue
|
||
|
|
raise
|
||
|
|
if mode == Mode.GENAI_TOOLS:
|
||
|
|
fc = chunk.candidates[0].content.parts[0].function_call.args
|
||
|
|
yield json.dumps(fc)
|
||
|
|
if mode == Mode.GEMINI_JSON:
|
||
|
|
try:
|
||
|
|
yield chunk.text
|
||
|
|
except ValueError as e:
|
||
|
|
if "valid `Part`" in str(e):
|
||
|
|
# Skip chunk with invalid Part (e.g., due to finish_reason=1 token limit)
|
||
|
|
continue
|
||
|
|
raise
|
||
|
|
if mode == Mode.GEMINI_TOOLS:
|
||
|
|
resp = chunk.candidates[0].content.parts[0].function_call
|
||
|
|
resp_dict = type(resp).to_dict(resp) # type:ignore
|
||
|
|
if "args" in resp_dict:
|
||
|
|
yield json.dumps(resp_dict["args"])
|
||
|
|
|
||
|
|
if mode in {
|
||
|
|
Mode.RESPONSES_TOOLS,
|
||
|
|
Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS,
|
||
|
|
}:
|
||
|
|
from openai.types.responses import (
|
||
|
|
ResponseFunctionCallArgumentsDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDoneEvent,
|
||
|
|
)
|
||
|
|
|
||
|
|
if isinstance(chunk, ResponseFunctionCallArgumentsDeltaEvent):
|
||
|
|
yield chunk.delta
|
||
|
|
elif on_event is not None and isinstance(
|
||
|
|
chunk,
|
||
|
|
(
|
||
|
|
ResponseReasoningSummaryTextDeltaEvent,
|
||
|
|
ResponseReasoningSummaryTextDoneEvent,
|
||
|
|
),
|
||
|
|
):
|
||
|
|
if asyncio.iscoroutinefunction(on_event):
|
||
|
|
await on_event(chunk)
|
||
|
|
else:
|
||
|
|
on_event(chunk)
|
||
|
|
elif chunk.choices:
|
||
|
|
if mode == Mode.FUNCTIONS:
|
||
|
|
Mode.warn_mode_functions_deprecation()
|
||
|
|
if json_chunk := chunk.choices[0].delta.function_call.arguments:
|
||
|
|
yield json_chunk
|
||
|
|
elif mode in {
|
||
|
|
Mode.JSON,
|
||
|
|
Mode.MD_JSON,
|
||
|
|
Mode.JSON_SCHEMA,
|
||
|
|
Mode.CEREBRAS_JSON,
|
||
|
|
Mode.FIREWORKS_JSON,
|
||
|
|
Mode.PERPLEXITY_JSON,
|
||
|
|
Mode.WRITER_JSON,
|
||
|
|
}:
|
||
|
|
if json_chunk := chunk.choices[0].delta.content:
|
||
|
|
yield json_chunk
|
||
|
|
elif mode in {
|
||
|
|
Mode.TOOLS,
|
||
|
|
Mode.TOOLS_STRICT,
|
||
|
|
Mode.FIREWORKS_TOOLS,
|
||
|
|
Mode.WRITER_TOOLS,
|
||
|
|
}:
|
||
|
|
if json_chunk := chunk.choices[0].delta.tool_calls:
|
||
|
|
if json_chunk[0].function.arguments:
|
||
|
|
yield json_chunk[0].function.arguments
|
||
|
|
else:
|
||
|
|
raise NotImplementedError(
|
||
|
|
f"Mode {mode} is not supported for MultiTask streaming"
|
||
|
|
)
|
||
|
|
except AttributeError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class Partial(Generic[T_Model]):
|
||
|
|
"""Generate a new class which has PartialBase as a base class.
|
||
|
|
|
||
|
|
Notes:
|
||
|
|
This will enable partial validation of the model while streaming.
|
||
|
|
|
||
|
|
Example:
|
||
|
|
Partial[SomeModel]
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __new__(
|
||
|
|
cls,
|
||
|
|
*args: object, # noqa
|
||
|
|
**kwargs: object, # noqa
|
||
|
|
) -> Partial[T_Model]:
|
||
|
|
"""Cannot instantiate.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
TypeError: Direct instantiation not allowed.
|
||
|
|
"""
|
||
|
|
raise TypeError("Cannot instantiate abstract Partial class.")
|
||
|
|
|
||
|
|
def __init_subclass__(
|
||
|
|
cls,
|
||
|
|
*args: object,
|
||
|
|
**kwargs: object,
|
||
|
|
) -> NoReturn:
|
||
|
|
"""Cannot subclass.
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
TypeError: Subclassing not allowed.
|
||
|
|
"""
|
||
|
|
raise TypeError(f"Cannot subclass {cls.__module__}.Partial")
|
||
|
|
|
||
|
|
def __class_getitem__(
|
||
|
|
cls,
|
||
|
|
wrapped_class: type[T_Model] | tuple[type[T_Model], type[MakeFieldsOptional]],
|
||
|
|
) -> type[T_Model]:
|
||
|
|
"""Convert model to one that inherits from PartialBase.
|
||
|
|
|
||
|
|
We don't make the fields optional at this point, we just wrap them with `Partial` so the names of the nested models will be
|
||
|
|
`Partial{ModelName}`. We want the output of `model_json_schema()` to
|
||
|
|
reflect the name change, but everything else should be the same as the
|
||
|
|
original model. During validation, we'll generate a true partial model
|
||
|
|
to support partially defined fields.
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
make_fields_optional = None
|
||
|
|
if isinstance(wrapped_class, tuple):
|
||
|
|
wrapped_class, make_fields_optional = wrapped_class
|
||
|
|
|
||
|
|
def _wrap_models(field: FieldInfo) -> tuple[object, FieldInfo]:
|
||
|
|
tmp_field = deepcopy(field)
|
||
|
|
|
||
|
|
annotation = field.annotation
|
||
|
|
|
||
|
|
# Handle generics (like List, Dict, etc.)
|
||
|
|
if get_origin(annotation) is not None:
|
||
|
|
# Get the generic base (like List, Dict) and its arguments (like User in List[User])
|
||
|
|
generic_base = get_origin(annotation)
|
||
|
|
generic_args = get_args(annotation)
|
||
|
|
|
||
|
|
modified_args = tuple(_process_generic_arg(arg) for arg in generic_args)
|
||
|
|
|
||
|
|
# Reconstruct the generic type with modified arguments
|
||
|
|
tmp_field.annotation = (
|
||
|
|
_rebuild_generic_annotation(generic_base, modified_args)
|
||
|
|
if generic_base
|
||
|
|
else None
|
||
|
|
)
|
||
|
|
# If the field is a BaseModel, then recursively convert it's
|
||
|
|
# attributes to optionals.
|
||
|
|
elif isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||
|
|
# Prevent infinite recursion for self-referential models
|
||
|
|
if annotation in _processing_models:
|
||
|
|
tmp_field.annotation = (
|
||
|
|
annotation # Already processing, keep unwrapped
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_processing_models.add(annotation)
|
||
|
|
try:
|
||
|
|
tmp_field.annotation = Partial[annotation]
|
||
|
|
finally:
|
||
|
|
_processing_models.discard(annotation)
|
||
|
|
return tmp_field.annotation, tmp_field
|
||
|
|
|
||
|
|
model_name = (
|
||
|
|
wrapped_class.__name__
|
||
|
|
if wrapped_class.__name__.startswith("Partial")
|
||
|
|
else f"Partial{wrapped_class.__name__}"
|
||
|
|
)
|
||
|
|
|
||
|
|
partial_model = create_model(
|
||
|
|
model_name,
|
||
|
|
__base__=(wrapped_class, PartialBase), # type: ignore
|
||
|
|
__module__=wrapped_class.__module__,
|
||
|
|
**{
|
||
|
|
field_name: (
|
||
|
|
_make_field_optional(field_info)
|
||
|
|
if make_fields_optional is not None
|
||
|
|
else _wrap_models(field_info)
|
||
|
|
)
|
||
|
|
for field_name, field_info in wrapped_class.model_fields.items()
|
||
|
|
}, # type: ignore
|
||
|
|
)
|
||
|
|
|
||
|
|
# Store reference to original model for final validation
|
||
|
|
partial_model._original_model = wrapped_class # type: ignore[attr-defined]
|
||
|
|
|
||
|
|
return partial_model
|