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

fix linter

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-05 12:05:40 -08:00
parent 5695c0b969
commit 565e460115
1 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/providers/impls/meta_reference/memory/faiss.py
View file

@ -16,6 +16,7 @@ from llama_models.llama3.api.datatypes import * # noqa: F403
from llama_stack.apis.memory import * # noqa: F403 from llama_stack.apis.memory import * # noqa: F403
from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
from llama_stack.providers.utils.kvstore import kvstore_impl
from llama_stack.providers.utils.memory.vector_store import ( from llama_stack.providers.utils.memory.vector_store import (
ALL_MINILM_L6_V2_DIMENSION, ALL_MINILM_L6_V2_DIMENSION,
@ -25,7 +26,6 @@ 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__)
@ -73,19 +73,18 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
self.config = config self.config = config
self.cache = {} self.cache = {}
self.kvstore = None self.kvstore = None
async def initialize(self) -> None: async def initialize(self) -> None:
self.kvstore = await kvstore_impl(self.config.kvstore) self.kvstore = await kvstore_impl(self.config.kvstore)
# Load existing banks from kvstore # Load existing banks from kvstore
start_key = MEMORY_BANKS_PREFIX start_key = MEMORY_BANKS_PREFIX
end_key = f"{MEMORY_BANKS_PREFIX}\xff" end_key = f"{MEMORY_BANKS_PREFIX}\xff"
stored_banks = await self.kvstore.range(start_key, end_key) stored_banks = await self.kvstore.range(start_key, end_key)
for bank_data in stored_banks: for bank_data in stored_banks:
bank = VectorMemoryBankDef.model_validate_json(bank_data) bank = VectorMemoryBankDef.model_validate_json(bank_data)
index = BankWithIndex( index = BankWithIndex(
bank=bank, bank=bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
) )
self.cache[bank.identifier] = index self.cache[bank.identifier] = index
@ -110,8 +109,7 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
# Store in cache # Store in cache
index = BankWithIndex( index = BankWithIndex(
bank=memory_bank, bank=memory_bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
) )
self.cache[memory_bank.identifier] = index self.cache[memory_bank.identifier] = index