Skip to content

Feat/add litellm provider - #384

Open
RheagalFire wants to merge 2 commits into
droidrun:mainfrom
RheagalFire:feat/add-litellm-provider
Open

Feat/add litellm provider#384
RheagalFire wants to merge 2 commits into
droidrun:mainfrom
RheagalFire:feat/add-litellm-provider

Conversation

@RheagalFire

Copy link
Copy Markdown

Summary

  • Adds LiteLLM as a first-class AI gateway provider, giving Mobilerun users access to 100+ LLM providers (OpenAI, Anthropic, Azure, AWS Bedrock, Google Vertex, Groq, Mistral, etc.) through a single unified interface.
  • LiteLLM is installed as an optional extra (pip install mobilerun[litellm]), keeping the base install unaffected.
  • Uses litellm SDK directly (import litellm / litellm.completion()) with a custom MobilerunLiteLLM wrapper extending llama-index's CustomLLM.

Changes

  • pyproject.toml - Added litellm optional dependency group with litellm>=1.83.0,<2.0.0
  • mobilerun/agent/providers/registry.py - Added LiteLLM provider family with API key auth mode and env key slot
  • mobilerun/agent/utils/llm_picker.py - Added MobilerunLiteLLM class (extends CustomLLM, calls litellm.completion() / litellm.acompletion() directly with drop_params=True), registered in SUPPORTED_PROVIDERS, PROVIDER_ALIASES, and load_llm() dispatch
  • mobilerun/agent/usage.py - Added LiteLLM to supported providers list, alias map, and usage extraction (prompt_tokens, completion_tokens, total_tokens)
  • mobilerun/cli/main.py - Updated CLI --provider help text to include LiteLLM
  • README.md - Added LiteLLM to feature list and installation instructions
  • tests/test_llm_picker.py - 9 tests: alias resolution, class loading, base_url translation, provider/model format, bare model name, API key forwarding, temperature forwarding, max_tokens forwarding, unsupported provider error
  • tests/test_provider_registry.py - 1 test: provider family variant resolution
  • tests/test_usage.py - 6 tests: usage extraction, tracker support, null usage fallback, dict raw response, total computation from parts, no-raw-response error

Tests

1. Unit tests (71/71 passing):

tests/test_llm_picker.py::test_litellm_loads_correct_class PASSED
tests/test_llm_picker.py::test_litellm_alias_resolves PASSED
tests/test_llm_picker.py::test_litellm_base_url_translates_to_api_base PASSED
tests/test_llm_picker.py::test_litellm_provider_slash_model_format PASSED
tests/test_llm_picker.py::test_litellm_bare_model_name_preserved PASSED
tests/test_llm_picker.py::test_litellm_api_key_forwarded PASSED
tests/test_llm_picker.py::test_litellm_temperature_forwarded PASSED
tests/test_llm_picker.py::test_litellm_max_tokens_forwarded PASSED
tests/test_llm_picker.py::test_litellm_unsupported_provider_raises PASSED
tests/test_provider_registry.py::test_litellm_family_resolves_api_key_variant PASSED
tests/test_usage.py::test_litellm_extracts_usage_from_response PASSED
tests/test_usage.py::test_track_usage_supports_litellm PASSED
tests/test_usage.py::test_litellm_usage_with_null_usage_returns_zeros PASSED
tests/test_usage.py::test_litellm_usage_with_dict_raw_response PASSED
tests/test_usage.py::test_litellm_usage_computes_total_when_missing PASSED
tests/test_usage.py::test_litellm_no_raw_response_raises PASSED
============================== 71 passed in 1.08s ==============================

2. Lint: ruff check mobilerun/ tests/ -> All checks passed.

3. Live E2E against Anthropic (via Azure Foundry):

from mobilerun.agent.utils.llm_picker import load_llm
from mobilerun.agent.usage import get_usage_from_response
from llama_index.core.base.llms.types import ChatMessage

llm = load_llm('LiteLLM', model='anthropic/claude-sonnet-4-6',
               api_base='https://amanrai-test-resource.services.ai.azure.com/anthropic')
messages = [ChatMessage(role='user', content='What is 2+2? Reply with just the number.')]
response = llm.chat(messages)
usage = get_usage_from_response('LiteLLM', response)
Provider class: MobilerunLiteLLM
Model: anthropic/claude-sonnet-4-6
Response: 4
Usage: prompt_tokens=20, completion_tokens=5, total_tokens=25

4. Pin resolution verified: pip install -e '.[litellm]' resolves cleanly against litellm>=1.83.0,<2.0.0.

Risk / Compatibility

  • Additive only. No existing providers touched.
  • litellm is an optional extra - base install and all existing providers remain unaffected.
  • Uses litellm SDK directly (import litellm) - no additional wrapper packages needed.
  • LiteLLM appears in mobilerun configure wizard automatically via the provider family registry.

Example usage

CLI:

# Install with LiteLLM support
pip install "mobilerun[litellm]"

# Use with any LiteLLM-supported provider
mobilerun "Open Settings" --provider LiteLLM --model anthropic/claude-sonnet-4-6

# Or with Azure OpenAI
AZURE_API_KEY=... AZURE_API_BASE=... \
mobilerun "Take a screenshot" --provider LiteLLM --model azure/gpt-4o

# Or with AWS Bedrock
AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... \
mobilerun "Open Chrome" --provider LiteLLM --model bedrock/anthropic.claude-v2

Python SDK:

from mobilerun.agent.utils.llm_picker import load_llm

llm = load_llm("LiteLLM", model="anthropic/claude-sonnet-4-6")
# Use with MobileAgent as usual

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant