fix: preserve generation kwargs in relation extraction (#1213)

* fix: preserve generation kwargs in relation extraction

* fix: include generation params in extraction cache keys

* fix: cover provider-specific generation params in extraction cache key

_GENERATION_CACHE_KEYS only covered the common OpenAI-shaped generation
params, so calls that differed only in Anthropic's system/stop_sequences,
Gemini's candidate_count, or Ollama's repeat_penalty/num_ctx/context_window
could still return a stale cached result generated under different settings.

Add these provider-specific keys to the cache key and add regression tests
covering system prompt, stop_sequences, and repeat_penalty.

---------

Co-authored-by: KaifAhmad1 <kaifahmad087@gmail.com>

semantica/semantic_extract/methods.py

@@ -140,6 +140,47 @@ _result_cache = ExtractionCache(
 if not config.get("cache_enabled", True):
     _result_cache.enabled = False
 
+# Generation kwargs that affect provider output and must therefore be part of
+# the cache key. This is the union of every generation-affecting parameter
+# read across providers.py, including params picked up outside _add_if_set
+# (e.g. AnthropicProvider's manual pass-through loop). Sensitive values
+# (api_key, token, etc.) are already filtered out by
+# ExtractionCache._generate_key, so they need not be excluded here.
+_GENERATION_CACHE_KEYS = frozenset({
+    "max_tokens",
+    "max_completion_tokens",
+    "temperature",
+    "top_p",
+    "top_k",
+    "seed",
+    "frequency_penalty",
+    "presence_penalty",
+    "stop",
+    "stop_sequences",  # Anthropic/Gemini spelling of "stop"
+    "logit_bias",
+    "user",
+    "system",  # Anthropic system prompt
+    "metadata",  # Anthropic request metadata
+    "candidate_count",  # Gemini
+    "repeat_penalty",  # Ollama
+    "num_ctx",  # Ollama
+    "context_window",  # Ollama alias for num_ctx
+})
+
+
+def _generation_cache_params(kwargs: dict) -> dict:
+    """Return the subset of *kwargs* that affects generation output.
+
+    Only keys listed in ``_GENERATION_CACHE_KEYS`` are included so that
+    irrelevant or sensitive caller kwargs do not pollute the cache key.
+    Values that are ``None`` are omitted; a caller passing
+    ``temperature=None`` is equivalent to not passing it at all.
+    """
+    return {
+        k: v for k, v in kwargs.items()
+        if k in _GENERATION_CACHE_KEYS and v is not None
+    }
+
 # Try to import spaCy
 from ..utils.helpers import safe_import
 
@@ -957,6 +998,7 @@ def extract_entities_llm(
         "max_text_length": max_text_length,
         "structured_output_mode": structured_output_mode,
         "entity_types": kwargs.get("entity_types"),
+        **_generation_cache_params(kwargs),
     }
     cached_result = _result_cache.get("entities", text, **cache_params)
     if cached_result is not None:
@@ -1706,7 +1748,8 @@ def extract_relations_llm(
         "relation_types": kwargs.get("relation_types"),
         "extract_temporal_bounds": extract_temporal_bounds,
         # Include entities hash/str in cache key implicitly via **cache_params
-        "entities_hash": hash(tuple(sorted([e.text for e in entities]))) if entities else 0
+        "entities_hash": hash(tuple(sorted([e.text for e in entities]))) if entities else 0,
+        **_generation_cache_params(kwargs),
     }
     cached_result = _result_cache.get("relations", text, **cache_params)
     if cached_result is not None:
@@ -1906,13 +1949,10 @@ Entities found in text: {entities_str}"""
                 "[methods.extract_relations_llm] Calling llm.generate_typed (%s/%s)...",
                 provider, model,
             )
-        # Only forward minimal, safe parameters to provider calls
-        call_kwargs = {}
-        if "temperature" in kwargs:
-            call_kwargs["temperature"] = kwargs["temperature"]
-        if "verbose" in kwargs:
-            call_kwargs["verbose"] = kwargs["verbose"]
-
+        # Forward all caller-supplied generation kwargs so they reach
+        # generate_typed and the underlying provider API. max_retries is
+        # always set from the explicit parameter.
+        call_kwargs = kwargs.copy()
         call_kwargs["max_retries"] = max_retries
 
         # Select schema based on whether temporal extraction is requested
@@ -2364,7 +2404,8 @@ def extract_triplets_llm(
         "triplet_types": kwargs.get("triplet_types"),
         # Include entities/relations hash in cache key implicitly via **cache_params
         "entities_hash": hash(tuple(sorted([e.text for e in entities]))) if entities else 0,
-        "relations_hash": hash(tuple(sorted([str(r) for r in relations]))) if relations else 0
+        "relations_hash": hash(tuple(sorted([str(r) for r in relations]))) if relations else 0,
+        **_generation_cache_params(kwargs),
     }
     cached_result = _result_cache.get("triplets", text, **cache_params)
     if cached_result is not None:

tests/reproduce_issue_176.py

@@ -96,5 +96,239 @@ class TestMaxTokensPropagation(unittest.TestCase):
         self.assertIn("max_tokens", kwargs)
         self.assertEqual(kwargs["max_tokens"], 128000)
 
+
+class TestCacheKeyIncludesGenerationParams(unittest.TestCase):
+    """Regression tests for the cache-key bug: two calls with identical extraction
+    inputs but different generation settings must NOT share a cached result.
+
+    Before the fix, extract_relations_llm (and entities/triplets) built
+    cache_params without generation kwargs, so max_tokens=4096 and
+    max_tokens=128000 hashed to the same key. The second call would return the
+    first cached result without ever running generate_typed again.
+    """
+
+    def _make_mock_llm(self, relations=None, entities=None, triplets=None):
+        mock_llm = MagicMock()
+        mock_llm.is_available.return_value = True
+        resp = MagicMock()
+        resp.relations = relations if relations is not None else []
+        resp.entities = entities if entities is not None else []
+        resp.triplets = triplets if triplets is not None else []
+        mock_llm.generate_typed.return_value = resp
+        return mock_llm
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_different_max_tokens_bypass_cache(self, mock_create_provider):
+        """Two relation extraction calls with the same text/entities but different
+        max_tokens must each call generate_typed (2 calls total), not reuse the
+        first cached result."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="openai", model="gpt-4", max_tokens=4096
+        )
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="openai", model="gpt-4", max_tokens=128000
+        )
+
+        # generate_typed must have been called twice — once per unique key
+        self.assertEqual(
+            mock_llm.generate_typed.call_count, 2,
+            "Different max_tokens values must produce different cache keys; "
+            "second call must not reuse the first cached result."
+        )
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_same_max_tokens_uses_cache(self, mock_create_provider):
+        """Two identical calls must reuse the cache (generate_typed called once)."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="openai", model="gpt-4", max_tokens=4096
+        )
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="openai", model="gpt-4", max_tokens=4096
+        )
+
+        self.assertEqual(
+            mock_llm.generate_typed.call_count, 1,
+            "Identical calls must reuse the cache."
+        )
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_different_temperature_bypass_cache(self, mock_create_provider):
+        """Different temperature values must also produce different cache keys."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Bar", label="PERSON", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="other text", entities=entities,
+            provider="openai", model="gpt-4", temperature=0.0
+        )
+        extract_relations_llm(
+            text="other text", entities=entities,
+            provider="openai", model="gpt-4", temperature=1.0
+        )
+
+        self.assertEqual(mock_llm.generate_typed.call_count, 2)
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_entities_different_max_tokens_bypass_cache(self, mock_create_provider):
+        """extract_entities_llm: different max_tokens must bypass cache."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("entities")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        extract_entities_llm(
+            text="some entity text", provider="openai", model="gpt-4",
+            max_tokens=4096
+        )
+        extract_entities_llm(
+            text="some entity text", provider="openai", model="gpt-4",
+            max_tokens=128000
+        )
+
+        self.assertEqual(mock_llm.generate_typed.call_count, 2)
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_triplets_different_max_tokens_bypass_cache(self, mock_create_provider):
+        """extract_triplets_llm: different max_tokens must bypass cache."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("triplets")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        extract_triplets_llm(
+            text="some triplet text", provider="openai", model="gpt-4",
+            max_tokens=4096
+        )
+        extract_triplets_llm(
+            text="some triplet text", provider="openai", model="gpt-4",
+            max_tokens=128000
+        )
+
+        self.assertEqual(mock_llm.generate_typed.call_count, 2)
+
+
+class TestCacheKeyIncludesProviderSpecificGenerationParams(unittest.TestCase):
+    """Regression tests for provider-specific generation params that aren't part
+    of the common OpenAI-shaped kwargs (max_tokens, temperature, etc.) but still
+    change provider output and must therefore also change the cache key.
+
+    See providers.py: AnthropicProvider.generate/generate_structured read
+    'system' and 'stop_sequences' via a manual pass-through loop (not
+    _add_if_set); GeminiProvider.generate reads 'candidate_count' and
+    'stop_sequences'; OllamaProvider._build_options reads 'repeat_penalty' and
+    'num_ctx'/'context_window'.
+    """
+
+    def _make_mock_llm(self):
+        mock_llm = MagicMock()
+        mock_llm.is_available.return_value = True
+        resp = MagicMock()
+        resp.relations = []
+        mock_llm.generate_typed.return_value = resp
+        return mock_llm
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_different_system_prompt_bypass_cache(self, mock_create_provider):
+        """Anthropic 'system' prompt changes output; must not share a cache entry."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="anthropic", model="claude-3-sonnet-20240229",
+            system="Extract only ORG relations."
+        )
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="anthropic", model="claude-3-sonnet-20240229",
+            system="Extract only PERSON relations."
+        )
+
+        self.assertEqual(
+            mock_llm.generate_typed.call_count, 2,
+            "Different 'system' prompts must produce different cache keys."
+        )
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_different_stop_sequences_bypass_cache(self, mock_create_provider):
+        """Anthropic/Gemini 'stop_sequences' must also be part of the cache key."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="anthropic", model="claude-3-sonnet-20240229",
+            stop_sequences=["\n\n"]
+        )
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="anthropic", model="claude-3-sonnet-20240229",
+            stop_sequences=["STOP"]
+        )
+
+        self.assertEqual(mock_llm.generate_typed.call_count, 2)
+
+    @patch("semantica.semantic_extract.methods.create_provider")
+    def test_relations_different_repeat_penalty_bypass_cache(self, mock_create_provider):
+        """Ollama 'repeat_penalty' must also be part of the cache key."""
+        from semantica.semantic_extract.methods import _result_cache
+        _result_cache.clear("relations")
+
+        mock_llm = self._make_mock_llm()
+        mock_create_provider.return_value = mock_llm
+
+        entities = [Entity(text="Foo", label="ORG", start_char=0, end_char=3)]
+
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="ollama", model="llama2",
+            repeat_penalty=1.0
+        )
+        extract_relations_llm(
+            text="some text", entities=entities,
+            provider="ollama", model="llama2",
+            repeat_penalty=1.5
+        )
+
+        self.assertEqual(mock_llm.generate_typed.call_count, 2)
+
+
 if __name__ == "__main__":
     unittest.main()