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, diff --git a/langextract/providers/gemini.py b/langextract/providers/gemini.py index 18416820..b07a58c2 100644 --- a/langextract/providers/gemini.py +++ b/langextract/providers/gemini.py @@ -352,6 +352,8 @@ def _process_single_prompt( self, prompt: str, config: dict ) -> core_types.ScoredOutput: """Run one Gemini request with per-chunk retries for transient failures.""" + if not isinstance(prompt, str): + prompt = str(prompt) delay = self.retry_delay for attempt in range(self.max_retries + 1): try: diff --git a/langextract/providers/gemini_batch.py b/langextract/providers/gemini_batch.py index 4aa6eb6d..a9dd1111 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_config: 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_config: Optional provider schema config for structured output, as produced by GeminiSchema.to_provider_config(). Supports response_json_schema (JSON Schema) and response_schema @@ -282,12 +287,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}]} @@ -406,9 +423,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, diff --git a/langextract/providers/ollama.py b/langextract/providers/ollama.py index 2e3470a8..5a94a383 100644 --- a/langextract/providers/ollama.py +++ b/langextract/providers/ollama.py @@ -606,6 +606,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 86974a49..ddbdc1ed 100644 --- a/langextract/providers/openai.py +++ b/langextract/providers/openai.py @@ -245,6 +245,8 @@ def _process_single_prompt( self, prompt: str, config: dict ) -> core_types.ScoredOutput: """Sends one prompt while preserving provider-specific error types.""" + if not isinstance(prompt, str): + prompt = str(prompt) try: api_params = self._build_chat_completions_params(prompt, config) response = self._client.chat.completions.create(**api_params) 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 001caccd..8e03aa30 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()