registry.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. from __future__ import annotations
  2. from collections.abc import Iterable
  3. from copy import deepcopy
  4. from agent_lab.application.events.models import EventDefinition
  5. class EventRegistry:
  6. def __init__(self, definitions: Iterable[EventDefinition] = ()) -> None:
  7. self._definitions: dict[str, EventDefinition] = {}
  8. for definition in definitions:
  9. self.register(definition)
  10. def register(self, definition: EventDefinition) -> None:
  11. if definition.name in self._definitions:
  12. raise ValueError(f"duplicate event definition: {definition.name}")
  13. self._definitions[definition.name] = definition
  14. def definition(self, name: str) -> EventDefinition | None:
  15. return self._definitions.get(name)
  16. def catalog(self, enabled_names: Iterable[str] | None = None) -> list[dict]:
  17. enabled = set(enabled_names) if enabled_names is not None else None
  18. return [
  19. {
  20. "name": definition.name,
  21. "description": definition.description,
  22. "parameters": deepcopy(definition.parameters),
  23. }
  24. for definition in self._definitions.values()
  25. if enabled is None or definition.name in enabled
  26. ]
  27. def tool_schema(self, name: str) -> dict | None:
  28. definition = self.definition(name)
  29. if definition is None:
  30. return None
  31. return {
  32. "type": "function",
  33. "function": {
  34. "name": definition.name,
  35. "description": definition.description,
  36. "parameters": deepcopy(definition.parameters),
  37. },
  38. }