mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-04 12:07:34 +00:00
fix(inference_store): on duplicate chat completion IDs, replace (#3408)
Some checks failed
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 2s
Test External Providers Installed via Module / test-external-providers-from-module (venv) (push) Has been skipped
Python Package Build Test / build (3.13) (push) Failing after 2s
SqlStore Integration Tests / test-postgres (3.13) (push) Failing after 9s
Vector IO Integration Tests / test-matrix (push) Failing after 4s
API Conformance Tests / check-schema-compatibility (push) Successful in 7s
Unit Tests / unit-tests (3.12) (push) Failing after 4s
Integration Tests (Replay) / Integration Tests (, , , client=, ) (push) Failing after 3s
Python Package Build Test / build (3.12) (push) Failing after 3s
SqlStore Integration Tests / test-postgres (3.12) (push) Failing after 8s
Unit Tests / unit-tests (3.13) (push) Failing after 5s
Update ReadTheDocs / update-readthedocs (push) Failing after 23s
Test External API and Providers / test-external (venv) (push) Failing after 30s
UI Tests / ui-tests (22) (push) Successful in 35s
Pre-commit / pre-commit (push) Successful in 1m45s
Some checks failed
Integration Auth Tests / test-matrix (oauth2_token) (push) Failing after 2s
Test External Providers Installed via Module / test-external-providers-from-module (venv) (push) Has been skipped
Python Package Build Test / build (3.13) (push) Failing after 2s
SqlStore Integration Tests / test-postgres (3.13) (push) Failing after 9s
Vector IO Integration Tests / test-matrix (push) Failing after 4s
API Conformance Tests / check-schema-compatibility (push) Successful in 7s
Unit Tests / unit-tests (3.12) (push) Failing after 4s
Integration Tests (Replay) / Integration Tests (, , , client=, ) (push) Failing after 3s
Python Package Build Test / build (3.12) (push) Failing after 3s
SqlStore Integration Tests / test-postgres (3.12) (push) Failing after 8s
Unit Tests / unit-tests (3.13) (push) Failing after 5s
Update ReadTheDocs / update-readthedocs (push) Failing after 23s
Test External API and Providers / test-external (venv) (push) Failing after 30s
UI Tests / ui-tests (22) (push) Successful in 35s
Pre-commit / pre-commit (push) Successful in 1m45s
# What does this PR do? Duplicate chat completion IDs can be generated during tests especially if they are replaying recorded responses across different tests. No need to warn or error under those circumstances. In the wild, this is not likely to happen at all (no evidence) so we aren't really hiding any problem.
This commit is contained in:
parent
c04f1c1e8c
commit
0c7f49490c
2 changed files with 53 additions and 9 deletions
|
@ -6,6 +6,8 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from sqlalchemy.exc import IntegrityError
|
||||||
|
|
||||||
from llama_stack.apis.inference import (
|
from llama_stack.apis.inference import (
|
||||||
ListOpenAIChatCompletionResponse,
|
ListOpenAIChatCompletionResponse,
|
||||||
OpenAIChatCompletion,
|
OpenAIChatCompletion,
|
||||||
|
@ -129,16 +131,44 @@ class InferenceStore:
|
||||||
raise ValueError("Inference store is not initialized")
|
raise ValueError("Inference store is not initialized")
|
||||||
|
|
||||||
data = chat_completion.model_dump()
|
data = chat_completion.model_dump()
|
||||||
|
record_data = {
|
||||||
|
"id": data["id"],
|
||||||
|
"created": data["created"],
|
||||||
|
"model": data["model"],
|
||||||
|
"choices": data["choices"],
|
||||||
|
"input_messages": [message.model_dump() for message in input_messages],
|
||||||
|
}
|
||||||
|
|
||||||
await self.sql_store.insert(
|
try:
|
||||||
table="chat_completions",
|
await self.sql_store.insert(
|
||||||
data={
|
table="chat_completions",
|
||||||
"id": data["id"],
|
data=record_data,
|
||||||
"created": data["created"],
|
)
|
||||||
"model": data["model"],
|
except IntegrityError as e:
|
||||||
"choices": data["choices"],
|
# Duplicate chat completion IDs can be generated during tests especially if they are replaying
|
||||||
"input_messages": [message.model_dump() for message in input_messages],
|
# recorded responses across different tests. No need to warn or error under those circumstances.
|
||||||
},
|
# In the wild, this is not likely to happen at all (no evidence) so we aren't really hiding any problem.
|
||||||
|
|
||||||
|
# Check if it's a unique constraint violation
|
||||||
|
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"]})
|
||||||
|
else:
|
||||||
|
# Re-raise if it's not a unique constraint error
|
||||||
|
raise
|
||||||
|
|
||||||
|
def _is_unique_constraint_error(self, error_message: str) -> bool:
|
||||||
|
"""Check if the error is specifically a unique constraint violation."""
|
||||||
|
error_lower = error_message.lower()
|
||||||
|
return any(
|
||||||
|
indicator in error_lower
|
||||||
|
for indicator in [
|
||||||
|
"unique constraint failed", # SQLite
|
||||||
|
"duplicate key", # PostgreSQL
|
||||||
|
"unique violation", # PostgreSQL alternative
|
||||||
|
"duplicate entry", # MySQL
|
||||||
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
async def list_chat_completions(
|
async def list_chat_completions(
|
||||||
|
|
|
@ -172,6 +172,20 @@ class AuthorizedSqlStore:
|
||||||
|
|
||||||
return results.data[0] if results.data else None
|
return results.data[0] if results.data else None
|
||||||
|
|
||||||
|
async def update(self, table: str, data: Mapping[str, Any], where: Mapping[str, Any]) -> None:
|
||||||
|
"""Update rows with automatic access control attribute capture."""
|
||||||
|
enhanced_data = dict(data)
|
||||||
|
|
||||||
|
current_user = get_authenticated_user()
|
||||||
|
if current_user:
|
||||||
|
enhanced_data["owner_principal"] = current_user.principal
|
||||||
|
enhanced_data["access_attributes"] = current_user.attributes
|
||||||
|
else:
|
||||||
|
enhanced_data["owner_principal"] = None
|
||||||
|
enhanced_data["access_attributes"] = None
|
||||||
|
|
||||||
|
await self.sql_store.update(table, enhanced_data, where)
|
||||||
|
|
||||||
async def delete(self, table: str, where: Mapping[str, Any]) -> None:
|
async def delete(self, table: str, where: Mapping[str, Any]) -> None:
|
||||||
"""Delete rows with automatic access control filtering."""
|
"""Delete rows with automatic access control filtering."""
|
||||||
await self.sql_store.delete(table, where)
|
await self.sql_store.delete(table, where)
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue