Skip to content

base

Module for handling client routing.

Serves as a workaround for Spark serialization issues.

ModelClass

Bases: str, Enum

Enumeration of available model classes.

Source code in spark_instructor/client/base.py 
163
164
165
166
167
168
169
class ModelClass(str, Enum):
    """Enumeration of available model classes."""

    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    DATABRICKS = "databricks"
    OLLAMA = "ollama"

get_anthropic_aclient(mode=instructor.Mode.ANTHROPIC_JSON, base_url=None, api_key=None, enable_prompt_caching=False) cached

Get the async Anthropic client.

Ensure that ANTHROPIC_API_KEY is set.

Source code in spark_instructor/client/base.py 
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
@lru_cache
def get_anthropic_aclient(
    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    enable_prompt_caching: bool = False,
) -> instructor.AsyncInstructor:
    """Get the async Anthropic client.

    Ensure that ``ANTHROPIC_API_KEY`` is set.
    """
    try:
        from anthropic import AsyncAnthropic
    except ImportError:
        raise ImportError(
            "Please install ``anthropic`` by running ``pip install anthropic`` "
            "or ``pip install spark-instructor[anthropic]``"
        )

    if not api_key:
        assert_env_is_set("ANTHROPIC_API_KEY")
        return instructor.from_anthropic(AsyncAnthropic(), mode=mode, enable_prompt_caching=enable_prompt_caching)
    return instructor.from_anthropic(
        AsyncAnthropic(api_key=api_key, base_url=base_url), mode=mode, enable_prompt_caching=enable_prompt_caching
    )

get_anthropic_client(mode=instructor.Mode.ANTHROPIC_JSON, base_url=None, api_key=None, enable_prompt_caching=False) cached

Get the Anthropic client.

Ensure that ANTHROPIC_API_KEY is set.

Source code in spark_instructor/client/base.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
@lru_cache
def get_anthropic_client(
    mode: instructor.Mode = instructor.Mode.ANTHROPIC_JSON,
    base_url: Optional[str] = None,
    api_key: Optional[str] = None,
    enable_prompt_caching: bool = False,
) -> instructor.Instructor:
    """Get the Anthropic client.

    Ensure that ``ANTHROPIC_API_KEY`` is set.
    """
    try:
        from anthropic import Anthropic
    except ImportError:
        raise ImportError(
            "Please install ``anthropic`` by running ``pip install anthropic`` "
            "or ``pip install spark-instructor[anthropic]``"
        )
    if not api_key:
        assert_env_is_set("ANTHROPIC_API_KEY")
        return instructor.from_anthropic(Anthropic(), mode=mode, enable_prompt_caching=enable_prompt_caching)
    return instructor.from_anthropic(
        Anthropic(api_key=api_key, base_url=base_url), mode=mode, enable_prompt_caching=enable_prompt_caching
    )

get_async_instructor(model_class=None, mode=None, api_key=None, base_url=None)

Get the instructor client based on the model class and mode.

Source code in spark_instructor/client/base.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
def get_async_instructor(
    model_class: Optional[ModelClass] = None,
    mode: Optional[instructor.Mode] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
) -> instructor.AsyncInstructor:
    """Get the instructor client based on the model class and mode."""
    if model_class is None:
        # Use OpenAI by default
        model_class = ModelClass.OPENAI
    kwargs: Dict[str, Any] = dict(api_key=api_key, base_url=base_url)
    if mode:
        kwargs |= dict(mode=mode)
    return MODEL_CLASS_ROUTE_ASYNC[model_class](**kwargs)

get_databricks_aclient(mode=instructor.Mode.MD_JSON, base_url=None, api_key=None) cached

Get the async databricks client.

Unless passed as arguments, ensure that the DATABRICKS_HOST and DATABRICKS_TOKEN environment variables are set.

Source code in spark_instructor/client/base.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
@lru_cache
def get_databricks_aclient(
    mode: instructor.Mode = instructor.Mode.MD_JSON, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.AsyncInstructor:
    """Get the async databricks client.

    Unless passed as arguments,
    ensure that the ``DATABRICKS_HOST`` and ``DATABRICKS_TOKEN`` environment variables are set.
    """
    if base_url is None:
        base_url = f"{get_env_variable('DATABRICKS_HOST')}/serving-endpoints"
    if api_key is None:
        api_key = get_env_variable("DATABRICKS_TOKEN")
    return instructor.from_openai(
        AsyncOpenAI(api_key=api_key, base_url=base_url),
        mode=mode,
    )

get_databricks_client(mode=instructor.Mode.MD_JSON, base_url=None, api_key=None) cached

Get the databricks client.

Unless passed as arguments, ensure that the DATABRICKS_HOST and DATABRICKS_TOKEN environment variables are set.

Source code in spark_instructor/client/base.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
@lru_cache
def get_databricks_client(
    mode: instructor.Mode = instructor.Mode.MD_JSON, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.Instructor:
    """Get the databricks client.

    Unless passed as arguments,
    ensure that the ``DATABRICKS_HOST`` and ``DATABRICKS_TOKEN`` environment variables are set.
    """
    if base_url is None:
        base_url = f"{get_env_variable('DATABRICKS_HOST')}/serving-endpoints"
    if api_key is None:
        api_key = get_env_variable("DATABRICKS_TOKEN")
    return instructor.from_openai(
        OpenAI(api_key=api_key, base_url=base_url),
        mode=mode,
    )

get_instructor(model_class=None, mode=None, api_key=None, base_url=None)

Get the instructor client based on the model class and mode.

Source code in spark_instructor/client/base.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def get_instructor(
    model_class: Optional[ModelClass] = None,
    mode: Optional[instructor.Mode] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
) -> instructor.Instructor:
    """Get the instructor client based on the model class and mode."""
    if model_class is None:
        # Use OpenAI by default
        model_class = ModelClass.OPENAI
    kwargs: Dict[str, Any] = dict(api_key=api_key, base_url=base_url)
    if mode:
        kwargs |= dict(mode=mode)
    return MODEL_CLASS_ROUTE[model_class](**kwargs)

get_ollama_aclient(mode=instructor.Mode.JSON, base_url=None, api_key=None) cached

Get the async Ollama client.

Source code in spark_instructor/client/base.py 
149
150
151
152
153
154
155
156
157
158
159
160
@lru_cache
def get_ollama_aclient(
    mode: instructor.Mode = instructor.Mode.JSON, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.AsyncInstructor:
    """Get the async Ollama client."""
    if not base_url:
        host = "http://localhost:11434"
        base_url = f"{host}/v1"

    return instructor.from_openai(
        AsyncOpenAI(base_url=base_url, api_key="ollama" if not api_key else api_key), mode=mode
    )

get_ollama_client(mode=instructor.Mode.JSON, base_url=None, api_key=None) cached

Get the Ollama client.

Source code in spark_instructor/client/base.py 
137
138
139
140
141
142
143
144
145
146
@lru_cache
def get_ollama_client(
    mode: instructor.Mode = instructor.Mode.JSON, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.Instructor:
    """Get the Ollama client."""
    if not base_url:
        host = "http://localhost:11434"
        base_url = f"{host}/v1"

    return instructor.from_openai(OpenAI(base_url=base_url, api_key="ollama" if not api_key else api_key), mode=mode)

get_openai_aclient(mode=instructor.Mode.TOOLS, base_url=None, api_key=None) cached

Get the async OpenAI client.

Unless passes as an argument, Ensure that OPENAI_API_KEY is set.

Source code in spark_instructor/client/base.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
@lru_cache
def get_openai_aclient(
    mode: instructor.Mode = instructor.Mode.TOOLS, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.AsyncInstructor:
    """Get the async OpenAI client.

    Unless passes as an argument,
    Ensure that ``OPENAI_API_KEY`` is set.
    """
    if not api_key:
        assert_env_is_set("OPENAI_API_KEY")
        return instructor.from_openai(AsyncOpenAI(base_url=base_url), mode=mode)
    return instructor.from_openai(AsyncOpenAI(base_url=base_url, api_key=api_key), mode=mode)

get_openai_client(mode=instructor.Mode.TOOLS, base_url=None, api_key=None) cached

Get the OpenAI client.

Unless passes as an argument, Ensure that OPENAI_API_KEY is set.

Source code in spark_instructor/client/base.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
@lru_cache
def get_openai_client(
    mode: instructor.Mode = instructor.Mode.TOOLS, base_url: Optional[str] = None, api_key: Optional[str] = None
) -> instructor.Instructor:
    """Get the OpenAI client.

    Unless passes as an argument,
    Ensure that ``OPENAI_API_KEY`` is set.
    """
    if not api_key:
        assert_env_is_set("OPENAI_API_KEY")
        return instructor.from_openai(OpenAI(base_url=base_url), mode=mode)
    return instructor.from_openai(OpenAI(base_url=base_url, api_key=api_key), mode=mode)

infer_model_class(model_name)

Attempt to infer the model class from the model name.

Source code in spark_instructor/client/base.py 
172
173
174
175
176
177
178
179
180
181
182
def infer_model_class(model_name: str) -> ModelClass:
    """Attempt to infer the model class from the model name."""
    if "databricks" in model_name:
        return ModelClass.DATABRICKS
    elif "gpt" in model_name:
        return ModelClass.OPENAI
    elif "claude" in model_name:
        return ModelClass.ANTHROPIC
    elif "llama" in model_name:
        return ModelClass.OLLAMA
    raise ValueError(f"Model name `{model_name}` does not match any of the available model classes.")