Credentials: Custom Provider Input #166

backend/app/alembic/versions/904ed70e7dab_added_provider_column_to_the_credential_.py

@@ -0,0 +1,68 @@
+"""Added provider column to the credential table
+
+Revision ID: 904ed70e7dab
+Revises: 543f97951bd0
+Create Date: 2025-05-10 11:13:17.868238
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+
+# revision identifiers, used by Alembic.
+revision = "904ed70e7dab"
+down_revision = "f23675767ed2"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "credential",
+        sa.Column("provider", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+    )
+    op.create_index(
+        op.f("ix_credential_provider"), "credential", ["provider"], unique=False
+    )
+    op.drop_constraint(
+        "credential_organization_id_fkey", "credential", type_="foreignkey"
+    )
+    op.create_foreign_key(
+        "credential_organization_id_fkey",
+        "credential",
+        "organization",
+        ["organization_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.drop_constraint("project_organization_id_fkey", "project", type_="foreignkey")
+    op.create_foreign_key(None, "project", "organization", ["organization_id"], ["id"])
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, "project", type_="foreignkey")
+    op.create_foreign_key(
+        "project_organization_id_fkey",
+        "project",
+        "organization",
+        ["organization_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.drop_constraint(
+        "credential_organization_id_fkey", "credential", type_="foreignkey"
+    )
+    op.create_foreign_key(
+        "credential_organization_id_fkey",
+        "credential",
+        "organization",
+        ["organization_id"],
+        ["id"],
+    )
+    op.drop_index(op.f("ix_credential_provider"), table_name="credential")
+    op.drop_column("credential", "provider")
+    # ### end Alembic commands ###

backend/app/api/routes/credentials.py

@@ -2,126 +2,208 @@
 from app.api.deps import SessionDep, get_current_active_superuser
 from app.crud.credentials import (
     get_creds_by_org,
-    get_key_by_org,
+    get_provider_credential,
     remove_creds_for_org,
     set_creds_for_org,
     update_creds_for_org,
+    remove_provider_credential,
 )
 from app.models import CredsCreate, CredsPublic, CredsUpdate
 from app.utils import APIResponse
 from datetime import datetime
+from app.core.providers import validate_provider
+from typing import List
+from sqlalchemy.exc import IntegrityError
+from app.models.organization import Organization
 
 router = APIRouter(prefix="/credentials", tags=["credentials"])
 
 
 @router.post(
     "/",
     dependencies=[Depends(get_current_active_superuser)],
-    response_model=APIResponse[CredsPublic],
+    response_model=APIResponse[List[CredsPublic]],
+    summary="Create new credentials for an organization",
+    description="Creates new credentials for a specific organization. This endpoint requires superuser privileges. If credentials already exist for the organization, it will return an error.",
 )
 def create_new_credential(*, session: SessionDep, creds_in: CredsCreate):
-    new_creds = None
     try:
         existing_creds = get_creds_by_org(
             session=session, org_id=creds_in.organization_id
         )
-        if not existing_creds:
-            new_creds = set_creds_for_org(session=session, creds_add=creds_in)
+        if existing_creds:
+            raise HTTPException(
+                status_code=400,
+                detail="Credentials already exist for this organization",
+            )
+
+        new_creds = set_creds_for_org(session=session, creds_add=creds_in)
+        if not new_creds:
+            raise HTTPException(status_code=500, detail="Failed to create credentials")
+
+        # Return all created credentials
+        return APIResponse.success_response(new_creds)
+
+    except ValueError as e:
+        if "Unsupported provider" in str(e):
+            raise HTTPException(status_code=400, detail=str(e))
+        raise HTTPException(status_code=404, detail=str(e))
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
-    # Ensure inserted_at is set during creation
-    new_creds.inserted_at = datetime.utcnow()
-
-    return APIResponse.success_response(new_creds)
-
 
 @router.get(
     "/{org_id}",
     dependencies=[Depends(get_current_active_superuser)],
-    response_model=APIResponse[CredsPublic],
+    response_model=APIResponse[List[CredsPublic]],
+    summary="Get all credentials for an organization",
+    description="Retrieves all provider credentials associated with a specific organization. This endpoint requires superuser privileges.",
 )
 def read_credential(*, session: SessionDep, org_id: int):
     try:
         creds = get_creds_by_org(session=session, org_id=org_id)
+        if not creds:
+            raise HTTPException(status_code=404, detail="Credentials not found")
+        return APIResponse.success_response(creds)
+    except HTTPException as e:
+        raise e  # Ensure HTTPException is not wrapped again
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
-    if creds is None:
-        raise HTTPException(status_code=404, detail="Credentials not found")
-
-    return APIResponse.success_response(creds)
-
 
 @router.get(
-    "/{org_id}/api-key",
+    "/{org_id}/{provider}",
     dependencies=[Depends(get_current_active_superuser)],
     response_model=APIResponse[dict],
+    summary="Get specific provider credentials",
+    description="Retrieves credentials for a specific provider (e.g., 'openai', 'anthropic') for a given organization. This endpoint requires superuser privileges.",
 )
-def read_api_key(*, session: SessionDep, org_id: int):
+def read_provider_credential(*, session: SessionDep, org_id: int, provider: str):
     try:
-        api_key = get_key_by_org(session=session, org_id=org_id)
+        provider_enum = validate_provider(provider)
+        provider_creds = get_provider_credential(
+            session=session, org_id=org_id, provider=provider_enum
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
-    if api_key is None:
-        raise HTTPException(status_code=404, detail="API key not found")
+    if provider_creds is None:
+        raise HTTPException(status_code=404, detail="Provider credentials not found")
 
-    return APIResponse.success_response({"api_key": api_key})
+    return APIResponse.success_response(provider_creds)
 
 
 @router.patch(
     "/{org_id}",
     dependencies=[Depends(get_current_active_superuser)],
-    response_model=APIResponse[CredsPublic],
+    response_model=APIResponse[List[CredsPublic]],
+    summary="Update organization credentials",
+    description="Updates credentials for a specific organization. Can update specific provider credentials or add new providers. This endpoint requires superuser privileges.",
 )
 def update_credential(*, session: SessionDep, org_id: int, creds_in: CredsUpdate):
     try:
+        # Validate incoming payload
+        if not creds_in or not creds_in.provider or not creds_in.credential:
+            raise HTTPException(
+                status_code=400, detail="Provider and credential must be provided"
+            )
+
+        # Defensive check to ensure organization exists
+        try:
+            organization = session.get(Organization, org_id)
+        except Exception as e:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to fetch organization: {str(e)}"
+            )
+
+        if not organization:
+            raise HTTPException(status_code=404, detail="Organization not found")
+
         updated_creds = update_creds_for_org(
             session=session, org_id=org_id, creds_in=creds_in
         )
-
-        updated_creds.updated_at = datetime.utcnow()
+        if not updated_creds:
+            raise HTTPException(status_code=404, detail="Failed to update credentials")
 
         return APIResponse.success_response(updated_creds)
+
+    except IntegrityError as e:
+        if "ForeignKeyViolation" in str(e):
+            raise HTTPException(
+                status_code=400,
+                detail="Invalid organization ID. Ensure the organization exists before updating credentials.",
+            )
+        raise HTTPException(
+            status_code=500, detail=f"An unexpected database error occurred: {str(e)}"
+        )
     except ValueError as e:
+        if "Unsupported provider" in str(e):
+            raise HTTPException(status_code=400, detail=str(e))
         raise HTTPException(status_code=404, detail=str(e))
+    except HTTPException:
+        raise
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
 
-from fastapi import HTTPException, Depends
-from app.crud.credentials import remove_creds_for_org
-from app.utils import APIResponse
-from app.api.deps import SessionDep, get_current_active_superuser
+@router.delete(
+    "/{org_id}/{provider}",
+    dependencies=[Depends(get_current_active_superuser)],
+    response_model=APIResponse[dict],
+    summary="Delete specific provider credentials",
+    description="Removes credentials for a specific provider while keeping other provider credentials intact. This endpoint requires superuser privileges.",
+)
+def delete_provider_credential(*, session: SessionDep, org_id: int, provider: str):
+    try:
+        provider_enum = validate_provider(provider)
+        updated_creds = remove_provider_credential(
+            session=session, org_id=org_id, provider=provider_enum
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=404, detail="Provider credentials not found"
+        )  # Updated to return 404
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"An unexpected error occurred: {str(e)}"
+        )
+
+    if not updated_creds:  # Ensure proper check for no credentials found
+        raise HTTPException(status_code=404, detail="Provider credentials not found")
+
+    return APIResponse.success_response(
+        {"message": "Provider credentials removed successfully"}
+    )
 
 
 @router.delete(
-    "/{org_id}/api-key",
+    "/{org_id}",
     dependencies=[Depends(get_current_active_superuser)],
     response_model=APIResponse[dict],
+    summary="Delete all organization credentials",
+    description="Removes all credentials for a specific organization. This is a soft delete operation that marks credentials as inactive. This endpoint requires superuser privileges.",
 )
-def delete_credential(*, session: SessionDep, org_id: int):
+def delete_all_credentials(*, session: SessionDep, org_id: int):
     try:
         creds = remove_creds_for_org(session=session, org_id=org_id)
     except Exception as e:
         raise HTTPException(
             status_code=500, detail=f"An unexpected error occurred: {str(e)}"
         )
 
-    if creds is None:
+    if not creds:  # Ensure proper check for no credentials found
         raise HTTPException(
             status_code=404, detail="Credentials for organization not found"
         )
 
-    # No need to manually set deleted_at and is_active if it's done in remove_creds_for_org
-    # Simply return the success response
     return APIResponse.success_response({"message": "Credentials deleted successfully"})

backend/app/core/providers.py

@@ -0,0 +1,58 @@
+from typing import Dict, List, Optional
+from enum import Enum
+
+
+class Provider(str, Enum):
+    """Enumeration of supported credential providers."""
+
+    OPENAI = "openai"
+    GEMINI = "gemini"
+    ANTHROPIC = "anthropic"
+    MISTRAL = "mistral"
+    COHERE = "cohere"
+    HUGGINGFACE = "huggingface"
+    AZURE = "azure"
+    AWS = "aws"
+    GOOGLE = "google"
+
+
+# Required fields for each provider's credentials
+PROVIDER_REQUIRED_FIELDS: Dict[str, List[str]] = {
+    Provider.OPENAI: ["api_key"],
+    Provider.GEMINI: ["api_key"],
+    Provider.ANTHROPIC: ["api_key"],
+    Provider.MISTRAL: ["api_key"],
+    Provider.COHERE: ["api_key"],
+    Provider.HUGGINGFACE: ["api_key"],
+    Provider.AZURE: ["api_key", "endpoint"],
+    Provider.AWS: ["access_key_id", "secret_access_key", "region"],
+    Provider.GOOGLE: ["api_key"],
+}
+
+
+def validate_provider(provider: str) -> Provider:
+    """Validate that the provider name is supported and return the Provider enum."""
+    try:
+        return Provider(provider.lower())
+    except ValueError:
+        supported = ", ".join([p.value for p in Provider])
+        raise ValueError(
+            f"Unsupported provider: {provider}. Supported providers are: {supported}"
+        )
+
+
+def validate_provider_credentials(provider: str, credentials: dict) -> None:
+    """Validate that the credentials contain all required fields for the provider."""
+    provider_enum = validate_provider(provider)
+    required_fields = PROVIDER_REQUIRED_FIELDS[provider_enum]
+
+    missing_fields = [field for field in required_fields if field not in credentials]
+    if missing_fields:
+        raise ValueError(
+            f"Missing required fields for {provider}: {', '.join(missing_fields)}"
+        )
+
+
+def get_supported_providers() -> List[str]:
+    """Return a list of all supported provider names."""
+    return [p.value for p in Provider]