chore: merge some RuntimeConfig methods

phil65 · Dec 5, 2024 · b4a49a1 · b4a49a1
1 parent 9aa5a53
commit b4a49a1
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 49 deletions.
diff --git a/docs/usage.md b/docs/usage.md
@@ -124,9 +124,9 @@ class ResourceObserver(RegistryEvents):
 
 async with RuntimeConfig.open("config.yml") as runtime:
     # Add observers
-    runtime.add_resource_observer(ResourceObserver())
-    runtime.add_prompt_observer(PromptObserver())
-    runtime.add_tool_observer(ToolObserver())
+    runtime.add_observer(ResourceObserver(), registry_type="resource")
+    runtime.add_observer(PromptObserver(), registry_type="prompt")
+    runtime.add_observer(ToolObserver(), registry_type="tool")
 ```
 
 ## Example: Agent Integration

diff --git a/src/llmling/config/runtime.py b/src/llmling/config/runtime.py
@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any, Literal, Self
 
 import depkit
 import logfire
@@ -44,6 +44,7 @@
 
 
 logger = get_logger(__name__)
+RegistryType = Literal["resource", "prompt", "tool"]
 
 
 class RuntimeConfig(EventEmitter):
@@ -607,55 +608,44 @@ async def process_content(
         )
 
     # Registry Observation
-    def add_resource_observer(self, observer: RegistryEvents[str, Resource]) -> None:
-        """Add observer for resource changes.
 
-        Args:
-            observer: Observer to add
-        """
-        self._resource_registry.add_observer(observer)
-
-    def add_prompt_observer(self, observer: RegistryEvents[str, BasePrompt]) -> None:
-        """Add observer for prompt changes.
-
-        Args:
-            observer: Observer to add
-        """
-        self._prompt_registry.add_observer(observer)
-
-    def add_tool_observer(self, observer: RegistryEvents[str, LLMCallableTool]) -> None:
-        """Add observer for tool changes.
+    def add_observer(
+        self,
+        observer: RegistryEvents[str, Any],
+        registry_type: RegistryType,
+    ) -> None:
+        """Add an observer for registry changes.
 
         Args:
             observer: Observer to add
+            registry_type: Type of registry to observe
         """
-        self._tool_registry.add_observer(observer)
-
-    def remove_resource_observer(self, observer: RegistryEvents[str, Resource]) -> None:
-        """Remove resource observer.
+        match registry_type:
+            case "resource":
+                self._resource_registry.add_observer(observer)
+            case "prompt":
+                self._prompt_registry.add_observer(observer)
+            case "tool":
+                self._tool_registry.add_observer(observer)
 
-        Args:
-            observer: Observer to remove
-        """
-        self._resource_registry.remove_observer(observer)
-
-    def remove_prompt_observer(self, observer: RegistryEvents[str, BasePrompt]) -> None:
-        """Remove prompt observer.
-
-        Args:
-            observer: Observer to remove
-        """
-        self._prompt_registry.remove_observer(observer)
-
-    def remove_tool_observer(
-        self, observer: RegistryEvents[str, LLMCallableTool]
+    def remove_observer(
+        self,
+        observer: RegistryEvents[str, Any],
+        registry_type: RegistryType,
     ) -> None:
-        """Remove tool observer.
+        """Remove a registry observer.
 
         Args:
             observer: Observer to remove
-        """
-        self._tool_registry.remove_observer(observer)
+            registry_type: Type of registry to remove from
+        """
+        match registry_type:
+            case "resource":
+                self._resource_registry.remove_observer(observer)
+            case "prompt":
+                self._prompt_registry.remove_observer(observer)
+            case "tool":
+                self._tool_registry.remove_observer(observer)
 
     @property
     def original_config(self) -> Config:

diff --git a/src/llmling/server/server.py b/src/llmling/server/server.py
@@ -124,9 +124,9 @@ def _setup_observers(self) -> None:
         self.prompt_observer = PromptObserver(self)
         self.tool_observer = ToolObserver(self)
 
-        self.runtime.add_resource_observer(self.resource_observer.events)
-        self.runtime.add_prompt_observer(self.prompt_observer.events)
-        self.runtime.add_tool_observer(self.tool_observer.events)
+        self.runtime.add_observer(self.resource_observer.events, "resource")
+        self.runtime.add_observer(self.prompt_observer.events, "prompt")
+        self.runtime.add_observer(self.tool_observer.events, "tool")
 
     async def start(self, *, raise_exceptions: bool = False) -> None:
         """Start the server."""
@@ -149,9 +149,9 @@ async def shutdown(self) -> None:
                 self._tasks.clear()
 
             # Remove observers
-            self.runtime.remove_resource_observer(self.resource_observer.events)
-            self.runtime.remove_prompt_observer(self.prompt_observer.events)
-            self.runtime.remove_tool_observer(self.tool_observer.events)
+            self.runtime.remove_observer(self.resource_observer.events, "resource")
+            self.runtime.remove_observer(self.prompt_observer.events, "prompt")
+            self.runtime.remove_observer(self.tool_observer.events, "tool")
 
             # Shutdown runtime
             await self.runtime.shutdown()