diff --git a/instructor/v2/dsl/citation.py b/instructor/v2/dsl/citation.py index c5ccd4f04..3aa044f00 100644 --- a/instructor/v2/dsl/citation.py +++ b/instructor/v2/dsl/citation.py @@ -69,6 +69,9 @@ def validate_sources(self, info: ValidationInfo) -> "CitationMixin": # Get the context from the info text_chunks = info.context.get("context", None) + if text_chunks is None: + return self + # Get the spans of the substring_phrase in the context spans = list(self.get_spans(text_chunks)) # Replace the substring_phrase with the actual substring diff --git a/instructor/v2/dsl/partial.py b/instructor/v2/dsl/partial.py index 464d32ca8..0cc66ae51 100644 --- a/instructor/v2/dsl/partial.py +++ b/instructor/v2/dsl/partial.py @@ -9,6 +9,7 @@ from __future__ import annotations import re +import threading import types import warnings from collections.abc import AsyncGenerator, Callable, Generator, Iterable @@ -42,6 +43,7 @@ # 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() +_processing_models_lock = threading.Lock() def _unwrap_optional_base_model(annotation: Any) -> type[BaseModel] | None: @@ -269,9 +271,10 @@ 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: - return arg # Already processing this model, return unwrapped - _processing_models.add(arg) + with _processing_models_lock: + if arg in _processing_models: + return arg # Already processing this model, return unwrapped + _processing_models.add(arg) try: return ( _make_partial_type(arg, make_fields_optional=True) @@ -279,7 +282,8 @@ def _process_generic_arg( else Partial[arg] ) finally: - _processing_models.discard(arg) + with _processing_models_lock: + _processing_models.discard(arg) else: return arg @@ -636,15 +640,17 @@ 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: - tmp_field.annotation = ( - annotation # Already processing, keep unwrapped - ) - else: - _processing_models.add(annotation) - try: - tmp_field.annotation = Partial[annotation] - finally: + with _processing_models_lock: + 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: + with _processing_models_lock: _processing_models.discard(annotation) return tmp_field.annotation, tmp_field diff --git a/tests/dsl/test_citation.py b/tests/dsl/test_citation.py new file mode 100644 index 000000000..bbf9fdefb --- /dev/null +++ b/tests/dsl/test_citation.py @@ -0,0 +1,43 @@ +"""Regression test for CitationMixin crash when context missing 'context' key. + +See: https://github.com/567-labs/instructor/issues/2459 +""" +import pytest +from pydantic import BaseModel, Field + +from instructor import CitationMixin + + +class CitationModel(CitationMixin, BaseModel): + answer: str = Field(description="The answer") + substring_quotes: list[str] = Field( + default_factory=list, + description="Quotes supporting the answer", + ) + + +class TestCitationMixinNoneContext: + def test_context_missing_context_key_does_not_crash(self): + """validate_sources should not crash when context dict lacks 'context' key.""" + model = CitationModel.model_validate( + {"answer": "test", "substring_quotes": ["test"]}, + context={"foo": "bar"}, + ) + assert model.answer == "test" + + def test_context_is_none_does_not_crash(self): + """validate_sources should handle None context gracefully.""" + model = CitationModel.model_validate( + {"answer": "test", "substring_quotes": ["test"]}, + ) + assert model.answer == "test" + + def test_context_with_context_key_works_normally(self): + """validate_sources should work normally when context key is present.""" + context = "Jason was a student. Jason is 20 years old." + model = CitationModel.model_validate( + {"answer": "Jason", "substring_quotes": ["Jason"]}, + context={"context": context}, + ) + assert model.answer == "Jason" + assert len(model.substring_quotes) >= 1