diff --git a/semantica/semantic_extract/methods.py b/semantica/semantic_extract/methods.py index 39144d7c..00482df1 100644 --- a/semantica/semantic_extract/methods.py +++ b/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: diff --git a/tests/reproduce_issue_176.py b/tests/reproduce_issue_176.py index a24ba7c8..79b24076 100644 --- a/tests/reproduce_issue_176.py +++ b/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()