Skip to content

Feat/add litellm provider#329

Merged
huangyz0918 merged 2 commits into
MLSysOps:mainfrom
RheagalFire:feat/add-litellm-provider
Jul 10, 2026
Merged

Feat/add litellm provider#329
huangyz0918 merged 2 commits into
MLSysOps:mainfrom
RheagalFire:feat/add-litellm-provider

Conversation

@RheagalFire

Copy link
Copy Markdown
Contributor

What has been done to verify that this works as intended?

15 unit tests pass:

$ pytest tests/test_litellm_model.py -v
15 passed in 6.31s

Covers: init with explicit params / env vars / defaults, missing openai import, query, streaming, None content, auth error (401), rate limit (429), timeout, empty choices, tool-to-function conversion.

6 live E2E tests pass (LiteLLM proxy -> Azure Foundry / Anthropic Claude Sonnet 4.6):

Test Input Response Result
Basic query "What is 2+2?" "4" PASS
Multi-turn conversation system + user + assistant + followup "What is my name?" "Your name is Alice." PASS
Streaming longer response "Count from 1 to 5" "1, 2, 3, 4, 5" (streamed in chunks) PASS
Temperature 0 deterministic "Capital of Japan?" x2 "Tokyo" both times PASS
System prompt respected "Respond in exactly 3 words. Describe the ocean." "Vast, blue, deep." PASS
Empty user message "" Graceful response, no crash PASS

Why is this the best possible solution? Were any other approaches considered?

Follows the exact same pattern as the existing DeepSeek and vLLM providers - uses the openai SDK pointed at a LiteLLM proxy URL. Zero new dependencies. The alternative was adding litellm as a direct Python SDK dependency, but since MLE-agent already uses the openai SDK for multiple providers, pointing it at a LiteLLM proxy is simpler, lighter, and consistent with the vLLM provider pattern.

How does this change affect users? Describe intentional changes to behavior and behavior that could have accidentally been affected by code changes. In other words, what are the regression risks?

Additive only - no existing code modified except registering the new provider in __init__.py. Existing providers and workflows are completely unaffected. Users opt in by setting platform: LiteLLM in .mle/project.yml.

Changes:

  • mle/model/litellm.py - new LiteLLMModel class extending Model with query() and stream()
  • mle/model/__init__.py - added MODEL_LITELLM = 'LiteLLM', import, and load_model() branch
  • tests/test_litellm_model.py - 15 unit tests

Usage:

# .mle/project.yml
platform: LiteLLM
base_url: http://litellm-proxy:4000/v1
api_key: sk-litellm-key  # optional if proxy has no auth

Do we need any specific form for testing your changes? If so, please attach one.

No. A running LiteLLM proxy is needed for live testing. Unit tests run without one (fully mocked).

Does this change require updates to documentation? If so, please file an issue here and include the link below.

No external documentation updates required. The provider follows the same configuration pattern as all existing providers.

Before submitting this PR, please make sure you have:

  • confirmed all checks still pass OR confirm CI build passes.
  • verified that any code or assets from external sources are properly credited in comments and/or in
    the credit file.

@dosubot dosubot Bot added size:L This PR changes 100-499 lines, ignoring generated files. enhancement New feature or request labels Jul 3, 2026
@RheagalFire

Copy link
Copy Markdown
Contributor Author

cc @huangyz0918 @leeeizhang

@huangyz0918 huangyz0918 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thank you @RheagalFire

@dosubot dosubot Bot added the lgtm This PR has been approved by a maintainer label Jul 7, 2026
@leeeizhang

Copy link
Copy Markdown
Collaborator

many thanks! @RheagalFire

@huangyz0918
huangyz0918 merged commit a29287e into MLSysOps:main Jul 10, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request lgtm This PR has been approved by a maintainer size:L This PR changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants