mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-30 23:51:00 +00:00
init
This commit is contained in:
parent
8de845a96d
commit
2f51af1bb7
2 changed files with 35 additions and 6 deletions
|
@ -5,9 +5,10 @@
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
from llama_models.schema_utils import json_schema_type
|
from llama_models.schema_utils import json_schema_type
|
||||||
|
from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
||||||
@json_schema_type
|
@json_schema_type
|
||||||
class FaissImplConfig(BaseModel): ...
|
class FaissImplConfig(BaseModel):
|
||||||
|
kvstore: KVStoreConfig = SqliteKVStoreConfig() # Uses default SQLite config
|
||||||
|
|
|
@ -25,9 +25,12 @@ from llama_stack.providers.utils.memory.vector_store import (
|
||||||
from llama_stack.providers.utils.telemetry import tracing
|
from llama_stack.providers.utils.telemetry import tracing
|
||||||
|
|
||||||
from .config import FaissImplConfig
|
from .config import FaissImplConfig
|
||||||
|
from llama_stack.providers.utils.kvstore import kvstore_impl
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
MEMORY_BANKS_PREFIX = "memory_banks:"
|
||||||
|
|
||||||
|
|
||||||
class FaissIndex(EmbeddingIndex):
|
class FaissIndex(EmbeddingIndex):
|
||||||
id_by_index: Dict[int, str]
|
id_by_index: Dict[int, str]
|
||||||
|
@ -69,10 +72,26 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
|
||||||
def __init__(self, config: FaissImplConfig) -> None:
|
def __init__(self, config: FaissImplConfig) -> None:
|
||||||
self.config = config
|
self.config = config
|
||||||
self.cache = {}
|
self.cache = {}
|
||||||
|
self.kvstore = None
|
||||||
|
|
||||||
|
async def initialize(self) -> None:
|
||||||
|
self.kvstore = await kvstore_impl(self.config.kvstore)
|
||||||
|
# Load existing banks from kvstore
|
||||||
|
start_key = MEMORY_BANKS_PREFIX
|
||||||
|
end_key = f"{MEMORY_BANKS_PREFIX}\xff"
|
||||||
|
stored_banks = await self.kvstore.range(start_key, end_key)
|
||||||
|
|
||||||
|
for bank_data in stored_banks:
|
||||||
|
bank = VectorMemoryBankDef.model_validate_json(bank_data)
|
||||||
|
index = BankWithIndex(
|
||||||
|
bank=bank,
|
||||||
|
index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
|
||||||
|
)
|
||||||
|
self.cache[bank.identifier] = index
|
||||||
|
|
||||||
async def initialize(self) -> None: ...
|
async def shutdown(self) -> None:
|
||||||
|
# Cleanup if needed
|
||||||
async def shutdown(self) -> None: ...
|
pass
|
||||||
|
|
||||||
async def register_memory_bank(
|
async def register_memory_bank(
|
||||||
self,
|
self,
|
||||||
|
@ -82,8 +101,17 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
|
||||||
memory_bank.type == MemoryBankType.vector.value
|
memory_bank.type == MemoryBankType.vector.value
|
||||||
), f"Only vector banks are supported {memory_bank.type}"
|
), f"Only vector banks are supported {memory_bank.type}"
|
||||||
|
|
||||||
|
# Store in kvstore
|
||||||
|
key = f"{MEMORY_BANKS_PREFIX}{memory_bank.identifier}"
|
||||||
|
await self.kvstore.set(
|
||||||
|
key=key,
|
||||||
|
value=memory_bank.json(),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Store in cache
|
||||||
index = BankWithIndex(
|
index = BankWithIndex(
|
||||||
bank=memory_bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
|
bank=memory_bank,
|
||||||
|
index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
|
||||||
)
|
)
|
||||||
self.cache[memory_bank.identifier] = index
|
self.cache[memory_bank.identifier] = index
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue