Pass memory bank API to agent impl

This commit is contained in:
Ashwin Bharambe 2024-10-09 21:16:57 -07:00
parent 6788173ffc
commit 2d94ca71a9
3 changed files with 10 additions and 1 deletions

View file

@ -11,6 +11,7 @@ from typing import AsyncGenerator
from llama_stack.apis.inference import Inference
from llama_stack.apis.memory import Memory
from llama_stack.apis.memory_banks import MemoryBanks
from llama_stack.apis.safety import Safety
from llama_stack.apis.agents import * # noqa: F403
@ -30,11 +31,14 @@ class MetaReferenceAgentsImpl(Agents):
inference_api: Inference,
memory_api: Memory,
safety_api: Safety,
memory_banks_api: MemoryBanks,
):
self.config = config
self.inference_api = inference_api
self.memory_api = memory_api
self.safety_api = safety_api
self.memory_banks_api = memory_banks_api
self.in_memory_store = InmemoryKVStoreImpl()
async def initialize(self) -> None:
@ -81,6 +85,7 @@ class MetaReferenceAgentsImpl(Agents):
inference_api=self.inference_api,
safety_api=self.safety_api,
memory_api=self.memory_api,
memory_banks_api=self.memory_banks_api,
persistence_store=(
self.persistence_store
if agent_config.enable_session_persistence