feat: Expand LiteLlm supported models and add registry tests

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 869873309
This commit is contained in:
Anmol Jaiswal
2026-02-13 13:47:13 -08:00
committed by Copybara-Service
co-authored by George Weale
parent fbe9eccd05
commit d5332f4434
2 changed files with 34 additions and 0 deletions
+26
View File
@@ -2038,8 +2038,34 @@ class LiteLlm(BaseLlm):
return [
# For OpenAI models (e.g., "openai/gpt-4o")
r"openai/.*",
# For Azure OpenAI models (e.g., "azure/gpt-4o")
r"azure/.*",
# For Azure AI models (e.g., "azure_ai/command-r-plus")
r"azure_ai/.*",
# For Groq models via Groq API (e.g., "groq/llama3-70b-8192")
r"groq/.*",
# For Anthropic models (e.g., "anthropic/claude-3-opus-20240229")
r"anthropic/.*",
# For AWS Bedrock models (e.g., "bedrock/anthropic.claude-3-sonnet")
r"bedrock/.*",
# For Ollama models excluding Gemma3 (handled by Gemma3Ollama)
r"ollama/(?!gemma3).*",
# For Ollama chat models (e.g., "ollama_chat/llama3")
r"ollama_chat/.*",
# For Together AI models (e.g., "together_ai/meta-llama/Llama-3-70b")
r"together_ai/.*",
# For Vertex AI non-Gemini models (e.g., "vertex_ai/claude-3-sonnet")
r"vertex_ai/.*",
# For Mistral AI models (e.g., "mistral/mistral-large-latest")
r"mistral/.*",
# For DeepSeek models (e.g., "deepseek/deepseek-chat")
r"deepseek/.*",
# For Fireworks AI models (e.g., "fireworks_ai/llama-v3-70b")
r"fireworks_ai/.*",
# For Cohere models (e.g., "cohere/command-r-plus")
r"cohere/.*",
# For Databricks models (e.g., "databricks/dbrx-instruct")
r"databricks/.*",
# For AI21 models (e.g., "ai21/jamba-1.5-large")
r"ai21/.*",
]
+8
View File
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from google.adk import models
from google.adk.models.gemma_llm import Gemma
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
@@ -509,10 +510,17 @@ def test_process_response_last_json_object():
# Tests for Gemma3Ollama (only run when LiteLLM is installed)
try:
from google.adk.models.gemma_llm import Gemma3Ollama
from google.adk.models.lite_llm import LiteLlm
def test_gemma3_ollama_supported_models():
assert Gemma3Ollama.supported_models() == [r"ollama/gemma3.*"]
def test_gemma3_ollama_registry_resolution():
assert models.LLMRegistry.resolve("ollama/gemma3:12b") is Gemma3Ollama
def test_non_gemma_ollama_registry_resolution():
assert models.LLMRegistry.resolve("ollama/llama3.2") is LiteLlm
@pytest.mark.parametrize(
"model_arg,expected_model",
[