feat: exception handling for api keys and models (#197)

* feat: exception handling for api keys and models

* fix: type error in calling agentmodelnotvalid

* fix: exclude `id` when getting tools to match openai tools schema

* refactor: Lint agents-api (CI)

---------

Co-authored-by: Diwank Singh Tomer <diwank.singh@gmail.com>
Co-authored-by: creatorrr <creatorrr@users.noreply.github.com>

agents-api/agents_api/common/exceptions/agents.py

@@ -2,7 +2,6 @@
 
 from uuid import UUID
 from . import BaseCommonException
-from agents_api.model_registry import ALL_AVAILABLE_MODELS
 
 
 class BaseAgentException(BaseCommonException):
@@ -43,9 +42,17 @@ def __init__(self, agent_id: UUID | str, doc_id: UUID | str):
 
 
 class AgentModelNotValid(BaseAgentException):
-    def __init__(self, model: str):
+    def __init__(self, model: str, all_models: list[str]):
         super().__init__(
             f"Unknown model: {model}. Please provide a valid model name."
-            "Known models are: " + ", ".join(ALL_AVAILABLE_MODELS.keys()),
+            "Available models: " + ", ".join(all_models),
+            http_code=400,
+        )
+
+
+class MissingAgentModelAPIKeyError(BaseAgentException):
+    def __init__(self, model: str):
+        super().__init__(
+            f"API key missing for model: {model}. Please provide a valid API key in the configuration",
             http_code=400,
         )

agents-api/agents_api/env.py

@@ -25,7 +25,7 @@
 )
 model_api_key: str = env.str("MODEL_API_KEY", default=None)
 model_inference_url: str = env.str("MODEL_INFERENCE_URL", default=None)
-openai_api_key: str = env.str("OPENAI_API_KEY", default=None)
+openai_api_key: str = env.str("OPENAI_API_KEY", default="")
 summarization_ratio_threshold: float = env.float(
     "MAX_TOKENS_RATIO_TO_SUMMARIZE", default=0.5
 )
@@ -78,8 +78,12 @@
     truncate_embed_text=truncate_embed_text,
     temporal_worker_url=temporal_worker_url,
     temporal_namespace=temporal_namespace,
+    openai_api_key=openai_api_key,
 )
 
+if openai_api_key == "":
+    print("OpenAI API key not found. OpenAI API will not be enabled.")
+
 if debug:
     # Print the loaded environment variables for debugging purposes.
     print("Environment variables:")

agents-api/agents_api/model_registry.py

@@ -4,6 +4,10 @@
 
 from typing import Dict
 from agents_api.clients.model import julep_client, openai_client
+from agents_api.common.exceptions.agents import (
+    AgentModelNotValid,
+    MissingAgentModelAPIKeyError,
+)
 from openai import AsyncOpenAI
 
 
@@ -106,20 +110,15 @@
 }
 
 
-# TODO: implement
-def validate_configuration():
+def validate_configuration(model: str):
     """
-    function that validates the config based on the model
+    Validates the model specified in the request
     """
-    pass
-
-
-# TODO: implement
-def validate_request():
-    """
-    function that validates the config based on the model
-    """
-    pass
+    if model not in ALL_AVAILABLE_MODELS:
+        raise AgentModelNotValid(model, list(ALL_AVAILABLE_MODELS.keys()))
+    model_client = get_model_client(model)
+    if model_client.api_key == "":
+        raise MissingAgentModelAPIKeyError(model)
 
 
 def get_model_client(model: str) -> AsyncOpenAI:

agents-api/agents_api/models/agent/create_agent.py

@@ -3,15 +3,13 @@
 It includes functions to construct and execute datalog queries for inserting new agent records.
 """
 
-from agents_api.common.exceptions.agents import AgentModelNotValid
 from uuid import UUID
 
 import pandas as pd
 from pycozo.client import Client as CozoClient
 
 from ...clients.cozo import client
 from ...common.utils.cozo import cozo_process_mutate_data
-from ...model_registry import ALL_AVAILABLE_MODELS
 
 
 """
@@ -44,8 +42,6 @@ def create_agent_query(
     default_settings: dict = {},
     client: CozoClient = client,
 ) -> pd.DataFrame:
-    if model not in ALL_AVAILABLE_MODELS.keys():
-        raise AgentModelNotValid(model)
 
     settings_cols, settings_vals = cozo_process_mutate_data(
         {

agents-api/agents_api/routers/agents/routers.py

@@ -3,6 +3,7 @@
 from typing import Annotated
 from uuid import uuid4
 
+from agents_api.model_registry import validate_configuration
 from fastapi import APIRouter, HTTPException, status, Depends
 import pandas as pd
 from pycozo.client import QueryException
@@ -210,6 +211,7 @@ async def create_agent(
     request: CreateAgentRequest,
     x_developer_id: Annotated[UUID4, Depends(get_developer_id)],
 ) -> ResourceCreatedResponse:
+    validate_configuration(request.model)
     resp = create_agent_query(
         agent_id=uuid4(),
         developer_id=x_developer_id,

agents-api/agents_api/routers/sessions/session.py

@@ -204,7 +204,7 @@ async def generate(
         ]
         tools = None
         if settings.tools:
-            tools = [tool.model_dump(mode="json") for tool in settings.tools]
+            tools = [(tool.model_dump(exclude="id")) for tool in settings.tools]
         model_client = get_model_client(settings.model)
         extra_body = (
             dict(