mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
co-authored by
George Weale
parent
fbe9eccd05
commit
d5332f4434
@@ -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/.*",
|
||||
]
|
||||
|
||||
@@ -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",
|
||||
[
|
||||
|
||||
Reference in New Issue
Block a user