wehub-resource-sync
83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import cast
|
|
|
|
from agents.items import TResponseInputItem
|
|
from agents.memory.session import Session
|
|
from agents.memory.session_settings import SessionSettings
|
|
|
|
|
|
class SimpleListSession(Session):
|
|
"""A minimal in-memory session implementation for tests."""
|
|
|
|
session_settings: SessionSettings | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str = "test",
|
|
history: list[TResponseInputItem] | None = None,
|
|
) -> None:
|
|
self.session_id = session_id
|
|
self._items: list[TResponseInputItem] = list(history) if history else []
|
|
# Some session implementations strip IDs on write; tests can opt-in via attribute.
|
|
self._ignore_ids_for_matching = False
|
|
# Mirror saved_items used by some tests for inspection.
|
|
self.saved_items: list[TResponseInputItem] = self._items
|
|
|
|
async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]:
|
|
if limit is None:
|
|
return list(self._items)
|
|
if limit <= 0:
|
|
return []
|
|
return self._items[-limit:]
|
|
|
|
async def add_items(self, items: list[TResponseInputItem]) -> None:
|
|
self._items.extend(items)
|
|
|
|
async def pop_item(self) -> TResponseInputItem | None:
|
|
if not self._items:
|
|
return None
|
|
return self._items.pop()
|
|
|
|
async def clear_session(self) -> None:
|
|
self._items.clear()
|
|
|
|
|
|
class CountingSession(SimpleListSession):
|
|
"""Session that tracks how many times pop_item is invoked (for rewind tests)."""
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str = "test",
|
|
history: list[TResponseInputItem] | None = None,
|
|
) -> None:
|
|
super().__init__(session_id=session_id, history=history)
|
|
self.pop_calls = 0
|
|
|
|
async def pop_item(self) -> TResponseInputItem | None:
|
|
self.pop_calls += 1
|
|
return await super().pop_item()
|
|
|
|
|
|
class IdStrippingSession(CountingSession):
|
|
"""Session that strips IDs on add to mimic hosted stores that reassign IDs."""
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str = "test",
|
|
history: list[TResponseInputItem] | None = None,
|
|
) -> None:
|
|
super().__init__(session_id=session_id, history=history)
|
|
self._ignore_ids_for_matching = True
|
|
|
|
async def add_items(self, items: list[TResponseInputItem]) -> None:
|
|
sanitized: list[TResponseInputItem] = []
|
|
for item in items:
|
|
if isinstance(item, dict):
|
|
clean = dict(item)
|
|
clean.pop("id", None)
|
|
sanitized.append(cast(TResponseInputItem, clean))
|
|
else:
|
|
sanitized.append(item)
|
|
await super().add_items(sanitized)
|