phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-12 20:12:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

simplified some, walked back some decisions

Browse source
This commit is contained in:
Ashwin Bharambe 2025-10-17 10:05:07 -07:00
parent af7472cdb0
commit 636764c2a1
90 changed files with 887 additions and 570 deletions
Download patch file Download diff file Expand all files Collapse all files

8
llama_stack/providers/inline/agents/meta_reference/config.py
View file

@ -8,14 +8,14 @@ from typing import Any
from pydantic import BaseModel
from llama_stack.core.storage.datatypes import KVStoreReference, SqlStoreReference
from llama_stack.core.storage.datatypes import KVStoreReference, ResponsesStoreReference
class AgentPersistenceConfig(BaseModel):
"""Nested persistence configuration for agents."""
agent_state: KVStoreReference
responses: SqlStoreReference
responses: ResponsesStoreReference
class MetaReferenceAgentsImplConfig(BaseModel):
@ -26,9 +26,11 @@ class MetaReferenceAgentsImplConfig(BaseModel):
return {
"persistence": {
"agent_state": KVStoreReference(
backend="kv_default",
namespace="agents",
).model_dump(exclude_none=True),
"responses": SqlStoreReference(
"responses": ResponsesStoreReference(
backend="sql_default",
table_name="responses",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/batches/reference/config.py
View file

@ -34,6 +34,7 @@ class ReferenceBatchesImplConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str) -> dict:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="batches",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/datasetio/localfs/config.py
View file

@ -17,6 +17,7 @@ class LocalFSDatasetIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="datasetio::localfs",
).model_dump(exclude_none=True)
}

1
llama_stack/providers/inline/eval/meta_reference/config.py
View file

@ -17,6 +17,7 @@ class MetaReferenceEvalConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="eval",
).model_dump(exclude_none=True)
}

1
llama_stack/providers/inline/files/localfs/config.py
View file

@ -25,6 +25,7 @@ class LocalfsFilesImplConfig(BaseModel):
return {
"storage_dir": "${env.FILES_STORAGE_DIR:=" + __distro_dir__ + "/files}",
"metadata_store": SqlStoreReference(
backend="sql_default",
table_name="files_metadata",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/vector_io/chroma/config.py
View file

@ -24,6 +24,7 @@ class ChromaVectorIOConfig(BaseModel):
return {
"db_path": db_path,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::chroma",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/vector_io/faiss/config.py
View file

@ -20,6 +20,7 @@ class FaissVectorIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::faiss",
).model_dump(exclude_none=True)
}

1
llama_stack/providers/inline/vector_io/milvus/config.py
View file

@ -23,6 +23,7 @@ class MilvusVectorIOConfig(BaseModel):
return {
"db_path": "${env.MILVUS_DB_PATH:=" + __distro_dir__ + "}/" + "milvus.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::milvus",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/vector_io/qdrant/config.py
View file

@ -23,6 +23,7 @@ class QdrantVectorIOConfig(BaseModel):
return {
"path": "${env.QDRANT_PATH:=~/.llama/" + __distro_dir__ + "}/" + "qdrant.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::qdrant",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/inline/vector_io/sqlite_vec/config.py
View file

@ -20,6 +20,7 @@ class SQLiteVectorIOConfig(BaseModel):
return {
"db_path": "${env.SQLITE_STORE_DIR:=" + __distro_dir__ + "}/" + "sqlite_vec.db",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::sqlite_vec",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/remote/datasetio/huggingface/config.py
View file

@ -17,6 +17,7 @@ class HuggingfaceDatasetIOConfig(BaseModel):
def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
return {
"kvstore": KVStoreReference(
backend="kv_default",
namespace="datasetio::huggingface",
).model_dump(exclude_none=True)
}

1
llama_stack/providers/remote/files/s3/config.py
View file

@ -36,6 +36,7 @@ class S3FilesImplConfig(BaseModel):
"endpoint_url": "${env.S3_ENDPOINT_URL:=}",
"auto_create_bucket": "${env.S3_AUTO_CREATE_BUCKET:=false}",
"metadata_store": SqlStoreReference(
backend="sql_default",
table_name="s3_files_metadata",
).model_dump(exclude_none=True),
}

2
llama_stack/providers/remote/files/s3/files.py
View file

@ -166,7 +166,7 @@ class S3FilesImpl(Files):
self._client = _create_s3_client(self._config)
await _create_bucket_if_not_exists(self._client, self._config)
self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.persistence), self.policy)
self._sql_store = AuthorizedSqlStore(sqlstore_impl(self._config.metadata_store), self.policy)
await self._sql_store.create_table(
"openai_files",
{

1
llama_stack/providers/remote/vector_io/chroma/config.py
View file

@ -22,6 +22,7 @@ class ChromaVectorIOConfig(BaseModel):
return {
"url": url,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::chroma_remote",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/remote/vector_io/milvus/config.py
View file

@ -29,6 +29,7 @@ class MilvusVectorIOConfig(BaseModel):
"uri": "${env.MILVUS_ENDPOINT}",
"token": "${env.MILVUS_TOKEN}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::milvus_remote",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/remote/vector_io/pgvector/config.py
View file

@ -41,6 +41,7 @@ class PGVectorVectorIOConfig(BaseModel):
"user": user,
"password": password,
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::pgvector",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/remote/vector_io/qdrant/config.py
View file

@ -31,6 +31,7 @@ class QdrantVectorIOConfig(BaseModel):
return {
"api_key": "${env.QDRANT_API_KEY:=}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::qdrant_remote",
).model_dump(exclude_none=True),
}

1
llama_stack/providers/remote/vector_io/weaviate/config.py
View file

@ -30,6 +30,7 @@ class WeaviateVectorIOConfig(BaseModel):
"weaviate_api_key": None,
"weaviate_cluster_url": "${env.WEAVIATE_CLUSTER_URL:=localhost:8080}",
"kvstore": KVStoreReference(
backend="kv_default",
namespace="vector_io::weaviate",
).model_dump(exclude_none=True),
}

13
llama_stack/providers/utils/inference/inference_store.py
View file

@ -48,12 +48,13 @@ class InferenceStore:
self.sql_store = AuthorizedSqlStore(base_store, self.policy)
# Disable write queue for SQLite to avoid concurrency issues
backend_name = self.reference.backend or list(_SQLSTORE_BACKENDS.keys())[0] if _SQLSTORE_BACKENDS else None
if backend_name:
backend_config = _SQLSTORE_BACKENDS.get(backend_name)
self.enable_write_queue = backend_config.type != StorageBackendType.SQL_SQLITE
else:
self.enable_write_queue = True
backend_name = self.reference.backend
backend_config = _SQLSTORE_BACKENDS.get(backend_name)
if backend_config is None:
raise ValueError(
f"Unregistered SQL backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
self.enable_write_queue = backend_config.type != StorageBackendType.SQL_SQLITE
await self.sql_store.create_table(
"chat_completions",
{

25
llama_stack/providers/utils/kvstore/kvstore.py
View file

@ -1,3 +1,9 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the terms described in the LICENSE file in
@ -47,36 +53,19 @@ class InmemoryKVStoreImpl(KVStore):
_KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {}
_KVSTORE_DEFAULT_BACKEND: str | None = None
def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None:
"""Register the set of available KV store backends for reference resolution."""
global _KVSTORE_BACKENDS
def _set_default_backend(name: str) -> None:
global _KVSTORE_DEFAULT_BACKEND
if _KVSTORE_DEFAULT_BACKEND and _KVSTORE_DEFAULT_BACKEND != name:
raise ValueError(
f"Multiple KVStore backends marked as default: '{_KVSTORE_DEFAULT_BACKEND}' and '{name}'. "
"Only one backend can be the default."
)
_KVSTORE_DEFAULT_BACKEND = name
_KVSTORE_BACKENDS.clear()
for name, cfg in backends.items():
if cfg.default:
_set_default_backend(name)
_KVSTORE_BACKENDS[name] = cfg
async def kvstore_impl(reference: KVStoreReference) -> KVStore:
backend_name = reference.backend or _KVSTORE_DEFAULT_BACKEND
if not backend_name:
raise ValueError(
"KVStore reference did not specify a backend and no default backend is configured. "
f"Available backends: {sorted(_KVSTORE_BACKENDS)}"
)
backend_name = reference.backend
backend_config = _KVSTORE_BACKENDS.get(backend_name)
if backend_config is None:

45
llama_stack/providers/utils/responses/responses_store.py
View file

@ -18,13 +18,13 @@ from llama_stack.apis.agents.openai_responses import (
OpenAIResponseObjectWithInput,
)
from llama_stack.apis.inference import OpenAIMessageParam
from llama_stack.core.datatypes import AccessRule, ResponsesStoreConfig
from llama_stack.core.utils.config_dirs import RUNTIME_BASE_DIR
from llama_stack.core.datatypes import AccessRule
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference, StorageBackendType
from llama_stack.log import get_logger
from ..sqlstore.api import ColumnDefinition, ColumnType
from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
from ..sqlstore.sqlstore import SqliteSqlStoreConfig, SqlStoreConfig, SqlStoreType, sqlstore_impl
from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
logger = get_logger(name=__name__, category="openai_responses")
@ -45,39 +45,38 @@ class _OpenAIResponseObjectWithInputAndMessages(OpenAIResponseObjectWithInput):
class ResponsesStore:
def __init__(
self,
config: ResponsesStoreConfig | SqlStoreConfig,
reference: ResponsesStoreReference | SqlStoreReference,
policy: list[AccessRule],
):
# Handle backward compatibility
if not isinstance(config, ResponsesStoreConfig):
# Legacy: SqlStoreConfig passed directly as config
config = ResponsesStoreConfig(
sql_store_config=config,
)
if isinstance(reference, ResponsesStoreReference):
self.reference = reference
else:
self.reference = ResponsesStoreReference(**reference.model_dump())
self.config = config
self.sql_store_config = config.sql_store_config
if not self.sql_store_config:
self.sql_store_config = SqliteSqlStoreConfig(
db_path=(RUNTIME_BASE_DIR / "sqlstore.db").as_posix(),
)
self.sql_store = None
self.policy = policy
# Disable write queue for SQLite to avoid concurrency issues
self.enable_write_queue = self.sql_store_config.type != SqlStoreType.sqlite
self.sql_store = None
self.enable_write_queue = True
# Async write queue and worker control
self._queue: (
asyncio.Queue[tuple[OpenAIResponseObject, list[OpenAIResponseInput], list[OpenAIMessageParam]]] | None
) = None
self._worker_tasks: list[asyncio.Task[Any]] = []
self._max_write_queue_size: int = config.max_write_queue_size
self._num_writers: int = max(1, config.num_writers)
self._max_write_queue_size: int = self.reference.max_write_queue_size
self._num_writers: int = max(1, self.reference.num_writers)
async def initialize(self):
"""Create the necessary tables if they don't exist."""
self.sql_store = AuthorizedSqlStore(sqlstore_impl(self.sql_store_config), self.policy)
base_store = sqlstore_impl(self.reference)
self.sql_store = AuthorizedSqlStore(base_store, self.policy)
backend_config = _SQLSTORE_BACKENDS.get(self.reference.backend)
if backend_config is None:
raise ValueError(
f"Unregistered SQL backend '{self.reference.backend}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
if backend_config.type == StorageBackendType.SQL_SQLITE:
self.enable_write_queue = False
await self.sql_store.create_table(
"openai_responses",
{

2
llama_stack/providers/utils/sqlstore/authorized_sqlstore.py
View file

@ -12,10 +12,10 @@ from llama_stack.core.access_control.conditions import ProtectedResource
from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
from llama_stack.core.datatypes import User
from llama_stack.core.request_headers import get_authenticated_user
from llama_stack.core.storage.datatypes import StorageBackendType
from llama_stack.log import get_logger
from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
from .sqlstore import StorageBackendType
logger = get_logger(name=__name__, category="providers::utils")

2
llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py
View file

@ -26,10 +26,10 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine
from sqlalchemy.sql.elements import ColumnElement
from llama_stack.apis.common.responses import PaginatedResponse
from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
from llama_stack.log import get_logger
from .api import ColumnDefinition, ColumnType, SqlStore
from .sqlstore import SqlAlchemySqlStoreConfig
logger = get_logger(name=__name__, category="providers::utils")

32
llama_stack/providers/utils/sqlstore/sqlstore.py
View file

@ -4,7 +4,7 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from typing import Annotated
from typing import Annotated, cast
from pydantic import Field
@ -21,7 +21,6 @@ from .api import SqlStore
sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
_SQLSTORE_BACKENDS: dict[str, StorageBackendConfig] = {}
_SQLSTORE_DEFAULT_BACKEND: str | None = None
SqlStoreConfig = Annotated[
@ -45,19 +44,18 @@ def get_pip_packages(store_config: dict | SqlStoreConfig) -> list[str]:
def sqlstore_impl(reference: SqlStoreReference) -> SqlStore:
backend_name = reference.backend or _SQLSTORE_DEFAULT_BACKEND
if not backend_name:
raise ValueError(
"SQL store reference did not specify a backend and no default backend is configured. "
f"Available backends: {sorted(_SQLSTORE_BACKENDS)}"
)
backend_name = reference.backend
backend_config = _SQLSTORE_BACKENDS.get(backend_name)
if backend_config.type in [StorageBackendType.SQL_SQLITE, StorageBackendType.SQL_POSTGRES]:
if backend_config is None:
raise ValueError(
f"Unknown SQL store backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
)
if isinstance(backend_config, SqliteSqlStoreConfig | PostgresSqlStoreConfig):
from .sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
config = backend_config.model_copy()
config.table_name = reference.table_name
config = cast(SqliteSqlStoreConfig | PostgresSqlStoreConfig, backend_config).model_copy()
return SqlAlchemySqlStoreImpl(config)
else:
raise ValueError(f"Unknown sqlstore type {backend_config.type}")
@ -67,18 +65,6 @@ def register_sqlstore_backends(backends: dict[str, StorageBackendConfig]) -> Non
"""Register the set of available SQL store backends for reference resolution."""
global _SQLSTORE_BACKENDS
def _set_default_backend(name: str) -> None:
global _SQLSTORE_DEFAULT_BACKEND
if _SQLSTORE_DEFAULT_BACKEND and _SQLSTORE_DEFAULT_BACKEND != name:
raise ValueError(
f"Multiple SQL store backends marked as default: '{_SQLSTORE_DEFAULT_BACKEND}' and '{name}'. "
"Only one backend can be the default."
)
_SQLSTORE_DEFAULT_BACKEND = name
_SQLSTORE_BACKENDS.clear()
for name, cfg in backends.items():
if cfg.default:
_set_default_backend(name)
_SQLSTORE_BACKENDS[name] = cfg