fix: add threading lock for _processing_models set in partial.py

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

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,17 +271,19 @@ 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)
                 if make_fields_optional
                 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
 

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