From a76a4a3a191a9245b90a35f86c0efff70ccd7311 Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Fri, 17 Apr 2026 19:02:58 +0000 Subject: [PATCH 1/7] perf(prompting): add PromptParts to share few-shot preamble across prompts Add a frozen PromptParts dataclass that splits rendered prompts into prefix (description + context), examples (large, shared by reference), and suffix (question + answer prefix). QAPromptGenerator caches the formatted examples text in __post_init__ and exposes render_parts() which returns a PromptParts whose examples field is always the same string object. render() is reimplemented as str(render_parts(...)). PromptBuilder.build_prompt() and ContextAwarePromptBuilder.build_prompt() now return PromptParts instead of str, so downstream consumers receive structured prompts that share the large examples allocation. Closes #446 --- langextract/prompting.py | 93 +++++++++++++++++++++++++++++++--------- 1 file changed, 73 insertions(+), 20 deletions(-) diff --git a/langextract/prompting.py b/langextract/prompting.py index cdf1c827..d75c0ec1 100644 --- a/langextract/prompting.py +++ b/langextract/prompting.py @@ -28,6 +28,31 @@ from langextract.core import format_handler +@dataclasses.dataclass(frozen=True, slots=True) +class PromptParts: + """Structured prompt with shared and unique parts for memory efficiency. + + When many prompts share the same few-shot examples text, storing it as a + shared reference avoids O(N * examples_size) memory for N prompts. + + Attributes: + prefix: Description and optional context (small, unique per prompt). + examples: Formatted few-shot examples (large, shared across prompts). + suffix: Question text and answer prefix (small, unique per prompt). + """ + + prefix: str + examples: str + suffix: str + + def __str__(self) -> str: + parts = [self.prefix] + if self.examples: + parts.append(self.examples) + parts.append(self.suffix) + return "\n".join(parts) + + class PromptBuilderError(exceptions.LangExtractError): """Failure to build prompt.""" @@ -91,6 +116,15 @@ class QAPromptGenerator: question_prefix: str = "Q: " answer_prefix: str = "A: " + def __post_init__(self): + if self.template.examples: + lines = [self.examples_heading] + for ex in self.template.examples: + lines.append(self.format_example_as_text(ex)) + self._examples_text = "\n".join(lines) + else: + self._examples_text = "" + def __str__(self) -> str: """Returns a string representation of the prompt with an empty question.""" return self.render("") @@ -112,8 +146,13 @@ def format_example_as_text(self, example: data.ExampleData) -> str: f"{self.answer_prefix}{answer}\n", ]) - def render(self, question: str, additional_context: str | None = None) -> str: - """Generate a text representation of the prompt. + def render_parts( + self, question: str, additional_context: str | None = None + ) -> PromptParts: + """Generate a structured prompt split into shared and unique parts. + + The examples text is cached and shared across all calls, avoiding + O(N * examples_size) memory when building many prompts. Args: question: That will be presented to the model. @@ -121,21 +160,35 @@ def render(self, question: str, additional_context: str | None = None) -> str: string is ignored. Returns: - Text prompt with a question to be presented to a language model. + PromptParts with prefix, examples, and suffix. """ - prompt_lines: list[str] = [f"{self.template.description}\n"] - + prefix_lines: list[str] = [f"{self.template.description}\n"] if additional_context: - prompt_lines.append(f"{additional_context}\n") + prefix_lines.append(f"{additional_context}\n") - if self.template.examples: - prompt_lines.append(self.examples_heading) - for ex in self.template.examples: - prompt_lines.append(self.format_example_as_text(ex)) + suffix = "\n".join([ + f"{self.question_prefix}{question}", + self.answer_prefix, + ]) + + return PromptParts( + prefix="\n".join(prefix_lines), + examples=self._examples_text, + suffix=suffix, + ) + + def render(self, question: str, additional_context: str | None = None) -> str: + """Generate a text representation of the prompt. - prompt_lines.append(f"{self.question_prefix}{question}") - prompt_lines.append(self.answer_prefix) - return "\n".join(prompt_lines) + Args: + question: That will be presented to the model. + additional_context: Additional context to include in the prompt. An empty + string is ignored. + + Returns: + Text prompt with a question to be presented to a language model. + """ + return str(self.render_parts(question, additional_context)) class PromptBuilder: @@ -158,7 +211,7 @@ def build_prompt( chunk_text: str, document_id: str, additional_context: str | None = None, - ) -> str: + ) -> PromptParts: """Builds a prompt for the given chunk. Args: @@ -167,10 +220,10 @@ def build_prompt( additional_context: Optional additional context from the document. Returns: - The rendered prompt string ready for the language model. + PromptParts with shared examples and unique prefix/suffix. """ del document_id # Unused in base class. - return self._generator.render( + return self._generator.render_parts( question=chunk_text, additional_context=additional_context, ) @@ -217,7 +270,7 @@ def build_prompt( chunk_text: str, document_id: str, additional_context: str | None = None, - ) -> str: + ) -> PromptParts: """Builds a prompt, injecting previous chunk context if enabled. Args: @@ -227,17 +280,17 @@ def build_prompt( additional_context: Optional additional context from the document. Returns: - The rendered prompt string ready for the language model. + PromptParts with shared examples and unique prefix/suffix. """ effective_context = self._build_effective_context( document_id, additional_context ) - prompt = self._generator.render( + prompt_parts = self._generator.render_parts( question=chunk_text, additional_context=effective_context, ) self._update_state(document_id, chunk_text) - return prompt + return prompt_parts def _build_effective_context( self, From ba28efc05bc712ef9502811c7f87aa1d4fbb8d63 Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Fri, 17 Apr 2026 19:03:14 +0000 Subject: [PATCH 2/7] perf(providers): use multi-part contents for PromptParts in batch API When _build_request receives a PromptParts with non-empty examples, emit three text parts in contents[0].parts instead of one concatenated string. The middle part holds the shared examples reference, so 10,000 requests share one ~300 KB string instead of duplicating it per request. Gemini single-prompt, OpenAI, and Ollama providers convert PromptParts to str at their entry points; since they process prompts one at a time (or in small thread pools), the temporary string has negligible memory impact. --- langextract/providers/gemini.py | 2 ++ langextract/providers/gemini_batch.py | 25 +++++++++++++++++++++---- langextract/providers/ollama.py | 2 ++ langextract/providers/openai.py | 2 ++ 4 files changed, 27 insertions(+), 4 deletions(-) diff --git a/langextract/providers/gemini.py b/langextract/providers/gemini.py index a82afe1e..91fb8cb2 100644 --- a/langextract/providers/gemini.py +++ b/langextract/providers/gemini.py @@ -203,6 +203,8 @@ def _process_single_prompt( self, prompt: str, config: dict ) -> core_types.ScoredOutput: """Process a single prompt and return a ScoredOutput.""" + if not isinstance(prompt, str): + prompt = str(prompt) try: # Apply stored kwargs that weren't already set in config for key, value in self._extra_kwargs.items(): diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index 220a262d..de8d120b 100644 --- a/langextract/providers/gemini_batch.py +++ b/langextract/providers/gemini_batch.py @@ -44,6 +44,7 @@ from google.cloud import storage from langextract.core import exceptions +from langextract.prompting import PromptParts _MIME_TYPE_JSON = "application/json" _DEFAULT_LOCATION = "us-central1" @@ -257,7 +258,7 @@ def _ensure_bucket_lifecycle( def _build_request( - prompt: str, + prompt: str | PromptParts, schema_dict: dict | None, gen_config: dict | None, system_instruction: str | None = None, @@ -270,8 +271,12 @@ def _build_request( can include its own generationConfig with schema and generation parameters, as well as top-level systemInstruction and safetySettings. + When prompt is a PromptParts with non-empty examples, uses multi-part + contents so the large examples string can be shared by reference across + requests, avoiding O(N * examples_size) memory. + Args: - prompt: The text prompt to send to the model. + prompt: The text prompt or structured PromptParts. schema_dict: Optional JSON schema for structured output. gen_config: Optional generation configuration parameters. system_instruction: Optional system instruction text. @@ -279,12 +284,24 @@ def _build_request( Returns: A dictionary formatted for REST API file-based submission, containing: - * contents: The prompt content. + * contents: The prompt content (multi-part when PromptParts). * systemInstruction: Optional system instructions. * safetySettings: Optional safety settings. * generationConfig: Optional generation configuration and schema. """ - request = {"contents": [{"role": "user", "parts": [{"text": prompt}]}]} + if isinstance(prompt, PromptParts) and prompt.examples: + # Multi-part content: keeps the large examples string as a shared + # Python reference across all requests in the batch. The \n placement + # matches the separator that PromptParts.__str__ uses via "\n".join(). + parts = [ + {"text": prompt.prefix + "\n"}, + {"text": prompt.examples}, + {"text": "\n" + prompt.suffix}, + ] + request = {"contents": [{"role": "user", "parts": parts}]} + else: + text = str(prompt) if not isinstance(prompt, str) else prompt + request = {"contents": [{"role": "user", "parts": [{"text": text}]}]} if system_instruction: request["systemInstruction"] = {"parts": [{"text": system_instruction}]} diff --git a/langextract/providers/ollama.py b/langextract/providers/ollama.py index c6be9379..c145c605 100644 --- a/langextract/providers/ollama.py +++ b/langextract/providers/ollama.py @@ -340,6 +340,8 @@ def _ollama_query( InferenceRuntimeError: For any other HTTP errors, timeouts, or request exceptions. """ + if not isinstance(prompt, str): + prompt = str(prompt) model = model or self._model model_url = model_url or self._model_url if structured_output_format is None and self.format_type is not None: diff --git a/langextract/providers/openai.py b/langextract/providers/openai.py index 3dde40cf..ea187e27 100644 --- a/langextract/providers/openai.py +++ b/langextract/providers/openai.py @@ -118,6 +118,8 @@ def _process_single_prompt( self, prompt: str, config: dict ) -> core_types.ScoredOutput: """Process a single prompt and return a ScoredOutput.""" + if not isinstance(prompt, str): + prompt = str(prompt) try: normalized_config = config.copy() From 8d73b08e201110cf5422dfe5468a184d09eaae5e Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Fri, 17 Apr 2026 19:03:38 +0000 Subject: [PATCH 3/7] test: update tests for PromptParts return type Update PromptBuilder, ContextAwarePromptBuilder, Annotator, and extract() tests to work with PromptParts instead of plain strings. Add test_build_prompt_shares_examples_reference and test_context_aware_shares_examples_reference to verify the memory- sharing invariant (all prompts from the same generator share a single examples string object via `assertIs`). --- tests/annotation_test.py | 67 ++++++++++++++++++++++++++-------- tests/init_test.py | 2 +- tests/prompting_test.py | 79 +++++++++++++++++++++++++++++----------- 3 files changed, 110 insertions(+), 38 deletions(-) diff --git a/tests/annotation_test.py b/tests/annotation_test.py index f0bd9f50..4f25e3a0 100644 --- a/tests/annotation_test.py +++ b/tests/annotation_test.py @@ -200,7 +200,13 @@ def test_annotate_text_single_chunk(self): text, actual_annotated_text.extractions ) self.mock_language_model.infer.assert_called_once_with( - batch_prompts=[f"\n\nQ: {text}\nA: "], + batch_prompts=[ + prompting.PromptParts( + prefix="\n", + examples="", + suffix=f"Q: {text}\nA: ", + ) + ], ) def test_annotate_text_without_index_suffix(self): @@ -319,7 +325,13 @@ def test_annotate_text_without_index_suffix(self): text, actual_annotated_text.extractions ) self.mock_language_model.infer.assert_called_once_with( - batch_prompts=[f"\n\nQ: {text}\nA: "], + batch_prompts=[ + prompting.PromptParts( + prefix="\n", + examples="", + suffix=f"Q: {text}\nA: ", + ) + ], ) def test_annotate_text_with_attributes_suffix(self): @@ -463,7 +475,13 @@ def test_annotate_text_with_attributes_suffix(self): text, actual_annotated_text.extractions ) self.mock_language_model.infer.assert_called_once_with( - batch_prompts=[f"\n\nQ: {text}\nA: "], + batch_prompts=[ + prompting.PromptParts( + prefix="\n", + examples="", + suffix=f"Q: {text}\nA: ", + ) + ], ) def test_annotate_text_multiple_chunks(self): @@ -556,12 +574,22 @@ def test_annotate_text_multiple_chunks(self): self.mock_language_model.infer.assert_has_calls([ mock.call( batch_prompts=[ - "\n\nQ: Patient takes one Aspirin for headaches.\nA: " + prompting.PromptParts( + prefix="\n", + examples="", + suffix="Q: Patient takes one Aspirin for headaches.\nA: ", + ) ], enable_fuzzy_alignment=False, ), mock.call( - batch_prompts=["\n\nQ: Pt has fever.\nA: "], + batch_prompts=[ + prompting.PromptParts( + prefix="\n", + examples="", + suffix="Q: Pt has fever.\nA: ", + ) + ], enable_fuzzy_alignment=False, ), ]) @@ -588,7 +616,13 @@ def test_annotate_text_no_extractions(self): ) self.assertDataclassEqual(expected_annotated_text, actual_annotated_text) self.mock_language_model.infer.assert_called_once_with( - batch_prompts=[f"\n\nQ: {text}\nA: "], + batch_prompts=[ + prompting.PromptParts( + prefix="\n", + examples="", + suffix=f"Q: {text}\nA: ", + ) + ], ) @@ -1132,14 +1166,15 @@ def setUp(self): def mock_infer(batch_prompts, **_): """Return medication extractions based on prompt content.""" for prompt in batch_prompts: - if "Ibuprofen" in prompt: + prompt_str = str(prompt) + if "Ibuprofen" in prompt_str: text = textwrap.dedent(f"""\ ```yaml {data.EXTRACTIONS_KEY}: - medication: "Ibuprofen" medication_index: 4 ```""") - elif "Cefazolin" in prompt: + elif "Cefazolin" in prompt_str: text = textwrap.dedent(f"""\ ```yaml {data.EXTRACTIONS_KEY}: @@ -1260,11 +1295,11 @@ def test_context_window_includes_previous_chunk_text(self): calls = self.mock_language_model.infer.call_args_list self.assertLen(calls, 2) - first_prompt = calls[0].kwargs["batch_prompts"][0] + first_prompt = str(calls[0].kwargs["batch_prompts"][0]) context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX self.assertNotIn(context_prefix, first_prompt) - second_prompt = calls[1].kwargs["batch_prompts"][0] + second_prompt = str(calls[1].kwargs["batch_prompts"][0]) self.assertIn(context_prefix, second_prompt) self.assertIn("cardiologist", second_prompt) @@ -1300,8 +1335,8 @@ def test_no_context_included_when_disabled(self): self.assertLen(calls, 2) context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX - first_prompt = calls[0].kwargs["batch_prompts"][0] - second_prompt = calls[1].kwargs["batch_prompts"][0] + first_prompt = str(calls[0].kwargs["batch_prompts"][0]) + second_prompt = str(calls[1].kwargs["batch_prompts"][0]) self.assertNotIn(context_prefix, first_prompt) self.assertNotIn(context_prefix, second_prompt) @@ -1341,10 +1376,10 @@ def test_context_window_per_document_isolation(self): context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX # Extract prompts in order: doc1_chunk1, doc1_chunk2, doc2_chunk1, doc2_chunk2 - doc1_chunk1_prompt = calls[0].kwargs["batch_prompts"][0] - doc1_chunk2_prompt = calls[1].kwargs["batch_prompts"][0] - doc2_chunk1_prompt = calls[2].kwargs["batch_prompts"][0] - doc2_chunk2_prompt = calls[3].kwargs["batch_prompts"][0] + doc1_chunk1_prompt = str(calls[0].kwargs["batch_prompts"][0]) + doc1_chunk2_prompt = str(calls[1].kwargs["batch_prompts"][0]) + doc2_chunk1_prompt = str(calls[2].kwargs["batch_prompts"][0]) + doc2_chunk2_prompt = str(calls[3].kwargs["batch_prompts"][0]) # First chunks of each document should NOT have context prefix self.assertNotIn(context_prefix, doc1_chunk1_prompt) diff --git a/tests/init_test.py b/tests/init_test.py index a1e124fe..578e3db1 100644 --- a/tests/init_test.py +++ b/tests/init_test.py @@ -149,7 +149,7 @@ def test_lang_extract_as_lx_extract( mock_gemini_schema.assert_not_called() mock_create_model.assert_called_once() mock_model.infer.assert_called_once_with( - batch_prompts=[prompt_generator.render(input_text)], + batch_prompts=[prompt_generator.render_parts(input_text)], max_workers=10, ) diff --git a/tests/prompting_test.py b/tests/prompting_test.py index 37e03ab9..688d6c0a 100644 --- a/tests/prompting_test.py +++ b/tests/prompting_test.py @@ -439,8 +439,10 @@ def test_build_prompt_renders_chunk_text(self): document_id="doc1", ) - self.assertIn("Test input text.", prompt) - self.assertIn("Extract entities.", prompt) + self.assertIsInstance(prompt, prompting.PromptParts) + prompt_str = str(prompt) + self.assertIn("Test input text.", prompt_str) + self.assertIn("Extract entities.", prompt_str) def test_build_prompt_includes_additional_context(self): """Verifies build_prompt passes additional_context to renderer.""" @@ -453,7 +455,18 @@ def test_build_prompt_includes_additional_context(self): additional_context="Important context here.", ) - self.assertIn("Important context here.", prompt) + prompt_str = str(prompt) + self.assertIn("Important context here.", prompt_str) + + def test_build_prompt_shares_examples_reference(self): + """Verifies all prompts share the same examples string object.""" + generator = self._create_generator() + builder = prompting.PromptBuilder(generator) + + p1 = builder.build_prompt(chunk_text="Chunk 1.", document_id="doc1") + p2 = builder.build_prompt(chunk_text="Chunk 2.", document_id="doc1") + + self.assertIs(p1.examples, p2.examples) class ContextAwarePromptBuilderTest(absltest.TestCase): @@ -511,8 +524,10 @@ def test_first_chunk_has_no_previous_context(self): document_id="doc1", ) - self.assertNotIn(context_prefix, prompt) - self.assertIn("First chunk text.", prompt) + self.assertIsInstance(prompt, prompting.PromptParts) + prompt_str = str(prompt) + self.assertNotIn(context_prefix, prompt_str) + self.assertIn("First chunk text.", prompt_str) def test_second_chunk_includes_previous_context(self): """Verifies the second chunk includes text from the first chunk.""" @@ -523,9 +538,11 @@ def test_second_chunk_includes_previous_context(self): context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX builder.build_prompt(chunk_text="First chunk ending.", document_id="doc1") - second_prompt = builder.build_prompt( - chunk_text="Second chunk text.", - document_id="doc1", + second_prompt = str( + builder.build_prompt( + chunk_text="Second chunk text.", + document_id="doc1", + ) ) self.assertIn(context_prefix, second_prompt) @@ -540,9 +557,11 @@ def test_context_disabled_when_none(self): context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX builder.build_prompt(chunk_text="First chunk.", document_id="doc1") - second_prompt = builder.build_prompt( - chunk_text="Second chunk.", - document_id="doc1", + second_prompt = str( + builder.build_prompt( + chunk_text="Second chunk.", + document_id="doc1", + ) ) self.assertNotIn(context_prefix, second_prompt) @@ -557,13 +576,17 @@ def test_context_isolated_per_document(self): builder.build_prompt(chunk_text="Doc A chunk one.", document_id="docA") builder.build_prompt(chunk_text="Doc B chunk one.", document_id="docB") - prompt_a2 = builder.build_prompt( - chunk_text="Doc A chunk two.", - document_id="docA", + prompt_a2 = str( + builder.build_prompt( + chunk_text="Doc A chunk two.", + document_id="docA", + ) ) - prompt_b2 = builder.build_prompt( - chunk_text="Doc B chunk two.", - document_id="docB", + prompt_b2 = str( + builder.build_prompt( + chunk_text="Doc B chunk two.", + document_id="docB", + ) ) self.assertIn("Doc A chunk one", prompt_a2) @@ -580,16 +603,30 @@ def test_combines_previous_context_with_additional_context(self): context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX builder.build_prompt(chunk_text="Previous chunk text.", document_id="doc1") - prompt = builder.build_prompt( - chunk_text="Current chunk.", - document_id="doc1", - additional_context="Extra info here.", + prompt = str( + builder.build_prompt( + chunk_text="Current chunk.", + document_id="doc1", + additional_context="Extra info here.", + ) ) self.assertIn(context_prefix, prompt) self.assertIn("Previous chunk text.", prompt) self.assertIn("Extra info here.", prompt) + def test_context_aware_shares_examples_reference(self): + """Verifies context-aware builder shares examples across prompts.""" + generator = self._create_generator() + builder = prompting.ContextAwarePromptBuilder( + generator, context_window_chars=50 + ) + + p1 = builder.build_prompt(chunk_text="Chunk 1.", document_id="doc1") + p2 = builder.build_prompt(chunk_text="Chunk 2.", document_id="doc1") + + self.assertIs(p1.examples, p2.examples) + if __name__ == "__main__": absltest.main() From 221917c799edcb185ce346666d7ac856f3c96814 Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Fri, 17 Apr 2026 19:30:38 +0000 Subject: [PATCH 4/7] fix(gemini_batch): preserve cache key compatibility with PromptParts Convert PromptParts to str before inserting into cache key_data dicts so that the SHA256 hash matches the old string-based format. This avoids a full cache miss on upgrade. The str() call creates one temporary string per prompt, processed sequentially, so peak memory is unchanged. --- langextract/providers/gemini_batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index de8d120b..2ad7462c 100644 --- a/langextract/providers/gemini_batch.py +++ b/langextract/providers/gemini_batch.py @@ -799,7 +799,7 @@ def infer_batch( for prompt in prompts: key_data_list.append({ "model_id": model_id, - "prompt": prompt, + "prompt": str(prompt) if not isinstance(prompt, str) else prompt, "system_instruction": system_instruction, "gen_config": gen_config, "safety_settings": safety_settings, @@ -890,7 +890,7 @@ def _process_batch( prompt = prompts[idx] key_data = { "model_id": model_id, - "prompt": prompt, + "prompt": str(prompt) if not isinstance(prompt, str) else prompt, "system_instruction": system_instruction, "gen_config": gen_config, "safety_settings": safety_settings, From 1d16e5859af439f0ed0257299e85502fbbeee709 Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Wed, 22 Apr 2026 19:09:59 +0000 Subject: [PATCH 5/7] Revert "fix(gemini_batch): preserve cache key compatibility with PromptParts" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The str(prompt) conversion in key_data dicts negates the memory optimization from PromptParts by materializing 10,000 × ~640 KB concatenated strings in key_data_list (6.4 GB). Without it, PromptParts serializes via dataclasses.asdict in _json_default, keeping the shared examples reference intact. Cache keys will differ from pre-PromptParts entries, but those expire via GCS lifecycle (retention_days) anyway. This reverts commit b9c1238. --- langextract/providers/gemini_batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index 2ad7462c..de8d120b 100644 --- a/langextract/providers/gemini_batch.py +++ b/langextract/providers/gemini_batch.py @@ -799,7 +799,7 @@ def infer_batch( for prompt in prompts: key_data_list.append({ "model_id": model_id, - "prompt": str(prompt) if not isinstance(prompt, str) else prompt, + "prompt": prompt, "system_instruction": system_instruction, "gen_config": gen_config, "safety_settings": safety_settings, @@ -890,7 +890,7 @@ def _process_batch( prompt = prompts[idx] key_data = { "model_id": model_id, - "prompt": str(prompt) if not isinstance(prompt, str) else prompt, + "prompt": prompt, "system_instruction": system_instruction, "gen_config": gen_config, "safety_settings": safety_settings, From b672754bc7fbd9bfbc33a4bfa16876480265892d Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Wed, 22 Apr 2026 19:46:38 +0000 Subject: [PATCH 6/7] fix(gemini_batch): resolve non-primitive key_data values in _compute_hash MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Convert non-primitive values (e.g. PromptParts) to str inside _compute_hash rather than at key_data construction time. This keeps PromptParts references in key_data_list (shared examples, ~0.4 MB) while producing hashes identical to the old string-based format (cache compat preserved). Only one transient str copy exists at a time during sequential hash computation. Replaces the reverted str(prompt) approach which materialized all prompts upfront in key_data_list, negating the PromptParts memory optimization (10,000 × 640 KB = 6.4 GB). --- langextract/providers/gemini_batch.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index de8d120b..7dfbe195 100644 --- a/langextract/providers/gemini_batch.py +++ b/langextract/providers/gemini_batch.py @@ -413,9 +413,20 @@ def __init__(self, bucket_name: str, project: str | None = None): self._bucket = self._client.bucket(bucket_name) def _compute_hash(self, key_data: dict) -> str: - """Compute SHA256 hash of the canonicalized request data.""" + """Compute SHA256 hash of the canonicalized request data. + + Non-primitive values (e.g. PromptParts) are converted to str before + serialization so the hash matches the format produced when prompts + were plain strings. The conversion is transient — only one + stringified copy exists at a time — so it does not affect peak memory. + """ + _PRIMITIVES = (str, int, float, bool, type(None), dict, list) + resolved = { + k: str(v) if not isinstance(v, _PRIMITIVES) else v + for k, v in key_data.items() + } canonical_json = json.dumps( - key_data, + resolved, sort_keys=True, ensure_ascii=False, default=_json_default, From 9529d127381397bf2eb732abc8684739e2697e90 Mon Sep 17 00:00:00 2001 From: Dan Robert Date: Tue, 21 Jul 2026 15:22:22 +0000 Subject: [PATCH 7/7] fix(gemini_batch): rename _PRIMITIVES to lowercase for pylint compliance Local variable inside method was using module-level constant naming convention, triggering C0103 invalid-name lint error. --- langextract/providers/gemini_batch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index 937135e4..a9dd1111 100644 --- a/langextract/providers/gemini_batch.py +++ b/langextract/providers/gemini_batch.py @@ -430,9 +430,9 @@ def _compute_hash(self, key_data: dict) -> str: were plain strings. The conversion is transient — only one stringified copy exists at a time — so it does not affect peak memory. """ - _PRIMITIVES = (str, int, float, bool, type(None), dict, list) + primitives = (str, int, float, bool, type(None), dict, list) resolved = { - k: str(v) if not isinstance(v, _PRIMITIVES) else v + k: str(v) if not isinstance(v, primitives) else v for k, v in key_data.items() } canonical_json = json.dumps(