162 lines
4.8 KiB
Python
162 lines
4.8 KiB
Python
"""
|
|
Model registry for managing available models/agents.
|
|
|
|
This registry maintains the list of available models and their capabilities.
|
|
It provides a central place to:
|
|
- List all models for the /v1/models endpoint
|
|
- Instantiate agents for specific models
|
|
- Check model capabilities
|
|
"""
|
|
|
|
import time
|
|
from typing import Type
|
|
|
|
from src.agents.base import AgentInterface
|
|
from src.agents.lorem_tester import LoremTesterAgent
|
|
from src.agents.tatlock import TatlockAgent
|
|
from src.core.exceptions import ModelNotFoundError
|
|
|
|
|
|
class ModelRegistry:
|
|
"""
|
|
Central registry for all available models.
|
|
|
|
Each model maps to an agent implementation. The registry provides:
|
|
- Model listing (for /v1/models endpoint)
|
|
- Agent instantiation (for request handling)
|
|
- Capability information (for client discovery)
|
|
"""
|
|
|
|
# Model configurations
|
|
# Add new models here as they're implemented
|
|
MODELS: dict[str, dict] = {
|
|
"lorem-tester": {
|
|
"agent_class": LoremTesterAgent,
|
|
"description": "Testing agent with mock Responses API features (Lorem Ipsum)",
|
|
"created": 1733529600, # 2025-12-06
|
|
"owned_by": "tatlock",
|
|
# Capabilities are retrieved from agent instance
|
|
},
|
|
"Tatlock": {
|
|
"agent_class": TatlockAgent,
|
|
"description": "Tatlock - Your homelab butler (British household coordinator)",
|
|
"created": 1733529600, # 2025-12-06
|
|
"owned_by": "tatlock",
|
|
# Capabilities are retrieved from agent instance
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def get_agent(cls, model_id: str) -> AgentInterface:
|
|
"""
|
|
Instantiate agent for given model ID.
|
|
|
|
Args:
|
|
model_id: Model identifier (e.g., "lorem-tester", "tatlock")
|
|
|
|
Returns:
|
|
AgentInterface: Instance of the agent
|
|
|
|
Raises:
|
|
ModelNotFoundError: If model_id not found in registry
|
|
"""
|
|
if model_id not in cls.MODELS:
|
|
raise ModelNotFoundError(model_id)
|
|
|
|
agent_class: Type[AgentInterface] = cls.MODELS[model_id]["agent_class"]
|
|
return agent_class()
|
|
|
|
@classmethod
|
|
async def list_models(cls) -> list[dict]:
|
|
"""
|
|
Return all models in OpenAI-compatible format.
|
|
|
|
Returns:
|
|
list[dict]: List of model objects with:
|
|
- id: Model identifier
|
|
- object: Always "model"
|
|
- created: Unix timestamp
|
|
- owned_by: Owner identifier
|
|
- capabilities: Dict of capabilities
|
|
- description: Human-readable description
|
|
|
|
Example:
|
|
[
|
|
{
|
|
"id": "lorem-tester",
|
|
"object": "model",
|
|
"created": 1733529600,
|
|
"owned_by": "tatlock",
|
|
"capabilities": {
|
|
"streaming": True,
|
|
"reasoning": True,
|
|
"tools": True,
|
|
"vision": False,
|
|
"audio": False
|
|
},
|
|
"description": "Testing agent..."
|
|
},
|
|
...
|
|
]
|
|
"""
|
|
models = []
|
|
|
|
for model_id, config in cls.MODELS.items():
|
|
# Instantiate agent to get capabilities
|
|
agent = cls.get_agent(model_id)
|
|
capabilities = await agent.get_capabilities()
|
|
|
|
models.append({
|
|
"id": model_id,
|
|
"object": "model",
|
|
"created": config["created"],
|
|
"owned_by": config["owned_by"],
|
|
"capabilities": capabilities,
|
|
"description": config["description"],
|
|
})
|
|
|
|
return models
|
|
|
|
@classmethod
|
|
def model_exists(cls, model_id: str) -> bool:
|
|
"""
|
|
Check if a model exists in the registry.
|
|
|
|
Args:
|
|
model_id: Model identifier
|
|
|
|
Returns:
|
|
bool: True if model exists
|
|
"""
|
|
return model_id in cls.MODELS
|
|
|
|
@classmethod
|
|
async def get_model_info(cls, model_id: str) -> dict:
|
|
"""
|
|
Get detailed information about a specific model.
|
|
|
|
Args:
|
|
model_id: Model identifier
|
|
|
|
Returns:
|
|
dict: Model information
|
|
|
|
Raises:
|
|
ModelNotFoundError: If model not found
|
|
"""
|
|
if not cls.model_exists(model_id):
|
|
raise ModelNotFoundError(model_id)
|
|
|
|
config = cls.MODELS[model_id]
|
|
agent = cls.get_agent(model_id)
|
|
capabilities = await agent.get_capabilities()
|
|
|
|
return {
|
|
"id": model_id,
|
|
"object": "model",
|
|
"created": config["created"],
|
|
"owned_by": config["owned_by"],
|
|
"capabilities": capabilities,
|
|
"description": config["description"],
|
|
}
|