diff --git a/CHANGELOG.md b/CHANGELOG.md index 0fcbd6f1b..ab0daf2f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ Versioning: [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ## [Unreleased] ### Fixed +- **Partial models**: Isolate recursive-model guards per request context so concurrent partial-model construction cannot skip nested model conversion. - **v2 message handling**: Preserve caller-owned message lists and nested content across request preparation and retries for OpenAI-compatible, Cohere, Mistral, OpenRouter, Writer, and xAI handlers. ([#2417](https://github.com/567-labs/instructor/issues/2417), [#2428](https://github.com/567-labs/instructor/issues/2428)) - **v2 JSON extraction**: Prefer the final complete top-level JSON value in text responses and retain every JSON object when multiple objects arrive in one streaming chunk. - **v2 schemas**: Treat fields with Pydantic `default_factory` values as optional in generated OpenAI tool schemas. diff --git a/instructor/v2/dsl/partial.py b/instructor/v2/dsl/partial.py index 464d32ca8..1e2ee432f 100644 --- a/instructor/v2/dsl/partial.py +++ b/instructor/v2/dsl/partial.py @@ -12,6 +12,7 @@ import types import warnings from collections.abc import AsyncGenerator, Callable, Generator, Iterable +from contextvars import ContextVar from copy import deepcopy from functools import cache from functools import reduce @@ -40,8 +41,11 @@ UNION_ORIGINS = (Union, UNION_TYPE) if UNION_TYPE is not None else (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() +# with self-referential models (e.g., TreeNode with children: List["TreeNode"]). +# Each top-level partial-model construction receives an isolated guard. +_processing_models: ContextVar[set[type] | None] = ContextVar( + "processing_models", default=None +) def _unwrap_optional_base_model(annotation: Any) -> type[BaseModel] | None: @@ -250,6 +254,15 @@ def _process_generic_arg( arg: Any, make_fields_optional: bool = False, ) -> Any: + if _processing_models.get() is None: + token = _processing_models.set(set()) + try: + return _process_generic_arg(arg, make_fields_optional=make_fields_optional) + finally: + _processing_models.reset(token) + + processing_models = _processing_models.get() + assert processing_models is not None arg_origin = get_origin(arg) if arg_origin is not None: @@ -269,9 +282,9 @@ def _process_generic_arg( return arg_origin[modified_nested_args] if isinstance(arg, type) and issubclass(arg, BaseModel): # Prevent infinite recursion for self-referential models - if arg in _processing_models: + if arg in processing_models: return arg # Already processing this model, return unwrapped - _processing_models.add(arg) + processing_models.add(arg) try: return ( _make_partial_type(arg, make_fields_optional=True) @@ -279,7 +292,7 @@ def _process_generic_arg( else Partial[arg] ) finally: - _processing_models.discard(arg) + processing_models.discard(arg) else: return arg @@ -607,6 +620,15 @@ def __class_getitem__( to support partially defined fields. """ + if _processing_models.get() is None: + token = _processing_models.set(set()) + try: + return cls.__class_getitem__(wrapped_class) + finally: + _processing_models.reset(token) + + processing_models = _processing_models.get() + assert processing_models is not None make_fields_optional = None if isinstance(wrapped_class, tuple): @@ -636,16 +658,16 @@ def _wrap_models(field: FieldInfo) -> tuple[object, FieldInfo]: # attributes to optionals. elif isinstance(annotation, type) and issubclass(annotation, BaseModel): # Prevent infinite recursion for self-referential models - if annotation in _processing_models: + if annotation in processing_models: tmp_field.annotation = ( annotation # Already processing, keep unwrapped ) else: - _processing_models.add(annotation) + processing_models.add(annotation) try: tmp_field.annotation = Partial[annotation] finally: - _processing_models.discard(annotation) + processing_models.discard(annotation) return tmp_field.annotation, tmp_field model_name = ( diff --git a/tests/coverage/test_dsl_partial_coverage.py b/tests/coverage/test_dsl_partial_coverage.py index 98ed7c553..78dc89406 100644 --- a/tests/coverage/test_dsl_partial_coverage.py +++ b/tests/coverage/test_dsl_partial_coverage.py @@ -1,5 +1,7 @@ from collections.abc import AsyncGenerator, Callable, Generator, Iterable +from concurrent.futures import ThreadPoolExecutor from enum import Enum +from threading import Event, Lock from typing import Any, Optional, Protocol, Union, cast import typing @@ -444,3 +446,64 @@ class RecursiveNode(BaseModel): assert schema["$defs"]["RecursiveNode"]["properties"]["child"] == { "$ref": "#/$defs/RecursiveNode" } + + +def test_partial_processing_isolated_between_threads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class Address(BaseModel): + street: str + + class Home(BaseModel): + addresses: list[Address] + + class Office(BaseModel): + addresses: list[Address] + + first_entered = Event() + second_entered = Event() + release_first = Event() + call_lock = Lock() + calls = 0 + original_make_partial_type = partial_module._make_partial_type + + def block_first_address_conversion( + annotation: type[BaseModel], *, make_fields_optional: bool = False + ) -> type[BaseModel]: + nonlocal calls + if annotation is Address: + with call_lock: + calls += 1 + call_number = calls + + if call_number == 1: + first_entered.set() + assert release_first.wait(timeout=1) + elif call_number == 2: + second_entered.set() + + return original_make_partial_type( + annotation, make_fields_optional=make_fields_optional + ) + + monkeypatch.setattr( + partial_module, "_make_partial_type", block_first_address_conversion + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit( + partial_module._make_partial_type, Home, make_fields_optional=True + ) + assert first_entered.wait(timeout=1) + + second = executor.submit( + partial_module._make_partial_type, Office, make_fields_optional=True + ) + assert second_entered.wait(timeout=1) + release_first.set() + + home_partial = first.result() + office_partial = second.result() + + assert "PartialAddress" in home_partial.model_json_schema()["$defs"] + assert "PartialAddress" in office_partial.model_json_schema()["$defs"]