From d5332f44347f44d60360e14205a2342a0c990d66 Mon Sep 17 00:00:00 2001
From: Anmol Jaiswal <anmolg1997@users.noreply.github.com>
Date: Fri, 13 Feb 2026 13:46:47 -0800
Subject: [PATCH] feat: Expand LiteLlm supported models and add registry tests

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 869873309
---
 src/google/adk/models/lite_llm.py        | 26 ++++++++++++++++++++++++
 tests/unittests/models/test_gemma_llm.py |  8 ++++++++
 2 files changed, 34 insertions(+)

diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py
index cb6e0230..b954d8a0 100644
--- a/src/google/adk/models/lite_llm.py
+++ b/src/google/adk/models/lite_llm.py
@@ -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/.*",
     ]
diff --git a/tests/unittests/models/test_gemma_llm.py b/tests/unittests/models/test_gemma_llm.py
index a8c2254d..07b4b092 100644
--- a/tests/unittests/models/test_gemma_llm.py
+++ b/tests/unittests/models/test_gemma_llm.py
@@ -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",
       [