Add get_api_key_from_model()

agenta-cli/agenta/sdk/assets.py

@@ -1,23 +1,9 @@
 supported_llm_models = {
-    "Mistral AI": [
-        "mistral/mistral-tiny",
-        "mistral/mistral-small",
-        "mistral/mistral-medium",
-        "mistral/mistral-large-latest",
-    ],
-    "Open AI": [
-        "gpt-3.5-turbo-1106",
-        "gpt-3.5-turbo",
-        "gpt-4",
-        "gpt-4o",
-        "gpt-4o-mini",
-        "gpt-4-1106-preview",
-    ],
-    "Gemini": ["gemini/gemini-1.5-pro-latest", "gemini/gemini-1.5-flash"],
-    "Cohere": [
-        "cohere/command-light",
-        "cohere/command-r-plus",
-        "cohere/command-nightly",
+    "Aleph Alpha": [
+        "luminous-base",
+        "luminous-base-control",
+        "luminous-extended-control",
+        "luminous-supreme",
     ],
     "Anthropic": [
         "anthropic/claude-3-5-sonnet-20240620",
@@ -33,11 +19,10 @@
         "anyscale/meta-llama/Llama-2-13b-chat-hf",
         "anyscale/meta-llama/Llama-2-70b-chat-hf",
     ],
-    "Perplexity AI": [
-        "perplexity/pplx-7b-chat",
-        "perplexity/pplx-70b-chat",
-        "perplexity/pplx-7b-online",
-        "perplexity/pplx-70b-online",
+    "Cohere": [
+        "cohere/command-light",
+        "cohere/command-r-plus",
+        "cohere/command-nightly",
     ],
     "DeepInfra": [
         "deepinfra/meta-llama/Llama-2-70b-chat-hf",
@@ -46,6 +31,46 @@
         "deepinfra/mistralai/Mistral-7B-Instruct-v0.1",
         "deepinfra/jondurbin/airoboros-l2-70b-gpt4-1.4.1",
     ],
+    "Gemini": [
+        "gemini/gemini-1.5-pro-latest",
+        "gemini/gemini-1.5-flash",
+    ],
+    "Groq": [
+        "groq/llama3-8b-8192",
+        "groq/llama3-70b-8192",
+        "groq/llama2-70b-4096",
+        "groq/mixtral-8x7b-32768",
+        "groq/gemma-7b-it",
+    ],
+    "Mistral AI": [
+        "mistral/mistral-tiny",
+        "mistral/mistral-small",
+        "mistral/mistral-medium",
+        "mistral/mistral-large-latest",
+    ],
+    "Open AI": [
+        "gpt-3.5-turbo-1106",
+        "gpt-3.5-turbo",
+        "gpt-4",
+        "gpt-4o",
+        "gpt-4o-mini",
+        "gpt-4-1106-preview",
+    ],
+    "OpenRouter": [
+        "openrouter/openai/gpt-3.5-turbo",
+        "openrouter/openai/gpt-3.5-turbo-16k",
+        "openrouter/anthropic/claude-instant-v1",
+        "openrouter/google/palm-2-chat-bison",
+        "openrouter/google/palm-2-codechat-bison",
+        "openrouter/meta-llama/llama-2-13b-chat",
+        "openrouter/meta-llama/llama-2-70b-chat",
+    ],
+    "Perplexity AI": [
+        "perplexity/pplx-7b-chat",
+        "perplexity/pplx-70b-chat",
+        "perplexity/pplx-7b-online",
+        "perplexity/pplx-70b-online",
+    ],
     "Together AI": [
         "together_ai/togethercomputer/llama-2-70b-chat",
         "together_ai/togethercomputer/llama-2-70b",
@@ -59,26 +84,12 @@
         "together_ai/NousResearch/Nous-Hermes-Llama2-13b",
         "together_ai/Austism/chronos-hermes-13b",
     ],
-    "Aleph Alpha": [
-        "luminous-base",
-        "luminous-base-control",
-        "luminous-extended-control",
-        "luminous-supreme",
-    ],
-    "OpenRouter": [
-        "openrouter/openai/gpt-3.5-turbo",
-        "openrouter/openai/gpt-3.5-turbo-16k",
-        "openrouter/anthropic/claude-instant-v1",
-        "openrouter/google/palm-2-chat-bison",
-        "openrouter/google/palm-2-codechat-bison",
-        "openrouter/meta-llama/llama-2-13b-chat",
-        "openrouter/meta-llama/llama-2-70b-chat",
-    ],
-    "Groq": [
-        "groq/llama3-8b-8192",
-        "groq/llama3-70b-8192",
-        "groq/llama2-70b-4096",
-        "groq/mixtral-8x7b-32768",
-        "groq/gemma-7b-it",
-    ],
+}
+
+providers_list = list(supported_llm_models.keys())
+
+model_to_provider_mapping = {
+    model: provider
+    for provider, models in supported_llm_models.items()
+    for model in models
 }

agenta-cli/agenta/sdk/managers/vault.py

@@ -2,6 +2,8 @@
 
 from agenta.sdk.context.routing import routing_context
 
+from agenta.sdk.assets import model_to_provider_mapping
+
 
 class VaultManager:
     @staticmethod
@@ -14,3 +16,23 @@ def get_from_route() -> Optional[Dict[str, Any]]:
             return None
 
         return secrets
+
+    @staticmethod
+    def get_api_key_for_model(model: str) -> str:
+        secrets = VaultManager.get_from_route()
+
+        if not secrets:
+            return None
+
+        provider = model_to_provider_mapping.get(model)
+
+        if not provider:
+            return None
+
+        provider = provider.lower().replace(" ", "")
+
+        for secret in secrets:
+            if secret["data"]["provider"] == provider:
+                return secret["data"]["key"]
+
+        return None

agenta-cli/agenta/sdk/middleware/vault.py

@@ -1,4 +1,4 @@
-from typing import Callable, Dict, Optional
+from typing import Callable, Dict, Optional, List
 
 from enum import Enum
 from os import getenv
@@ -24,18 +24,18 @@ class SecretKind(str, Enum):
 
 # TODO: Move to backend client types
 class ProviderKind(str, Enum):
-    OPENAI = "openai"
-    COHERE = "cohere"
+    ALEPHALPHA = "alephalpha"
+    ANTHROPIC = "anthropic"
     ANYSCALE = "anyscale"
+    COHERE = "cohere"
     DEEPINFRA = "deepinfra"
-    ALEPHALPHA = "alephalpha"
+    GEMINI = "gemini"
     GROQ = "groq"
     MISTRALAI = "mistralai"
-    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    OPENROUTER = "openrouter"
     PERPLEXITYAI = "perplexityai"
     TOGETHERAI = "togetherai"
-    OPENROUTER = "openrouter"
-    GEMINI = "gemini"
 
 
 # TODO: Move to backend client types
@@ -97,7 +97,7 @@ async def _get_secrets(self, request: Request) -> Optional[Dict]:
 
                 return secrets
 
-        local_secrets = []
+        local_secrets: List[SecretDTO] = []
 
         try:
             for provider_kind in ProviderKind:
@@ -120,7 +120,7 @@ async def _get_secrets(self, request: Request) -> Optional[Dict]:
         except:  # pylint: disable=bare-except
             display_exception("Vault: Local Secrets Exception")
 
-        vault_secrets = []
+        vault_secrets: List[SecretDTO] = []
 
         try:
             async with httpx.AsyncClient() as client:
@@ -135,7 +135,8 @@ async def _get_secrets(self, request: Request) -> Optional[Dict]:
                 else:
                     vault = response.json()
 
-                    vault_secrets = vault.get("secrets")
+                    vault_secrets = [secret.secret.model_dump() for secret in vault]
+
         except:  # pylint: disable=bare-except
             display_exception("Vault: Vault Secrets Exception")