mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
Merge remote-tracking branch 'upstream/main' into url
This commit is contained in:
commit
b6aca4d366
79 changed files with 424 additions and 348 deletions
|
|
@ -6723,9 +6723,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -7125,6 +7126,11 @@ components:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
- type: 'null'
|
- type: 'null'
|
||||||
|
parallel_tool_calls:
|
||||||
|
anyOf:
|
||||||
|
- type: boolean
|
||||||
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -7251,9 +7257,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
|
||||||
19
docs/static/deprecated-llama-stack-spec.yaml
vendored
19
docs/static/deprecated-llama-stack-spec.yaml
vendored
|
|
@ -3566,9 +3566,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -3968,6 +3969,11 @@ components:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
- type: 'null'
|
- type: 'null'
|
||||||
|
parallel_tool_calls:
|
||||||
|
anyOf:
|
||||||
|
- type: boolean
|
||||||
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -4094,9 +4100,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
|
||||||
14
docs/static/experimental-llama-stack-spec.yaml
vendored
14
docs/static/experimental-llama-stack-spec.yaml
vendored
|
|
@ -3263,9 +3263,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -3662,9 +3663,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
|
||||||
19
docs/static/llama-stack-spec.yaml
vendored
19
docs/static/llama-stack-spec.yaml
vendored
|
|
@ -5744,9 +5744,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -6146,6 +6147,11 @@ components:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
- type: 'null'
|
- type: 'null'
|
||||||
|
parallel_tool_calls:
|
||||||
|
anyOf:
|
||||||
|
- type: boolean
|
||||||
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -6272,9 +6278,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
|
||||||
19
docs/static/stainless-llama-stack-spec.yaml
vendored
19
docs/static/stainless-llama-stack-spec.yaml
vendored
|
|
@ -6723,9 +6723,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -7125,6 +7126,11 @@ components:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
- type: 'null'
|
- type: 'null'
|
||||||
|
parallel_tool_calls:
|
||||||
|
anyOf:
|
||||||
|
- type: boolean
|
||||||
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
@ -7251,9 +7257,10 @@ components:
|
||||||
type: array
|
type: array
|
||||||
title: Output
|
title: Output
|
||||||
parallel_tool_calls:
|
parallel_tool_calls:
|
||||||
type: boolean
|
anyOf:
|
||||||
title: Parallel Tool Calls
|
- type: boolean
|
||||||
default: false
|
- type: 'null'
|
||||||
|
default: true
|
||||||
previous_response_id:
|
previous_response_id:
|
||||||
anyOf:
|
anyOf:
|
||||||
- type: string
|
- type: string
|
||||||
|
|
|
||||||
|
|
@ -356,6 +356,10 @@ exclude = [
|
||||||
module = [
|
module = [
|
||||||
"yaml",
|
"yaml",
|
||||||
"fire",
|
"fire",
|
||||||
|
"redis.asyncio",
|
||||||
|
"psycopg2",
|
||||||
|
"psycopg2.extras",
|
||||||
|
"psycopg2.extensions",
|
||||||
"torchtune.*",
|
"torchtune.*",
|
||||||
"fairscale.*",
|
"fairscale.*",
|
||||||
"torchvision.*",
|
"torchvision.*",
|
||||||
|
|
|
||||||
|
|
@ -11,10 +11,9 @@ from typing import Any, Literal
|
||||||
from pydantic import BaseModel, TypeAdapter
|
from pydantic import BaseModel, TypeAdapter
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule, StackRunConfig
|
from llama_stack.core.datatypes import AccessRule, StackRunConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
Conversation,
|
Conversation,
|
||||||
ConversationDeletedResource,
|
ConversationDeletedResource,
|
||||||
|
|
@ -25,6 +24,7 @@ from llama_stack_api import (
|
||||||
Conversations,
|
Conversations,
|
||||||
Metadata,
|
Metadata,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="openai_conversations")
|
logger = get_logger(name=__name__, category="openai_conversations")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,7 +10,7 @@ from typing import Any
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
from llama_stack.core.datatypes import StackRunConfig
|
from llama_stack.core.datatypes import StackRunConfig
|
||||||
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
|
from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
|
||||||
from llama_stack_api import ListPromptsResponse, Prompt, Prompts
|
from llama_stack_api import ListPromptsResponse, Prompt, Prompts
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -11,9 +11,9 @@ from datetime import UTC, datetime, timedelta
|
||||||
from starlette.types import ASGIApp, Receive, Scope, Send
|
from starlette.types import ASGIApp, Receive, Scope, Send
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendType
|
from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendType
|
||||||
|
from llama_stack.core.storage.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
from llama_stack.providers.utils.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
|
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="core::server")
|
logger = get_logger(name=__name__, category="core::server")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -385,8 +385,8 @@ def _initialize_storage(run_config: StackRunConfig):
|
||||||
else:
|
else:
|
||||||
raise ValueError(f"Unknown storage backend type: {type}")
|
raise ValueError(f"Unknown storage backend type: {type}")
|
||||||
|
|
||||||
from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends
|
from llama_stack.core.storage.kvstore.kvstore import register_kvstore_backends
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
|
|
||||||
register_kvstore_backends(kv_backends)
|
register_kvstore_backends(kv_backends)
|
||||||
register_sqlstore_backends(sql_backends)
|
register_sqlstore_backends(sql_backends)
|
||||||
|
|
|
||||||
|
|
@ -12,6 +12,8 @@ from typing import Annotated, Literal
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, field_validator
|
from pydantic import BaseModel, Field, field_validator
|
||||||
|
|
||||||
|
from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
|
||||||
|
|
||||||
|
|
||||||
class StorageBackendType(StrEnum):
|
class StorageBackendType(StrEnum):
|
||||||
KV_REDIS = "kv_redis"
|
KV_REDIS = "kv_redis"
|
||||||
|
|
@ -256,15 +258,24 @@ class ResponsesStoreReference(InferenceStoreReference):
|
||||||
|
|
||||||
class ServerStoresConfig(BaseModel):
|
class ServerStoresConfig(BaseModel):
|
||||||
metadata: KVStoreReference | None = Field(
|
metadata: KVStoreReference | None = Field(
|
||||||
default=None,
|
default=KVStoreReference(
|
||||||
|
backend="kv_default",
|
||||||
|
namespace="registry",
|
||||||
|
),
|
||||||
description="Metadata store configuration (uses KV backend)",
|
description="Metadata store configuration (uses KV backend)",
|
||||||
)
|
)
|
||||||
inference: InferenceStoreReference | None = Field(
|
inference: InferenceStoreReference | None = Field(
|
||||||
default=None,
|
default=InferenceStoreReference(
|
||||||
|
backend="sql_default",
|
||||||
|
table_name="inference_store",
|
||||||
|
),
|
||||||
description="Inference store configuration (uses SQL backend)",
|
description="Inference store configuration (uses SQL backend)",
|
||||||
)
|
)
|
||||||
conversations: SqlStoreReference | None = Field(
|
conversations: SqlStoreReference | None = Field(
|
||||||
default=None,
|
default=SqlStoreReference(
|
||||||
|
backend="sql_default",
|
||||||
|
table_name="openai_conversations",
|
||||||
|
),
|
||||||
description="Conversations store configuration (uses SQL backend)",
|
description="Conversations store configuration (uses SQL backend)",
|
||||||
)
|
)
|
||||||
responses: ResponsesStoreReference | None = Field(
|
responses: ResponsesStoreReference | None = Field(
|
||||||
|
|
@ -272,13 +283,21 @@ class ServerStoresConfig(BaseModel):
|
||||||
description="Responses store configuration (uses SQL backend)",
|
description="Responses store configuration (uses SQL backend)",
|
||||||
)
|
)
|
||||||
prompts: KVStoreReference | None = Field(
|
prompts: KVStoreReference | None = Field(
|
||||||
default=None,
|
default=KVStoreReference(backend="kv_default", namespace="prompts"),
|
||||||
description="Prompts store configuration (uses KV backend)",
|
description="Prompts store configuration (uses KV backend)",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class StorageConfig(BaseModel):
|
class StorageConfig(BaseModel):
|
||||||
backends: dict[str, StorageBackendConfig] = Field(
|
backends: dict[str, StorageBackendConfig] = Field(
|
||||||
|
default={
|
||||||
|
"kv_default": SqliteKVStoreConfig(
|
||||||
|
db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/kvstore.db",
|
||||||
|
),
|
||||||
|
"sql_default": SqliteSqlStoreConfig(
|
||||||
|
db_path=f"${{env.SQLITE_STORE_DIR:={DISTRIBS_BASE_DIR}}}/sql_store.db",
|
||||||
|
),
|
||||||
|
},
|
||||||
description="Named backend configurations (e.g., 'default', 'cache')",
|
description="Named backend configurations (e.g., 'default', 'cache')",
|
||||||
)
|
)
|
||||||
stores: ServerStoresConfig = Field(
|
stores: ServerStoresConfig = Field(
|
||||||
|
|
|
||||||
|
|
@ -4,4 +4,6 @@
|
||||||
# This source code is licensed under the terms described in the LICENSE file in
|
# This source code is licensed under the terms described in the LICENSE file in
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore as KVStore
|
||||||
|
|
||||||
from .kvstore import * # noqa: F401, F403
|
from .kvstore import * # noqa: F401, F403
|
||||||
|
|
@ -13,11 +13,19 @@ from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
|
from datetime import datetime
|
||||||
|
from typing import cast
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType
|
from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .api import KVStore
|
from .config import (
|
||||||
from .config import KVStoreConfig
|
KVStoreConfig,
|
||||||
|
MongoDBKVStoreConfig,
|
||||||
|
PostgresKVStoreConfig,
|
||||||
|
RedisKVStoreConfig,
|
||||||
|
SqliteKVStoreConfig,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def kvstore_dependencies():
|
def kvstore_dependencies():
|
||||||
|
|
@ -33,7 +41,7 @@ def kvstore_dependencies():
|
||||||
|
|
||||||
class InmemoryKVStoreImpl(KVStore):
|
class InmemoryKVStoreImpl(KVStore):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self._store = {}
|
self._store: dict[str, str] = {}
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
pass
|
pass
|
||||||
|
|
@ -41,7 +49,7 @@ class InmemoryKVStoreImpl(KVStore):
|
||||||
async def get(self, key: str) -> str | None:
|
async def get(self, key: str) -> str | None:
|
||||||
return self._store.get(key)
|
return self._store.get(key)
|
||||||
|
|
||||||
async def set(self, key: str, value: str) -> None:
|
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
|
||||||
self._store[key] = value
|
self._store[key] = value
|
||||||
|
|
||||||
async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
|
async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
|
||||||
|
|
@ -70,7 +78,8 @@ def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None
|
||||||
_KVSTORE_INSTANCES.clear()
|
_KVSTORE_INSTANCES.clear()
|
||||||
_KVSTORE_LOCKS.clear()
|
_KVSTORE_LOCKS.clear()
|
||||||
for name, cfg in backends.items():
|
for name, cfg in backends.items():
|
||||||
_KVSTORE_BACKENDS[name] = cfg
|
typed_cfg = cast(KVStoreConfig, cfg)
|
||||||
|
_KVSTORE_BACKENDS[name] = typed_cfg
|
||||||
|
|
||||||
|
|
||||||
async def kvstore_impl(reference: KVStoreReference) -> KVStore:
|
async def kvstore_impl(reference: KVStoreReference) -> KVStore:
|
||||||
|
|
@ -94,19 +103,20 @@ async def kvstore_impl(reference: KVStoreReference) -> KVStore:
|
||||||
config = backend_config.model_copy()
|
config = backend_config.model_copy()
|
||||||
config.namespace = reference.namespace
|
config.namespace = reference.namespace
|
||||||
|
|
||||||
if config.type == StorageBackendType.KV_REDIS.value:
|
impl: KVStore
|
||||||
|
if isinstance(config, RedisKVStoreConfig):
|
||||||
from .redis import RedisKVStoreImpl
|
from .redis import RedisKVStoreImpl
|
||||||
|
|
||||||
impl = RedisKVStoreImpl(config)
|
impl = RedisKVStoreImpl(config)
|
||||||
elif config.type == StorageBackendType.KV_SQLITE.value:
|
elif isinstance(config, SqliteKVStoreConfig):
|
||||||
from .sqlite import SqliteKVStoreImpl
|
from .sqlite import SqliteKVStoreImpl
|
||||||
|
|
||||||
impl = SqliteKVStoreImpl(config)
|
impl = SqliteKVStoreImpl(config)
|
||||||
elif config.type == StorageBackendType.KV_POSTGRES.value:
|
elif isinstance(config, PostgresKVStoreConfig):
|
||||||
from .postgres import PostgresKVStoreImpl
|
from .postgres import PostgresKVStoreImpl
|
||||||
|
|
||||||
impl = PostgresKVStoreImpl(config)
|
impl = PostgresKVStoreImpl(config)
|
||||||
elif config.type == StorageBackendType.KV_MONGODB.value:
|
elif isinstance(config, MongoDBKVStoreConfig):
|
||||||
from .mongodb import MongoDBKVStoreImpl
|
from .mongodb import MongoDBKVStoreImpl
|
||||||
|
|
||||||
impl = MongoDBKVStoreImpl(config)
|
impl = MongoDBKVStoreImpl(config)
|
||||||
|
|
@ -9,8 +9,8 @@ from datetime import datetime
|
||||||
from pymongo import AsyncMongoClient
|
from pymongo import AsyncMongoClient
|
||||||
from pymongo.asynchronous.collection import AsyncCollection
|
from pymongo.asynchronous.collection import AsyncCollection
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import KVStore
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import KVStore
|
|
||||||
|
|
||||||
from ..config import MongoDBKVStoreConfig
|
from ..config import MongoDBKVStoreConfig
|
||||||
|
|
||||||
|
|
@ -6,12 +6,13 @@
|
||||||
|
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
import psycopg2
|
import psycopg2 # type: ignore[import-not-found]
|
||||||
from psycopg2.extras import DictCursor
|
from psycopg2.extensions import connection as PGConnection # type: ignore[import-not-found]
|
||||||
|
from psycopg2.extras import DictCursor # type: ignore[import-not-found]
|
||||||
|
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from ..api import KVStore
|
|
||||||
from ..config import PostgresKVStoreConfig
|
from ..config import PostgresKVStoreConfig
|
||||||
|
|
||||||
log = get_logger(name=__name__, category="providers::utils")
|
log = get_logger(name=__name__, category="providers::utils")
|
||||||
|
|
@ -20,12 +21,12 @@ log = get_logger(name=__name__, category="providers::utils")
|
||||||
class PostgresKVStoreImpl(KVStore):
|
class PostgresKVStoreImpl(KVStore):
|
||||||
def __init__(self, config: PostgresKVStoreConfig):
|
def __init__(self, config: PostgresKVStoreConfig):
|
||||||
self.config = config
|
self.config = config
|
||||||
self.conn = None
|
self._conn: PGConnection | None = None
|
||||||
self.cursor = None
|
self._cursor: DictCursor | None = None
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
try:
|
try:
|
||||||
self.conn = psycopg2.connect(
|
self._conn = psycopg2.connect(
|
||||||
host=self.config.host,
|
host=self.config.host,
|
||||||
port=self.config.port,
|
port=self.config.port,
|
||||||
database=self.config.db,
|
database=self.config.db,
|
||||||
|
|
@ -34,11 +35,11 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
sslmode=self.config.ssl_mode,
|
sslmode=self.config.ssl_mode,
|
||||||
sslrootcert=self.config.ca_cert_path,
|
sslrootcert=self.config.ca_cert_path,
|
||||||
)
|
)
|
||||||
self.conn.autocommit = True
|
self._conn.autocommit = True
|
||||||
self.cursor = self.conn.cursor(cursor_factory=DictCursor)
|
self._cursor = self._conn.cursor(cursor_factory=DictCursor)
|
||||||
|
|
||||||
# Create table if it doesn't exist
|
# Create table if it doesn't exist
|
||||||
self.cursor.execute(
|
self._cursor.execute(
|
||||||
f"""
|
f"""
|
||||||
CREATE TABLE IF NOT EXISTS {self.config.table_name} (
|
CREATE TABLE IF NOT EXISTS {self.config.table_name} (
|
||||||
key TEXT PRIMARY KEY,
|
key TEXT PRIMARY KEY,
|
||||||
|
|
@ -51,6 +52,11 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
log.exception("Could not connect to PostgreSQL database server")
|
log.exception("Could not connect to PostgreSQL database server")
|
||||||
raise RuntimeError("Could not connect to PostgreSQL database server") from e
|
raise RuntimeError("Could not connect to PostgreSQL database server") from e
|
||||||
|
|
||||||
|
def _cursor_or_raise(self) -> DictCursor:
|
||||||
|
if self._cursor is None:
|
||||||
|
raise RuntimeError("Postgres client not initialized")
|
||||||
|
return self._cursor
|
||||||
|
|
||||||
def _namespaced_key(self, key: str) -> str:
|
def _namespaced_key(self, key: str) -> str:
|
||||||
if not self.config.namespace:
|
if not self.config.namespace:
|
||||||
return key
|
return key
|
||||||
|
|
@ -58,7 +64,8 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
|
|
||||||
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
|
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
self.cursor.execute(
|
cursor = self._cursor_or_raise()
|
||||||
|
cursor.execute(
|
||||||
f"""
|
f"""
|
||||||
INSERT INTO {self.config.table_name} (key, value, expiration)
|
INSERT INTO {self.config.table_name} (key, value, expiration)
|
||||||
VALUES (%s, %s, %s)
|
VALUES (%s, %s, %s)
|
||||||
|
|
@ -70,7 +77,8 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
|
|
||||||
async def get(self, key: str) -> str | None:
|
async def get(self, key: str) -> str | None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
self.cursor.execute(
|
cursor = self._cursor_or_raise()
|
||||||
|
cursor.execute(
|
||||||
f"""
|
f"""
|
||||||
SELECT value FROM {self.config.table_name}
|
SELECT value FROM {self.config.table_name}
|
||||||
WHERE key = %s
|
WHERE key = %s
|
||||||
|
|
@ -78,12 +86,13 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
""",
|
""",
|
||||||
(key,),
|
(key,),
|
||||||
)
|
)
|
||||||
result = self.cursor.fetchone()
|
result = cursor.fetchone()
|
||||||
return result[0] if result else None
|
return result[0] if result else None
|
||||||
|
|
||||||
async def delete(self, key: str) -> None:
|
async def delete(self, key: str) -> None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
self.cursor.execute(
|
cursor = self._cursor_or_raise()
|
||||||
|
cursor.execute(
|
||||||
f"DELETE FROM {self.config.table_name} WHERE key = %s",
|
f"DELETE FROM {self.config.table_name} WHERE key = %s",
|
||||||
(key,),
|
(key,),
|
||||||
)
|
)
|
||||||
|
|
@ -92,7 +101,8 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
start_key = self._namespaced_key(start_key)
|
start_key = self._namespaced_key(start_key)
|
||||||
end_key = self._namespaced_key(end_key)
|
end_key = self._namespaced_key(end_key)
|
||||||
|
|
||||||
self.cursor.execute(
|
cursor = self._cursor_or_raise()
|
||||||
|
cursor.execute(
|
||||||
f"""
|
f"""
|
||||||
SELECT value FROM {self.config.table_name}
|
SELECT value FROM {self.config.table_name}
|
||||||
WHERE key >= %s AND key < %s
|
WHERE key >= %s AND key < %s
|
||||||
|
|
@ -101,14 +111,15 @@ class PostgresKVStoreImpl(KVStore):
|
||||||
""",
|
""",
|
||||||
(start_key, end_key),
|
(start_key, end_key),
|
||||||
)
|
)
|
||||||
return [row[0] for row in self.cursor.fetchall()]
|
return [row[0] for row in cursor.fetchall()]
|
||||||
|
|
||||||
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
|
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
|
||||||
start_key = self._namespaced_key(start_key)
|
start_key = self._namespaced_key(start_key)
|
||||||
end_key = self._namespaced_key(end_key)
|
end_key = self._namespaced_key(end_key)
|
||||||
|
|
||||||
self.cursor.execute(
|
cursor = self._cursor_or_raise()
|
||||||
|
cursor.execute(
|
||||||
f"SELECT key FROM {self.config.table_name} WHERE key >= %s AND key < %s",
|
f"SELECT key FROM {self.config.table_name} WHERE key >= %s AND key < %s",
|
||||||
(start_key, end_key),
|
(start_key, end_key),
|
||||||
)
|
)
|
||||||
return [row[0] for row in self.cursor.fetchall()]
|
return [row[0] for row in cursor.fetchall()]
|
||||||
|
|
@ -6,18 +6,25 @@
|
||||||
|
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
|
|
||||||
from redis.asyncio import Redis
|
from redis.asyncio import Redis # type: ignore[import-not-found]
|
||||||
|
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from ..api import KVStore
|
|
||||||
from ..config import RedisKVStoreConfig
|
from ..config import RedisKVStoreConfig
|
||||||
|
|
||||||
|
|
||||||
class RedisKVStoreImpl(KVStore):
|
class RedisKVStoreImpl(KVStore):
|
||||||
def __init__(self, config: RedisKVStoreConfig):
|
def __init__(self, config: RedisKVStoreConfig):
|
||||||
self.config = config
|
self.config = config
|
||||||
|
self._redis: Redis | None = None
|
||||||
|
|
||||||
async def initialize(self) -> None:
|
async def initialize(self) -> None:
|
||||||
self.redis = Redis.from_url(self.config.url)
|
self._redis = Redis.from_url(self.config.url)
|
||||||
|
|
||||||
|
def _client(self) -> Redis:
|
||||||
|
if self._redis is None:
|
||||||
|
raise RuntimeError("Redis client not initialized")
|
||||||
|
return self._redis
|
||||||
|
|
||||||
def _namespaced_key(self, key: str) -> str:
|
def _namespaced_key(self, key: str) -> str:
|
||||||
if not self.config.namespace:
|
if not self.config.namespace:
|
||||||
|
|
@ -26,30 +33,37 @@ class RedisKVStoreImpl(KVStore):
|
||||||
|
|
||||||
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
|
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
await self.redis.set(key, value)
|
client = self._client()
|
||||||
|
await client.set(key, value)
|
||||||
if expiration:
|
if expiration:
|
||||||
await self.redis.expireat(key, expiration)
|
await client.expireat(key, expiration)
|
||||||
|
|
||||||
async def get(self, key: str) -> str | None:
|
async def get(self, key: str) -> str | None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
value = await self.redis.get(key)
|
client = self._client()
|
||||||
|
value = await client.get(key)
|
||||||
if value is None:
|
if value is None:
|
||||||
return None
|
return None
|
||||||
await self.redis.ttl(key)
|
await client.ttl(key)
|
||||||
return value
|
if isinstance(value, bytes):
|
||||||
|
return value.decode("utf-8")
|
||||||
|
if isinstance(value, str):
|
||||||
|
return value
|
||||||
|
return str(value)
|
||||||
|
|
||||||
async def delete(self, key: str) -> None:
|
async def delete(self, key: str) -> None:
|
||||||
key = self._namespaced_key(key)
|
key = self._namespaced_key(key)
|
||||||
await self.redis.delete(key)
|
await self._client().delete(key)
|
||||||
|
|
||||||
async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
|
async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
|
||||||
start_key = self._namespaced_key(start_key)
|
start_key = self._namespaced_key(start_key)
|
||||||
end_key = self._namespaced_key(end_key)
|
end_key = self._namespaced_key(end_key)
|
||||||
|
client = self._client()
|
||||||
cursor = 0
|
cursor = 0
|
||||||
pattern = start_key + "*" # Match all keys starting with start_key prefix
|
pattern = start_key + "*" # Match all keys starting with start_key prefix
|
||||||
matching_keys = []
|
matching_keys: list[str | bytes] = []
|
||||||
while True:
|
while True:
|
||||||
cursor, keys = await self.redis.scan(cursor, match=pattern, count=1000)
|
cursor, keys = await client.scan(cursor, match=pattern, count=1000)
|
||||||
|
|
||||||
for key in keys:
|
for key in keys:
|
||||||
key_str = key.decode("utf-8") if isinstance(key, bytes) else key
|
key_str = key.decode("utf-8") if isinstance(key, bytes) else key
|
||||||
|
|
@ -61,7 +75,7 @@ class RedisKVStoreImpl(KVStore):
|
||||||
|
|
||||||
# Then fetch all values in a single MGET call
|
# Then fetch all values in a single MGET call
|
||||||
if matching_keys:
|
if matching_keys:
|
||||||
values = await self.redis.mget(matching_keys)
|
values = await client.mget(matching_keys)
|
||||||
return [
|
return [
|
||||||
value.decode("utf-8") if isinstance(value, bytes) else value for value in values if value is not None
|
value.decode("utf-8") if isinstance(value, bytes) else value for value in values if value is not None
|
||||||
]
|
]
|
||||||
|
|
@ -70,7 +84,18 @@ class RedisKVStoreImpl(KVStore):
|
||||||
|
|
||||||
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
|
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
|
||||||
"""Get all keys in the given range."""
|
"""Get all keys in the given range."""
|
||||||
matching_keys = await self.redis.zrangebylex(self.namespace, f"[{start_key}", f"[{end_key}")
|
start_key = self._namespaced_key(start_key)
|
||||||
if not matching_keys:
|
end_key = self._namespaced_key(end_key)
|
||||||
return []
|
client = self._client()
|
||||||
return [k.decode("utf-8") for k in matching_keys]
|
cursor = 0
|
||||||
|
pattern = start_key + "*"
|
||||||
|
result: list[str] = []
|
||||||
|
while True:
|
||||||
|
cursor, keys = await client.scan(cursor, match=pattern, count=1000)
|
||||||
|
for key in keys:
|
||||||
|
key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
|
||||||
|
if start_key <= key_str <= end_key:
|
||||||
|
result.append(key_str)
|
||||||
|
if cursor == 0:
|
||||||
|
break
|
||||||
|
return result
|
||||||
|
|
@ -10,8 +10,8 @@ from datetime import datetime
|
||||||
import aiosqlite
|
import aiosqlite
|
||||||
|
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from ..api import KVStore
|
|
||||||
from ..config import SqliteKVStoreConfig
|
from ..config import SqliteKVStoreConfig
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="providers::utils")
|
logger = get_logger(name=__name__, category="providers::utils")
|
||||||
17
src/llama_stack/core/storage/sqlstore/__init__.py
Normal file
17
src/llama_stack/core/storage/sqlstore/__init__.py
Normal file
|
|
@ -0,0 +1,17 @@
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from llama_stack_api.internal.sqlstore import (
|
||||||
|
ColumnDefinition as ColumnDefinition,
|
||||||
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import (
|
||||||
|
ColumnType as ColumnType,
|
||||||
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import (
|
||||||
|
SqlStore as SqlStore,
|
||||||
|
)
|
||||||
|
|
||||||
|
from .sqlstore import * # noqa: F401,F403
|
||||||
|
|
@ -14,8 +14,8 @@ from llama_stack.core.datatypes import User
|
||||||
from llama_stack.core.request_headers import get_authenticated_user
|
from llama_stack.core.request_headers import get_authenticated_user
|
||||||
from llama_stack.core.storage.datatypes import StorageBackendType
|
from llama_stack.core.storage.datatypes import StorageBackendType
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
|
from llama_stack_api import PaginatedResponse
|
||||||
from .api import ColumnDefinition, ColumnType, PaginatedResponse, SqlStore
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="providers::utils")
|
logger = get_logger(name=__name__, category="providers::utils")
|
||||||
|
|
||||||
|
|
@ -29,8 +29,7 @@ from sqlalchemy.sql.elements import ColumnElement
|
||||||
from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
|
from llama_stack.core.storage.datatypes import SqlAlchemySqlStoreConfig
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack_api import PaginatedResponse
|
from llama_stack_api import PaginatedResponse
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType, SqlStore
|
||||||
from .api import ColumnDefinition, ColumnType, SqlStore
|
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="providers::utils")
|
logger = get_logger(name=__name__, category="providers::utils")
|
||||||
|
|
||||||
|
|
@ -16,8 +16,7 @@ from llama_stack.core.storage.datatypes import (
|
||||||
StorageBackendConfig,
|
StorageBackendConfig,
|
||||||
StorageBackendType,
|
StorageBackendType,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import SqlStore
|
||||||
from .api import SqlStore
|
|
||||||
|
|
||||||
sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
|
sql_store_pip_packages = ["sqlalchemy[asyncio]", "aiosqlite", "asyncpg"]
|
||||||
|
|
||||||
|
|
@ -12,8 +12,8 @@ import pydantic
|
||||||
|
|
||||||
from llama_stack.core.datatypes import RoutableObjectWithProvider
|
from llama_stack.core.datatypes import RoutableObjectWithProvider
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference
|
from llama_stack.core.storage.datatypes import KVStoreReference
|
||||||
|
from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
|
|
||||||
|
|
||||||
logger = get_logger(__name__, category="core::registry")
|
logger = get_logger(__name__, category="core::registry")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,6 +17,8 @@ from llama_stack.core.datatypes import (
|
||||||
ToolGroupInput,
|
ToolGroupInput,
|
||||||
VectorStoresConfig,
|
VectorStoresConfig,
|
||||||
)
|
)
|
||||||
|
from llama_stack.core.storage.kvstore.config import PostgresKVStoreConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import PostgresSqlStoreConfig
|
||||||
from llama_stack.core.utils.dynamic import instantiate_class_type
|
from llama_stack.core.utils.dynamic import instantiate_class_type
|
||||||
from llama_stack.distributions.template import DistributionTemplate, RunConfigSettings
|
from llama_stack.distributions.template import DistributionTemplate, RunConfigSettings
|
||||||
from llama_stack.providers.inline.files.localfs.config import LocalfsFilesImplConfig
|
from llama_stack.providers.inline.files.localfs.config import LocalfsFilesImplConfig
|
||||||
|
|
@ -35,8 +37,6 @@ from llama_stack.providers.remote.vector_io.pgvector.config import (
|
||||||
)
|
)
|
||||||
from llama_stack.providers.remote.vector_io.qdrant.config import QdrantVectorIOConfig
|
from llama_stack.providers.remote.vector_io.qdrant.config import QdrantVectorIOConfig
|
||||||
from llama_stack.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig
|
from llama_stack.providers.remote.vector_io.weaviate.config import WeaviateVectorIOConfig
|
||||||
from llama_stack.providers.utils.kvstore.config import PostgresKVStoreConfig
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig
|
|
||||||
from llama_stack_api import RemoteProviderSpec
|
from llama_stack_api import RemoteProviderSpec
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -35,13 +35,13 @@ from llama_stack.core.storage.datatypes import (
|
||||||
SqlStoreReference,
|
SqlStoreReference,
|
||||||
StorageBackendType,
|
StorageBackendType,
|
||||||
)
|
)
|
||||||
|
from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
|
||||||
|
from llama_stack.core.storage.kvstore.config import get_pip_packages as get_kv_pip_packages
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import get_pip_packages as get_sql_pip_packages
|
||||||
from llama_stack.core.utils.dynamic import instantiate_class_type
|
from llama_stack.core.utils.dynamic import instantiate_class_type
|
||||||
from llama_stack.core.utils.image_types import LlamaStackImageType
|
from llama_stack.core.utils.image_types import LlamaStackImageType
|
||||||
from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
|
from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
|
||||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
|
||||||
from llama_stack.providers.utils.kvstore.config import get_pip_packages as get_kv_pip_packages
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import get_pip_packages as get_sql_pip_packages
|
|
||||||
from llama_stack_api import DatasetPurpose, ModelType
|
from llama_stack_api import DatasetPurpose, ModelType
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,8 +6,8 @@
|
||||||
|
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
|
from llama_stack.core.storage.kvstore import InmemoryKVStoreImpl, kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
|
|
||||||
from llama_stack.providers.utils.responses.responses_store import ResponsesStore
|
from llama_stack.providers.utils.responses.responses_store import ResponsesStore
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
Agents,
|
Agents,
|
||||||
|
|
@ -92,6 +92,7 @@ class MetaReferenceAgentsImpl(Agents):
|
||||||
model: str,
|
model: str,
|
||||||
prompt: OpenAIResponsePrompt | None = None,
|
prompt: OpenAIResponsePrompt | None = None,
|
||||||
instructions: str | None = None,
|
instructions: str | None = None,
|
||||||
|
parallel_tool_calls: bool | None = True,
|
||||||
previous_response_id: str | None = None,
|
previous_response_id: str | None = None,
|
||||||
conversation: str | None = None,
|
conversation: str | None = None,
|
||||||
store: bool | None = True,
|
store: bool | None = True,
|
||||||
|
|
@ -120,6 +121,7 @@ class MetaReferenceAgentsImpl(Agents):
|
||||||
include,
|
include,
|
||||||
max_infer_iters,
|
max_infer_iters,
|
||||||
guardrails,
|
guardrails,
|
||||||
|
parallel_tool_calls,
|
||||||
max_tool_calls,
|
max_tool_calls,
|
||||||
)
|
)
|
||||||
return result # type: ignore[no-any-return]
|
return result # type: ignore[no-any-return]
|
||||||
|
|
|
||||||
|
|
@ -252,6 +252,7 @@ class OpenAIResponsesImpl:
|
||||||
include: list[str] | None = None,
|
include: list[str] | None = None,
|
||||||
max_infer_iters: int | None = 10,
|
max_infer_iters: int | None = 10,
|
||||||
guardrails: list[str | ResponseGuardrailSpec] | None = None,
|
guardrails: list[str | ResponseGuardrailSpec] | None = None,
|
||||||
|
parallel_tool_calls: bool | None = None,
|
||||||
max_tool_calls: int | None = None,
|
max_tool_calls: int | None = None,
|
||||||
):
|
):
|
||||||
stream = bool(stream)
|
stream = bool(stream)
|
||||||
|
|
@ -296,6 +297,7 @@ class OpenAIResponsesImpl:
|
||||||
tools=tools,
|
tools=tools,
|
||||||
max_infer_iters=max_infer_iters,
|
max_infer_iters=max_infer_iters,
|
||||||
guardrail_ids=guardrail_ids,
|
guardrail_ids=guardrail_ids,
|
||||||
|
parallel_tool_calls=parallel_tool_calls,
|
||||||
max_tool_calls=max_tool_calls,
|
max_tool_calls=max_tool_calls,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -346,6 +348,7 @@ class OpenAIResponsesImpl:
|
||||||
tools: list[OpenAIResponseInputTool] | None = None,
|
tools: list[OpenAIResponseInputTool] | None = None,
|
||||||
max_infer_iters: int | None = 10,
|
max_infer_iters: int | None = 10,
|
||||||
guardrail_ids: list[str] | None = None,
|
guardrail_ids: list[str] | None = None,
|
||||||
|
parallel_tool_calls: bool | None = True,
|
||||||
max_tool_calls: int | None = None,
|
max_tool_calls: int | None = None,
|
||||||
) -> AsyncIterator[OpenAIResponseObjectStream]:
|
) -> AsyncIterator[OpenAIResponseObjectStream]:
|
||||||
# These should never be None when called from create_openai_response (which sets defaults)
|
# These should never be None when called from create_openai_response (which sets defaults)
|
||||||
|
|
@ -385,6 +388,7 @@ class OpenAIResponsesImpl:
|
||||||
created_at=created_at,
|
created_at=created_at,
|
||||||
text=text,
|
text=text,
|
||||||
max_infer_iters=max_infer_iters,
|
max_infer_iters=max_infer_iters,
|
||||||
|
parallel_tool_calls=parallel_tool_calls,
|
||||||
tool_executor=self.tool_executor,
|
tool_executor=self.tool_executor,
|
||||||
safety_api=self.safety_api,
|
safety_api=self.safety_api,
|
||||||
guardrail_ids=guardrail_ids,
|
guardrail_ids=guardrail_ids,
|
||||||
|
|
|
||||||
|
|
@ -114,6 +114,7 @@ class StreamingResponseOrchestrator:
|
||||||
safety_api,
|
safety_api,
|
||||||
guardrail_ids: list[str] | None = None,
|
guardrail_ids: list[str] | None = None,
|
||||||
prompt: OpenAIResponsePrompt | None = None,
|
prompt: OpenAIResponsePrompt | None = None,
|
||||||
|
parallel_tool_calls: bool | None = None,
|
||||||
max_tool_calls: int | None = None,
|
max_tool_calls: int | None = None,
|
||||||
):
|
):
|
||||||
self.inference_api = inference_api
|
self.inference_api = inference_api
|
||||||
|
|
@ -128,6 +129,8 @@ class StreamingResponseOrchestrator:
|
||||||
self.prompt = prompt
|
self.prompt = prompt
|
||||||
# System message that is inserted into the model's context
|
# System message that is inserted into the model's context
|
||||||
self.instructions = instructions
|
self.instructions = instructions
|
||||||
|
# Whether to allow more than one function tool call generated per turn.
|
||||||
|
self.parallel_tool_calls = parallel_tool_calls
|
||||||
# Max number of total calls to built-in tools that can be processed in a response
|
# Max number of total calls to built-in tools that can be processed in a response
|
||||||
self.max_tool_calls = max_tool_calls
|
self.max_tool_calls = max_tool_calls
|
||||||
self.sequence_number = 0
|
self.sequence_number = 0
|
||||||
|
|
@ -190,6 +193,7 @@ class StreamingResponseOrchestrator:
|
||||||
usage=self.accumulated_usage,
|
usage=self.accumulated_usage,
|
||||||
instructions=self.instructions,
|
instructions=self.instructions,
|
||||||
prompt=self.prompt,
|
prompt=self.prompt,
|
||||||
|
parallel_tool_calls=self.parallel_tool_calls,
|
||||||
max_tool_calls=self.max_tool_calls,
|
max_tool_calls=self.max_tool_calls,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,7 +7,7 @@
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule, Api
|
from llama_stack.core.datatypes import AccessRule, Api
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack_api import Files, Inference, Models
|
from llama_stack_api import Files, Inference, Models
|
||||||
|
|
||||||
from .batches import ReferenceBatchesImpl
|
from .batches import ReferenceBatchesImpl
|
||||||
|
|
|
||||||
|
|
@ -16,8 +16,8 @@ from typing import Any, Literal
|
||||||
from openai.types.batch import BatchError, Errors
|
from openai.types.batch import BatchError, Errors
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import KVStore
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import KVStore
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
Batches,
|
Batches,
|
||||||
BatchObject,
|
BatchObject,
|
||||||
|
|
|
||||||
|
|
@ -5,8 +5,8 @@
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_uri
|
from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_uri
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.pagination import paginate_records
|
from llama_stack.providers.utils.pagination import paginate_records
|
||||||
from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
|
from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -8,8 +8,8 @@ from typing import Any
|
||||||
|
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.providers.utils.common.data_schema_validator import ColumnName
|
from llama_stack.providers.utils.common.data_schema_validator import ColumnName
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
Agents,
|
Agents,
|
||||||
Benchmark,
|
Benchmark,
|
||||||
|
|
|
||||||
|
|
@ -13,11 +13,10 @@ from fastapi import Depends, File, Form, Response, UploadFile
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
from llama_stack.core.id_generation import generate_object_id
|
from llama_stack.core.id_generation import generate_object_id
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
ExpiresAfter,
|
ExpiresAfter,
|
||||||
Files,
|
Files,
|
||||||
|
|
@ -28,6 +27,7 @@ from llama_stack_api import (
|
||||||
Order,
|
Order,
|
||||||
ResourceNotFoundError,
|
ResourceNotFoundError,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
|
|
||||||
from .config import LocalfsFilesImplConfig
|
from .config import LocalfsFilesImplConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,9 +14,8 @@ import faiss # type: ignore[import-untyped]
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from numpy.typing import NDArray
|
from numpy.typing import NDArray
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
|
|
@ -32,6 +31,7 @@ from llama_stack_api import (
|
||||||
VectorStoreNotFoundError,
|
VectorStoreNotFoundError,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .config import FaissVectorIOConfig
|
from .config import FaissVectorIOConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,9 +14,8 @@ import numpy as np
|
||||||
import sqlite_vec # type: ignore[import-untyped]
|
import sqlite_vec # type: ignore[import-untyped]
|
||||||
from numpy.typing import NDArray
|
from numpy.typing import NDArray
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import (
|
from llama_stack.providers.utils.memory.vector_store import (
|
||||||
RERANKER_TYPE_RRF,
|
RERANKER_TYPE_RRF,
|
||||||
|
|
@ -35,6 +34,7 @@ from llama_stack_api import (
|
||||||
VectorStoreNotFoundError,
|
VectorStoreNotFoundError,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="vector_io")
|
logger = get_logger(name=__name__, category="vector_io")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,7 @@
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
|
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_dependencies
|
from llama_stack.core.storage.kvstore import kvstore_dependencies
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
Api,
|
Api,
|
||||||
InlineProviderSpec,
|
InlineProviderSpec,
|
||||||
|
|
|
||||||
|
|
@ -4,7 +4,7 @@
|
||||||
# This source code is licensed under the terms described in the LICENSE file in
|
# This source code is licensed under the terms described in the LICENSE file in
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import sql_store_pip_packages
|
from llama_stack.core.storage.sqlstore.sqlstore import sql_store_pip_packages
|
||||||
from llama_stack_api import Api, InlineProviderSpec, ProviderSpec, RemoteProviderSpec
|
from llama_stack_api import Api, InlineProviderSpec, ProviderSpec, RemoteProviderSpec
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,7 @@
|
||||||
from typing import Any
|
from typing import Any
|
||||||
from urllib.parse import parse_qs, urlparse
|
from urllib.parse import parse_qs, urlparse
|
||||||
|
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.providers.utils.pagination import paginate_records
|
from llama_stack.providers.utils.pagination import paginate_records
|
||||||
from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
|
from llama_stack_api import Dataset, DatasetIO, DatasetsProtocolPrivate, PaginatedResponse
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,10 +10,9 @@ from typing import Annotated, Any
|
||||||
from fastapi import Depends, File, Form, Response, UploadFile
|
from fastapi import Depends, File, Form, Response, UploadFile
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
|
||||||
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
ExpiresAfter,
|
ExpiresAfter,
|
||||||
Files,
|
Files,
|
||||||
|
|
@ -24,6 +23,7 @@ from llama_stack_api import (
|
||||||
Order,
|
Order,
|
||||||
ResourceNotFoundError,
|
ResourceNotFoundError,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
|
|
||||||
from .config import OpenAIFilesImplConfig
|
from .config import OpenAIFilesImplConfig
|
||||||
|
|
|
||||||
|
|
@ -19,10 +19,9 @@ if TYPE_CHECKING:
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
from llama_stack.core.id_generation import generate_object_id
|
from llama_stack.core.id_generation import generate_object_id
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
|
||||||
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
from llama_stack.providers.utils.files.form_data import parse_expires_after
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
ExpiresAfter,
|
ExpiresAfter,
|
||||||
Files,
|
Files,
|
||||||
|
|
@ -33,6 +32,7 @@ from llama_stack_api import (
|
||||||
Order,
|
Order,
|
||||||
ResourceNotFoundError,
|
ResourceNotFoundError,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
|
|
||||||
from .config import S3FilesImplConfig
|
from .config import S3FilesImplConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -11,10 +11,9 @@ from urllib.parse import urlparse
|
||||||
import chromadb
|
import chromadb
|
||||||
from numpy.typing import NDArray
|
from numpy.typing import NDArray
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
|
from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
|
|
@ -27,6 +26,7 @@ from llama_stack_api import (
|
||||||
VectorStore,
|
VectorStore,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
|
from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -11,10 +11,9 @@ from typing import Any
|
||||||
from numpy.typing import NDArray
|
from numpy.typing import NDArray
|
||||||
from pymilvus import AnnSearchRequest, DataType, Function, FunctionType, MilvusClient, RRFRanker, WeightedRanker
|
from pymilvus import AnnSearchRequest, DataType, Function, FunctionType, MilvusClient, RRFRanker, WeightedRanker
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
|
from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import (
|
from llama_stack.providers.utils.memory.vector_store import (
|
||||||
RERANKER_TYPE_WEIGHTED,
|
RERANKER_TYPE_WEIGHTED,
|
||||||
|
|
@ -34,6 +33,7 @@ from llama_stack_api import (
|
||||||
VectorStoreNotFoundError,
|
VectorStoreNotFoundError,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
|
from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -13,10 +13,9 @@ from psycopg2 import sql
|
||||||
from psycopg2.extras import Json, execute_values
|
from psycopg2.extras import Json, execute_values
|
||||||
from pydantic import BaseModel, TypeAdapter
|
from pydantic import BaseModel, TypeAdapter
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
|
from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
||||||
from llama_stack.providers.utils.vector_io.vector_utils import WeightedInMemoryAggregator, sanitize_collection_name
|
from llama_stack.providers.utils.vector_io.vector_utils import WeightedInMemoryAggregator, sanitize_collection_name
|
||||||
|
|
@ -31,6 +30,7 @@ from llama_stack_api import (
|
||||||
VectorStoreNotFoundError,
|
VectorStoreNotFoundError,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .config import PGVectorVectorIOConfig
|
from .config import PGVectorVectorIOConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -13,9 +13,9 @@ from numpy.typing import NDArray
|
||||||
from qdrant_client import AsyncQdrantClient, models
|
from qdrant_client import AsyncQdrantClient, models
|
||||||
from qdrant_client.models import PointStruct
|
from qdrant_client.models import PointStruct
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
|
from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
|
|
|
||||||
|
|
@ -13,9 +13,8 @@ from weaviate.classes.init import Auth
|
||||||
from weaviate.classes.query import Filter, HybridFusion
|
from weaviate.classes.query import Filter, HybridFusion
|
||||||
|
|
||||||
from llama_stack.core.request_headers import NeedsRequestProviderData
|
from llama_stack.core.request_headers import NeedsRequestProviderData
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl
|
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
|
||||||
from llama_stack.providers.utils.memory.vector_store import (
|
from llama_stack.providers.utils.memory.vector_store import (
|
||||||
RERANKER_TYPE_RRF,
|
RERANKER_TYPE_RRF,
|
||||||
|
|
@ -35,6 +34,7 @@ from llama_stack_api import (
|
||||||
VectorStoreNotFoundError,
|
VectorStoreNotFoundError,
|
||||||
VectorStoresProtocolPrivate,
|
VectorStoresProtocolPrivate,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
from .config import WeaviateVectorIOConfig
|
from .config import WeaviateVectorIOConfig
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,8 @@ from sqlalchemy.exc import IntegrityError
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
|
from llama_stack.core.storage.datatypes import InferenceStoreReference, StorageBackendType
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
ListOpenAIChatCompletionResponse,
|
ListOpenAIChatCompletionResponse,
|
||||||
|
|
@ -18,10 +20,7 @@ from llama_stack_api import (
|
||||||
OpenAIMessageParam,
|
OpenAIMessageParam,
|
||||||
Order,
|
Order,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
from ..sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from ..sqlstore.sqlstore import _SQLSTORE_BACKENDS, sqlstore_impl
|
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="inference")
|
logger = get_logger(name=__name__, category="inference")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,20 +0,0 @@
|
||||||
# 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.
|
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
|
||||||
|
|
||||||
from llama_stack_api import json_schema_type
|
|
||||||
|
|
||||||
|
|
||||||
@json_schema_type
|
|
||||||
class SqliteControlPlaneConfig(BaseModel):
|
|
||||||
db_path: str = Field(
|
|
||||||
description="File path for the sqlite database",
|
|
||||||
)
|
|
||||||
table_name: str = Field(
|
|
||||||
default="llamastack_control_plane",
|
|
||||||
description="Table into which all the keys will be placed",
|
|
||||||
)
|
|
||||||
|
|
@ -17,7 +17,6 @@ from pydantic import TypeAdapter
|
||||||
|
|
||||||
from llama_stack.core.id_generation import generate_object_id
|
from llama_stack.core.id_generation import generate_object_id
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack.providers.utils.kvstore.api import KVStore
|
|
||||||
from llama_stack.providers.utils.memory.vector_store import (
|
from llama_stack.providers.utils.memory.vector_store import (
|
||||||
ChunkForDeletion,
|
ChunkForDeletion,
|
||||||
content_from_data_and_mime_type,
|
content_from_data_and_mime_type,
|
||||||
|
|
@ -53,6 +52,7 @@ from llama_stack_api import (
|
||||||
VectorStoreSearchResponse,
|
VectorStoreSearchResponse,
|
||||||
VectorStoreSearchResponsePage,
|
VectorStoreSearchResponsePage,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.kvstore import KVStore
|
||||||
|
|
||||||
EMBEDDING_DIMENSION = 768
|
EMBEDDING_DIMENSION = 768
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,6 +6,8 @@
|
||||||
|
|
||||||
from llama_stack.core.datatypes import AccessRule
|
from llama_stack.core.datatypes import AccessRule
|
||||||
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference
|
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqlStoreReference
|
||||||
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
|
||||||
from llama_stack.log import get_logger
|
from llama_stack.log import get_logger
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
ListOpenAIResponseInputItem,
|
ListOpenAIResponseInputItem,
|
||||||
|
|
@ -17,10 +19,7 @@ from llama_stack_api import (
|
||||||
OpenAIResponseObjectWithInput,
|
OpenAIResponseObjectWithInput,
|
||||||
Order,
|
Order,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
from ..sqlstore.api import ColumnDefinition, ColumnType
|
|
||||||
from ..sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
|
||||||
from ..sqlstore.sqlstore import sqlstore_impl
|
|
||||||
|
|
||||||
logger = get_logger(name=__name__, category="openai_responses")
|
logger = get_logger(name=__name__, category="openai_responses")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,140 +0,0 @@
|
||||||
# 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.
|
|
||||||
|
|
||||||
from collections.abc import Mapping, Sequence
|
|
||||||
from enum import Enum
|
|
||||||
from typing import Any, Literal, Protocol
|
|
||||||
|
|
||||||
from pydantic import BaseModel
|
|
||||||
|
|
||||||
from llama_stack_api import PaginatedResponse
|
|
||||||
|
|
||||||
|
|
||||||
class ColumnType(Enum):
|
|
||||||
INTEGER = "INTEGER"
|
|
||||||
STRING = "STRING"
|
|
||||||
TEXT = "TEXT"
|
|
||||||
FLOAT = "FLOAT"
|
|
||||||
BOOLEAN = "BOOLEAN"
|
|
||||||
JSON = "JSON"
|
|
||||||
DATETIME = "DATETIME"
|
|
||||||
|
|
||||||
|
|
||||||
class ColumnDefinition(BaseModel):
|
|
||||||
type: ColumnType
|
|
||||||
primary_key: bool = False
|
|
||||||
nullable: bool = True
|
|
||||||
default: Any = None
|
|
||||||
|
|
||||||
|
|
||||||
class SqlStore(Protocol):
|
|
||||||
"""
|
|
||||||
A protocol for a SQL store.
|
|
||||||
"""
|
|
||||||
|
|
||||||
async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None:
|
|
||||||
"""
|
|
||||||
Create a table.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None:
|
|
||||||
"""
|
|
||||||
Insert a row or batch of rows into a table.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def upsert(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
data: Mapping[str, Any],
|
|
||||||
conflict_columns: list[str],
|
|
||||||
update_columns: list[str] | None = None,
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
Insert a row and update specified columns when conflicts occur.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def fetch_all(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
where: Mapping[str, Any] | None = None,
|
|
||||||
where_sql: str | None = None,
|
|
||||||
limit: int | None = None,
|
|
||||||
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
|
||||||
cursor: tuple[str, str] | None = None,
|
|
||||||
) -> PaginatedResponse:
|
|
||||||
"""
|
|
||||||
Fetch all rows from a table with optional cursor-based pagination.
|
|
||||||
|
|
||||||
:param table: The table name
|
|
||||||
:param where: Simple key-value WHERE conditions
|
|
||||||
:param where_sql: Raw SQL WHERE clause for complex queries
|
|
||||||
:param limit: Maximum number of records to return
|
|
||||||
:param order_by: List of (column, order) tuples for sorting
|
|
||||||
:param cursor: Tuple of (key_column, cursor_id) for pagination (None for first page)
|
|
||||||
Requires order_by with exactly one column when used
|
|
||||||
:return: PaginatedResult with data and has_more flag
|
|
||||||
|
|
||||||
Note: Cursor pagination only supports single-column ordering for simplicity.
|
|
||||||
Multi-column ordering is allowed without cursor but will raise an error with cursor.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def fetch_one(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
where: Mapping[str, Any] | None = None,
|
|
||||||
where_sql: str | None = None,
|
|
||||||
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
|
||||||
) -> dict[str, Any] | None:
|
|
||||||
"""
|
|
||||||
Fetch one row from a table.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def update(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
data: Mapping[str, Any],
|
|
||||||
where: Mapping[str, Any],
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
Update a row in a table.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def delete(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
where: Mapping[str, Any],
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
Delete a row from a table.
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def add_column_if_not_exists(
|
|
||||||
self,
|
|
||||||
table: str,
|
|
||||||
column_name: str,
|
|
||||||
column_type: ColumnType,
|
|
||||||
nullable: bool = True,
|
|
||||||
) -> None:
|
|
||||||
"""
|
|
||||||
Add a column to an existing table if the column doesn't already exist.
|
|
||||||
|
|
||||||
This is useful for table migrations when adding new functionality.
|
|
||||||
If the table doesn't exist, this method should do nothing.
|
|
||||||
If the column already exists, this method should do nothing.
|
|
||||||
|
|
||||||
:param table: Table name
|
|
||||||
:param column_name: Name of the column to add
|
|
||||||
:param column_type: Type of the column to add
|
|
||||||
:param nullable: Whether the column should be nullable (default: True)
|
|
||||||
"""
|
|
||||||
pass
|
|
||||||
|
|
@ -72,6 +72,7 @@ class Agents(Protocol):
|
||||||
model: str,
|
model: str,
|
||||||
prompt: OpenAIResponsePrompt | None = None,
|
prompt: OpenAIResponsePrompt | None = None,
|
||||||
instructions: str | None = None,
|
instructions: str | None = None,
|
||||||
|
parallel_tool_calls: bool | None = True,
|
||||||
previous_response_id: str | None = None,
|
previous_response_id: str | None = None,
|
||||||
conversation: str | None = None,
|
conversation: str | None = None,
|
||||||
store: bool | None = True,
|
store: bool | None = True,
|
||||||
|
|
|
||||||
|
|
@ -3,3 +3,7 @@
|
||||||
#
|
#
|
||||||
# This source code is licensed under the terms described in the LICENSE file in
|
# This source code is licensed under the terms described in the LICENSE file in
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
|
# Internal subpackage for shared interfaces that are not part of the public API.
|
||||||
|
|
||||||
|
__all__: list[str] = []
|
||||||
|
|
@ -9,6 +9,8 @@ from typing import Protocol
|
||||||
|
|
||||||
|
|
||||||
class KVStore(Protocol):
|
class KVStore(Protocol):
|
||||||
|
"""Protocol for simple key/value storage backends."""
|
||||||
|
|
||||||
# TODO: make the value type bytes instead of str
|
# TODO: make the value type bytes instead of str
|
||||||
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: ...
|
async def set(self, key: str, value: str, expiration: datetime | None = None) -> None: ...
|
||||||
|
|
||||||
|
|
@ -19,3 +21,6 @@ class KVStore(Protocol):
|
||||||
async def values_in_range(self, start_key: str, end_key: str) -> list[str]: ...
|
async def values_in_range(self, start_key: str, end_key: str) -> list[str]: ...
|
||||||
|
|
||||||
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
|
async def keys_in_range(self, start_key: str, end_key: str) -> list[str]: ...
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["KVStore"]
|
||||||
79
src/llama_stack_api/internal/sqlstore.py
Normal file
79
src/llama_stack_api/internal/sqlstore.py
Normal file
|
|
@ -0,0 +1,79 @@
|
||||||
|
# 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.
|
||||||
|
|
||||||
|
from collections.abc import Mapping, Sequence
|
||||||
|
from enum import Enum
|
||||||
|
from typing import Any, Literal, Protocol
|
||||||
|
|
||||||
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
from llama_stack_api import PaginatedResponse
|
||||||
|
|
||||||
|
|
||||||
|
class ColumnType(Enum):
|
||||||
|
INTEGER = "INTEGER"
|
||||||
|
STRING = "STRING"
|
||||||
|
TEXT = "TEXT"
|
||||||
|
FLOAT = "FLOAT"
|
||||||
|
BOOLEAN = "BOOLEAN"
|
||||||
|
JSON = "JSON"
|
||||||
|
DATETIME = "DATETIME"
|
||||||
|
|
||||||
|
|
||||||
|
class ColumnDefinition(BaseModel):
|
||||||
|
type: ColumnType
|
||||||
|
primary_key: bool = False
|
||||||
|
nullable: bool = True
|
||||||
|
default: Any = None
|
||||||
|
|
||||||
|
|
||||||
|
class SqlStore(Protocol):
|
||||||
|
"""Protocol for common SQL-store functionality."""
|
||||||
|
|
||||||
|
async def create_table(self, table: str, schema: Mapping[str, ColumnType | ColumnDefinition]) -> None: ...
|
||||||
|
|
||||||
|
async def insert(self, table: str, data: Mapping[str, Any] | Sequence[Mapping[str, Any]]) -> None: ...
|
||||||
|
|
||||||
|
async def upsert(
|
||||||
|
self,
|
||||||
|
table: str,
|
||||||
|
data: Mapping[str, Any],
|
||||||
|
conflict_columns: list[str],
|
||||||
|
update_columns: list[str] | None = None,
|
||||||
|
) -> None: ...
|
||||||
|
|
||||||
|
async def fetch_all(
|
||||||
|
self,
|
||||||
|
table: str,
|
||||||
|
where: Mapping[str, Any] | None = None,
|
||||||
|
where_sql: str | None = None,
|
||||||
|
limit: int | None = None,
|
||||||
|
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
||||||
|
cursor: tuple[str, str] | None = None,
|
||||||
|
) -> PaginatedResponse: ...
|
||||||
|
|
||||||
|
async def fetch_one(
|
||||||
|
self,
|
||||||
|
table: str,
|
||||||
|
where: Mapping[str, Any] | None = None,
|
||||||
|
where_sql: str | None = None,
|
||||||
|
order_by: list[tuple[str, Literal["asc", "desc"]]] | None = None,
|
||||||
|
) -> dict[str, Any] | None: ...
|
||||||
|
|
||||||
|
async def update(self, table: str, data: Mapping[str, Any], where: Mapping[str, Any]) -> None: ...
|
||||||
|
|
||||||
|
async def delete(self, table: str, where: Mapping[str, Any]) -> None: ...
|
||||||
|
|
||||||
|
async def add_column_if_not_exists(
|
||||||
|
self,
|
||||||
|
table: str,
|
||||||
|
column_name: str,
|
||||||
|
column_type: ColumnType,
|
||||||
|
nullable: bool = True,
|
||||||
|
) -> None: ...
|
||||||
|
|
||||||
|
|
||||||
|
__all__ = ["ColumnDefinition", "ColumnType", "SqlStore"]
|
||||||
|
|
@ -585,7 +585,7 @@ class OpenAIResponseObject(BaseModel):
|
||||||
:param model: Model identifier used for generation
|
:param model: Model identifier used for generation
|
||||||
:param object: Object type identifier, always "response"
|
:param object: Object type identifier, always "response"
|
||||||
:param output: List of generated output items (messages, tool calls, etc.)
|
:param output: List of generated output items (messages, tool calls, etc.)
|
||||||
:param parallel_tool_calls: Whether tool calls can be executed in parallel
|
:param parallel_tool_calls: (Optional) Whether to allow more than one function tool call generated per turn.
|
||||||
:param previous_response_id: (Optional) ID of the previous response in a conversation
|
:param previous_response_id: (Optional) ID of the previous response in a conversation
|
||||||
:param prompt: (Optional) Reference to a prompt template and its variables.
|
:param prompt: (Optional) Reference to a prompt template and its variables.
|
||||||
:param status: Current status of the response generation
|
:param status: Current status of the response generation
|
||||||
|
|
@ -605,7 +605,7 @@ class OpenAIResponseObject(BaseModel):
|
||||||
model: str
|
model: str
|
||||||
object: Literal["response"] = "response"
|
object: Literal["response"] = "response"
|
||||||
output: Sequence[OpenAIResponseOutput]
|
output: Sequence[OpenAIResponseOutput]
|
||||||
parallel_tool_calls: bool = False
|
parallel_tool_calls: bool | None = True
|
||||||
previous_response_id: str | None = None
|
previous_response_id: str | None = None
|
||||||
prompt: OpenAIResponsePrompt | None = None
|
prompt: OpenAIResponsePrompt | None = None
|
||||||
status: str
|
status: str
|
||||||
|
|
|
||||||
|
|
@ -175,7 +175,7 @@ def test_expires_after_requests(openai_client):
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.xfail(message="User isolation broken for current providers, must be fixed.")
|
@pytest.mark.xfail(message="User isolation broken for current providers, must be fixed.")
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
def test_files_authentication_isolation(mock_get_authenticated_user, llama_stack_client):
|
def test_files_authentication_isolation(mock_get_authenticated_user, llama_stack_client):
|
||||||
"""Test that users can only access their own files."""
|
"""Test that users can only access their own files."""
|
||||||
from llama_stack_client import NotFoundError
|
from llama_stack_client import NotFoundError
|
||||||
|
|
@ -275,7 +275,7 @@ def test_files_authentication_isolation(mock_get_authenticated_user, llama_stack
|
||||||
raise e
|
raise e
|
||||||
|
|
||||||
|
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
def test_files_authentication_shared_attributes(
|
def test_files_authentication_shared_attributes(
|
||||||
mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
|
mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
|
||||||
):
|
):
|
||||||
|
|
@ -335,7 +335,7 @@ def test_files_authentication_shared_attributes(
|
||||||
raise e
|
raise e
|
||||||
|
|
||||||
|
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
def test_files_authentication_anonymous_access(
|
def test_files_authentication_anonymous_access(
|
||||||
mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
|
mock_get_authenticated_user, llama_stack_client, provider_type_is_openai
|
||||||
):
|
):
|
||||||
|
|
|
||||||
|
|
@ -13,14 +13,14 @@ import pytest
|
||||||
from llama_stack.core.access_control.access_control import default_policy
|
from llama_stack.core.access_control.access_control import default_policy
|
||||||
from llama_stack.core.datatypes import User
|
from llama_stack.core.datatypes import User
|
||||||
from llama_stack.core.storage.datatypes import SqlStoreReference
|
from llama_stack.core.storage.datatypes import SqlStoreReference
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnType
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
|
from llama_stack.core.storage.sqlstore.sqlstore import (
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import (
|
|
||||||
PostgresSqlStoreConfig,
|
PostgresSqlStoreConfig,
|
||||||
SqliteSqlStoreConfig,
|
SqliteSqlStoreConfig,
|
||||||
register_sqlstore_backends,
|
register_sqlstore_backends,
|
||||||
sqlstore_impl,
|
sqlstore_impl,
|
||||||
)
|
)
|
||||||
|
from llama_stack_api.internal.sqlstore import ColumnType
|
||||||
|
|
||||||
|
|
||||||
def get_postgres_config():
|
def get_postgres_config():
|
||||||
|
|
@ -96,7 +96,7 @@ async def cleanup_records(sql_store, table_name, record_ids):
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
|
@pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
async def test_authorized_store_attributes(mock_get_authenticated_user, authorized_store, request):
|
async def test_authorized_store_attributes(mock_get_authenticated_user, authorized_store, request):
|
||||||
"""Test that JSON column comparisons work correctly for both PostgreSQL and SQLite"""
|
"""Test that JSON column comparisons work correctly for both PostgreSQL and SQLite"""
|
||||||
backend_name = request.node.callspec.id
|
backend_name = request.node.callspec.id
|
||||||
|
|
@ -190,7 +190,7 @@ async def test_authorized_store_attributes(mock_get_authenticated_user, authoriz
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
|
@pytest.mark.parametrize("backend_config", BACKEND_CONFIGS)
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
async def test_user_ownership_policy(mock_get_authenticated_user, authorized_store, request):
|
async def test_user_ownership_policy(mock_get_authenticated_user, authorized_store, request):
|
||||||
"""Test that 'user is owner' policies work correctly with record ownership"""
|
"""Test that 'user is owner' policies work correctly with record ownership"""
|
||||||
from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
|
from llama_stack.core.access_control.datatypes import AccessRule, Action, Scope
|
||||||
|
|
|
||||||
|
|
@ -23,7 +23,7 @@ from llama_stack.core.storage.datatypes import (
|
||||||
SqlStoreReference,
|
SqlStoreReference,
|
||||||
StorageConfig,
|
StorageConfig,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack_api import OpenAIResponseInputMessageContentText, OpenAIResponseMessage
|
from llama_stack_api import OpenAIResponseInputMessageContentText, OpenAIResponseMessage
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -38,6 +38,9 @@ async def service():
|
||||||
},
|
},
|
||||||
stores=ServerStoresConfig(
|
stores=ServerStoresConfig(
|
||||||
conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
|
conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
|
||||||
|
metadata=None,
|
||||||
|
inference=None,
|
||||||
|
prompts=None,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})
|
register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})
|
||||||
|
|
@ -142,6 +145,9 @@ async def test_policy_configuration():
|
||||||
},
|
},
|
||||||
stores=ServerStoresConfig(
|
stores=ServerStoresConfig(
|
||||||
conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
|
conversations=SqlStoreReference(backend="sql_test", table_name="openai_conversations"),
|
||||||
|
metadata=None,
|
||||||
|
inference=None,
|
||||||
|
prompts=None,
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})
|
register_sqlstore_backends({"sql_test": storage.backends["sql_test"]})
|
||||||
|
|
|
||||||
|
|
@ -10,8 +10,9 @@ from unittest.mock import AsyncMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.datatypes import QualifiedModel, SafetyConfig, StackRunConfig, StorageConfig, VectorStoresConfig
|
from llama_stack.core.datatypes import QualifiedModel, SafetyConfig, StackRunConfig, VectorStoresConfig
|
||||||
from llama_stack.core.stack import validate_safety_config, validate_vector_stores_config
|
from llama_stack.core.stack import validate_safety_config, validate_vector_stores_config
|
||||||
|
from llama_stack.core.storage.datatypes import ServerStoresConfig, StorageConfig
|
||||||
from llama_stack_api import Api, ListModelsResponse, ListShieldsResponse, Model, ModelType, Shield
|
from llama_stack_api import Api, ListModelsResponse, ListShieldsResponse, Model, ModelType, Shield
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -21,7 +22,15 @@ class TestVectorStoresValidation:
|
||||||
run_config = StackRunConfig(
|
run_config = StackRunConfig(
|
||||||
image_name="test",
|
image_name="test",
|
||||||
providers={},
|
providers={},
|
||||||
storage=StorageConfig(backends={}, stores={}),
|
storage=StorageConfig(
|
||||||
|
backends={},
|
||||||
|
stores=ServerStoresConfig(
|
||||||
|
metadata=None,
|
||||||
|
inference=None,
|
||||||
|
conversations=None,
|
||||||
|
prompts=None,
|
||||||
|
),
|
||||||
|
),
|
||||||
vector_stores=VectorStoresConfig(
|
vector_stores=VectorStoresConfig(
|
||||||
default_provider_id="faiss",
|
default_provider_id="faiss",
|
||||||
default_embedding_model=QualifiedModel(
|
default_embedding_model=QualifiedModel(
|
||||||
|
|
@ -41,7 +50,15 @@ class TestVectorStoresValidation:
|
||||||
run_config = StackRunConfig(
|
run_config = StackRunConfig(
|
||||||
image_name="test",
|
image_name="test",
|
||||||
providers={},
|
providers={},
|
||||||
storage=StorageConfig(backends={}, stores={}),
|
storage=StorageConfig(
|
||||||
|
backends={},
|
||||||
|
stores=ServerStoresConfig(
|
||||||
|
metadata=None,
|
||||||
|
inference=None,
|
||||||
|
conversations=None,
|
||||||
|
prompts=None,
|
||||||
|
),
|
||||||
|
),
|
||||||
vector_stores=VectorStoresConfig(
|
vector_stores=VectorStoresConfig(
|
||||||
default_provider_id="faiss",
|
default_provider_id="faiss",
|
||||||
default_embedding_model=QualifiedModel(
|
default_embedding_model=QualifiedModel(
|
||||||
|
|
|
||||||
|
|
@ -9,11 +9,11 @@ import pytest
|
||||||
|
|
||||||
from llama_stack.core.access_control.access_control import default_policy
|
from llama_stack.core.access_control.access_control import default_policy
|
||||||
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
|
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack.providers.inline.files.localfs import (
|
from llama_stack.providers.inline.files.localfs import (
|
||||||
LocalfsFilesImpl,
|
LocalfsFilesImpl,
|
||||||
LocalfsFilesImplConfig,
|
LocalfsFilesImplConfig,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
|
||||||
from llama_stack_api import OpenAIFilePurpose, Order, ResourceNotFoundError
|
from llama_stack_api import OpenAIFilePurpose, Order, ResourceNotFoundError
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,9 +6,9 @@
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
|
||||||
|
from llama_stack.core.storage.kvstore.sqlite import SqliteKVStoreImpl
|
||||||
from llama_stack.core.store.registry import CachedDiskDistributionRegistry, DiskDistributionRegistry
|
from llama_stack.core.store.registry import CachedDiskDistributionRegistry, DiskDistributionRegistry
|
||||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
|
||||||
from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="function")
|
@pytest.fixture(scope="function")
|
||||||
|
|
|
||||||
|
|
@ -18,7 +18,7 @@ from llama_stack.core.storage.datatypes import (
|
||||||
SqlStoreReference,
|
SqlStoreReference,
|
||||||
StorageConfig,
|
StorageConfig,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.kvstore import register_kvstore_backends
|
from llama_stack.core.storage.kvstore import register_kvstore_backends
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -17,6 +17,7 @@ from openai.types.chat.chat_completion_chunk import (
|
||||||
|
|
||||||
from llama_stack.core.access_control.access_control import default_policy
|
from llama_stack.core.access_control.access_control import default_policy
|
||||||
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
|
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack.providers.inline.agents.meta_reference.responses.openai_responses import (
|
from llama_stack.providers.inline.agents.meta_reference.responses.openai_responses import (
|
||||||
OpenAIResponsesImpl,
|
OpenAIResponsesImpl,
|
||||||
)
|
)
|
||||||
|
|
@ -24,7 +25,6 @@ from llama_stack.providers.utils.responses.responses_store import (
|
||||||
ResponsesStore,
|
ResponsesStore,
|
||||||
_OpenAIResponseObjectWithInputAndMessages,
|
_OpenAIResponseObjectWithInputAndMessages,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
|
||||||
from llama_stack_api.agents import Order
|
from llama_stack_api.agents import Order
|
||||||
from llama_stack_api.inference import (
|
from llama_stack_api.inference import (
|
||||||
OpenAIAssistantMessageParam,
|
OpenAIAssistantMessageParam,
|
||||||
|
|
|
||||||
|
|
@ -13,9 +13,9 @@ from unittest.mock import AsyncMock
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl, register_kvstore_backends
|
||||||
from llama_stack.providers.inline.batches.reference.batches import ReferenceBatchesImpl
|
from llama_stack.providers.inline.batches.reference.batches import ReferenceBatchesImpl
|
||||||
from llama_stack.providers.inline.batches.reference.config import ReferenceBatchesImplConfig
|
from llama_stack.providers.inline.batches.reference.config import ReferenceBatchesImplConfig
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -9,8 +9,8 @@ import pytest
|
||||||
from moto import mock_aws
|
from moto import mock_aws
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
|
from llama_stack.core.storage.datatypes import SqliteSqlStoreConfig, SqlStoreReference
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack.providers.remote.files.s3 import S3FilesImplConfig, get_adapter_impl
|
from llama_stack.providers.remote.files.s3 import S3FilesImplConfig, get_adapter_impl
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
|
||||||
|
|
||||||
|
|
||||||
class MockUploadFile:
|
class MockUploadFile:
|
||||||
|
|
|
||||||
|
|
@ -18,11 +18,11 @@ async def test_listing_hides_other_users_file(s3_provider, sample_text_file):
|
||||||
user_a = User("user-a", {"roles": ["team-a"]})
|
user_a = User("user-a", {"roles": ["team-a"]})
|
||||||
user_b = User("user-b", {"roles": ["team-b"]})
|
user_b = User("user-b", {"roles": ["team-b"]})
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_a
|
mock_get_user.return_value = user_a
|
||||||
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_b
|
mock_get_user.return_value = user_b
|
||||||
listed = await s3_provider.openai_list_files()
|
listed = await s3_provider.openai_list_files()
|
||||||
assert all(f.id != uploaded.id for f in listed.data)
|
assert all(f.id != uploaded.id for f in listed.data)
|
||||||
|
|
@ -41,11 +41,11 @@ async def test_cannot_access_other_user_file(s3_provider, sample_text_file, op):
|
||||||
user_a = User("user-a", {"roles": ["team-a"]})
|
user_a = User("user-a", {"roles": ["team-a"]})
|
||||||
user_b = User("user-b", {"roles": ["team-b"]})
|
user_b = User("user-b", {"roles": ["team-b"]})
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_a
|
mock_get_user.return_value = user_a
|
||||||
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_b
|
mock_get_user.return_value = user_b
|
||||||
with pytest.raises(ResourceNotFoundError):
|
with pytest.raises(ResourceNotFoundError):
|
||||||
await op(s3_provider, uploaded.id)
|
await op(s3_provider, uploaded.id)
|
||||||
|
|
@ -56,11 +56,11 @@ async def test_shared_role_allows_listing(s3_provider, sample_text_file):
|
||||||
user_a = User("user-a", {"roles": ["shared-role"]})
|
user_a = User("user-a", {"roles": ["shared-role"]})
|
||||||
user_b = User("user-b", {"roles": ["shared-role"]})
|
user_b = User("user-b", {"roles": ["shared-role"]})
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_a
|
mock_get_user.return_value = user_a
|
||||||
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_b
|
mock_get_user.return_value = user_b
|
||||||
listed = await s3_provider.openai_list_files()
|
listed = await s3_provider.openai_list_files()
|
||||||
assert any(f.id == uploaded.id for f in listed.data)
|
assert any(f.id == uploaded.id for f in listed.data)
|
||||||
|
|
@ -79,10 +79,10 @@ async def test_shared_role_allows_access(s3_provider, sample_text_file, op):
|
||||||
user_x = User("user-x", {"roles": ["shared-role"]})
|
user_x = User("user-x", {"roles": ["shared-role"]})
|
||||||
user_y = User("user-y", {"roles": ["shared-role"]})
|
user_y = User("user-y", {"roles": ["shared-role"]})
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_x
|
mock_get_user.return_value = user_x
|
||||||
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
uploaded = await s3_provider.openai_upload_file(file=sample_text_file, purpose=OpenAIFilePurpose.ASSISTANTS)
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
with patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user") as mock_get_user:
|
||||||
mock_get_user.return_value = user_y
|
mock_get_user.return_value = user_y
|
||||||
await op(s3_provider, uploaded.id)
|
await op(s3_provider, uploaded.id)
|
||||||
|
|
|
||||||
|
|
@ -11,13 +11,13 @@ import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
||||||
|
from llama_stack.core.storage.kvstore import register_kvstore_backends
|
||||||
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
|
from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
|
||||||
from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter
|
from llama_stack.providers.inline.vector_io.faiss.faiss import FaissIndex, FaissVectorIOAdapter
|
||||||
from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig
|
from llama_stack.providers.inline.vector_io.sqlite_vec import SQLiteVectorIOConfig
|
||||||
from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter
|
from llama_stack.providers.inline.vector_io.sqlite_vec.sqlite_vec import SQLiteVecIndex, SQLiteVecVectorIOAdapter
|
||||||
from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig
|
from llama_stack.providers.remote.vector_io.pgvector.config import PGVectorVectorIOConfig
|
||||||
from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter
|
from llama_stack.providers.remote.vector_io.pgvector.pgvector import PGVectorIndex, PGVectorVectorIOAdapter
|
||||||
from llama_stack.providers.utils.kvstore import register_kvstore_backends
|
|
||||||
from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, VectorStore
|
from llama_stack_api import Chunk, ChunkMetadata, QueryChunksResponse, VectorStore
|
||||||
|
|
||||||
EMBEDDING_DIMENSION = 768
|
EMBEDDING_DIMENSION = 768
|
||||||
|
|
@ -279,7 +279,7 @@ async def pgvector_vec_adapter(unique_kvstore_config, mock_inference_api, embedd
|
||||||
) as mock_check_version:
|
) as mock_check_version:
|
||||||
mock_check_version.return_value = "0.5.1"
|
mock_check_version.return_value = "0.5.1"
|
||||||
|
|
||||||
with patch("llama_stack.providers.utils.kvstore.kvstore_impl") as mock_kvstore_impl:
|
with patch("llama_stack.core.storage.kvstore.kvstore_impl") as mock_kvstore_impl:
|
||||||
mock_kvstore = AsyncMock()
|
mock_kvstore = AsyncMock()
|
||||||
mock_kvstore_impl.return_value = mock_kvstore
|
mock_kvstore_impl.return_value = mock_kvstore
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -9,12 +9,12 @@ import pytest
|
||||||
|
|
||||||
from llama_stack.core.datatypes import VectorStoreWithOwner
|
from llama_stack.core.datatypes import VectorStoreWithOwner
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
||||||
|
from llama_stack.core.storage.kvstore import kvstore_impl, register_kvstore_backends
|
||||||
from llama_stack.core.store.registry import (
|
from llama_stack.core.store.registry import (
|
||||||
KEY_FORMAT,
|
KEY_FORMAT,
|
||||||
CachedDiskDistributionRegistry,
|
CachedDiskDistributionRegistry,
|
||||||
DiskDistributionRegistry,
|
DiskDistributionRegistry,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.kvstore import kvstore_impl, register_kvstore_backends
|
|
||||||
from llama_stack_api import Model, VectorStore
|
from llama_stack_api import Model, VectorStore
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -15,7 +15,7 @@ from starlette.middleware.base import BaseHTTPMiddleware
|
||||||
from llama_stack.core.datatypes import QuotaConfig, QuotaPeriod
|
from llama_stack.core.datatypes import QuotaConfig, QuotaPeriod
|
||||||
from llama_stack.core.server.quota import QuotaMiddleware
|
from llama_stack.core.server.quota import QuotaMiddleware
|
||||||
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
from llama_stack.core.storage.datatypes import KVStoreReference, SqliteKVStoreConfig
|
||||||
from llama_stack.providers.utils.kvstore import register_kvstore_backends
|
from llama_stack.core.storage.kvstore import register_kvstore_backends
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -24,8 +24,8 @@ from llama_stack.core.storage.datatypes import (
|
||||||
SqlStoreReference,
|
SqlStoreReference,
|
||||||
StorageConfig,
|
StorageConfig,
|
||||||
)
|
)
|
||||||
from llama_stack.providers.utils.kvstore import register_kvstore_backends
|
from llama_stack.core.storage.kvstore import register_kvstore_backends
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack_api import Inference, InlineProviderSpec, ProviderSpec
|
from llama_stack_api import Inference, InlineProviderSpec, ProviderSpec
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -9,8 +9,8 @@ import time
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig
|
from llama_stack.core.storage.datatypes import InferenceStoreReference, SqliteSqlStoreConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack.providers.utils.inference.inference_store import InferenceStore
|
from llama_stack.providers.utils.inference.inference_store import InferenceStore
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
|
||||||
from llama_stack_api import (
|
from llama_stack_api import (
|
||||||
OpenAIAssistantMessageParam,
|
OpenAIAssistantMessageParam,
|
||||||
OpenAIChatCompletion,
|
OpenAIChatCompletion,
|
||||||
|
|
|
||||||
|
|
@ -5,8 +5,8 @@
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
|
|
||||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
from llama_stack.core.storage.kvstore.config import SqliteKVStoreConfig
|
||||||
from llama_stack.providers.utils.kvstore.sqlite.sqlite import SqliteKVStoreImpl
|
from llama_stack.core.storage.kvstore.sqlite.sqlite import SqliteKVStoreImpl
|
||||||
|
|
||||||
|
|
||||||
async def test_memory_kvstore_persistence_behavior():
|
async def test_memory_kvstore_persistence_behavior():
|
||||||
|
|
|
||||||
|
|
@ -11,8 +11,8 @@ from uuid import uuid4
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
|
from llama_stack.core.storage.datatypes import ResponsesStoreReference, SqliteSqlStoreConfig
|
||||||
|
from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
|
||||||
from llama_stack.providers.utils.responses.responses_store import ResponsesStore
|
from llama_stack.providers.utils.responses.responses_store import ResponsesStore
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
|
|
||||||
from llama_stack_api import OpenAIMessageParam, OpenAIResponseInput, OpenAIResponseObject, OpenAIUserMessageParam, Order
|
from llama_stack_api import OpenAIMessageParam, OpenAIResponseInput, OpenAIResponseObject, OpenAIUserMessageParam, Order
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -9,9 +9,9 @@ from tempfile import TemporaryDirectory
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
|
from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
|
||||||
from llama_stack.providers.utils.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
|
from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
|
from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
|
||||||
|
|
||||||
|
|
||||||
async def test_sqlite_sqlstore():
|
async def test_sqlite_sqlstore():
|
||||||
|
|
|
||||||
|
|
@ -10,13 +10,13 @@ from unittest.mock import patch
|
||||||
from llama_stack.core.access_control.access_control import default_policy, is_action_allowed
|
from llama_stack.core.access_control.access_control import default_policy, is_action_allowed
|
||||||
from llama_stack.core.access_control.datatypes import Action
|
from llama_stack.core.access_control.datatypes import Action
|
||||||
from llama_stack.core.datatypes import User
|
from llama_stack.core.datatypes import User
|
||||||
from llama_stack.providers.utils.sqlstore.api import ColumnType
|
from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore, SqlRecord
|
||||||
from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore, SqlRecord
|
from llama_stack.core.storage.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
|
||||||
from llama_stack.providers.utils.sqlstore.sqlalchemy_sqlstore import SqlAlchemySqlStoreImpl
|
from llama_stack.core.storage.sqlstore.sqlstore import SqliteSqlStoreConfig
|
||||||
from llama_stack.providers.utils.sqlstore.sqlstore import SqliteSqlStoreConfig
|
from llama_stack_api.internal.sqlstore import ColumnType
|
||||||
|
|
||||||
|
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
async def test_authorized_fetch_with_where_sql_access_control(mock_get_authenticated_user):
|
async def test_authorized_fetch_with_where_sql_access_control(mock_get_authenticated_user):
|
||||||
"""Test that fetch_all works correctly with where_sql for access control"""
|
"""Test that fetch_all works correctly with where_sql for access control"""
|
||||||
with TemporaryDirectory() as tmp_dir:
|
with TemporaryDirectory() as tmp_dir:
|
||||||
|
|
@ -78,7 +78,7 @@ async def test_authorized_fetch_with_where_sql_access_control(mock_get_authentic
|
||||||
assert row["title"] == "User Document"
|
assert row["title"] == "User Document"
|
||||||
|
|
||||||
|
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
async def test_sql_policy_consistency(mock_get_authenticated_user):
|
async def test_sql_policy_consistency(mock_get_authenticated_user):
|
||||||
"""Test that SQL WHERE clause logic exactly matches is_action_allowed policy logic"""
|
"""Test that SQL WHERE clause logic exactly matches is_action_allowed policy logic"""
|
||||||
with TemporaryDirectory() as tmp_dir:
|
with TemporaryDirectory() as tmp_dir:
|
||||||
|
|
@ -164,7 +164,7 @@ async def test_sql_policy_consistency(mock_get_authenticated_user):
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
|
@patch("llama_stack.core.storage.sqlstore.authorized_sqlstore.get_authenticated_user")
|
||||||
async def test_authorized_store_user_attribute_capture(mock_get_authenticated_user):
|
async def test_authorized_store_user_attribute_capture(mock_get_authenticated_user):
|
||||||
"""Test that user attributes are properly captured during insert"""
|
"""Test that user attributes are properly captured during insert"""
|
||||||
with TemporaryDirectory() as tmp_dir:
|
with TemporaryDirectory() as tmp_dir:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue