mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-04 04:04:14 +00:00
fix precommit errors
This commit is contained in:
parent
d32d8ec94b
commit
82ca72bde6
2 changed files with 9 additions and 7 deletions
|
@ -21,7 +21,7 @@ class SqliteKVStoreImpl(KVStore):
|
||||||
def __init__(self, config: SqliteKVStoreConfig):
|
def __init__(self, config: SqliteKVStoreConfig):
|
||||||
self.db_path = config.db_path
|
self.db_path = config.db_path
|
||||||
self.table_name = "kvstore"
|
self.table_name = "kvstore"
|
||||||
self._conn = None
|
self._conn: aiosqlite.Connection | None = None
|
||||||
|
|
||||||
def __str__(self):
|
def __str__(self):
|
||||||
return f"SqliteKVStoreImpl(db_path={self.db_path}, table_name={self.table_name})"
|
return f"SqliteKVStoreImpl(db_path={self.db_path}, table_name={self.table_name})"
|
||||||
|
@ -57,6 +57,7 @@ class SqliteKVStoreImpl(KVStore):
|
||||||
self._conn = None
|
self._conn = None
|
||||||
|
|
||||||
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:
|
||||||
|
assert self._conn is not None, "Connection not initialized. Call initialize() first."
|
||||||
await self._conn.execute(
|
await self._conn.execute(
|
||||||
f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
|
f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
|
||||||
(key, value, expiration),
|
(key, value, expiration),
|
||||||
|
@ -64,7 +65,10 @@ class SqliteKVStoreImpl(KVStore):
|
||||||
await self._conn.commit()
|
await self._conn.commit()
|
||||||
|
|
||||||
async def get(self, key: str) -> str | None:
|
async def get(self, key: str) -> str | None:
|
||||||
async with self._conn.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
|
assert self._conn is not None, "Connection not initialized. Call initialize() first."
|
||||||
|
async with self._conn.execute(
|
||||||
|
f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)
|
||||||
|
) as cursor:
|
||||||
row = await cursor.fetchone()
|
row = await cursor.fetchone()
|
||||||
if row is None:
|
if row is None:
|
||||||
return None
|
return None
|
||||||
|
@ -75,10 +79,12 @@ class SqliteKVStoreImpl(KVStore):
|
||||||
return value
|
return value
|
||||||
|
|
||||||
async def delete(self, key: str) -> None:
|
async def delete(self, key: str) -> None:
|
||||||
|
assert self._conn is not None, "Connection not initialized. Call initialize() first."
|
||||||
await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
|
await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
|
||||||
await self._conn.commit()
|
await self._conn.commit()
|
||||||
|
|
||||||
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]:
|
||||||
|
assert self._conn is not None, "Connection not initialized. Call initialize() first."
|
||||||
async with self._conn.execute(
|
async with self._conn.execute(
|
||||||
f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
|
f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
|
||||||
(start_key, end_key),
|
(start_key, end_key),
|
||||||
|
@ -91,6 +97,7 @@ class SqliteKVStoreImpl(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."""
|
||||||
|
assert self._conn is not None, "Connection not initialized. Call initialize() first."
|
||||||
cursor = await self._conn.execute(
|
cursor = await self._conn.execute(
|
||||||
f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
|
f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
|
||||||
(start_key, end_key),
|
(start_key, end_key),
|
||||||
|
|
|
@ -4,13 +4,11 @@
|
||||||
# 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.
|
||||||
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
|
||||||
from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
|
from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_memory_kvstore_basic_operations():
|
async def test_memory_kvstore_basic_operations():
|
||||||
"""Test basic CRUD operations with :memory: database."""
|
"""Test basic CRUD operations with :memory: database."""
|
||||||
config = SqliteKVStoreConfig(db_path=":memory:")
|
config = SqliteKVStoreConfig(db_path=":memory:")
|
||||||
|
@ -35,7 +33,6 @@ async def test_memory_kvstore_basic_operations():
|
||||||
await kvstore.close()
|
await kvstore.close()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_memory_kvstore_range_operations():
|
async def test_memory_kvstore_range_operations():
|
||||||
"""Test range operations with :memory: database."""
|
"""Test range operations with :memory: database."""
|
||||||
config = SqliteKVStoreConfig(db_path=":memory:")
|
config = SqliteKVStoreConfig(db_path=":memory:")
|
||||||
|
@ -66,7 +63,6 @@ async def test_memory_kvstore_range_operations():
|
||||||
await kvstore.close()
|
await kvstore.close()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_memory_kvstore_multiple_instances():
|
async def test_memory_kvstore_multiple_instances():
|
||||||
"""Test that multiple :memory: instances are independent."""
|
"""Test that multiple :memory: instances are independent."""
|
||||||
config1 = SqliteKVStoreConfig(db_path=":memory:")
|
config1 = SqliteKVStoreConfig(db_path=":memory:")
|
||||||
|
@ -98,7 +94,6 @@ async def test_memory_kvstore_multiple_instances():
|
||||||
await kvstore2.close()
|
await kvstore2.close()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_memory_kvstore_persistence_behavior():
|
async def test_memory_kvstore_persistence_behavior():
|
||||||
"""Test that :memory: database doesn't persist across instances."""
|
"""Test that :memory: database doesn't persist across instances."""
|
||||||
config = SqliteKVStoreConfig(db_path=":memory:")
|
config = SqliteKVStoreConfig(db_path=":memory:")
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue