From 470fe55e874974c9dd3413939d23df6254df956e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Thu, 11 Dec 2025 14:50:23 +0100 Subject: [PATCH] fix(inference): respect table_name config in InferenceStore (#4371) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # What does this PR do? The InferenceStore class was ignoring the table_name field from InferenceStoreReference and always using the hardcoded value "chat_completions". This meant that any custom table_name configured in the run config (e.g., "inference_store" in run-with-postgres-store.yaml) was silently ignored. This change updates all SQL operations in InferenceStore to use self.reference.table_name instead of the hardcoded string, ensuring the configured table name is properly respected. A new test has been added to verify that custom table names work correctly for storing, retrieving, and listing chat completions. ## Test Plan CI Signed-off-by: Sébastien Han --- .../utils/inference/inference_store.py | 10 +++---- .../utils/inference/test_inference_store.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/llama_stack/providers/utils/inference/inference_store.py b/src/llama_stack/providers/utils/inference/inference_store.py index a8a0cace4..78327573b 100644 --- a/src/llama_stack/providers/utils/inference/inference_store.py +++ b/src/llama_stack/providers/utils/inference/inference_store.py @@ -55,7 +55,7 @@ class InferenceStore: logger.debug("Write queue disabled for SQLite (WAL mode handles concurrency)") await self.sql_store.create_table( - "chat_completions", + self.reference.table_name, { "id": ColumnDefinition(type=ColumnType.STRING, primary_key=True), "created": ColumnType.INTEGER, @@ -152,7 +152,7 @@ class InferenceStore: try: await self.sql_store.insert( - table="chat_completions", + table=self.reference.table_name, data=record_data, ) except IntegrityError as e: @@ -164,7 +164,7 @@ class InferenceStore: error_message = str(e.orig) if e.orig else str(e) if self._is_unique_constraint_error(error_message): # Update the existing record instead - await self.sql_store.update(table="chat_completions", data=record_data, where={"id": data["id"]}) + await self.sql_store.update(table=self.reference.table_name, data=record_data, where={"id": data["id"]}) else: # Re-raise if it's not a unique constraint error raise @@ -208,7 +208,7 @@ class InferenceStore: where_conditions["model"] = model paginated_result = await self.sql_store.fetch_all( - table="chat_completions", + table=self.reference.table_name, where=where_conditions if where_conditions else None, order_by=[("created", order.value)], cursor=("id", after) if after else None, @@ -237,7 +237,7 @@ class InferenceStore: raise ValueError("Inference store is not initialized") row = await self.sql_store.fetch_one( - table="chat_completions", + table=self.reference.table_name, where={"id": completion_id}, ) diff --git a/tests/unit/utils/inference/test_inference_store.py b/tests/unit/utils/inference/test_inference_store.py index 22d4ec1e5..c6bdc5c29 100644 --- a/tests/unit/utils/inference/test_inference_store.py +++ b/tests/unit/utils/inference/test_inference_store.py @@ -210,3 +210,32 @@ async def test_inference_store_pagination_no_limit(): assert result.data[0].id == "beta-second" # Most recent first assert result.data[1].id == "omega-first" assert result.has_more is False + + +async def test_inference_store_custom_table_name(): + """Test that the table_name from config is respected.""" + custom_table_name = "custom_inference_store" + reference = InferenceStoreReference(backend="sql_default", table_name=custom_table_name) + store = InferenceStore(reference, policy=[]) + await store.initialize() + + # Create and store a test chat completion + base_time = int(time.time()) + completion = create_test_chat_completion("custom-table-test", base_time) + input_messages = [OpenAIUserMessageParam(role="user", content="Test custom table")] + await store.store_chat_completion(completion, input_messages) + await store.flush() + + # Verify we can retrieve the completion + result = await store.get_chat_completion("custom-table-test") + assert result.id == "custom-table-test" + assert result.model == "test-model" + + # Verify listing works + list_result = await store.list_chat_completions() + assert len(list_result.data) == 1 + assert list_result.data[0].id == "custom-table-test" + + # Verify the error message uses the custom table name + with pytest.raises(ValueError, match=f"Record with id='non-existent' not found in table '{custom_table_name}'"): + await store.list_chat_completions(after="non-existent", limit=2)