From 5d18e3b180f47a87f46bf9c6baf5ee7b53780260 Mon Sep 17 00:00:00 2001
From: Abreham Melese <truecallerabreham@gmail.com>
Date: Sat, 18 Jul 2026 11:09:32 -0700
Subject: [PATCH] fix: prevent CitationMixin crash when context missing
 'context' key

Closes #2459
---
 instructor/v2/dsl/citation.py |  3 +++
 tests/dsl/test_citation.py    | 43 +++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 tests/dsl/test_citation.py

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/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