diff --git a/alembic/versions/36beb9b73c64_update_department_model.py b/alembic/versions/36beb9b73c64_update_department_model.py new file mode 100644 index 00000000..9b511dfd --- /dev/null +++ b/alembic/versions/36beb9b73c64_update_department_model.py @@ -0,0 +1,34 @@ +"""update department model + +Revision ID: 36beb9b73c64 +Revises: 0aeaf9df35a6 +Create Date: 2024-02-21 15:12:07.840057 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '36beb9b73c64' +down_revision: Union[str, None] = '0aeaf9df35a6' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('departments', sa.Column('total_users', sa.Integer(), nullable=True)) + op.add_column('departments', sa.Column('total_documents', sa.Integer(), nullable=True)) + # op.create_unique_constraint('unique_user_role', 'user_roles', ['user_id', 'role_id', 'company_id']) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + # op.drop_constraint('unique_user_role', 'user_roles', type_='unique') + op.drop_column('departments', 'total_documents') + op.drop_column('departments', 'total_users') + # ### end Alembic commands ### diff --git a/private_gpt/server/completions/completions_router.py b/private_gpt/server/completions/completions_router.py index 0a79d987..17fa76e5 100644 --- a/private_gpt/server/completions/completions_router.py +++ b/private_gpt/server/completions/completions_router.py @@ -16,8 +16,8 @@ from private_gpt.open_ai.openai_models import ( ) from private_gpt.server.chat.chat_router import ChatBody, chat_completion from private_gpt.server.utils.auth import authenticated -from private_gpt.users import crud, models, schemas from private_gpt.users.api import deps +from private_gpt.users import crud, models, schemas completions_router = APIRouter(prefix="/v1", dependencies=[Depends(authenticated)]) @@ -44,53 +44,53 @@ class CompletionsBody(BaseModel): } -@completions_router.post( - "/completions", - response_model=None, - summary="Completion", - responses={200: {"model": OpenAICompletion}}, - tags=["Contextual Completions"], -) -def prompt_completion( - request: Request, body: CompletionsBody -) -> OpenAICompletion | StreamingResponse: - """We recommend most users use our Chat completions API. +# @completions_router.post( +# "/completions", +# response_model=None, +# summary="Completion", +# responses={200: {"model": OpenAICompletion}}, +# tags=["Contextual Completions"], +# ) +# def prompt_completion( +# request: Request, body: CompletionsBody +# ) -> OpenAICompletion | StreamingResponse: +# """We recommend most users use our Chat completions API. - Given a prompt, the model will return one predicted completion. +# Given a prompt, the model will return one predicted completion. - Optionally include a `system_prompt` to influence the way the LLM answers. +# Optionally include a `system_prompt` to influence the way the LLM answers. - If `use_context` - is set to `true`, the model will use context coming from the ingested documents - to create the response. The documents being used can be filtered using the - `context_filter` and passing the document IDs to be used. Ingested documents IDs - can be found using `/ingest/list` endpoint. If you want all ingested documents to - be used, remove `context_filter` altogether. +# If `use_context` +# is set to `true`, the model will use context coming from the ingested documents +# to create the response. The documents being used can be filtered using the +# `context_filter` and passing the document IDs to be used. Ingested documents IDs +# can be found using `/ingest/list` endpoint. If you want all ingested documents to +# be used, remove `context_filter` altogether. - When using `'include_sources': true`, the API will return the source Chunks used - to create the response, which come from the context provided. +# When using `'include_sources': true`, the API will return the source Chunks used +# to create the response, which come from the context provided. - When using `'stream': true`, the API will return data chunks following [OpenAI's - streaming model](https://platform.openai.com/docs/api-reference/chat/streaming): - ``` - {"id":"12345","object":"completion.chunk","created":1694268190, - "model":"private-gpt","choices":[{"index":0,"delta":{"content":"Hello"}, - "finish_reason":null}]} - ``` - """ - messages = [OpenAIMessage(content=body.prompt, role="user")] - # If system prompt is passed, create a fake message with the system prompt. - if body.system_prompt: - messages.insert(0, OpenAIMessage(content=body.system_prompt, role="system")) +# When using `'stream': true`, the API will return data chunks following [OpenAI's +# streaming model](https://platform.openai.com/docs/api-reference/chat/streaming): +# ``` +# {"id":"12345","object":"completion.chunk","created":1694268190, +# "model":"private-gpt","choices":[{"index":0,"delta":{"content":"Hello"}, +# "finish_reason":null}]} +# ``` +# """ +# messages = [OpenAIMessage(content=body.prompt, role="user")] +# # If system prompt is passed, create a fake message with the system prompt. +# if body.system_prompt: +# messages.insert(0, OpenAIMessage(content=body.system_prompt, role="system")) - chat_body = ChatBody( - messages=messages, - use_context=body.use_context, - stream=body.stream, - include_sources=body.include_sources, - context_filter=body.context_filter, - ) - return chat_completion(request, chat_body) +# chat_body = ChatBody( +# messages=messages, +# use_context=body.use_context, +# stream=body.stream, +# include_sources=body.include_sources, +# context_filter=body.context_filter, +# ) +# return chat_completion(request, chat_body) @completions_router.post( @@ -101,7 +101,7 @@ def prompt_completion( tags=["Contextual Completions"], ) def prompt_completion( - request: Request, + request: Request, body: CompletionsBody, db: Session = Depends(deps.get_db), current_user: models.User = Security( @@ -111,15 +111,17 @@ def prompt_completion( try: service = request.state.injector.get(IngestService) - department = crud.department.get_by_id(db, id=current_user.department_id) + department = crud.department.get_by_id( + db, id=current_user.department_id) if not department: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"No department assigned to you") - documents = crud.documents.get_multi_documents(db, department_id=department.id) + documents = crud.documents.get_multi_documents( + db, department_id=department.id) if not documents: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"No documents uploaded for your department.") - docs_list = [document.filename for document in documents] + docs_list = [document.filename for document in documents] docs_ids = [] for filename in docs_list: doc_id = service.get_doc_ids_by_filename(filename) diff --git a/private_gpt/users/api/v1/routers/auth.py b/private_gpt/users/api/v1/routers/auth.py index a067a895..0873b981 100644 --- a/private_gpt/users/api/v1/routers/auth.py +++ b/private_gpt/users/api/v1/routers/auth.py @@ -34,13 +34,13 @@ def register_user( """ Register a new user in the database. """ - logging.info(f"User : {email} Password: {password} company_id: {company.id} deparment_id: {department.id}") + logging.info(f"User : {email} Password: {password} company_id: {company.id} deparment_id: {department}") user_in = schemas.UserCreate( email=email, password=password, fullname=fullname, company_id=company.id, - department_id=department.id, + department_id=department, ) try: send_registration_email(fullname, email, password) @@ -204,8 +204,8 @@ def register( # password: str = Body(...), company_id: int = Body(None, title="Company ID", description="Company ID for the user (if applicable)"), - department_name: str = Body(None, title="Department Name", - description="Department name for the user (if applicable)"), + department_id: int = Body(None, title="Department Id", + description="Department Id for the user (if applicable)"), role_name: str = Body(None, title="Role Name", description="User role name (if applicable)"), current_user: models.User = Security( @@ -232,17 +232,15 @@ def register( status_code=404, detail="Company not found.", ) - if department_name: - department = crud.department.get_by_department_name( - db=db, name=department_name) - if not department: - raise HTTPException( - status_code=404, - detail="Department not found.", - ) - logging.info(f"Department is {department}") + + if not department_id: + raise HTTPException( + status_code=404, + detail="Department not found.", + ) + logging.info(f"Department is {department_id}") user = register_user( - db, email, fullname, random_password, company, department + db, email, fullname, random_password, company, department_id ) user_role_name = role_name or Role.GUEST["name"] user_role = create_user_role(db, user, user_role_name, company) diff --git a/private_gpt/users/api/v1/routers/departments.py b/private_gpt/users/api/v1/routers/departments.py index 4b6b10bf..3aebd50c 100644 --- a/private_gpt/users/api/v1/routers/departments.py +++ b/private_gpt/users/api/v1/routers/departments.py @@ -56,7 +56,7 @@ def create_deparment( ) -@router.get("/{department_id}", response_model=schemas.Department) +@router.post("/read", response_model=schemas.Department) def read_department( department_id: int, db: Session = Depends(deps.get_db), @@ -74,9 +74,8 @@ def read_department( return department -@router.put("/{department_id}", response_model=schemas.Department) +@router.post("/update", response_model=schemas.Department) def update_department( - department_id: int, department_in: schemas.DepartmentUpdate, db: Session = Depends(deps.get_db), current_user: models.User = Security( @@ -87,7 +86,7 @@ def update_department( """ Update a Department by ID """ - department = crud.department.get_by_id(db, id=department_id) + department = crud.department.get_by_id(db, id=department_in.id) if department is None: raise HTTPException(status_code=404, detail="department not found") @@ -103,9 +102,9 @@ def update_department( ) -@router.delete("/{department_id}", response_model=schemas.Department) +@router.post("/delete", response_model=schemas.Department) def delete_department( - department_id: int, + department_in: schemas.DepartmentDelete, db: Session = Depends(deps.get_db), current_user: models.User = Security( deps.get_current_user, @@ -115,6 +114,10 @@ def delete_department( """ Delete a Department by ID """ + department_id = department_in.id + department = crud.department.get(db, id=department_id) + if department is None: + raise HTTPException(status_code=404, detail="User not found") department = crud.department.remove(db=db, id=department_id) if department is None: diff --git a/private_gpt/users/api/v1/routers/documents.py b/private_gpt/users/api/v1/routers/documents.py index fe5a6839..cae510af 100644 --- a/private_gpt/users/api/v1/routers/documents.py +++ b/private_gpt/users/api/v1/routers/documents.py @@ -15,6 +15,7 @@ logger = logging.getLogger(__name__) router = APIRouter(prefix='/documents', tags=['Documents']) + @router.get("", response_model=List[schemas.Document]) def list_files( request: Request, @@ -26,18 +27,36 @@ def list_files( scopes=[Role.SUPER_ADMIN["name"]], ) ): + def get_department_name(db, id): + dep = crud.department.get_by_id(db=db, id=id) + return dep.name + + def get_username(db, id): + user = crud.user.get_by_id(db=db, id=id) + return user.fullname try: docs = crud.documents.get_multi(db, skip=skip, limit=limit) + docs = [ + schemas.Document( + id=doc.id, + filename=doc.filename, + uploaded_at=doc.uploaded_at, + uploaded_by=get_username(db, doc.uploaded_by), + department=get_department_name(db, doc.department_id) + ) + for doc in docs + ] return docs except Exception as e: print(traceback.format_exc()) logger.error(f"There was an error listing the file(s).") raise HTTPException( status_code=500, - detail="Internal Server Error: Unable to ingest file.", + detail="Internal Server Error", ) + @router.get('{department_id}', response_model=List[schemas.Document]) def list_files_by_department( request: Request, @@ -62,7 +81,7 @@ def list_files_by_department( logger.error(f"There was an error listing the file(s).") raise HTTPException( status_code=500, - detail="Internal Server Error: Unable to ingest file.", + detail="Internal Server Error.", ) @@ -90,5 +109,5 @@ def list_files_by_department( logger.error(f"There was an error listing the file(s).") raise HTTPException( status_code=500, - detail="Internal Server Error: Unable to ingest file.", + detail="Internal Server Error.", ) diff --git a/private_gpt/users/crud/department_crud.py b/private_gpt/users/crud/department_crud.py index 0c026716..3f5586fb 100644 --- a/private_gpt/users/crud/department_crud.py +++ b/private_gpt/users/crud/department_crud.py @@ -6,7 +6,7 @@ from typing import Optional class CRUDDepartments(CRUDBase[Department, DepartmentCreate, DepartmentUpdate]): - def get_by_id(self, db: Session, *, id: str) -> Optional[Department]: + def get_by_id(self, db: Session, *, id: int) -> Optional[Department]: return db.query(self.model).filter(Department.id == id).first() def get_by_department_name(self, db: Session, *, name: str) -> Optional[Department]: diff --git a/private_gpt/users/crud/document_crud.py b/private_gpt/users/crud/document_crud.py index f7bf668c..7e297c69 100644 --- a/private_gpt/users/crud/document_crud.py +++ b/private_gpt/users/crud/document_crud.py @@ -6,7 +6,7 @@ from typing import Optional, List class CRUDDocuments(CRUDBase[Document, DocumentCreate, DocumentUpdate]): - def get_by_id(self, db: Session, *, id: str) -> Optional[Document]: + def get_by_id(self, db: Session, *, id: int) -> Optional[Document]: return db.query(self.model).filter(Document.id == id).first() def get_by_filename(self, db: Session, *, file_name: str) -> Optional[Document]: diff --git a/private_gpt/users/crud/user_crud.py b/private_gpt/users/crud/user_crud.py index 48529623..c416991d 100644 --- a/private_gpt/users/crud/user_crud.py +++ b/private_gpt/users/crud/user_crud.py @@ -104,4 +104,8 @@ class CRUDUser(CRUDBase[User, UserCreate, UserUpdate]): .limit(limit) .all() ) + + def get_by_id(self, db: Session, *, id: int) -> Optional[User]: + return db.query(self.model).filter(User.id == id).first() + user = CRUDUser(User) diff --git a/private_gpt/users/models/department.py b/private_gpt/users/models/department.py index 595d0f95..fb54fbe9 100644 --- a/private_gpt/users/models/department.py +++ b/private_gpt/users/models/department.py @@ -1,8 +1,10 @@ -from sqlalchemy import ForeignKey -from sqlalchemy.orm import relationship +from sqlalchemy import ForeignKey, event +from sqlalchemy.orm import relationship, Session from sqlalchemy import Column, Integer, String from private_gpt.users.db.base_class import Base +from private_gpt.users.models.document import Document +from private_gpt.users.models.user import User class Department(Base): @@ -17,4 +19,27 @@ class Department(Base): company = relationship("Company", back_populates="departments") users = relationship("User", back_populates="department") - documents = relationship("Document", back_populates="department") \ No newline at end of file + documents = relationship("Document", back_populates="department") + + total_users = Column(Integer, default=0) + total_documents = Column(Integer, default=0) + + + +def update_total_users(mapper, connection, target): + session = Session(bind=connection) + target.total_users = session.query(User).filter_by( + department_id=target.id).count() + + +def update_total_documents(mapper, connection, target): + session = Session(bind=connection) + target.total_documents = session.query( + Document).filter_by(department_id=target.id).count() + + +# Attach event listeners to Department model +event.listen(Department, 'after_insert', update_total_users) +event.listen(Department, 'after_update', update_total_users) +event.listen(Department, 'after_insert', update_total_documents) +event.listen(Department, 'after_update', update_total_documents) diff --git a/private_gpt/users/schemas/__init__.py b/private_gpt/users/schemas/__init__.py index b9d5ea25..dac90b6c 100644 --- a/private_gpt/users/schemas/__init__.py +++ b/private_gpt/users/schemas/__init__.py @@ -5,4 +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, DocumentList -from .department import Department, DepartmentCreate, DepartmentUpdate, DepartmentAdminCreate \ No newline at end of file +from .department import Department, DepartmentCreate, DepartmentUpdate, DepartmentAdminCreate, DepartmentDelete \ No newline at end of file diff --git a/private_gpt/users/schemas/department.py b/private_gpt/users/schemas/department.py index 9fdaaf73..9e7631a8 100644 --- a/private_gpt/users/schemas/department.py +++ b/private_gpt/users/schemas/department.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Optional from pydantic import BaseModel @@ -11,12 +11,18 @@ class DepartmentCreate(DepartmentBase): class DepartmentUpdate(DepartmentBase): - pass + id: int + +class DepartmentDelete(BaseModel): + id: int + class DepartmentInDB(DepartmentBase): id: int company_id: int + total_users: Optional[int] + total_documents: Optional[int] class Config: orm_mode = True @@ -29,3 +35,4 @@ class DepartmentAdminCreate(DepartmentBase): class Department(DepartmentInDB): pass + diff --git a/private_gpt/users/schemas/documents.py b/private_gpt/users/schemas/documents.py index e4426877..d1cf1284 100644 --- a/private_gpt/users/schemas/documents.py +++ b/private_gpt/users/schemas/documents.py @@ -23,9 +23,9 @@ class DocumentList(DocumentsBase): class Document(DocumentsBase): id: int - uploaded_by: int + uploaded_by: str uploaded_at: datetime - department_id: int + department: str class Config: orm_mode = True