superagent-ai · homanp · Feb 29, 2024 · Feb 27, 2024 · Feb 27, 2024 · Feb 27, 2024
diff --git a/libs/superagent/app/api/api_keys.py b/libs/superagent/app/api/api_keys.py
@@ -0,0 +1,141 @@
+import segment.analytics as analytics
+from decouple import config
+from fastapi import APIRouter, Depends
+from fastapi.responses import JSONResponse
+
+from app.models.request import ApiKey as ApiKeyRequest
+from app.models.response import ApiKey as ApiKeyResponse
+from app.models.response import ApiKeyCreate as ApiKeyCreateResponse
+from app.models.response import ApiKeyList as ApiKeyListResponse
+from app.utils.api import generate_jwt, get_current_api_user, handle_exception
+from app.utils.prisma import prisma
+
+SEGMENT_WRITE_KEY = config("SEGMENT_WRITE_KEY", None)
+
+router = APIRouter()
+analytics.write_key = SEGMENT_WRITE_KEY
+
+
+def get_display_api_key(key: str):
+    return f"{key[:3]}****{key[-4:]}"
+
+
+@router.post(
+    "/api-keys",
+    name="create",
+    description="Create a new API key",
+    response_model=ApiKeyCreateResponse,
+)
+async def create(body: ApiKeyRequest, api_user=Depends(get_current_api_user)):
+    """Endpoint for creating an agent"""
+    try:
+        api_key = generate_jwt(
+            {
+                "api_user_id": api_user.id,
+            }
+        )
+
+        display_api_key = get_display_api_key(api_key)
+
+        api_key_data = await prisma.apikey.create(
+            {
+                "name": body.name,
+                "displayApiKey": display_api_key,
+                "apiUserId": api_user.id,
+            }
+        )
+        if SEGMENT_WRITE_KEY:
+            analytics.track(api_user.id, "API Key Created", {"api_key": api_key})
+
+        return {
+            "success": True,
+            "data": {
+                **api_key_data.dict(),
+                "apiKey": api_key,
+            },
+        }
+    except Exception as e:
+        handle_exception(e)
+
+
+@router.delete(
+    "/api-keys/{id}",
+    name="delete",
+    description="Delete an API key",
+    response_model=ApiKeyResponse,
+)
+async def delete(id: str, api_user=Depends(get_current_api_user)):
+    """Endpoint for deleting an agent"""
+    try:
+        api_key = await prisma.apikey.find_unique(where={"id": id})
+        if not api_key:
+            raise Exception("API Key not found")
+
+        await prisma.apikey.delete(
+            where={
+                "id": id,
+                "apiUserId": api_user.id,
+            }
+        )
+        if SEGMENT_WRITE_KEY:
+            analytics.track(
+                api_user.id,
+                "API Key Deleted",
+            )
+
+        return {"success": True, "data": api_key}
+    except Exception as e:
+        handle_exception(e)
+
+
+@router.get(
+    "/api-keys",
+    name="list",
+    description="List API keys",
+    response_model=ApiKeyListResponse,
+)
+async def list(api_user=Depends(get_current_api_user)):
+    """Endpoint for listing all agents"""
+    try:
+        api_keys = await prisma.apikey.find_many(
+            where={
+                "apiUserId": api_user.id,
+            }
+        )
+
+        return {
+            "success": True,
+            "data": api_keys,
+        }
+    except Exception as e:
+        handle_exception(e)
+
+
+@router.patch(
+    "/api-keys/{id}",
+    name="update",
+    description="Update an API key",
+    response_model=ApiKeyResponse,
+)
+async def update(id: str, body: ApiKeyRequest, api_user=Depends(get_current_api_user)):
+    """Endpoint for updating an agent"""
+    try:
+        api_key = await prisma.apikey.find_unique(
+            where={"id": id, "apiUserId": api_user.id}
+        )
+        if not api_key:
+            return JSONResponse(
+                status_code=404,
+                content={"success": False, "error": {"message": "API Key not found"}},
+            )
+
+        api_key = await prisma.apikey.update(
+            where={"id": id},
+            data={
+                **body.dict(exclude_unset=True),
+            },
+        )
+
+        return {"success": True, "data": api_key}
+    except Exception as e:
+        handle_exception(e)
diff --git a/libs/superagent/app/models/request.py b/libs/superagent/app/models/request.py
@@ -14,6 +14,10 @@ class ApiUser(BaseModel):
     anonymousId: Optional[str]
 
 
+class ApiKey(BaseModel):
+    name: str
+
+
 class OpenAiAssistantParameters(BaseModel):
     metadata: Optional[Dict[str, Any]]
     fileIds: Optional[List[str]]

diff --git a/libs/superagent/app/models/response.py b/libs/superagent/app/models/response.py
@@ -14,6 +14,9 @@
 from prisma.models import (
     AgentTool as AgentToolModel,
 )
+from prisma.models import (
+    ApiKey as ApiKeyModel,
+)
 from prisma.models import (
     ApiUser as ApiUserModel,
 )
@@ -39,6 +42,25 @@ class ApiUser(BaseModel):
     data: Optional[ApiUserModel]
 
 
+class _ApiKeyCreateModel(ApiKeyModel):
+    apiKey: str
+
+
+class ApiKeyCreate(BaseModel):
+    success: bool
+    data: Optional[_ApiKeyCreateModel]
+
+
+class ApiKey(BaseModel):
+    success: bool
+    data: Optional[ApiKeyModel]
+
+
+class ApiKeyList(BaseModel):
+    success: bool
+    data: Optional[list[ApiKeyModel]]
+
+
 class Agent(BaseModel):
     success: bool
     data: Optional[AgentModel]

diff --git a/libs/superagent/app/routers.py b/libs/superagent/app/routers.py
@@ -2,6 +2,7 @@
 
 from app.api import (
     agents,
+    api_keys,
     api_user,
     datasources,
     llms,
@@ -17,6 +18,7 @@
 router.include_router(agents.router, tags=["Agent"], prefix=api_prefix)
 router.include_router(llms.router, tags=["LLM"], prefix=api_prefix)
 router.include_router(api_user.router, tags=["Api user"], prefix=api_prefix)
+router.include_router(api_keys.router, tags=["API key"], prefix=api_prefix)
 router.include_router(datasources.router, tags=["Datasource"], prefix=api_prefix)
 router.include_router(tools.router, tags=["Tool"], prefix=api_prefix)
 router.include_router(workflows.router, tags=["Workflow"], prefix=api_prefix)

diff --git a/libs/superagent/app/utils/api.py b/libs/superagent/app/utils/api.py
@@ -1,4 +1,5 @@
 import logging
+import uuid
 
 import jwt
 from decouple import config
@@ -19,6 +20,9 @@ def handle_exception(e):
 
 
 def generate_jwt(data: dict):
+    # for randomness
+    data.update({"jti": str(uuid.uuid4())})
+
     token = jwt.encode({**data}, config("JWT_SECRET"), algorithm="HS256")
     return token
 

diff --git a/libs/superagent/prisma/migrations/20240227075836_add_api_keys_table/migration.sql b/libs/superagent/prisma/migrations/20240227075836_add_api_keys_table/migration.sql
@@ -0,0 +1,17 @@
+-- CreateTable
+CREATE TABLE "ApiKey" (
+    "id" TEXT NOT NULL,
+    "name" TEXT NOT NULL,
+    "displayApiKey" TEXT NOT NULL,
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updatedAt" TIMESTAMP(3) NOT NULL,
+    "apiUserId" TEXT NOT NULL,
+
+    CONSTRAINT "ApiKey_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE INDEX "api_user_id" ON "ApiKey"("apiUserId");
+
+-- AddForeignKey
+ALTER TABLE "ApiKey" ADD CONSTRAINT "ApiKey_apiUserId_fkey" FOREIGN KEY ("apiUserId") REFERENCES "ApiUser"("id") ON DELETE RESTRICT ON UPDATE CASCADE;
diff --git a/libs/superagent/prisma/schema.prisma b/libs/superagent/prisma/schema.prisma
@@ -106,6 +106,19 @@ model ApiUser {
   workflows       Workflow[]
   vectorDb        VectorDb[]
   workflowConfigs WorkflowConfig[]
+  apiKeys         ApiKey[]
+}
+
+model ApiKey {
+  id            String   @id @default(uuid())
+  name          String
+  displayApiKey String
+  createdAt     DateTime @default(now())
+  updatedAt     DateTime @updatedAt
+  apiUserId     String
+  apiUser       ApiUser  @relation(fields: [apiUserId], references: [id])
+
+  @@index([apiUserId], name: "api_user_id")
 }
 
 model Agent {

diff --git a/libs/ui/app/onboarding/client-page.tsx b/libs/ui/app/onboarding/client-page.tsx
@@ -75,7 +75,7 @@ export default function OnboardingClientPage() {
     if (!profile.api_key) {
       const {
         data: { token: api_key },
-      } = await api.createApiKey({
+      } = await api.createApiUser({
         email: user.email,
         firstName: first_name,
         lastName: last_name,