chore: turn OpenAIMixin into a pydantic.BaseModel #3671

mattf · 2025-10-03T18:20:17Z

What does this PR do?

implement get_api_key instead of relying on LiteLLMOpenAIMixin.get_api_key
remove use of LiteLLMOpenAIMixin
add default initialize/shutdown methods to OpenAIMixin
remove __init__s to allow proper pydantic construction
remove dead code from vllm adapter and associated / duplicate unit tests
update vllm adapter to use openaimixin for model registration
remove ModelRegistryHelper from fireworks & together adapters
remove Inference from nvidia adapter
complete type hints on embedding_model_metadata
allow extra fields on OpenAIMixin, for model_store, provider_id, etc
new recordings for ollama
enhance the list models error handling
update cerebras (remove cerebras-cloud-sdk) and anthropic (custom model listing) inference adapters
parametrized test_inference_client_caching
remove cerebras, databricks, fireworks, together from blanket mypy exclude
removed unnecessary litellm deps

Test Plan

ci

…aking _model_cache details

…c.BaseModel - implement get_api_key instead of relying on LiteLLMOpenAIMixin.get_api_key - remove use of LiteLLMOpenAIMixin - add default initialize/shutdown methods to OpenAIMixin - remove __init__s to allow proper pydantic construction - remove dead code from vllm adapter and associated / duplicate unit tests - update vllm adapter to use openaimixin for model registration - remove ModelRegistryHelper from fireworks & together adapters - remove Inference from nvidia adapter - complete type hints on embedding_model_metadata - allow extra fields on OpenAIMixin, for model_store, __provider_id__, etc - new recordings for ollama - enhance the list models error handling w/ new tests - update cerebras (remove cerebras-cloud-sdk) and anthropic (custom model listing) inference adapters - parametrized test_inference_client_caching - remove cerebras, databricks, fireworks, together from blanket mypy exclude

mattf · 2025-10-06T14:13:01Z

this provides a foundation for #3517

franciscojavierarceo

this lgtm as you answered the open questions I had in the PR description but i'll wait for someone else who has spent more time on inference to do another pass. but it looks a rebase is needed.

leseb

missing 2 files (vllm.py and openai_mixin.py,)

leseb · 2025-10-06T13:52:03Z

llama_stack/providers/registry/inference.py

            adapter_type="anthropic",
            provider_type="remote::anthropic",
-            pip_packages=["litellm"],
+            pip_packages=["litellm", "anthropic"],


Do we still need litellm here? I only see OpenAIMixin being used?

good catch. i'll remove all the others too.

leseb · 2025-10-06T13:54:13Z

llama_stack/providers/remote/inference/azure/azure.py

 from urllib.parse import urljoin

-from llama_stack.apis.inference import ChatCompletionRequest
-from llama_stack.providers.utils.inference.litellm_openai_mixin import (


can we remove litellm from RemoteProviderSpec in the registry?

leseb · 2025-10-06T13:56:33Z

llama_stack/providers/remote/inference/databricks/config.py

    )
    api_token: SecretStr = Field(
-        default=SecretStr(None),
+        default=SecretStr(None),  # type: ignore[arg-type]


Why not set default=SecretStr("") and avoid ignore type comment? (yes I know, that same discussion again)

i hope to handle this later. i removed the blanket exclude so we can see the specific instances.

leseb · 2025-10-06T13:58:07Z

llama_stack/providers/remote/inference/cerebras/config.py

    )
    api_key: SecretStr = Field(
-        default=SecretStr(os.environ.get("CEREBRAS_API_KEY")),
+        default=SecretStr(os.environ.get("CEREBRAS_API_KEY")),  # type: ignore[arg-type]


Suggested change

default=SecretStr(os.environ.get("CEREBRAS_API_KEY")), # type: ignore[arg-type]

default_factory=lambda: SecretStr(os.getenv("CEREBRAS_API_KEY", ""))

And mypy will be happy

i removed the blanket excludes. i want to avoid rolling more things into this PR.

leseb · 2025-10-06T13:59:13Z

llama_stack/providers/remote/inference/fireworks/fireworks.py

-                    'Pass Fireworks API Key in the header X-LlamaStack-Provider-Data as { "fireworks_api_key": <your api key>}'
-                )
-            return provider_data.fireworks_api_key
+        return self.config.api_key.get_secret_value() if self.config.api_key else None  # type: ignore[return-value]


why do you need to type ignore comment?

i removed the blanket excludes to expose specific places that need improvement

leseb · 2025-10-06T14:01:06Z

llama_stack/providers/remote/inference/together/together.py

            response.usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)

-        return response
+        return response  # type: ignore[no-any-return]


can you leave a comment about why we are ignoring?

i removed the blanket exclude

mattf · 2025-10-06T14:43:55Z

missing 2 files (vllm.py and openai_mixin.py,)

@leseb what do you mean?

leseb · 2025-10-06T14:46:04Z

missing 2 files (vllm.py and openai_mixin.py,)

@leseb what do you mean?

Sorry, I typed too fast, I meant that I still haven't gone through those two files :)

leseb

HUGE! Thanks for this one!

A few questions:

can we remove llama_stack/providers/utils/inference/litellm_openai_mixin.py now (along with test files)?
can we remove litellm from pyproject after that too?

mattf · 2025-10-06T15:32:18Z

HUGE! Thanks for this one!

A few questions:

* can we remove llama_stack/providers/utils/inference/litellm_openai_mixin.py now (along with test files)?

* can we remove litellm from pyproject after that too?

unfortunately we cannot. the watsonx adapter is being updated and will rely on LiteLLMOpenAIMixin because watsonx does not provide an openai-compat endpoint.

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 3, 2025

chore: give OpenAIMixin subcalsses a change to list models without le…

c465472

…aking _model_cache details

mattf force-pushed the make-openaimix-pydantic branch 4 times, most recently from 241e77c to dc97160 Compare October 5, 2025 10:05

mattf force-pushed the make-openaimix-pydantic branch from dc97160 to fd06717 Compare October 5, 2025 11:31

mattf marked this pull request as ready for review October 5, 2025 11:37

mattf requested review from ashwinb, yanxi0830, hardikjshah, raghotham, ehhuang, terrytangyuan, leseb, bbrowning, reluctantfuturist, slekkala1 and franciscojavierarceo as code owners October 5, 2025 11:37

mattf added 5 commits October 6, 2025 06:37

get_models -> list_provider_model_ids

a765f1c

Merge branch 'refactor-list-models' into make-openaimix-pydantic

6fa23d8

Merge branch 'main' into make-openaimix-pydantic

d158596

new recordings

733e72c

update anthropic, databricks, tgi, together after get_models rename

8658949

franciscojavierarceo reviewed Oct 6, 2025

View reviewed changes

Merge branch 'main' into make-openaimix-pydantic

96232e2

leseb reviewed Oct 6, 2025

View reviewed changes

remove unnecessary litellm deps

7fe3dac

leseb approved these changes Oct 6, 2025

View reviewed changes

mattf merged commit d23ed26 into llamastack:main Oct 6, 2025
45 checks passed

	default=SecretStr(os.environ.get("CEREBRAS_API_KEY")), # type: ignore[arg-type]
	default_factory=lambda: SecretStr(os.getenv("CEREBRAS_API_KEY", ""))

chore: turn OpenAIMixin into a pydantic.BaseModel #3671

chore: turn OpenAIMixin into a pydantic.BaseModel #3671

Conversation

mattf commented Oct 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Test Plan

Uh oh!

mattf commented Oct 6, 2025

Uh oh!

franciscojavierarceo left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leseb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattf commented Oct 6, 2025

Uh oh!

leseb commented Oct 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

leseb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattf commented Oct 6, 2025

Uh oh!

Uh oh!

Uh oh!

mattf commented Oct 3, 2025 •

edited

Loading

franciscojavierarceo left a comment •

edited

Loading

leseb commented Oct 6, 2025 •

edited

Loading

leseb left a comment •

edited

Loading