Skip to content

registry

Module for defining registries of clients.

ClientRegistry dataclass

A registry for managing and accessing different client factories for various AI models.

This class serves as a central repository for client factories, allowing easy registration, retrieval, and mapping of model names to their respective factories. It provides a flexible way to manage multiple AI model providers in a single application.

Attributes:

Name Type Description
factories Dict[str, Type[ClientFactory]]

A dictionary mapping model class names to their corresponding ClientFactory classes.
model_map Dict[str, str]

A dictionary mapping model name patterns to their corresponding model class names.

The default configuration includes factories for Anthropic, OpenAI, Databricks, and Ollama, with mappings for common model name patterns.

Usage 
registry = ClientRegistry()
openai_factory = registry.get_factory_from_model("gpt-3.5-turbo")
Source code in spark_instructor/registry.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@dataclass
class ClientRegistry:
    """A registry for managing and accessing different client factories for various AI models.

    This class serves as a central repository for client factories, allowing easy registration,
    retrieval, and mapping of model names to their respective factories. It provides a flexible
    way to manage multiple AI model providers in a single application.

    Attributes:
        factories (Dict[str, Type[ClientFactory]]): A dictionary mapping model class names
            to their corresponding ClientFactory classes.
        model_map (Dict[str, str]): A dictionary mapping model name patterns to their
            corresponding model class names.

    The default configuration includes factories for Anthropic, OpenAI, Databricks, and Ollama,
    with mappings for common model name patterns.

    Usage:
        ```python
        registry = ClientRegistry()
        openai_factory = registry.get_factory_from_model("gpt-3.5-turbo")
        ```
    """

    factories: Dict[str, Type[ClientFactory]] = field(default_factory=get_default_factories)
    model_map: Dict[str, str] = field(
        default_factory=lambda: {
            "databricks": "databricks",
            "gpt": "openai",
            "claude": "anthropic",
            "llama": "ollama",
            "o1": "o1",
        }
    )

    def register(self, model_class: str, factory_class: Type[ClientFactory]):
        """Register a new factory for a given model class.

        This method allows adding new factory classes to the registry or overriding
        existing ones.

        Args:
            model_class (str): The name of the model class (e.g., "openai", "anthropic").
            factory_class (Type[ClientFactory]): The ClientFactory subclass to be registered.

        Note:
            The model_class is converted to lowercase before registration to ensure
            case-insensitive lookups.

        Example:
            ```python
            registry.register("custom_model", CustomModelFactory)
            ```
        """
        self.factories[model_class.lower()] = factory_class

    def get_factory(self, model_class: str) -> Type[ClientFactory]:
        """Retrieve a factory class for a given model class.

        This method looks up and returns the appropriate ClientFactory subclass
        for the specified model class.

        Args:
            model_class (str): The name of the model class to look up.

        Returns:
            Type[ClientFactory]: The corresponding ClientFactory subclass.

        Raises:
            ValueError: If no factory is registered for the given model class.

        Example:
            ```python
            openai_factory = registry.get_factory("openai")
            ```
        """
        factory_class = self.factories.get(model_class.lower())
        if not factory_class:
            raise ValueError(f"No factory registered for model class: {model_class}")
        return factory_class

    def get_factory_from_model(self, model: str) -> Type[ClientFactory]:
        """Retrieve a factory class based on a model name.

        This method uses the model_map to determine the appropriate model class
        for a given model name, then returns the corresponding factory.

        Args:
            model (str): The name of the model (e.g., "gpt-4o", "claude-3-5-sonnet-20240620").

        Returns:
            Type[ClientFactory]: The corresponding ClientFactory subclass.

        Raises:
            ValueError: If the model name doesn't match any known patterns in the model_map.

        Example:
            ```python
            claude_factory = registry.get_factory_from_model("claude-3-5-sonnet-20240620")
            ```

        Note:
            This method performs a partial match on the model name. For example,
            any model name containing "gpt" will be mapped to the OpenAI factory.
        """
        for key, val in self.model_map.items():
            if key in model:
                return self.get_factory(val)
        raise ValueError(f"`{model}` does not have a mapped factory. Consider updating the `model_map`.")

get_factory(model_class)

Retrieve a factory class for a given model class.

This method looks up and returns the appropriate ClientFactory subclass for the specified model class.

Parameters:

Name Type Description Default
model_class str

The name of the model class to look up.

 required

Returns:

Type Description
Type[ClientFactory]

Type[ClientFactory]: The corresponding ClientFactory subclass.

Raises:

Type Description
ValueError

If no factory is registered for the given model class.
Example 
openai_factory = registry.get_factory("openai")
Source code in spark_instructor/registry.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def get_factory(self, model_class: str) -> Type[ClientFactory]:
    """Retrieve a factory class for a given model class.

    This method looks up and returns the appropriate ClientFactory subclass
    for the specified model class.

    Args:
        model_class (str): The name of the model class to look up.

    Returns:
        Type[ClientFactory]: The corresponding ClientFactory subclass.

    Raises:
        ValueError: If no factory is registered for the given model class.

    Example:
        ```python
        openai_factory = registry.get_factory("openai")
        ```
    """
    factory_class = self.factories.get(model_class.lower())
    if not factory_class:
        raise ValueError(f"No factory registered for model class: {model_class}")
    return factory_class

get_factory_from_model(model)

Retrieve a factory class based on a model name.

This method uses the model_map to determine the appropriate model class for a given model name, then returns the corresponding factory.

Parameters:

Name Type Description Default
model str

The name of the model (e.g., "gpt-4o", "claude-3-5-sonnet-20240620").

 required

Returns:

Type Description
Type[ClientFactory]

Type[ClientFactory]: The corresponding ClientFactory subclass.

Raises:

Type Description
ValueError

If the model name doesn't match any known patterns in the model_map.
Example 
claude_factory = registry.get_factory_from_model("claude-3-5-sonnet-20240620")
Note

This method performs a partial match on the model name. For example, any model name containing "gpt" will be mapped to the OpenAI factory.

Source code in spark_instructor/registry.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def get_factory_from_model(self, model: str) -> Type[ClientFactory]:
    """Retrieve a factory class based on a model name.

    This method uses the model_map to determine the appropriate model class
    for a given model name, then returns the corresponding factory.

    Args:
        model (str): The name of the model (e.g., "gpt-4o", "claude-3-5-sonnet-20240620").

    Returns:
        Type[ClientFactory]: The corresponding ClientFactory subclass.

    Raises:
        ValueError: If the model name doesn't match any known patterns in the model_map.

    Example:
        ```python
        claude_factory = registry.get_factory_from_model("claude-3-5-sonnet-20240620")
        ```

    Note:
        This method performs a partial match on the model name. For example,
        any model name containing "gpt" will be mapped to the OpenAI factory.
    """
    for key, val in self.model_map.items():
        if key in model:
            return self.get_factory(val)
    raise ValueError(f"`{model}` does not have a mapped factory. Consider updating the `model_map`.")

register(model_class, factory_class)

Register a new factory for a given model class.

This method allows adding new factory classes to the registry or overriding existing ones.

Parameters:

Name Type Description Default
model_class str

The name of the model class (e.g., "openai", "anthropic").

 required
factory_class Type[ClientFactory]

The ClientFactory subclass to be registered.

 required
Note

The model_class is converted to lowercase before registration to ensure case-insensitive lookups.

Example 
registry.register("custom_model", CustomModelFactory)
Source code in spark_instructor/registry.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def register(self, model_class: str, factory_class: Type[ClientFactory]):
    """Register a new factory for a given model class.

    This method allows adding new factory classes to the registry or overriding
    existing ones.

    Args:
        model_class (str): The name of the model class (e.g., "openai", "anthropic").
        factory_class (Type[ClientFactory]): The ClientFactory subclass to be registered.

    Note:
        The model_class is converted to lowercase before registration to ensure
        case-insensitive lookups.

    Example:
        ```python
        registry.register("custom_model", CustomModelFactory)
        ```
    """
    self.factories[model_class.lower()] = factory_class

get_default_factories()

Get default factories based on whether anthropic is available.

Source code in spark_instructor/registry.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
def get_default_factories() -> Dict[str, Type[ClientFactory]]:
    """Get default factories based on whether anthropic is available."""
    if is_anthropic_available():
        from spark_instructor.factory import AnthropicFactory

        return {
            "anthropic": AnthropicFactory,
            "openai": OpenAIFactory,
            "databricks": DatabricksFactory,
            "ollama": OllamaFactory,
            "o1": O1Factory,
        }
    return {"openai": OpenAIFactory, "databricks": DatabricksFactory, "ollama": OllamaFactory, "o1": O1Factory}