71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
import logging
|
|
from collections import OrderedDict
|
|
|
|
try:
|
|
from fasta2a.schema import Artifact, Message, TaskState # type: ignore
|
|
from fasta2a.storage import InMemoryStorage, Storage # type: ignore
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"A2A support requires the 'a2a' extra. "
|
|
"Install with: uv pip install 'haiku.rag[a2a]'"
|
|
) from e
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LRUMemoryStorage(Storage[list[Message]]): # type: ignore
|
|
"""Storage wrapper with LRU eviction for contexts.
|
|
|
|
Enforces a maximum context limit using LRU (Least Recently Used) eviction.
|
|
"""
|
|
|
|
def __init__(self, storage: InMemoryStorage, max_contexts: int):
|
|
self.storage = storage
|
|
self.max_contexts = max_contexts
|
|
# Track context access order (LRU cache)
|
|
self.context_order: OrderedDict[str, None] = OrderedDict()
|
|
|
|
async def load_context(self, context_id: str) -> list[Message] | None:
|
|
"""Load context and update access order."""
|
|
result = await self.storage.load_context(context_id)
|
|
if result is not None:
|
|
# Move to end (most recently used)
|
|
self.context_order.pop(context_id, None)
|
|
self.context_order[context_id] = None
|
|
return result
|
|
|
|
async def update_context(self, context_id: str, context: list[Message]) -> None:
|
|
"""Update context and enforce LRU limit."""
|
|
await self.storage.update_context(context_id, context)
|
|
# Move to end (most recently used)
|
|
self.context_order.pop(context_id, None)
|
|
self.context_order[context_id] = None
|
|
|
|
# Enforce max contexts limit (LRU eviction)
|
|
while len(self.context_order) > self.max_contexts:
|
|
# Remove oldest (first item in OrderedDict)
|
|
oldest_context_id = next(iter(self.context_order))
|
|
self.context_order.pop(oldest_context_id)
|
|
logger.debug(
|
|
f"Evicted context {oldest_context_id} (LRU, limit={self.max_contexts})"
|
|
)
|
|
|
|
async def load_task(self, task_id: str, history_length: int | None = None):
|
|
"""Delegate to underlying storage."""
|
|
return await self.storage.load_task(task_id, history_length)
|
|
|
|
async def update_task(
|
|
self,
|
|
task_id: str,
|
|
state: TaskState,
|
|
new_artifacts: list[Artifact] | None = None,
|
|
new_messages: list[Message] | None = None,
|
|
):
|
|
"""Delegate to underlying storage."""
|
|
return await self.storage.update_task(
|
|
task_id, state, new_artifacts, new_messages
|
|
)
|
|
|
|
async def submit_task(self, context_id: str, message: Message):
|
|
"""Delegate to underlying storage."""
|
|
return await self.storage.submit_task(context_id, message)
|