llama-stack 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

llama_stack/providers/utils/inference/openai_mixin.py

@@ -10,11 +10,16 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterable
 from typing import Any
 
+import httpx
 from openai import AsyncOpenAI
 from pydantic import BaseModel, ConfigDict
 
 from llama_stack.core.request_headers import NeedsRequestProviderData
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.inference.http_client import (
+    _build_network_client_kwargs,
+    _merge_network_config_into_client,
+)
 from llama_stack.providers.utils.inference.model_registry import RemoteInferenceProviderConfig
 from llama_stack.providers.utils.inference.openai_compat import (
     get_stream_options_for_telemetry,
@@ -34,6 +39,7 @@ from llama_stack_api import (
     OpenAIEmbeddingsResponse,
     OpenAIEmbeddingUsage,
     OpenAIMessageParam,
+    validate_embeddings_input_is_text,
 )
 
 logger = get_logger(name=__name__, category="providers::utils")
@@ -82,6 +88,10 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
     # Set to False for providers that don't support stream_options (e.g., Ollama, vLLM)
     supports_stream_options: bool = True
 
+    # Allow subclasses to control whether the provider supports tokenized embeddings input
+    # Set to True for providers that support pre-tokenized input (list[int] and list[list[int]])
+    supports_tokenized_embeddings_input: bool = False
+
     # Embedding model metadata for this provider
     # Can be set by subclasses or instances to provide embedding models
     # Format: {"model_id": {"embedding_dimension": 1536, "context_length": 8192}}
@@ -121,7 +131,10 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         Get any extra parameters to pass to the AsyncOpenAI client.
 
         Child classes can override this method to provide additional parameters
-        such as timeout settings, proxies, etc.
+        such as custom http_client, timeout settings, proxies, etc.
+
+        Note: Network configuration from config.network is automatically applied
+        in the client property. This method is for provider-specific customizations.
 
         :return: A dictionary of extra parameters
         """
@@ -194,6 +207,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         Uses the abstract methods get_api_key() and get_base_url() which must be
         implemented by child classes.
 
+        Network configuration from config.network is automatically applied.
         Users can also provide the API key via the provider data header, which
         is used instead of any config API key.
         """
@@ -205,10 +219,30 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
                 message += f' Please provide a valid API key in the provider data header, e.g. x-llamastack-provider-data: {{"{self.provider_data_api_key_field}": "<API_KEY>"}}.'
             raise ValueError(message)
 
+        extra_params = self.get_extra_client_params()
+        network_kwargs = _build_network_client_kwargs(self.config.network)
+
+        # Handle http_client creation/merging:
+        # - If get_extra_client_params() provides an http_client (e.g., OCI with custom auth),
+        #   merge network config into it. The merge behavior:
+        #   * Preserves auth from get_extra_client_params() (provider-specific auth like OCI signer)
+        #   * Preserves headers from get_extra_client_params() as base
+        #   * Applies network config (TLS, proxy, timeout, headers) on top
+        #   * Network config headers take precedence over provider headers (allows override)
+        # - Otherwise, if network config exists, create http_client from it
+        # This allows providers with custom auth to still use standard network settings
+        if "http_client" in extra_params:
+            if network_kwargs:
+                extra_params["http_client"] = _merge_network_config_into_client(
+                    extra_params["http_client"], self.config.network
+                )
+        elif network_kwargs:
+            extra_params["http_client"] = httpx.AsyncClient(**network_kwargs)
+
         return AsyncOpenAI(
             api_key=api_key,
             base_url=self.get_base_url(),
-            **self.get_extra_client_params(),
+            **extra_params,
         )
 
     def _get_api_key_from_config_or_provider_data(self) -> str | None:
@@ -371,6 +405,8 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
             top_logprobs=params.top_logprobs,
             top_p=params.top_p,
             user=params.user,
+            safety_identifier=params.safety_identifier,
+            reasoning_effort=params.reasoning_effort,
         )
 
         if extra_body := params.model_extra:
@@ -386,6 +422,10 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         """
         Direct OpenAI embeddings API call.
         """
+        # Validate token array support if provider doesn't support it
+        if not self.supports_tokenized_embeddings_input:
+            validate_embeddings_input_is_text(params)
+
         provider_model_id = await self._get_provider_model_id(params.model)
         self._validate_model_allowed(provider_model_id)
 

llama_stack/providers/utils/inference/prompt_adapter.py

@@ -4,63 +4,24 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import asyncio
 import base64
-import io
-import json
 import re
 from typing import Any
 
 import httpx
-from PIL import Image as PIL_Image
 
 from llama_stack.log import get_logger
-from llama_stack.models.llama.datatypes import (
-    RawContent,
-    RawContentItem,
-    RawMediaItem,
-    RawMessage,
-    RawTextItem,
-    StopReason,
-    ToolCall,
-    ToolDefinition,
-    ToolPromptFormat,
-)
-from llama_stack.models.llama.llama3.chat_format import ChatFormat
-from llama_stack.models.llama.llama3.tokenizer import Tokenizer
-from llama_stack.models.llama.sku_list import resolve_model
-from llama_stack.models.llama.sku_types import ModelFamily, is_multimodal
 from llama_stack_api import (
-    CompletionRequest,
     ImageContentItem,
-    InterleavedContent,
-    InterleavedContentItem,
-    OpenAIAssistantMessageParam,
     OpenAIChatCompletionContentPartImageParam,
     OpenAIChatCompletionContentPartTextParam,
     OpenAIFile,
-    OpenAIMessageParam,
-    OpenAISystemMessageParam,
-    OpenAIToolMessageParam,
-    OpenAIUserMessageParam,
-    ResponseFormat,
-    ResponseFormatType,
     TextContentItem,
-    ToolChoice,
 )
 
 log = get_logger(name=__name__, category="providers::utils")
 
 
-class CompletionRequestWithRawContent(CompletionRequest):
-    content: RawContent
-
-
-def decode_assistant_message(content: str, stop_reason: StopReason) -> RawMessage:
-    formatter = ChatFormat(Tokenizer.get_instance())
-    return formatter.decode_assistant_message_from_content(content, stop_reason)
-
-
 def interleaved_content_as_str(
     content: Any,
     sep: str = " ",
@@ -86,92 +47,6 @@ def interleaved_content_as_str(
         return _process(content)
 
 
-async def interleaved_content_convert_to_raw(
-    content: InterleavedContent,
-) -> RawContent:
-    """Download content from URLs / files etc. so plain bytes can be sent to the model"""
-
-    async def _localize_single(c: str | InterleavedContentItem) -> str | RawContentItem:
-        if isinstance(c, str):
-            return RawTextItem(text=c)
-        elif isinstance(c, TextContentItem):
-            return RawTextItem(text=c.text)
-        elif isinstance(c, ImageContentItem):
-            image = c.image
-            if image.url:
-                # Load image bytes from URL
-                if image.url.uri.startswith("data"):
-                    match = re.match(r"data:image/(\w+);base64,(.+)", image.url.uri)
-                    if not match:
-                        raise ValueError(f"Invalid data URL format, {image.url.uri[:40]}...")
-                    _, image_data = match.groups()
-                    data = base64.b64decode(image_data)
-                elif image.url.uri.startswith("file://"):
-                    path = image.url.uri[len("file://") :]
-                    with open(path, "rb") as f:
-                        data = f.read()  # type: ignore
-                elif image.url.uri.startswith("http"):
-                    async with httpx.AsyncClient() as client:
-                        response = await client.get(image.url.uri)
-                        data = response.content
-                else:
-                    raise ValueError("Unsupported URL type")
-            elif image.data:
-                # data is a base64 encoded string, decode it to bytes for RawMediaItem
-                data = base64.b64decode(image.data)
-            else:
-                raise ValueError("No data or URL provided")
-
-            return RawMediaItem(data=data)
-        else:
-            raise ValueError(f"Unsupported content type: {type(c)}")
-
-    if isinstance(content, list):
-        return await asyncio.gather(*(_localize_single(c) for c in content))
-    else:
-        return await _localize_single(content)
-
-
-async def convert_openai_message_to_raw_message(message: OpenAIMessageParam) -> RawMessage:
-    """Convert OpenAI message format to RawMessage format used by Llama formatters."""
-    if isinstance(message, OpenAIUserMessageParam):
-        content = await interleaved_content_convert_to_raw(message.content)  # type: ignore[arg-type]
-        return RawMessage(role="user", content=content)
-    elif isinstance(message, OpenAISystemMessageParam):
-        content = await interleaved_content_convert_to_raw(message.content)  # type: ignore[arg-type]
-        return RawMessage(role="system", content=content)
-    elif isinstance(message, OpenAIAssistantMessageParam):
-        content = await interleaved_content_convert_to_raw(message.content or "")  # type: ignore[arg-type]
-        tool_calls = []
-        if message.tool_calls:
-            for tc in message.tool_calls:
-                if tc.function:
-                    tool_calls.append(
-                        ToolCall(
-                            call_id=tc.id or "",
-                            tool_name=tc.function.name or "",
-                            arguments=tc.function.arguments or "{}",
-                        )
-                    )
-        return RawMessage(role="assistant", content=content, tool_calls=tool_calls)
-    elif isinstance(message, OpenAIToolMessageParam):
-        content = await interleaved_content_convert_to_raw(message.content)  # type: ignore[arg-type]
-        return RawMessage(role="tool", content=content)
-    else:
-        # Handle OpenAIDeveloperMessageParam if needed
-        raise ValueError(f"Unsupported message type: {type(message)}")
-
-
-def content_has_media(content: InterleavedContent):
-    def _has_media_content(c):
-        return isinstance(c, ImageContentItem)
-
-    if isinstance(content, list):
-        return any(_has_media_content(c) for c in content)
-    else:
-        return _has_media_content(content)
-
-
 async def localize_image_content(uri: str) -> tuple[bytes, str] | None:
     if uri.startswith("http"):
         async with httpx.AsyncClient() as client:
@@ -194,87 +69,3 @@ async def localize_image_content(uri: str) -> tuple[bytes, str] | None:
         return content, fmt
     else:
         return None
-
-
-async def convert_image_content_to_url(
-    media: ImageContentItem, download: bool = False, include_format: bool = True
-) -> str:
-    image = media.image
-    if image.url and (not download or image.url.uri.startswith("data")):
-        return image.url.uri
-
-    if image.data:
-        # data is a base64 encoded string, decode it to bytes first
-        # TODO(mf): do this more efficiently, decode less
-        content = base64.b64decode(image.data)
-        pil_image = PIL_Image.open(io.BytesIO(content))
-        format = pil_image.format
-    else:
-        localize_result = await localize_image_content(image.url.uri)
-        if localize_result is None:
-            raise ValueError(f"Failed to localize image content from {image.url.uri}")
-        content, format = localize_result
-
-    if include_format:
-        return f"data:image/{format};base64," + base64.b64encode(content).decode("utf-8")
-    else:
-        return base64.b64encode(content).decode("utf-8")
-
-
-def augment_content_with_response_format_prompt(response_format, content):
-    if fmt_prompt := response_format_prompt(response_format):
-        if isinstance(content, list):
-            return content + [TextContentItem(text=fmt_prompt)]
-        elif isinstance(content, str):
-            return [TextContentItem(text=content), TextContentItem(text=fmt_prompt)]
-        else:
-            return [content, TextContentItem(text=fmt_prompt)]
-
-    return content
-
-
-def response_format_prompt(fmt: ResponseFormat | None):
-    if not fmt:
-        return None
-
-    if fmt.type == ResponseFormatType.json_schema.value:
-        return f"Please respond in JSON format with the schema: {json.dumps(fmt.json_schema)}"
-    elif fmt.type == ResponseFormatType.grammar.value:
-        raise NotImplementedError("Grammar response format not supported yet")
-    else:
-        raise ValueError(f"Unknown response format {fmt.type}")
-
-
-def _get_tool_choice_prompt(tool_choice: ToolChoice | str, tools: list[ToolDefinition]) -> str:
-    if tool_choice == ToolChoice.auto:
-        return ""
-    elif tool_choice == ToolChoice.required:
-        return "You MUST use one of the provided functions/tools to answer the user query."
-    elif tool_choice == ToolChoice.none:
-        # tools are already not passed in
-        return ""
-    else:
-        # specific tool
-        return f"You MUST use the tool `{tool_choice}` to answer the user query."
-
-
-def get_default_tool_prompt_format(model: str) -> ToolPromptFormat:
-    llama_model = resolve_model(model)
-    if llama_model is None:
-        log.warning(f"Could not resolve model {model}, defaulting to json tool prompt format")
-        return ToolPromptFormat.json
-
-    if llama_model.model_family == ModelFamily.llama3_1 or (
-        llama_model.model_family == ModelFamily.llama3_2 and is_multimodal(llama_model.core_model_id)
-    ):
-        # llama3.1 and llama3.2 multimodal models follow the same tool prompt format
-        return ToolPromptFormat.json
-    elif llama_model.model_family in (
-        ModelFamily.llama3_2,
-        ModelFamily.llama3_3,
-        ModelFamily.llama4,
-    ):
-        # llama3.2 and llama3.3 models follow the same tool prompt format
-        return ToolPromptFormat.python_list
-    else:
-        return ToolPromptFormat.json

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -671,6 +671,19 @@ class OpenAIVectorStoreMixin(ABC):
             search_query = query
 
         try:
+            # Validate neural ranker requires model parameter
+            if ranking_options is not None:
+                if getattr(ranking_options, "ranker", None) == "neural":
+                    model_value = getattr(ranking_options, "model", None)
+                    if model_value is None or (isinstance(model_value, str) and model_value.strip() == ""):
+                        # Return empty results when model is missing for neural ranker
+                        logger.warning("model parameter is required when ranker='neural', returning empty results")
+                        return VectorStoreSearchResponsePage(
+                            search_query=query if isinstance(query, list) else [query],
+                            data=[],
+                            has_more=False,
+                            next_page=None,
+                        )
             score_threshold = (
                 ranking_options.score_threshold
                 if ranking_options and ranking_options.score_threshold is not None
@@ -681,7 +694,10 @@ class OpenAIVectorStoreMixin(ABC):
                 "score_threshold": score_threshold,
                 "mode": search_mode,
             }
-            # TODO: Add support for ranking_options.ranker
+
+            # Use VectorStoresConfig defaults when ranking_options values are not provided
+            config = self.vector_stores_config or VectorStoresConfig()
+            params.update(self._build_reranker_params(ranking_options, config))
 
             response = await self.query_chunks(
                 vector_store_id=vector_store_id,
@@ -722,8 +738,8 @@ class OpenAIVectorStoreMixin(ABC):
             )
 
         except Exception as e:
+            # Log the error and return empty results
             logger.error(f"Error searching vector store {vector_store_id}: {e}")
-            # Return empty results on error
             return VectorStoreSearchResponsePage(
                 search_query=query if isinstance(query, list) else [query],
                 data=[],
@@ -731,6 +747,62 @@ class OpenAIVectorStoreMixin(ABC):
                 next_page=None,
             )
 
+    def _build_reranker_params(
+        self,
+        ranking_options: SearchRankingOptions | None,
+        config: VectorStoresConfig,
+    ) -> dict[str, Any]:
+        reranker_params: dict[str, Any] = {}
+        params: dict[str, Any] = {}
+
+        if ranking_options and ranking_options.ranker:
+            reranker_type = ranking_options.ranker
+
+            if ranking_options.ranker == "weighted":
+                alpha = ranking_options.alpha
+                if alpha is None:
+                    alpha = config.chunk_retrieval_params.weighted_search_alpha
+                reranker_params["alpha"] = alpha
+                if ranking_options.weights:
+                    reranker_params["weights"] = ranking_options.weights
+            elif ranking_options.ranker == "rrf":
+                # For RRF ranker, use impact_factor from request if provided, otherwise use VectorStoresConfig default
+                impact_factor = ranking_options.impact_factor
+                if impact_factor is None:
+                    impact_factor = config.chunk_retrieval_params.rrf_impact_factor
+                reranker_params["impact_factor"] = impact_factor
+                # If weights dict is provided (for neural combination), store it
+                if ranking_options.weights:
+                    reranker_params["weights"] = ranking_options.weights
+            elif ranking_options.ranker == "neural":
+                reranker_params["model"] = ranking_options.model
+            else:
+                logger.debug(f"Unknown ranker value: {ranking_options.ranker}, passing through")
+
+            params["reranker_type"] = reranker_type
+            params["reranker_params"] = reranker_params
+
+            # Store model and weights for neural reranking (TODO: implemented in Part II)
+            if ranking_options.model:
+                params["neural_model"] = ranking_options.model
+            if ranking_options.weights:
+                params["neural_weights"] = ranking_options.weights
+        elif ranking_options is None or ranking_options.ranker is None:
+            # No ranker specified in request - use VectorStoresConfig default
+            default_strategy = config.chunk_retrieval_params.default_reranker_strategy
+            if default_strategy in ("weighted", "rrf"):
+                params["reranker_type"] = default_strategy
+                reranker_params = {}
+
+                if default_strategy == "weighted":
+                    reranker_params["alpha"] = config.chunk_retrieval_params.weighted_search_alpha
+                elif default_strategy == "rrf":
+                    reranker_params["impact_factor"] = config.chunk_retrieval_params.rrf_impact_factor
+
+                params["reranker_params"] = reranker_params
+
+        return params
+
     def _matches_filters(self, metadata: dict[str, Any], filters: dict[str, Any]) -> bool:
         """Check if metadata matches the provided filters."""
         if not filters:
@@ -738,15 +810,29 @@ class OpenAIVectorStoreMixin(ABC):
 
         filter_type = filters.get("type")
 
+        if filter_type is None:
+            if "key" not in filters and "value" not in filters and "filters" not in filters:
+                for key, value in filters.items():
+                    if key not in metadata:
+                        return False
+                    if metadata[key] != value:
+                        return False
+                return True
+            else:
+                raise ValueError("Unsupported filter structure: missing 'type' field")
+
         if filter_type in ["eq", "ne", "gt", "gte", "lt", "lte"]:
             # Comparison filter
-            key = filters.get("key")
+            filter_key = filters.get("key")
             value = filters.get("value")
 
-            if key not in metadata:
+            if filter_key is None or not isinstance(filter_key, str):
+                return False
+
+            if filter_key not in metadata:
                 return False
 
-            metadata_value = metadata[key]
+            metadata_value = metadata[filter_key]
 
             if filter_type == "eq":
                 return bool(metadata_value == value)
@@ -901,6 +987,7 @@ class OpenAIVectorStoreMixin(ABC):
                 params = OpenAIEmbeddingsRequestWithExtraBody(
                     model=embedding_model,
                     input=[interleaved_content_as_str(c.content) for c in chunks],
+                    dimensions=embedding_dimension,
                 )
                 resp = await self.inference_api.openai_embeddings(params)
 

llama_stack/providers/utils/memory/vector_store.py

@@ -297,37 +297,64 @@ class VectorStoreWithIndex:
         mode = params.get("mode")
         score_threshold = params.get("score_threshold", 0.0)
 
-        ranker = params.get("ranker")
-        if ranker is None:
+        # Get reranker configuration from params (set by openai_vector_store_mixin)
+        # NOTE: Breaking change - removed support for old nested "ranker" format.
+        #       Now uses flattened format: reranker_type and reranker_params.
+        reranker_type = params.get("reranker_type")
+        reranker_params = params.get("reranker_params", {})
+
+        # If no ranker specified, use VectorStoresConfig default
+        if reranker_type is None:
             reranker_type = (
                 RERANKER_TYPE_RRF
                 if config.chunk_retrieval_params.default_reranker_strategy == "rrf"
                 else config.chunk_retrieval_params.default_reranker_strategy
             )
             reranker_params = {"impact_factor": config.chunk_retrieval_params.rrf_impact_factor}
+
+        # Normalize reranker_type to use constants
+        if reranker_type == "weighted":
+            reranker_type = RERANKER_TYPE_WEIGHTED
+            # Ensure alpha is set (use default if not provided)
+            if "alpha" not in reranker_params:
+                reranker_params["alpha"] = config.chunk_retrieval_params.weighted_search_alpha
+        elif reranker_type == "rrf":
+            reranker_type = RERANKER_TYPE_RRF
+            # Ensure impact_factor is set (use default if not provided)
+            if "impact_factor" not in reranker_params:
+                reranker_params["impact_factor"] = config.chunk_retrieval_params.rrf_impact_factor
+        elif reranker_type == "neural":
+            # TODO: Implement neural reranking
+            log.warning(
+                "TODO: Neural reranking for vector stores is not implemented yet; "
+                "using configured reranker params without algorithm fallback."
+            )
+        elif reranker_type == "normalized":
+            reranker_type = RERANKER_TYPE_NORMALIZED
         else:
-            strategy = ranker.get("strategy", config.chunk_retrieval_params.default_reranker_strategy)
-            if strategy == "weighted":
-                weights = ranker.get("params", {}).get("weights", [0.5, 0.5])
-                reranker_type = RERANKER_TYPE_WEIGHTED
-                reranker_params = {
-                    "alpha": weights[0] if len(weights) > 0 else config.chunk_retrieval_params.weighted_search_alpha
-                }
-            elif strategy == "normalized":
-                reranker_type = RERANKER_TYPE_NORMALIZED
-            else:
-                reranker_type = RERANKER_TYPE_RRF
-                k_value = ranker.get("params", {}).get("k", config.chunk_retrieval_params.rrf_impact_factor)
-                reranker_params = {"impact_factor": k_value}
+            # Default to RRF for unknown strategies
+            reranker_type = RERANKER_TYPE_RRF
+            if "impact_factor" not in reranker_params:
+                reranker_params["impact_factor"] = config.chunk_retrieval_params.rrf_impact_factor
+
+        # Store neural model and weights from params if provided (for future neural reranking in Part II)
+        if "neural_model" in params:
+            reranker_params["neural_model"] = params["neural_model"]
+        if "neural_weights" in params:
+            reranker_params["neural_weights"] = params["neural_weights"]
 
         query_string = interleaved_content_as_str(query)
         if mode == "keyword":
             return await self.index.query_keyword(query_string, k, score_threshold)
 
-        params = OpenAIEmbeddingsRequestWithExtraBody(
-            model=self.vector_store.embedding_model,
-            input=[query_string],
-        )
+        if "embedding_dimensions" in params:
+            params = OpenAIEmbeddingsRequestWithExtraBody(
+                model=self.vector_store.embedding_model,
+                input=[query_string],
+                dimensions=params.get("embedding_dimensions"),
+            )
+        else:
+            params = OpenAIEmbeddingsRequestWithExtraBody(model=self.vector_store.embedding_model, input=[query_string])
         embeddings_response = await self.inference_api.openai_embeddings(params)
         query_vector = np.array(embeddings_response.data[0].embedding, dtype=np.float32)
         if mode == "hybrid":

llama_stack/providers/utils/responses/responses_store.py

@@ -57,7 +57,7 @@ class ResponsesStore:
         self.sql_store = AuthorizedSqlStore(base_store, self.policy)
 
         await self.sql_store.create_table(
-            "openai_responses",
+            self.reference.table_name,
             {
                 "id": ColumnDefinition(type=ColumnType.STRING, primary_key=True),
                 "created_at": ColumnType.INTEGER,
@@ -89,6 +89,40 @@ class ResponsesStore:
     ) -> None:
         await self._write_response_object(response_object, input, messages)
 
+    async def upsert_response_object(
+        self,
+        response_object: OpenAIResponseObject,
+        input: list[OpenAIResponseInput],
+        messages: list[OpenAIMessageParam],
+    ) -> None:
+        """Upsert response object using INSERT on first call, UPDATE on subsequent calls.
+
+        This method enables incremental persistence during streaming, allowing clients
+        to poll GET /v1/responses/{response_id} and see in-progress turn state.
+
+        :param response_object: The response object to store/update.
+        :param input: The input items for the response.
+        :param messages: The chat completion messages (for conversation continuity).
+        """
+        if self.sql_store is None:
+            raise ValueError("Responses store is not initialized")
+
+        data = response_object.model_dump()
+        data["input"] = [input_item.model_dump() for input_item in input]
+        data["messages"] = [msg.model_dump() for msg in messages]
+
+        await self.sql_store.upsert(
+            table=self.reference.table_name,
+            data={
+                "id": data["id"],
+                "created_at": data["created_at"],
+                "model": data["model"],
+                "response_object": data,
+            },
+            conflict_columns=["id"],
+            update_columns=["response_object"],
+        )
+
     async def _write_response_object(
         self,
         response_object: OpenAIResponseObject,
@@ -103,7 +137,7 @@ class ResponsesStore:
         data["messages"] = [msg.model_dump() for msg in messages]
 
         await self.sql_store.insert(
-            "openai_responses",
+            self.reference.table_name,
             {
                 "id": data["id"],
                 "created_at": data["created_at"],
@@ -138,7 +172,7 @@ class ResponsesStore:
             where_conditions["model"] = model
 
         paginated_result = await self.sql_store.fetch_all(
-            table="openai_responses",
+            table=self.reference.table_name,
             where=where_conditions if where_conditions else None,
             order_by=[("created_at", order.value)],
             cursor=("id", after) if after else None,
@@ -161,7 +195,7 @@ class ResponsesStore:
             raise ValueError("Responses store is not initialized")
 
         row = await self.sql_store.fetch_one(
-            "openai_responses",
+            self.reference.table_name,
             where={"id": response_id},
         )
 
@@ -176,10 +210,10 @@ class ResponsesStore:
         if not self.sql_store:
             raise ValueError("Responses store is not initialized")
 
-        row = await self.sql_store.fetch_one("openai_responses", where={"id": response_id})
+        row = await self.sql_store.fetch_one(self.reference.table_name, where={"id": response_id})
         if not row:
             raise ValueError(f"Response with id {response_id} not found")
-        await self.sql_store.delete("openai_responses", where={"id": response_id})
+        await self.sql_store.delete(self.reference.table_name, where={"id": response_id})
         return OpenAIDeleteResponseObject(id=response_id)
 
     async def list_response_input_items(