[FEAT][Agent][api.py][v1/models]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms-cloud"
-version = "0.2.6"
+version = "0.2.7"
 description = "Swarms Cloud - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]

servers/agent/api.py

@@ -1,76 +1,13 @@
 import os
-from typing import List, Any
+from typing import List
 
 import tiktoken
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, model_validator
 from swarms import Agent, Anthropic, GPT4VisionAPI, OpenAIChat
 from swarms.utils.loguru_logger import logger
 from swarms_cloud.schema.cog_vlm_schemas import ChatCompletionResponse, UsageInfo
-
-
-# Define the input model using Pydantic
-class AgentInput(BaseModel):
-    agent_name: str = "Swarm Agent"
-    system_prompt: str = None
-    agent_description: str = None
-    model_name: str = "OpenAIChat"
-    max_loops: int = 1
-    autosave: bool = False
-    dynamic_temperature_enabled: bool = False
-    dashboard: bool = False
-    verbose: bool = False
-    streaming_on: bool = False
-    saved_state_path: str = "agent_saved_state.json"
-    sop: str = None
-    sop_list: List[str] = None
-    user_name: str = "User"
-    retry_attempts: int = 3
-    context_length: int = 8192
-    task: str = None
-    max_tokens: int = None
-    tool_schema: Any = None
-
-    @model_validator(mode="pre")
-    def check_required_fields(cls, values):
-        required_fields = [
-            "agent_name",
-            "system_prompt",
-            "task",
-            "max_loops",
-            "context_window",
-        ]
-
-        for field in required_fields:
-            if not values.get(field):
-
-                raise ValueError(f"{field} must not be empty or null")
-
-        if values["max_loops"] <= 0:
-            raise ValueError("max_loops must be greater than 0")
-
-        if values["context_window"] <= 0:
-            raise ValueError("context_window must be greater than 0")
-
-        return values
-
-
-# Define the output model using Pydantic
-class GenerationMetrics(BaseModel):
-    tokens_per_second: float = 0.0
-    completion_time: float = 0.0
-
-
-# Define the output model using Pydantic
-class AgentOutput(BaseModel):
-    agent: AgentInput
-    completions: ChatCompletionResponse
-    # metrics: GenerationMetrics
-
-
-# Define the available models
-AVAILABLE_MODELS = ["OpenAIChat", "GPT4VisionAPI", "Anthropic"]
+from swarms_cloud.schema.agent_api_schemas import AgentInput, AgentOutput, ModelList, ModelSchema
 
 
 async def count_tokens(text: str) -> int:
@@ -102,12 +39,14 @@ async def model_router(model_name: str):
 
     """
     # Logic to switch to the specified model
-    if model_name == "OpenAIChat":
+    if model_name == "gpt-4o":
         # Switch to OpenAIChat model
-        llm = OpenAIChat(max_tokens=4000)
-    elif model_name == "GPT4VisionAPI":
+        llm = OpenAIChat(max_tokens=4000, model_name="gpt-4o")
+    elif model_name == "gpt-4-vision-preview":
         # Switch to GPT4VisionAPI model
-        llm = GPT4VisionAPI()
+        llm = GPT4VisionAPI(
+            max_tokens=4000,
+        )
     elif model_name == "Anthropic":
         # Switch to Anthropic model
         llm = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))
@@ -119,7 +58,11 @@ async def model_router(model_name: str):
 
 
 # Create a FastAPI app
-app = FastAPI(debug=True)
+app = FastAPI(
+    debug=True,
+    title="Swarm Agent API",
+    version="0.1.0",
+)
 
 # Load the middleware to handle CORS
 app.add_middleware(
@@ -131,13 +74,23 @@ async def model_router(model_name: str):
 )
 
 
-@app.get("/v1/models", response_model=List[str])
+@app.get("/v1/models", response_model=ModelList)
 async def list_models() -> List[str]:
     """
     An endpoint to list available models. It returns a list of model names.
     This is useful for clients to query and understand what models are available for use.
     """
-    return AVAILABLE_MODELS
+    models = ModelList(
+        data = [
+            ModelSchema(id="gpt-4o", created_at="OpenAI"),
+            ModelSchema(id="gpt-4-vision-preview", created_at="OpenAI"),
+            ModelSchema(id="Anthropic", created_at="Anthropic"),
+            # ModelSchema(id="gpt-4o", created_at="OpenAI"),
+            ## Llama3.1
+        ]
+    )
+
+    return models
 
 
 @app.post("/v1/agent/completions", response_model=AgentOutput)
@@ -155,7 +108,7 @@ async def agent_completions(agent_input: AgentInput):
         # Model check
         model_name = agent_input.model_name
         model = await model_router(model_name)
-        
+
         # Initialize the agent
         agent = Agent(
             agent_name=agent_name,
@@ -181,7 +134,7 @@ async def agent_completions(agent_input: AgentInput):
         logger.info(f"Running agent with task: {task}")
         agent_history = agent.short_memory.return_history_as_string()
         completions = agent.run(task)
-        
+
         logger.info(f"Agent response: {completions}")
 
         # Costs calculation
@@ -209,7 +162,7 @@ async def agent_completions(agent_input: AgentInput):
                 usage_info=UsageInfo(
                     prompt_tokens=all_input_tokens,
                     completion_tokens=output_tokens,
-                    total_tokens=all_input_tokens + output_tokens,
+                    total_tokens=total_costs,
                 ),
             ),
         )

swarms_cloud/schema/__init__.py

@@ -0,0 +1,10 @@
+from swarms_cloud.schema.agent_api_schemas import AgentInput, ModelSchema, ModelList, GenerationMetrics, AgentOutput
+
+
+__all__ = [
+    "AgentInput",
+    "ModelSchema",
+    "ModelList",
+    "GenerationMetrics",
+    "AgentOutput",
+]

swarms_cloud/schema/agent_api_schemas.py

@@ -0,0 +1,79 @@
+import time
+from typing import List, Any
+
+from pydantic import BaseModel, model_validator
+
+from swarms_cloud.schema.cog_vlm_schemas import ChatCompletionResponse
+
+
+
+# Define the input model using Pydantic
+class AgentInput(BaseModel):
+    agent_name: str = "Swarm Agent"
+    system_prompt: str = None
+    agent_description: str = None
+    model_name: str = "OpenAIChat"
+    max_loops: int = 1
+    autosave: bool = False
+    dynamic_temperature_enabled: bool = False
+    dashboard: bool = False
+    verbose: bool = False
+    streaming_on: bool = False
+    saved_state_path: str = "agent_saved_state.json"
+    sop: str = None
+    sop_list: List[str] = None
+    user_name: str = "User"
+    retry_attempts: int = 3
+    context_length: int = 8192
+    task: str = None
+    max_tokens: int = None
+    tool_schema: Any = None
+
+    @model_validator(mode="pre")
+    def check_required_fields(cls, values):
+        required_fields = [
+            "agent_name",
+            "system_prompt",
+            "task",
+            "max_loops",
+            "context_window",
+        ]
+
+        for field in required_fields:
+            if not values.get(field):
+
+                raise ValueError(f"{field} must not be empty or null")
+
+        if values["max_loops"] <= 0:
+            raise ValueError("max_loops must be greater than 0")
+
+        if values["context_window"] <= 0:
+            raise ValueError("context_window must be greater than 0")
+
+        return values
+
+
+class ModelSchema(BaseModel):
+    id: str = None
+    object: str = "model"
+    created_at: int = time.time()
+    owned_by: str = "TGSC"
+
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data = List[ModelSchema] = []
+
+
+# Define the output model using Pydantic
+class GenerationMetrics(BaseModel):
+    tokens_per_second: float = 0.0
+    completion_time: float = 0.0
+
+
+# Define the output model using Pydantic
+class AgentOutput(BaseModel):
+    agent: AgentInput
+    completions: ChatCompletionResponse
+    # metrics: GenerationMetrics
+