litellm google mapping
frdel committed
Jul 4, 2025 at 11:06 UTC
223ef788e81ab0c267f9ae6ab46625ae3bd05964
1 file changed
+30
-16
models.py
+30
-16
@@ -34,11 +34,13 @@ from langchain_core.messages import (
34
from langchain.embeddings.base import Embeddings
35
from sentence_transformers import SentenceTransformer
36
37
+
38
# disable extra logging
39
def turn_off_logging():
39
- os.environ['LITELLM_LOG'] = "ERROR" # only errors
40
+ os.environ["LITELLM_LOG"] = "ERROR" # only errors
41
litellm.suppress_debug_info = True
42
43
+
44
# init
45
load_dotenv()
46
turn_off_logging()
@@ -85,9 +87,10 @@ def configure_litellm_environment():
87
"API_KEY_MISTRAL": "MISTRAL_API_KEY",
88
"API_KEY_OLLAMA": "OLLAMA_API_KEY",
89
"API_KEY_HUGGINGFACE": "HUGGINGFACE_API_KEY",
88
- "API_KEY_OPENAI_AZURE": "AZURE_API_KEY",
90
+ "API_KEY_OPENAI_AZURE": "AZURE_AI_API_KEY",
91
"API_KEY_DEEPSEEK": "DEEPSEEK_API_KEY",
92
"API_KEY_SAMBANOVA": "SAMBANOVA_API_KEY",
93
+ "API_KEY_GOOGLE": "GEMINI_API_KEY",
94
}
95
base_url_mappings = {
96
"OPENAI_BASE_URL": "OPENAI_API_BASE",
@@ -97,7 +100,7 @@ def configure_litellm_environment():
100
"MISTRAL_BASE_URL": "MISTRAL_API_BASE",
101
"OLLAMA_BASE_URL": "OLLAMA_API_BASE",
102
"HUGGINGFACE_BASE_URL": "HUGGINGFACE_API_BASE",
100
- "AZURE_BASE_URL": "AZURE_API_BASE",
103
+ "AZURE_BASE_URL": "AZURE_AI_API_BASE",
104
"DEEPSEEK_BASE_URL": "DEEPSEEK_API_BASE",
105
"SAMBANOVA_BASE_URL": "SAMBANOVA_API_BASE",
106
}
@@ -139,8 +142,8 @@ def _parse_chunk(chunk: Any) -> ChatChunk:
142
if isinstance(delta, dict)
143
else getattr(delta, "content", "")
144
) or (
142
- message.get("content", "")
143
- if isinstance(message, dict)
145
+ message.get("content", "")
146
+ if isinstance(message, dict)
147
else getattr(message, "content", "")
148
)
149
reasoning_delta = (
@@ -335,6 +338,7 @@ class BrowserCompatibleChatWrapper(LiteLLMChatWrapper):
338
A wrapper for browser agent that can filter/sanitize messages
339
before sending them to the LLM.
340
"""
341
+
342
def __init__(self, *args, **kwargs):
343
turn_off_logging()
344
super().__init__(*args, **kwargs)
@@ -373,14 +377,14 @@ class LiteLLMEmbeddingWrapper(Embeddings):
377
def embed_documents(self, texts: List[str]) -> List[List[float]]:
378
resp = embedding(model=self.model_name, input=texts, **self.kwargs)
379
return [
376
- item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore
377
- for item in resp.data # type: ignore
380
+ item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore
381
+ for item in resp.data # type: ignore
382
]
383
384
def embed_query(self, text: str) -> List[float]:
385
resp = embedding(model=self.model_name, input=[text], **self.kwargs)
382
- item = resp.data[0] # type: ignore
383
- return item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore
386
+ item = resp.data[0] # type: ignore
387
+ return item.get("embedding") if isinstance(item, dict) else item.embedding # type: ignore
388
389
390
class LocalSentenceTransformerWrapper(Embeddings):
@@ -395,15 +399,15 @@ class LocalSentenceTransformerWrapper(Embeddings):
399
self.model_name = model_name
400
401
def embed_documents(self, texts: List[str]) -> List[List[float]]:
398
- embeddings = self.model.encode(texts, convert_to_tensor=False) # type: ignore
399
- return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings # type: ignore
402
+ embeddings = self.model.encode(texts, convert_to_tensor=False) # type: ignore
403
+ return embeddings.tolist() if hasattr(embeddings, "tolist") else embeddings # type: ignore
404
405
def embed_query(self, text: str) -> List[float]:
402
- embedding = self.model.encode([text], convert_to_tensor=False) # type: ignore
406
+ embedding = self.model.encode([text], convert_to_tensor=False) # type: ignore
407
result = (
408
embedding[0].tolist() if hasattr(embedding[0], "tolist") else embedding[0]
409
)
406
- return result # type: ignore
410
+ return result # type: ignore
411
412
413
def _get_litellm_chat(
@@ -419,7 +423,7 @@ def _get_litellm_chat(
423
api_key = kwargs.pop("api_key", None) or get_api_key(provider_name)
424
425
# litellm will pick up base_url from env. We just need to control the api_key.
422
- base_url = dotenv.get_dotenv_value(f"{provider_name.upper()}_BASE_URL")
426
+ # base_url = dotenv.get_dotenv_value(f"{provider_name.upper()}_BASE_URL")
427
428
# If a base_url is set, ensure api_key is not passed to litellm
429
# > remove, this can be handled by api_key=None
@@ -479,7 +483,7 @@ def get_model(type: ModelType, provider: ModelProvider, name: str, **kwargs: Any
483
def get_chat_model(
484
provider: ModelProvider, name: str, **kwargs: Any
485
) -> LiteLLMChatWrapper:
482
- provider_name = provider.name.lower()
486
+ provider_name = _get_litellm_provider(provider)
487
kwargs = _normalize_chat_kwargs(kwargs)
488
model = _get_litellm_chat(LiteLLMChatWrapper, name, provider_name, **kwargs)
489
return model
@@ -510,4 +514,14 @@ def _normalize_chat_kwargs(kwargs: Any) -> Any:
514
515
516
def _normalize_embedding_kwargs(kwargs: Any) -> Any:
513
- return kwargs
\ No newline at end of file
517
+ return kwargs
518
+
519
+
520
+def _get_litellm_provider(provider: ModelProvider) -> str:
521
+ name = provider.name.lower()
522
+
523
+ # exceptions
524
+ if name == "google":
525
+ name = "gemini"
526
+
527
+ return name