mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-05 12:21:52 +00:00
add redis adapter + sqlite provider
This commit is contained in:
parent
b31f0f6e4e
commit
bbf0b59ae4
11 changed files with 251 additions and 0 deletions
15
llama_stack/providers/impls/sqlite/control_plane/__init__.py
Normal file
15
llama_stack/providers/impls/sqlite/control_plane/__init__.py
Normal file
|
@ -0,0 +1,15 @@
|
|||
# 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 .config import SqliteControlPlaneConfig
|
||||
|
||||
|
||||
async def get_provider_impl(config: SqliteControlPlaneConfig, _deps):
|
||||
from .control_plane import SqliteControlPlane
|
||||
|
||||
impl = SqliteControlPlane(config)
|
||||
await impl.initialize()
|
||||
return impl
|
19
llama_stack/providers/impls/sqlite/control_plane/config.py
Normal file
19
llama_stack/providers/impls/sqlite/control_plane/config.py
Normal file
|
@ -0,0 +1,19 @@
|
|||
# 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_models.schema_utils import json_schema_type
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
@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",
|
||||
)
|
|
@ -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.
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import aiosqlite
|
||||
|
||||
from llama_stack.apis.control_plane import * # noqa: F403
|
||||
|
||||
|
||||
from .config import SqliteControlPlaneConfig
|
||||
|
||||
|
||||
class SqliteControlPlane(ControlPlane):
|
||||
def __init__(self, config: SqliteControlPlaneConfig):
|
||||
self.db_path = config.db_path
|
||||
self.table_name = config.table_name
|
||||
|
||||
async def initialize(self):
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {self.table_name} (
|
||||
key TEXT PRIMARY KEY,
|
||||
value TEXT,
|
||||
expiration TIMESTAMP
|
||||
)
|
||||
"""
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def set(
|
||||
self, key: str, value: Any, expiration: Optional[datetime] = None
|
||||
) -> None:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
|
||||
(key, json.dumps(value), expiration),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get(self, key: str) -> Optional[ControlPlaneValue]:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
value, expiration = row
|
||||
return ControlPlaneValue(
|
||||
key=key, value=json.loads(value), expiration=expiration
|
||||
)
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
|
||||
await db.commit()
|
||||
|
||||
async def range(self, start_key: str, end_key: str) -> List[ControlPlaneValue]:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
|
||||
(start_key, end_key),
|
||||
) as cursor:
|
||||
result = []
|
||||
async for row in cursor:
|
||||
key, value, expiration = row
|
||||
result.append(
|
||||
ControlPlaneValue(
|
||||
key=key, value=json.loads(value), expiration=expiration
|
||||
)
|
||||
)
|
||||
return result
|
Loading…
Add table
Add a link
Reference in a new issue