perf: share few-shot preamble across prompts via PromptParts

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,30 +146,49 @@ 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.
       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.
+      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,

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:

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,

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:

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)