18class PromptRegistry:
19    """
20    Registry for managing prompt templates.
21    
22    Allows versioning, retrieval, and formatting of prompt templates.
23    """
24    
25    def __init__(self):
26        """Initialize prompt registry."""
27        self._prompts: Dict[str, PromptTemplate] = {}
28    
29    def register(
30        self,
31        name: str,
32        template: str,
33        version: str = "1.0",
34        variables: Optional[list[str]] = None,
35        description: str = ""
36    ) -> None:
37        """Register a new prompt template."""
38        key = f"{name}:{version}"
39        self._prompts[key] = PromptTemplate(
40            name=name,
41            template=template,
42            version=version,
43            variables=variables or [],
44            description=description
45        )
46    
47    def get(self, name: str, version: Optional[str] = None) -> Optional[PromptTemplate]:
48        """Get a prompt template."""
49        if version:
50            key = f"{name}:{version}"
51            return self._prompts.get(key)
52        
53        # Get latest version
54        prompts_with_name = [
55            p for k, p in self._prompts.items() if p.name == name
56        ]
57        if prompts_with_name:
58            return max(prompts_with_name, key=lambda p: p.version)
59        
60        return None
61    
62    def format(self, name: str, **variables) -> Optional[str]:
63        """Format a prompt template with variables."""
64        prompt = self.get(name)
65        if not prompt:
66            return None
67        
68        return prompt.template.format(**variables)

18class ToolRegistry:
19    """
20    Registry for managing tools available to agents.
21    
22    Tools can be functions, MCP servers, or other callable interfaces.
23    """
24    
25    def __init__(self):
26        """Initialize tool registry."""
27        self._tools: Dict[str, ToolInfo] = {}
28    
29    def register(
30        self,
31        name: str,
32        function: Callable,
33        description: str = "",
34        parameters: Optional[Dict[str, Any]] = None,
35        category: str = "general"
36    ) -> None:
37        """Register a new tool."""
38        self._tools[name] = ToolInfo(
39            name=name,
40            description=description,
41            function=function,
42            parameters=parameters or {},
43            category=category
44        )
45    
46    def get(self, name: str) -> Optional[ToolInfo]:
47        """Get tool information."""
48        return self._tools.get(name)
49    
50    def list_tools(self, category: Optional[str] = None) -> list[ToolInfo]:
51        """List all registered tools, optionally filtered by category."""
52        tools = list(self._tools.values())
53        if category:
54            tools = [t for t in tools if t.category == category]
55        return tools
56    
57    async def execute(self, name: str, **kwargs: Any) -> Any:
58        """Execute a tool by name."""
59        tool = self.get(name)
60        if not tool:
61            raise ValueError(f"Tool not found: {name}")
62        
63        return await tool.function(**kwargs)

 14class LLMProviderRegistry:
 15    """
 16    Registry for managing multiple LLM providers.
 17    
 18    Allows applications to:
 19    - Register multiple providers (Azure, Ollama, etc.)
 20    - Set a default provider
 21    - Retrieve providers by name
 22    - List available providers
 23    
 24    Example:
 25        >>> registry = LLMProviderRegistry()
 26        >>> registry.register("azure", azure_provider, is_default=True)
 27        >>> registry.register("ollama", ollama_provider)
 28        >>> 
 29        >>> # Get specific provider
 30        >>> provider = registry.get("ollama")
 31        >>> 
 32        >>> # Get default provider
 33        >>> default = registry.get_default()
 34    """
 35    
 36    def __init__(self):
 37        """Initialize the provider registry."""
 38        self._providers: Dict[str, BaseProvider] = {}
 39        self._default_provider: Optional[str] = None
 40        # Store models separately: { provider_name: { model_name: ModelInfo } }
 41        self._models: Dict[str, Dict[str, ModelInfo]] = {}
 42    
 43    def register(
 44        self,
 45        name: str,
 46        provider: BaseProvider,
 47        is_default: bool = False
 48    ) -> None:
 49        """
 50        Register a new LLM provider.
 51        
 52        Args:
 53            name: Unique identifier for the provider (e.g., "azure", "ollama")
 54            provider: The provider instance
 55            is_default: Whether this should be the default provider
 56            
 57        Raises:
 58            ValueError: If provider with this name already exists
 59            TypeError: If provider is not a BaseProvider instance
 60        """
 61        if not isinstance(provider, BaseProvider):
 62            raise TypeError(f"Provider must be an instance of BaseProvider, got {type(provider)}")
 63        
 64        if name in self._providers:
 65            raise ValueError(f"Provider '{name}' is already registered. Use unregister() first.")
 66        
 67        self._providers[name] = provider
 68        self._models[name] = {}  # Initialize empty model dict for this provider
 69        
 70        # Set as default if requested or if it's the first provider
 71        if is_default or self._default_provider is None:
 72            self._default_provider = name
 73    
 74    def unregister(self, name: str) -> None:
 75        """
 76        Unregister a provider.
 77        
 78        Args:
 79            name: Name of the provider to unregister
 80            
 81        Raises:
 82            KeyError: If provider doesn't exist
 83        """
 84        if name not in self._providers:
 85            raise KeyError(f"Provider '{name}' not found in registry")
 86        
 87        del self._providers[name]
 88        del self._models[name]  # Also remove all models for this provider
 89        
 90        # If we unregistered the default, set a new default if providers remain
 91        if self._default_provider == name:
 92            self._default_provider = next(iter(self._providers.keys())) if self._providers else None
 93    
 94    def get(self, name: str) -> BaseProvider:
 95        """
 96        Get a provider by name.
 97        
 98        Args:
 99            name: Name of the provider
100            
101        Returns:
102            The provider instance
103            
104        Raises:
105            KeyError: If provider doesn't exist
106        """
107        if name not in self._providers:
108            available = ", ".join(self._providers.keys()) if self._providers else "none"
109            raise KeyError(
110                f"Provider '{name}' not found. Available providers: {available}"
111            )
112        
113        return self._providers[name]
114    
115    def get_default(self) -> BaseProvider:
116        """
117        Get the default provider.
118        
119        Returns:
120            The default provider instance
121            
122        Raises:
123            RuntimeError: If no default provider is set
124        """
125        if self._default_provider is None:
126            raise RuntimeError("No default provider set. Register at least one provider.")
127        
128        return self._providers[self._default_provider]
129    
130    def set_default(self, name: str) -> None:
131        """
132        Set the default provider.
133        
134        Args:
135            name: Name of the provider to set as default
136            
137        Raises:
138            KeyError: If provider doesn't exist
139        """
140        if name not in self._providers:
141            raise KeyError(f"Provider '{name}' not found in registry")
142        
143        self._default_provider = name
144    
145    def has_provider(self, name: str) -> bool:
146        """
147        Check if a provider exists in the registry.
148        
149        Args:
150            name: Name of the provider
151            
152        Returns:
153            True if provider exists, False otherwise
154        """
155        return name in self._providers
156    
157    def list_providers(self) -> List[str]:
158        """
159        List all registered provider names.
160        
161        Returns:
162            List of provider names
163        """
164        return list(self._providers.keys())
165    
166    def get_default_name(self) -> Optional[str]:
167        """
168        Get the name of the default provider.
169        
170        Returns:
171            Name of default provider, or None if no providers registered
172        """
173        return self._default_provider
174    
175    def clear(self) -> None:
176        """
177        Remove all providers from the registry.
178        
179        Useful for testing or resetting state.
180        """
181        self._providers.clear()
182        self._models.clear()
183        self._default_provider = None
184    
185    # Model Management Methods
186    
187    def register_model(
188        self,
189        provider_name: str,
190        model_name: str,
191        capabilities: Optional[Dict] = None,
192        **metadata
193    ) -> None:
194        """
195        Register a model for a specific provider.
196        
197        Args:
198            provider_name: Name of the provider (e.g., "azure", "ollama")
199            model_name: Name of the model (e.g., "gpt-4", "llama3")
200            capabilities: Dict of model capabilities (max_tokens, context_window, etc.)
201            **metadata: Additional metadata about the model
202            
203        Raises:
204            KeyError: If provider doesn't exist
205            ValueError: If model already registered for this provider
206            
207        Example:
208            >>> registry.register("azure", azure_provider)
209            >>> registry.register_model(
210            ...     "azure",
211            ...     "gpt-4",
212            ...     capabilities={"max_tokens": 8192, "context_window": 128000},
213            ...     deployment="my-gpt4-deployment"
214            ... )
215        """
216        if provider_name not in self._providers:
217            raise KeyError(f"Provider '{provider_name}' not found. Register provider first.")
218        
219        if model_name in self._models[provider_name]:
220            raise ValueError(
221                f"Model '{model_name}' already registered for provider '{provider_name}'. "
222                f"Use unregister_model() first."
223            )
224        
225        provider = self._providers[provider_name]
226        model_info = ModelInfo(
227            name=model_name,
228            provider=provider.name,
229            capabilities=capabilities or {},
230            metadata=metadata
231        )
232        self._models[provider_name][model_name] = model_info
233    
234    def unregister_model(self, provider_name: str, model_name: str) -> None:
235        """
236        Unregister a model from a provider.
237        
238        Args:
239            provider_name: Name of the provider
240            model_name: Name of the model
241            
242        Raises:
243            KeyError: If provider or model doesn't exist
244        """
245        if provider_name not in self._providers:
246            raise KeyError(f"Provider '{provider_name}' not found")
247        
248        if model_name not in self._models[provider_name]:
249            raise KeyError(
250                f"Model '{model_name}' not registered for provider '{provider_name}'"
251            )
252        
253        del self._models[provider_name][model_name]
254    
255    def has_model(self, provider_name: str, model_name: str) -> bool:
256        """
257        Check if a model is registered for a provider.
258        
259        Args:
260            provider_name: Name of the provider
261            model_name: Name of the model
262            
263        Returns:
264            True if model is registered, False otherwise
265        """
266        if provider_name not in self._providers:
267            return False
268        return model_name in self._models[provider_name]
269    
270    def __len__(self) -> int:
271        """Return the number of registered providers."""
272        return len(self._providers)
273    
274    def __contains__(self, name: str) -> bool:
275        """Check if a provider exists (supports 'in' operator)."""
276        return name in self._providers
277    
278    # Model Discovery Methods
279    
280    async def list_all_models(self) -> List[ModelInfo]:
281        """
282        List all models from all registered providers.
283        
284        Returns:
285            List of ModelInfo objects from all providers
286            
287        Example:
288            >>> models = await registry.list_all_models()
289            >>> for model in models:
290            ...     print(f"{model.name} from {model.provider}")
291        """
292        all_models = []
293        for provider_models in self._models.values():
294            all_models.extend(provider_models.values())
295        return all_models
296    
297    async def list_models(self, provider_name: Optional[str] = None) -> List[ModelInfo]:
298        """
299        List models from a specific provider or the default provider.
300        
301        Args:
302            provider_name: Name of the provider, or None for default
303            
304        Returns:
305            List of ModelInfo objects from the specified provider
306            
307        Raises:
308            KeyError: If provider doesn't exist
309        """
310        if provider_name is None:
311            if self._default_provider is None:
312                return []
313            provider_name = self._default_provider
314        
315        if provider_name not in self._providers:
316            raise KeyError(f"Provider '{provider_name}' not found")
317        
318        return list(self._models[provider_name].values())
319    
320    async def get_model_info(self, model_name: str) -> Optional[ModelInfo]:
321        """
322        Find information about a specific model across all providers.
323        
324        Searches all registered providers for a model with the given name.
325        
326        Args:
327            model_name: Name of the model to find
328            
329        Returns:
330            ModelInfo object if found, None otherwise
331        """
332        for provider_models in self._models.values():
333            if model_name in provider_models:
334                return provider_models[model_name]
335        return None
336    
337    def __repr__(self) -> str:
338        """String representation of the registry."""
339        providers = ", ".join(self._providers.keys()) if self._providers else "none"
340        default = f" (default: {self._default_provider})" if self._default_provider else ""
341        return f"LLMProviderRegistry(providers=[{providers}]{default})"