feat(python): use router to call litellm

Author: SanjoDeundiak

implementations/python/python/ockam/models/model.py

@@ -1,6 +1,8 @@
 from typing import List, Optional
 
 import litellm
+from litellm import Router
+
 import os
 import boto3
 import threading
@@ -40,6 +42,14 @@
         "gemma3:27b": "ollama_chat/gemma3:27b",
         "nomic-embed-text": "ollama/nomic-embed-text",
     },
+    "ollama_chat": {
+        "deepseek-r1": "ollama_chat/deepseek-r1",
+        "llama3.2": "ollama_chat/llama3.2",
+        "llama3.3": "ollama_chat/llama3.3",
+        "gemma3": "ollama_chat/gemma3",
+        "gemma3:27b": "ollama_chat/gemma3:27b",
+        "nomic-embed-text": "ollama/nomic-embed-text",
+    },
     "bedrock": {
         "claude-3-5-haiku-v1": "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
         "claude-3-5-sonnet-v1": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
@@ -80,6 +90,16 @@
 init_lock = threading.Lock()
 
 
+# slightly modify the parameters to accommodate services
+litellm.modify_params = True
+
+model_list = []
+for model_name in ALL_PROVIDER_ALLOWED_FULL_NAMES:
+    model_list.append({"model_name": model_name, "litellm_params": {"model": model_name}})
+
+router = Router(model_list=model_list, default_max_parallel_requests=400)
+
+
 def construct_bedrock_arn(model_identifier: str) -> Optional[str]:
     global region, account_id, init_lock
     with init_lock:
@@ -192,17 +212,16 @@ async def complete_chat(self, messages: List[dict] | List[ConversationMessage],
                 if "tool_calls" in message:
                     del message["tool_calls"]
 
-        # slightly modify the parameters to accommodate services
-        litellm.modify_params = True
-
         # parameters provided in kwargs will override the default parameters
         kwargs = {**self.kwargs, **kwargs}
 
         # sometimes an empty tools list is interpreted as "please hallucinate tools",
         if "tools" in kwargs and len(kwargs["tools"]) == 0:
             del kwargs["tools"]
 
-        return await litellm.acompletion(self.name, messages=messages, stream=stream, **kwargs)
+        global router
+
+        return await router.acompletion(self.name, messages=messages, stream=stream, **kwargs)
 
     async def embeddings(self, text: List[str], **kwargs) -> List[List[float]]:
         # parameters provided in kwargs will override the default parameters