Skip to content

Commit 12bf9b6

Browse files
nishika26sourabhlodhaavirajsingh7AkhileshNegi
authored
OpenAI Key per Org (#82)
* Update README.md (#44) * changes (#45) * Readme update (#47) rename project and stack --------- Co-authored-by: sourabhlodha <sourabhlodha@Administrators-MacBook-Pro.local> * fix create_user endpoint (#62) * standard api response and http exception handling (#67) * Upgrade PostgreSQL to 16 & Fix CORS Configuration (#57) * use latest docker image * update envsample * Add Customizable Token Expiry Time in Login API (#70) * token expiry time can be customize * default to one day * Organization/project : Crud, Endpoint and Test Cases (#63) * trial * pushing all * models file * renaming * Rename Project.py to project.py * Rename oganization.py to organization.py * Update README.md (#44) * changes (#45) Co-authored-by: sourabhlodha <sourabhlodha@Administrators-MacBook-Pro.local> * Readme update (#47) rename project and stack --------- Co-authored-by: sourabhlodha <sourabhlodha@Administrators-MacBook-Pro.local> * fix create_user endpoint (#62) * standard api response and http exception handling (#67) * standardization and edits * small edits * small edits * small edits * fixed project post * trial * pushing all * models file * renaming * Rename Project.py to project.py * Rename oganization.py to organization.py * standardization and edits * small edits * small edits * small edits * fixed project post * remove these files since they were somehow pushed into this branch * re-push the docker file * re-push utils file * re-push the file * fixing test cases --------- Co-authored-by: Sourabh Lodha <sourabh_lodha@ymail.com> Co-authored-by: sourabhlodha <sourabhlodha@Administrators-MacBook-Pro.local> Co-authored-by: Aviraj Gour <100823015+avirajsingh7@users.noreply.github.com> Co-authored-by: Ishankoradia <ikoradia@umich.edu> * Add Project User Management (#65) * intial commit user project mapping and authorization * fix alembic migration * Use standard API response * add pagination * add index and use base model * Alembic: migration fixes for organization (#77) * fixing testcases and migrations * changes migration file name * remove old migration --------- Co-authored-by: Akhilesh Negi <akhileshnegi.an3@gmail.com> * Added Support of API Key Authentication (#76) * Intial setup api key * added Api key auth flow * support both api key and oauth --------- Co-authored-by: Sourabh Lodha <sourabh_lodha@ymail.com> * Main to stage code sync (#80) Back merge Production to staging code * added migration for api table (#81) * creds table * Refactor Authentication Logic and Testing Enhancements (#89) * fix authentication part * Modify test cases to compatible with new auth * Github: CI (#74) * issue CI * first stab at continuous integration * fixing testcases and migrations * syncing with master * moving to python version 3.11.7 * making copy of env * updating env * added migrations * added uv sync * updating working directory * added step to activate env * updating working directory * updating working directory for codecov upload * updating script to upload to codecov * remove working directory * added working directory for % check * clenaup * cleanup * activating env * update the issue template * update readme and env file * adding badges (#91) * OpenAI: Threads (#40) * getting threads up and running * added testcases and citation * removing ssl verify * using standardized APIResponse * getting rid of redundant files * refactor code after testing * refactor testcases * setting up init.py * fixing review comments * cleanup * cleanup * removed validate thread as it can be handled by default * fixing few code review suggestions * removed validation testcases for assistant ID * threads testcases fix (#93) * project router changes * endpoint,crud and migration file * models file * minor fix * fixes * fixes * test cases and fixes * alembic file * type checking * cleaner exception * fixing alembic revision heads * using crendentials * init module * running pre commit * running pre commit * final changes * migration file * Rename fa868aa8debd_add_credetial_table.py to fa868aa8debd_add_credential_table.py * test case change * test cases * test cases * test cases * removing duplicate lines * datetime columns addition * migration file --------- Co-authored-by: Sourabh Lodha <sourabh_lodha@ymail.com> Co-authored-by: Aviraj Gour <100823015+avirajsingh7@users.noreply.github.com> Co-authored-by: Akhilesh Negi <akhileshnegi.an3@gmail.com>
1 parent 945f676 commit 12bf9b6

File tree

9 files changed

+735
-0
lines changed

9 files changed

+735
-0
lines changed
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
"""add credetial table
2+
3+
Revision ID: 543f97951bd0
4+
Revises: 8d7a05fd0ad4
5+
Create Date: 2025-04-14 23:50:51.118373
6+
7+
"""
8+
from alembic import op
9+
import sqlalchemy as sa
10+
import sqlmodel.sql.sqltypes
11+
12+
13+
# revision identifiers, used by Alembic.
14+
revision = "543f97951bd0"
15+
down_revision = "8d7a05fd0ad4"
16+
branch_labels = None
17+
depends_on = None
18+
19+
20+
def upgrade():
21+
op.create_table(
22+
"credential",
23+
sa.Column("organization_id", sa.Integer(), nullable=False),
24+
sa.Column("is_active", sa.Boolean(), nullable=False),
25+
sa.Column("id", sa.Integer(), nullable=False),
26+
sa.Column("credential", sa.JSON(), nullable=True),
27+
sa.Column("inserted_at", sa.DateTime(), nullable=True),
28+
sa.Column("updated_at", sa.DateTime(), nullable=True),
29+
sa.Column("deleted_at", sa.DateTime(), nullable=True),
30+
sa.ForeignKeyConstraint(
31+
["organization_id"], ["organization.id"], ondelete="CASCADE"
32+
),
33+
sa.PrimaryKeyConstraint("id"),
34+
)
35+
36+
37+
def downgrade():
38+
op.drop_table("credential")

backend/app/api/main.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
threads,
1212
users,
1313
utils,
14+
credentials,
1415
)
1516
from app.core.config import settings
1617

@@ -24,6 +25,8 @@
2425
api_router.include_router(project.router)
2526
api_router.include_router(project_user.router)
2627
api_router.include_router(api_keys.router)
28+
api_router.include_router(credentials.router)
29+
2730

2831
if settings.ENVIRONMENT == "local":
2932
api_router.include_router(private.router)
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
from fastapi import APIRouter, Depends, HTTPException
2+
from app.api.deps import SessionDep, get_current_active_superuser
3+
from app.crud.credentials import (
4+
get_creds_by_org,
5+
get_key_by_org,
6+
remove_creds_for_org,
7+
set_creds_for_org,
8+
update_creds_for_org,
9+
)
10+
from app.models import CredsCreate, CredsPublic, CredsUpdate
11+
from app.utils import APIResponse
12+
from datetime import datetime
13+
14+
router = APIRouter(prefix="/credentials", tags=["credentials"])
15+
16+
17+
@router.post(
18+
"/",
19+
dependencies=[Depends(get_current_active_superuser)],
20+
response_model=APIResponse[CredsPublic],
21+
)
22+
def create_new_credential(*, session: SessionDep, creds_in: CredsCreate):
23+
new_creds = None
24+
try:
25+
existing_creds = get_creds_by_org(
26+
session=session, org_id=creds_in.organization_id
27+
)
28+
if not existing_creds:
29+
new_creds = set_creds_for_org(session=session, creds_add=creds_in)
30+
except Exception as e:
31+
raise HTTPException(
32+
status_code=500, detail=f"An unexpected error occurred: {str(e)}"
33+
)
34+
35+
# Ensure inserted_at is set during creation
36+
new_creds.inserted_at = datetime.utcnow()
37+
38+
return APIResponse.success_response(new_creds)
39+
40+
41+
@router.get(
42+
"/{org_id}",
43+
dependencies=[Depends(get_current_active_superuser)],
44+
response_model=APIResponse[CredsPublic],
45+
)
46+
def read_credential(*, session: SessionDep, org_id: int):
47+
try:
48+
creds = get_creds_by_org(session=session, org_id=org_id)
49+
except Exception as e:
50+
raise HTTPException(
51+
status_code=500, detail=f"An unexpected error occurred: {str(e)}"
52+
)
53+
54+
if creds is None:
55+
raise HTTPException(status_code=404, detail="Credentials not found")
56+
57+
return APIResponse.success_response(creds)
58+
59+
60+
@router.get(
61+
"/{org_id}/api-key",
62+
dependencies=[Depends(get_current_active_superuser)],
63+
response_model=APIResponse[dict],
64+
)
65+
def read_api_key(*, session: SessionDep, org_id: int):
66+
try:
67+
api_key = get_key_by_org(session=session, org_id=org_id)
68+
except Exception as e:
69+
raise HTTPException(
70+
status_code=500, detail=f"An unexpected error occurred: {str(e)}"
71+
)
72+
73+
if api_key is None:
74+
raise HTTPException(status_code=404, detail="API key not found")
75+
76+
return APIResponse.success_response({"api_key": api_key})
77+
78+
79+
@router.patch(
80+
"/{org_id}",
81+
dependencies=[Depends(get_current_active_superuser)],
82+
response_model=APIResponse[CredsPublic],
83+
)
84+
def update_credential(*, session: SessionDep, org_id: int, creds_in: CredsUpdate):
85+
try:
86+
updated_creds = update_creds_for_org(
87+
session=session, org_id=org_id, creds_in=creds_in
88+
)
89+
90+
updated_creds.updated_at = datetime.utcnow()
91+
92+
return APIResponse.success_response(updated_creds)
93+
except ValueError as e:
94+
raise HTTPException(status_code=404, detail=str(e))
95+
except Exception as e:
96+
raise HTTPException(
97+
status_code=500, detail=f"An unexpected error occurred: {str(e)}"
98+
)
99+
100+
101+
from fastapi import HTTPException, Depends
102+
from app.crud.credentials import remove_creds_for_org
103+
from app.utils import APIResponse
104+
from app.api.deps import SessionDep, get_current_active_superuser
105+
106+
107+
@router.delete(
108+
"/{org_id}/api-key",
109+
dependencies=[Depends(get_current_active_superuser)],
110+
response_model=APIResponse[dict],
111+
)
112+
def delete_credential(*, session: SessionDep, org_id: int):
113+
try:
114+
creds = remove_creds_for_org(session=session, org_id=org_id)
115+
except Exception as e:
116+
raise HTTPException(
117+
status_code=500, detail=f"An unexpected error occurred: {str(e)}"
118+
)
119+
120+
if creds is None:
121+
raise HTTPException(
122+
status_code=404, detail="Credentials for organization not found"
123+
)
124+
125+
# No need to manually set deleted_at and is_active if it's done in remove_creds_for_org
126+
# Simply return the success response
127+
return APIResponse.success_response({"message": "Credentials deleted successfully"})

backend/app/crud/credentials.py

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
from typing import Optional, Dict, Any
2+
from sqlmodel import Session, select
3+
from sqlalchemy.exc import IntegrityError
4+
from datetime import datetime
5+
6+
from app.models import Credential, CredsCreate, CredsUpdate
7+
8+
9+
def set_creds_for_org(*, session: Session, creds_add: CredsCreate) -> Credential:
10+
creds = Credential.model_validate(creds_add)
11+
12+
# Set the inserted_at timestamp (current UTC time)
13+
creds.inserted_at = datetime.utcnow()
14+
15+
try:
16+
session.add(creds)
17+
session.commit()
18+
session.refresh(creds)
19+
except IntegrityError as e:
20+
session.rollback() # Rollback the session if there's a unique constraint violation
21+
raise ValueError(f"Error while adding credentials: {str(e)}")
22+
23+
return creds
24+
25+
26+
def get_creds_by_org(*, session: Session, org_id: int) -> Optional[Credential]:
27+
"""Fetches the credentials for the given organization."""
28+
statement = select(Credential).where(Credential.organization_id == org_id)
29+
return session.exec(statement).first()
30+
31+
32+
def get_key_by_org(*, session: Session, org_id: int) -> Optional[str]:
33+
"""Fetches the API key from the credentials for the given organization."""
34+
statement = select(Credential).where(Credential.organization_id == org_id)
35+
creds = session.exec(statement).first()
36+
37+
# Check if creds exists and if the credential field contains the api_key
38+
if (
39+
creds
40+
and creds.credential
41+
and "openai" in creds.credential
42+
and "api_key" in creds.credential["openai"]
43+
):
44+
return creds.credential["openai"]["api_key"]
45+
46+
return None
47+
48+
49+
def update_creds_for_org(
50+
session: Session, org_id: int, creds_in: CredsUpdate
51+
) -> Credential:
52+
# Fetch the current credentials for the organization
53+
creds = session.exec(
54+
select(Credential).where(Credential.organization_id == org_id)
55+
).first()
56+
57+
if not creds:
58+
raise ValueError(f"Credentials not found")
59+
60+
# Update the credentials data with the provided values
61+
creds_data = creds_in.dict(exclude_unset=True)
62+
63+
# Directly update the fields on the original creds object instead of creating a new one
64+
for key, value in creds_data.items():
65+
setattr(creds, key, value)
66+
67+
# Set the updated_at timestamp (current UTC time)
68+
creds.updated_at = datetime.utcnow()
69+
70+
try:
71+
# Add the updated creds to the session and flush the changes to the database
72+
session.add(creds)
73+
session.flush() # This will flush the changes to the database but without committing
74+
session.commit() # Now we commit the changes to make them permanent
75+
except IntegrityError as e:
76+
# Rollback in case of any integrity errors (e.g., constraint violations)
77+
session.rollback()
78+
raise ValueError(f"Error while updating credentials: {str(e)}")
79+
80+
# Refresh the session to get the latest updated data
81+
session.refresh(creds)
82+
83+
return creds
84+
85+
86+
def remove_creds_for_org(*, session: Session, org_id: int) -> Optional[Credential]:
87+
"""Removes (soft deletes) the credentials for the given organization."""
88+
statement = select(Credential).where(Credential.organization_id == org_id)
89+
creds = session.exec(statement).first()
90+
91+
if creds:
92+
try:
93+
# Soft delete: Set is_active to False and set deleted_at timestamp
94+
creds.is_active = False
95+
creds.deleted_at = (
96+
datetime.utcnow()
97+
) # Set the current time as the deleted_at timestamp
98+
session.add(creds)
99+
session.commit()
100+
except IntegrityError as e:
101+
session.rollback() # Rollback in case of a failure during delete operation
102+
raise ValueError(f"Error while deleting credentials: {str(e)}")
103+
104+
return creds

backend/app/models/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,11 @@
4242
UsersPublic,
4343
UpdatePassword,
4444
)
45+
46+
from .credentials import (
47+
Credential,
48+
CredsBase,
49+
CredsCreate,
50+
CredsPublic,
51+
CredsUpdate,
52+
)

backend/app/models/credentials.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
from typing import Dict, Any, Optional
2+
import sqlalchemy as sa
3+
from sqlmodel import Field, Relationship, SQLModel
4+
from datetime import datetime
5+
6+
7+
class CredsBase(SQLModel):
8+
organization_id: int = Field(foreign_key="organization.id")
9+
is_active: bool = True
10+
11+
12+
class CredsCreate(CredsBase):
13+
credential: Dict[str, Any] = Field(default=None, sa_column=sa.Column(sa.JSON))
14+
15+
16+
class CredsUpdate(SQLModel):
17+
credential: Optional[Dict[str, Any]] = Field(
18+
default=None, sa_column=sa.Column(sa.JSON)
19+
)
20+
is_active: Optional[bool] = Field(default=None)
21+
22+
23+
class Credential(CredsBase, table=True):
24+
id: int = Field(default=None, primary_key=True)
25+
credential: Dict[str, Any] = Field(default=None, sa_column=sa.Column(sa.JSON))
26+
inserted_at: datetime = Field(
27+
default_factory=datetime.utcnow,
28+
sa_column=sa.Column(sa.DateTime, default=datetime.utcnow),
29+
)
30+
updated_at: datetime = Field(
31+
default_factory=datetime.utcnow,
32+
sa_column=sa.Column(sa.DateTime, onupdate=datetime.utcnow),
33+
)
34+
deleted_at: Optional[datetime] = Field(
35+
default=None, sa_column=sa.Column(sa.DateTime, nullable=True)
36+
)
37+
38+
organization: Optional["Organization"] = Relationship(back_populates="creds")
39+
40+
41+
class CredsPublic(CredsBase):
42+
id: int
43+
credential: Dict[str, Any]
44+
inserted_at: datetime
45+
updated_at: datetime
46+
deleted_at: Optional[datetime]

backend/app/models/organization.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,10 @@
1+
from typing import List, TYPE_CHECKING
12
from sqlmodel import Field, Relationship, SQLModel
3+
from sqlalchemy.orm import relationship
4+
5+
6+
if TYPE_CHECKING:
7+
from .credentials import Credential
28

39

410
# Shared properties for an Organization
@@ -22,7 +28,11 @@ class OrganizationUpdate(SQLModel):
2228
class Organization(OrganizationBase, table=True):
2329
id: int = Field(default=None, primary_key=True)
2430

31+
# Relationship back to Creds
2532
api_keys: list["APIKey"] = Relationship(back_populates="organization")
33+
creds: list["Credential"] = Relationship(
34+
back_populates="organization", sa_relationship_kwargs={"cascade": "all, delete"}
35+
)
2636

2737

2838
# Properties to return via API

0 commit comments

Comments
 (0)