mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-17 05:22:36 +00:00
e2d1b712e2
Summary: Implementing Memory provider fakes as discussed in this draft https://github.com/meta-llama/llama-stack/pull/490#issuecomment-2492877393. High level changes: * Fake provider is specified via the "fake" mark * Test config will setup a fake fixture for the run of the test * Test resolver checks fixtures and upon finding a fake provider it injects InlineProviderSpec for fake provider * Fake provider gets resolved through path/naming convention * Fake provider implementaion is contained to the tests/ directory and implements stubs and method fakes with minimal functionality to simulate real provider Instructins to creating a fake * Create the "fakes" module inside the provider test directory * Inside the module implement `get_provider_impl` that will return fake implementation object * Name the fake implementation class to match the fake provider id (e.g. memory_fake -> MemoryFakeImpl) * Same rule for the config (e.g. memory_fake -> MemoryFakeConfig) * Add fake fixture (in the fixtures.py) and setup methods stubs there Test Plan: Run memory tests ``` pytest llama_stack/providers/tests/memory/test_memory.py -m "fake" -v -s --tb=short ====================================================================================================== test session starts ====================================================================================================== platform darwin -- Python 3.11.10, pytest-8.3.3, pluggy-1.5.0 -- /opt/homebrew/Caskroom/miniconda/base/envs/llama-stack/bin/python cachedir: .pytest_cache rootdir: /Users/vivic/Code/llama-stack configfile: pyproject.toml plugins: asyncio-0.24.0, anyio-4.6.2.post1 asyncio: mode=Mode.STRICT, default_loop_scope=None collected 18 items / 15 deselected / 3 selected llama_stack/providers/tests/memory/test_memory.py::TestMemory::test_banks_list[fake] PASSED llama_stack/providers/tests/memory/test_memory.py::TestMemory::test_banks_register[fake] PASSED llama_stack/providers/tests/memory/test_memory.py::TestMemory::test_query_documents[fake] The scores are: [0.5] PASSED ========================================================================================= 3 passed, 15 deselected, 10 warnings in 0.46s ========================================================================================= ```
84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
# 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 typing import Any, Dict, List, Optional
|
|
|
|
from llama_models.schema_utils import json_schema_type
|
|
from pydantic import BaseModel
|
|
|
|
from llama_stack.apis.memory import * # noqa: F403
|
|
from llama_stack.apis.memory_banks import MemoryBank
|
|
from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
|
|
|
|
|
|
@json_schema_type
|
|
class InlineMemoryFakeConfig(BaseModel):
|
|
pass
|
|
|
|
|
|
class InlineMemoryFakeImpl(Memory, MemoryBanksProtocolPrivate):
|
|
method_stubs: Dict[str, Any] = {}
|
|
memory_banks: Dict[str, MemoryBank] = {}
|
|
|
|
@staticmethod
|
|
def stub_method(method_name: str, return_value_matchers: Dict[str, Any]) -> None:
|
|
if method_name in InlineMemoryFakeImpl.method_stubs:
|
|
InlineMemoryFakeImpl.method_stubs[method_name].update(return_value_matchers)
|
|
return
|
|
InlineMemoryFakeImpl.method_stubs[method_name] = return_value_matchers
|
|
|
|
def __init__(self, config: InlineMemoryFakeConfig) -> None:
|
|
pass
|
|
|
|
async def initialize(self) -> None:
|
|
pass
|
|
|
|
async def shutdown(self) -> None:
|
|
pass
|
|
|
|
async def register_memory_bank(
|
|
self,
|
|
memory_bank: MemoryBank,
|
|
) -> None:
|
|
InlineMemoryFakeImpl.memory_banks[memory_bank.memory_bank_id] = memory_bank
|
|
|
|
async def list_memory_banks(self) -> List[MemoryBank]:
|
|
return list(InlineMemoryFakeImpl.memory_banks.values())
|
|
|
|
async def unregister_memory_bank(self, memory_bank_id: str) -> None:
|
|
if memory_bank_id not in InlineMemoryFakeImpl.memory_banks:
|
|
raise ValueError(f"Bank {memory_bank_id} not found.")
|
|
del InlineMemoryFakeImpl.memory_banks[memory_bank_id]
|
|
|
|
async def insert_documents(
|
|
self,
|
|
bank_id: str,
|
|
documents: List[MemoryBankDocument],
|
|
ttl_seconds: Optional[int] = None,
|
|
) -> None:
|
|
pass
|
|
|
|
async def query_documents(
|
|
self,
|
|
bank_id: str,
|
|
query: InterleavedTextMedia,
|
|
params: Optional[Dict[str, Any]] = None,
|
|
) -> QueryDocumentsResponse:
|
|
if query in InlineMemoryFakeImpl.method_stubs["query_documents"]:
|
|
return InlineMemoryFakeImpl.method_stubs["query_documents"][query]
|
|
raise ValueError(
|
|
f"Stub for query '{query}' not found, please set up expected result"
|
|
)
|
|
|
|
|
|
async def get_provider_impl(config: InlineMemoryFakeConfig, _deps: Any):
|
|
assert isinstance(
|
|
config, InlineMemoryFakeConfig
|
|
), f"Unexpected config type: {type(config)}"
|
|
|
|
impl = InlineMemoryFakeImpl(config)
|
|
await impl.initialize()
|
|
return impl
|