Added department models

private_gpt/users/api/v1/routers/departments.py

@@ -0,0 +1,127 @@
+from typing import Any, List
+
+from sqlalchemy.orm import Session
+from fastapi.responses import JSONResponse
+from fastapi.encoders import jsonable_encoder
+from fastapi import APIRouter, Depends, HTTPException, status, Security
+
+from private_gpt.users.api import deps
+from private_gpt.users.constants.role import Role
+from private_gpt.users import crud, models, schemas
+
+
+router = APIRouter(prefix="/departments", tags=["Deparments"])
+
+
+@router.get("", response_model=List[schemas.Department])
+def list_deparments(
+    db: Session = Depends(deps.get_db),
+    skip: int = 0,
+    limit: int = 100,
+    current_user: models.User = Security(
+        deps.get_current_user,
+        scopes=[Role.SUPER_ADMIN["name"]],
+    ),
+) -> List[schemas.Department]:
+    """
+    Retrieve a list of companies with pagination support.
+    """
+    deparments = crud.deparment.get_multi(db, skip=skip, limit=limit)
+    return deparments
+
+
+@router.post("/create", response_model=schemas.Department)
+def create_deparment(
+    company_in: schemas.DepartmentCreate,
+    db: Session = Depends(deps.get_db),
+    current_user: models.User = Security(
+        deps.get_current_user,
+        scopes=[Role.SUPER_ADMIN["name"]],
+    ),
+) -> schemas.Company:
+    """
+    Create a new company
+    """
+    deparment = crud.deparment.create(db=db, obj_in=company_in)
+    deparment = jsonable_encoder(deparment)
+
+    return JSONResponse(
+        status_code=status.HTTP_201_CREATED,
+        content={
+            "message": "Department created successfully",
+            "department": deparment
+        },
+    )
+
+
+@router.get("/{deparment_id}", response_model=schemas.Department)
+def read_company(
+    deparment_id: int,
+    db: Session = Depends(deps.get_db),
+    current_user: models.User = Security(
+        deps.get_current_user,
+        scopes=[Role.SUPER_ADMIN["name"]],
+    ),
+) -> schemas.Company:
+    """
+    Read a company by ID
+    """
+    deparment = crud.deparment.get_by_id(db, id=deparment_id)
+    if deparment is None:
+        raise HTTPException(status_code=404, detail="Deparment not found")
+    return deparment
+
+
+@router.put("/{deparment_id}", response_model=schemas.Department)
+def update_company(
+    deparment_id: int,
+    deparment_in: schemas.DepartmentUpdate,
+    db: Session = Depends(deps.get_db),
+    current_user: models.User = Security(
+        deps.get_current_user,
+        scopes=[Role.SUPER_ADMIN["name"]],
+    ),
+) -> schemas.Company:
+    """
+    Update a company by ID
+    """
+    deparment = crud.deparment.get_by_id(db, id=deparment_id)
+    if deparment is None:
+        raise HTTPException(status_code=404, detail="Deparment not found")
+
+    updated_deparment = crud.deparment.update(
+        db=db, db_obj=deparment, obj_in=deparment_in)
+    updated_deparment = jsonable_encoder(updated_deparment)
+    return JSONResponse(
+        status_code=status.HTTP_200_OK,
+        content={
+            "message": f"{deparment_in} Deparment updated successfully",
+            "deparment": updated_deparment
+        },
+    )
+
+
+@router.delete("/{deparment_id}", response_model=schemas.Department)
+def delete_company(
+    deparment_id: int,
+    db: Session = Depends(deps.get_db),
+    current_user: models.User = Security(
+        deps.get_current_user,
+        scopes=[Role.SUPER_ADMIN["name"]],
+    ),
+) -> schemas.Company:
+    """
+    Delete a company by ID
+    """
+
+    deparment = crud.deparment.remove(db=db, id=deparment_id)
+    if deparment is None:
+        raise HTTPException(status_code=404, detail="Deparment not found")
+    deparment = jsonable_encoder(deparment)
+    return JSONResponse(
+        status_code=status.HTTP_200_OK,
+        content={
+            "message": "Deparment deleted successfully",
+            "deparment": deparment
+        },
+    )

private_gpt/users/crud/__init__.py

@@ -4,3 +4,4 @@ from .user_role_crud import user_role
 from .company_crud import company
 from .subscription_crud import subscription
 from .document_crud import documents
+from .department_crud import deparment

private_gpt/users/crud/department_crud.py

@@ -0,0 +1,16 @@
+from sqlalchemy.orm import Session
+from private_gpt.users.schemas.deparment import DepartmentCreate, DepartmentUpdate
+from private_gpt.users.models.department import Department
+from private_gpt.users.crud.base import CRUDBase
+from typing import Optional
+
+
+class CRUDDepartments(CRUDBase[Department, DepartmentCreate, DepartmentUpdate]):
+    def get_by_id(self, db: Session, *, id: str) -> Optional[Department]:
+        return db.query(self.model).filter(Department.id == id).first()
+
+    def get_by_deparment_name(self, db: Session, *, name: str) -> Optional[Department]:
+        return db.query(self.model).filter(Department.name == name).first()
+
+
+deparment = CRUDDepartments(Department)

private_gpt/users/models/__init__.py

@@ -4,3 +4,4 @@ from .user_role import UserRole
 from .role import Role
 from .documents import Documents
 from .subscription import Subscription
+from .department import Department

private_gpt/users/models/company.py

@@ -15,3 +15,4 @@ class Company(Base):
     subscriptions = relationship("Subscription", back_populates="company")
     users = relationship("User", back_populates="company") 
     user_roles = relationship("UserRole", back_populates="company")
+    departments = relationship("Department", back_populates="company")

private_gpt/users/models/department.py

@@ -0,0 +1,17 @@
+from sqlalchemy import ForeignKey
+from sqlalchemy.orm import relationship
+from sqlalchemy import Column, Integer, String
+
+from private_gpt.users.db.base_class import Base
+
+
+class Department(Base):
+    """Models a Department table."""
+
+    __tablename__ = "departments"
+
+    id = Column(Integer, primary_key=True, index=True)
+    name = Column(String, index=True, unique=True)
+
+    company_id = Column(Integer, ForeignKey('companies.id'))
+    company = relationship("Company", back_populates="departments")

private_gpt/users/models/documents.py

@@ -21,3 +21,8 @@ class Documents(Base):
     )
     uploaded_by_user = relationship(
         "User", back_populates="uploaded_documents")
+    department_id = Column(Integer, ForeignKey("departments.id"), nullable=True)
+    
+    uploaded_by_user = relationship(
+        "User", back_populates="uploaded_documents")
+    department = relationship("Department", back_populates="documents")

private_gpt/users/models/user.py

@@ -40,6 +40,10 @@ class User(Base):
     user_role = relationship(
         "UserRole", back_populates="user", uselist=False, cascade="all, delete-orphan")
     
+    department_id = Column(Integer, ForeignKey(
+        "departments.id"), nullable=True)
+    department = relationship("Department", back_populates="users")
+
     def __repr__(self):
         """Returns string representation of model instance"""
         return "<User {fullname!r}>".format(fullname=self.fullname)

private_gpt/users/schemas/__init__.py

@@ -5,3 +5,4 @@ from .user_role import UserRole, UserRoleCreate, UserRoleInDB, UserRoleUpdate
 from .subscription import Subscription, SubscriptionBase, SubscriptionCreate, SubscriptionUpdate
 from .company import Company, CompanyBase, CompanyCreate, CompanyUpdate
 from .documents import Document, DocumentCreate, DocumentsBase, DocumentUpdate
+from .deparment import Department, DepartmentCreate, DepartmentUpdate

private_gpt/users/schemas/company.py

@@ -1,19 +1,30 @@
 from typing import List
-from datetime import datetime
 from pydantic import BaseModel
 
+from private_gpt.users.schemas import Department, User
+
+
 class CompanyBase(BaseModel):
     name: str
 
+
 class CompanyCreate(CompanyBase):
     pass
 
+
 class CompanyUpdate(CompanyBase):
     pass
 
 
-class Company(CompanyBase):
+class CompanyInDB(CompanyBase):
     id: int
 
     class Config:
         orm_mode = True
+
+
+class Company(CompanyInDB):
+    subscriptions: List[str] = []
+    users: List[User] = []
+    user_roles: List[str] = []
+    departments: List[Department] = []

private_gpt/users/schemas/deparment.py

@@ -0,0 +1,26 @@
+from typing import List
+from pydantic import BaseModel
+
+
+class DepartmentBase(BaseModel):
+    name: str
+
+
+class DepartmentCreate(DepartmentBase):
+    pass
+
+
+class DepartmentUpdate(DepartmentBase):
+    pass
+
+
+class DepartmentInDB(DepartmentBase):
+    id: int
+    company_id: int
+
+    class Config:
+        orm_mode = True
+
+
+class Department(DepartmentInDB):
+    pass

private_gpt/users/schemas/user.py

@@ -47,14 +47,11 @@ class UserSchema(UserBaseSchema):
 	class Config:
 		orm_mode = True
 
-# Additional properties to return via API
 
 
 class User(UserSchema):
     pass
 
-# Additional properties stored in DB
-
 
 class UserInDB(UserSchema):
     hashed_password: str
@@ -72,3 +69,4 @@ class UserAdminUpdate(BaseModel):
 	id: int
 	fullname: str
 	role: str
+	department_id: Optional[int] = None