Update tests

airtai · Aug 23, 2024 · 46f37f2 · 46f37f2
1 parent d786137
commit 46f37f2
Show file tree

Hide file tree

Showing 14 changed files with 163 additions and 28 deletions.
diff --git a/fastagency/studio/models/agents/web_surfer.py b/fastagency/studio/models/agents/web_surfer.py
@@ -93,7 +93,8 @@ class WebSurferAgent(AgentBaseModel):
         int, Field(description="The viewport size of the browser")
     ] = 4096
     bing_api_key: Annotated[
-        Optional[BingAPIKeyRef], Field(title="Bing API Key", description="The Bing API key for the browser")
+        Optional[BingAPIKeyRef],
+        Field(title="Bing API Key", description="The Bing API key for the browser"),
     ] = None
 
     @classmethod

diff --git a/fastagency/studio/models/deployments/deployment.py b/fastagency/studio/models/deployments/deployment.py
@@ -56,11 +56,17 @@ class Deployment(Model):
     ]
     gh_token: Annotated[
         GitHubTokenRef,
-        Field(title="GH Token", description="The GitHub token to use for creating a new repository")
+        Field(
+            title="GH Token",
+            description="The GitHub token to use for creating a new repository",
+        ),
     ]
     fly_token: Annotated[
         FlyTokenRef,
-        Field(title="Fly Token", description="The Fly.io token to use for deploying the deployment")
+        Field(
+            title="Fly Token",
+            description="The Fly.io token to use for deploying the deployment",
+        ),
     ]
 
     @classmethod

diff --git a/fastagency/studio/models/llms/azure.py b/fastagency/studio/models/llms/azure.py
@@ -26,7 +26,9 @@
 
 @register("secret")
 class AzureOAIAPIKey(Model):
-    api_key: Annotated[str, Field(title="API Key", description="The API Key from Azure OpenAI")]
+    api_key: Annotated[
+        str, Field(title="API Key", description="The API Key from Azure OpenAI")
+    ]
 
     @classmethod
     async def create_autogen(cls, model_id: UUID, user_id: UUID, **kwargs: Any) -> str:
@@ -53,7 +55,7 @@ class AzureOAI(Model):
 
     api_key: Annotated[
         AzureOAIAPIKeyRef,
-        Field(title="API Key", description="The API Key from Azure OpenAI")
+        Field(title="API Key", description="The API Key from Azure OpenAI"),
     ]
 
     base_url: Annotated[
@@ -67,7 +69,10 @@ class AzureOAI(Model):
 
     api_version: Annotated[
         AzureApiVersionsLiteral,
-        Field(title="API Version", description="The version of the Azure OpenAI API, e.g. '2024-02-01'"),
+        Field(
+            title="API Version",
+            description="The version of the Azure OpenAI API, e.g. '2024-02-01'",
+        ),
     ] = "2024-02-01"
 
     temperature: Annotated[

diff --git a/fastagency/studio/models/llms/openai.py b/fastagency/studio/models/llms/openai.py
@@ -79,7 +79,7 @@ class OpenAI(Model):
         Field(
             title="API Key",
             description="The API Key from OpenAI",
-        )
+        ),
     ]
 
     base_url: Annotated[URL, Field(description="The base URL of the OpenAI API")] = URL(

diff --git a/fastagency/studio/models/llms/together.py b/fastagency/studio/models/llms/together.py
@@ -147,7 +147,7 @@ class TogetherAI(Model):
 
     api_key: Annotated[
         TogetherAIAPIKeyRef,
-        Field(title="API Key", description="The API Key from Together.ai")
+        Field(title="API Key", description="The API Key from Together.ai"),
     ]
 
     base_url: Annotated[URL, Field(description="The base URL of the OpenAI API")] = URL(

diff --git a/fastagency/studio/models/secrets/fly_token.py b/fastagency/studio/models/secrets/fly_token.py
@@ -12,7 +12,11 @@
 @register("secret")
 class FlyToken(Model):
     fly_token: Annotated[
-        str, Field(title="Fly Token", description="The Fly.io token to use for deploying the deployment")
+        str,
+        Field(
+            title="Fly Token",
+            description="The Fly.io token to use for deploying the deployment",
+        ),
     ]
 
     @classmethod

diff --git a/fastagency/studio/models/secrets/github_token.py b/fastagency/studio/models/secrets/github_token.py
@@ -12,7 +12,11 @@
 @register("secret")
 class GitHubToken(Model):
     gh_token: Annotated[
-        str, Field(title="GH Token", description="The GitHub token to use for creating a new repository")
+        str,
+        Field(
+            title="GH Token",
+            description="The GitHub token to use for creating a new repository",
+        ),
     ]
 
     @classmethod

diff --git a/tests/studio/models/agents/test_web_surfer.py b/tests/studio/models/agents/test_web_surfer.py
@@ -276,6 +276,7 @@ def test_web_surfer_model_schema(self) -> None:
                     "anyOf": [{"$ref": "#/$defs/BingAPIKeyRef"}, {"type": "null"}],
                     "default": None,
                     "description": "The Bing API key for the browser",
+                    "title": "Bing API Key",
                 },
             },
             "required": ["name", "llm", "summarizer_llm"],

diff --git a/tests/studio/models/deployments/test_deployment.py b/tests/studio/models/deployments/test_deployment.py
@@ -160,10 +160,18 @@ def test_deployment_model_schema(self) -> None:
                 "team": {
                     "allOf": [{"$ref": "#/$defs/TwoAgentTeamRef"}],
                     "description": "The team that is used in the deployment",
-                    "title": "Team name",
+                    "title": "Team Name",
+                },
+                "gh_token": {
+                    "allOf": [{"$ref": "#/$defs/GitHubTokenRef"}],
+                    "description": "The GitHub token to use for creating a new repository",
+                    "title": "GH Token",
+                },
+                "fly_token": {
+                    "allOf": [{"$ref": "#/$defs/FlyTokenRef"}],
+                    "description": "The Fly.io token to use for deploying the deployment",
+                    "title": "Fly Token",
                 },
-                "gh_token": {"$ref": "#/$defs/GitHubTokenRef"},
-                "fly_token": {"$ref": "#/$defs/FlyTokenRef"},
             },
             "required": [
                 "name",

diff --git a/tests/studio/models/llms/test_anthropic.py b/tests/studio/models/llms/test_anthropic.py
@@ -102,7 +102,11 @@ def test_anthropic_model_schema(self) -> None:
                     "title": "Model",
                     "type": "string",
                 },
-                "api_key": {"$ref": "#/$defs/AnthropicAPIKeyRef"},
+                "api_key": {
+                    "allOf": [{"$ref": "#/$defs/AnthropicAPIKeyRef"}],
+                    "description": "The API Key from Anthropic",
+                    "title": "API Key",
+                },
                 "base_url": {
                     "default": "https://api.anthropic.com/v1",
                     "description": "The base URL of the Anthropic API",

diff --git a/tests/studio/models/llms/test_azure.py b/tests/studio/models/llms/test_azure.py
@@ -30,7 +30,6 @@ async def test_azure_constructor(
         name = model.name
         api_key_uuid = model.api_key.uuid  # type: ignore [attr-defined]
         base_url = model.base_url  # type: ignore [attr-defined]
-
         expected = {
             "name": name,
             "model": "gpt-35-turbo-16k",
@@ -93,7 +92,11 @@ def test_azure_model_schema(self) -> None:
                     "title": "Model",
                     "type": "string",
                 },
-                "api_key": {"$ref": "#/$defs/AzureOAIAPIKeyRef"},
+                "api_key": {
+                    "allOf": [{"$ref": "#/$defs/AzureOAIAPIKeyRef"}],
+                    "description": "The API Key from Azure OpenAI",
+                    "title": "API Key",
+                },
                 "base_url": {
                     "default": "https://api.openai.com/v1",
                     "description": "The base URL of the Azure OpenAI API",
@@ -108,7 +111,7 @@ def test_azure_model_schema(self) -> None:
                     "default": "azure",
                     "description": "The type of the API, must be 'azure'",
                     "enum": ["azure"],
-                    "title": "API type",
+                    "title": "API Type",
                     "type": "string",
                 },
                 "api_version": {
@@ -124,14 +127,14 @@ def test_azure_model_schema(self) -> None:
                         "2024-05-01-preview",
                         "2024-02-01",
                     ],
-                    "title": "Api Version",
+                    "title": "API Version",
                     "type": "string",
                 },
                 "temperature": {
                     "default": 0.8,
                     "description": "The temperature to use for the model, must be between 0 and 2",
-                    "minimum": 0.0,
                     "maximum": 2.0,
+                    "minimum": 0.0,
                     "title": "Temperature",
                     "type": "number",
                 },

diff --git a/tests/studio/models/llms/test_openai.py b/tests/studio/models/llms/test_openai.py
@@ -141,7 +141,11 @@ def test_openai_schema(self) -> None:
                     "title": "Model",
                     "type": "string",
                 },
-                "api_key": {"$ref": "#/$defs/OpenAIAPIKeyRef"},
+                "api_key": {
+                    "allOf": [{"$ref": "#/$defs/OpenAIAPIKeyRef"}],
+                    "description": "The API Key from OpenAI",
+                    "title": "API Key",
+                },
                 "base_url": {
                     "default": "https://api.openai.com/v1",
                     "description": "The base URL of the OpenAI API",

diff --git a/tests/studio/models/llms/test_together.py b/tests/studio/models/llms/test_together.py
@@ -129,14 +129,109 @@ def test_togetherai_schema(self) -> None:
                 "model": {
                     "default": "Meta Llama 3 70B Reference",
                     "description": "The model to use for the Together API",
-                    #     "enum": [
-                    #         "Meta Llama 3 70B Reference",
-                    #         ...
-                    #     ],
+                    "enum": [
+                        "WizardLM v1.2 (13B)",
+                        "Code Llama Instruct (34B)",
+                        "Upstage SOLAR Instruct v1 (11B)",
+                        "Meta Llama 3 70B Reference",
+                        "OpenHermes-2-Mistral (7B)",
+                        "LLaMA-2-7B-32K-Instruct (7B)",
+                        "ReMM SLERP L2 (13B)",
+                        "Toppy M (7B)",
+                        "OpenChat 3.5",
+                        "Chronos Hermes (13B)",
+                        "Snorkel Mistral PairRM DPO (7B)",
+                        "Qwen 1.5 Chat (7B)",
+                        "Qwen 1.5 Chat (14B)",
+                        "Qwen 1.5 Chat (1.8B)",
+                        "Snowflake Arctic Instruct",
+                        "Nous Hermes 2 - Mixtral 8x7B-SFT",
+                        "Nous Hermes 2 - Mixtral 8x7B-DPO ",
+                        "Deepseek Coder Instruct (33B)",
+                        "Nous Hermes Llama-2 (13B)",
+                        "Vicuna v1.5 (13B)",
+                        "Qwen 1.5 Chat (0.5B)",
+                        "Code Llama Instruct (7B)",
+                        "Nous Hermes-2 Yi (34B)",
+                        "Code Llama Instruct (13B)",
+                        "Llama3 8B Chat HF INT4",
+                        "OpenHermes-2.5-Mistral (7B)",
+                        "Nous Capybara v1.9 (7B)",
+                        "Meta Llama 3.1 70B Instruct Turbo",
+                        "Nous Hermes 2 - Mistral DPO (7B)",
+                        "StripedHyena Nous (7B)",
+                        "Alpaca (7B)",
+                        "Platypus2 Instruct (70B)",
+                        "Gemma Instruct (2B)",
+                        "Gemma Instruct (7B)",
+                        "LLaMA-2 Chat (7B)",
+                        "OLMo Instruct (7B)",
+                        "Qwen 1.5 Chat (4B)",
+                        "MythoMax-L2 (13B)",
+                        "Mistral (7B) Instruct",
+                        "Mistral (7B) Instruct v0.2",
+                        "Meta Llama 3.1 8B Instruct Turbo",
+                        "OpenOrca Mistral (7B) 8K",
+                        "Nous Hermes LLaMA-2 (7B)",
+                        "Qwen 1.5 Chat (32B)",
+                        "Qwen 2 Instruct (72B)",
+                        "Qwen 1.5 Chat (72B)",
+                        "DeepSeek LLM Chat (67B)",
+                        "Vicuna v1.5 (7B)",
+                        "WizardLM-2 (8x22B)",
+                        "Togethercomputer Llama3 8B Instruct Int8",
+                        "Mistral (7B) Instruct v0.3",
+                        "Qwen 1.5 Chat (110B)",
+                        "LLaMA-2 Chat (13B)",
+                        "Gemma-2 Instruct (27B)",
+                        "01-ai Yi Chat (34B)",
+                        "Meta Llama 3 70B Instruct Turbo",
+                        "Meta Llama 3 70B Instruct Lite",
+                        "Gemma-2 Instruct (9B)",
+                        "Meta Llama 3 8B Reference",
+                        "Mixtral-8x7B Instruct v0.1",
+                        "Code Llama Instruct (70B)",
+                        "Meta Llama 3.1 405B Instruct Turbo",
+                        "DBRX Instruct",
+                        "Meta Llama 3.1 8B Instruct",
+                        "Meta Llama 3 8B Instruct Turbo",
+                        "Dolphin 2.5 Mixtral 8x7b",
+                        "Mixtral-8x22B Instruct v0.1",
+                        "Meta Llama 3 8B Instruct Lite",
+                        "LLaMA-2 Chat (70B)",
+                        "Koala (7B)",
+                        "Qwen 2 Instruct (1.5B)",
+                        "Qwen 2 Instruct (7B)",
+                        "Guanaco (65B) ",
+                        "Vicuna v1.3 (7B)",
+                        "Nous Hermes LLaMA-2 (70B)",
+                        "Vicuna v1.5 16K (13B)",
+                        "Zephyr-7B-ß",
+                        "Guanaco (13B) ",
+                        "Vicuna v1.3 (13B)",
+                        "Guanaco (33B) ",
+                        "Koala (13B)",
+                        "Upstage SOLAR Instruct v1 (11B)-Int4",
+                        "Guanaco (7B) ",
+                        "Meta Llama 3 8B Instruct",
+                        "Meta Llama 3 70B Instruct",
+                        "Hermes 2 Theta Llama-3 70B",
+                        "carson ml318bit",
+                        "carson ml31405bit",
+                        "carson ml3170bit",
+                        "carson ml318br",
+                        "Llama-3 70B Instruct Gradient 1048K",
+                        "Meta Llama 3.1 70B Instruct",
+                        "Meta Llama 3.1 70B",
+                    ],
                     "title": "Model",
                     "type": "string",
                 },
-                "api_key": {"$ref": "#/$defs/TogetherAIAPIKeyRef"},
+                "api_key": {
+                    "allOf": [{"$ref": "#/$defs/TogetherAIAPIKeyRef"}],
+                    "description": "The API Key from Together.ai",
+                    "title": "API Key",
+                },
                 "base_url": {
                     "default": "https://api.together.xyz/v1",
                     "description": "The base URL of the OpenAI API",

diff --git a/tests/studio/models/teams/test_two_agents_team.py b/tests/studio/models/teams/test_two_agents_team.py
@@ -155,7 +155,7 @@ def test_two_agents_team_schema(self) -> None:
                     "default": "ALWAYS",
                     "description": "Mode for human input",
                     "enum": ["ALWAYS", "TERMINATE", "NEVER"],
-                    "title": "Human input mode",
+                    "title": "Human Input Mode",
                     "type": "string",
                 },
                 "initial_agent": {
@@ -165,7 +165,7 @@ def test_two_agents_team_schema(self) -> None:
                         {"$ref": "#/$defs/WebSurferAgentRef"},
                     ],
                     "description": "Agent that starts the conversation",
-                    "title": "Initial agent",
+                    "title": "Initial Agent",
                 },
                 "secondary_agent": {
                     "anyOf": [
@@ -174,7 +174,7 @@ def test_two_agents_team_schema(self) -> None:
                         {"$ref": "#/$defs/WebSurferAgentRef"},
                     ],
                     "description": "Agent that continues the conversation",
-                    "title": "Secondary agent",
+                    "title": "Secondary Agent",
                 },
             },
             "required": ["name", "initial_agent", "secondary_agent"],