Merge pull request #123 from askui/fix/persistence-in-chat-api

Fix/persistence in chat api

Author: adi-wan-askui

mypy.ini

@@ -24,3 +24,6 @@ namespace_packages = true
 
 [mypy-jsonref.*]
 ignore_missing_imports = true
+
+[mypy-bson.*]
+ignore_missing_imports = true

pdm.lock

@@ -29,6 +29,7 @@ dependencies = [
     "filetype>=1.2.0",
     "markitdown[xls,xlsx,docx]>=0.1.2",
     "asyncer==0.0.8",
+    "bson>=0.5.10",
 ]
 requires-python = ">=3.10"
 readme = "README.md"
@@ -102,6 +103,7 @@ python_files = ["test_*.py"]
 python_functions = ["test_*"]
 testpaths = ["tests"]
 timeout = 60
+asyncio_default_fixture_loop_scope = "session"
 
 [tool.ruff]
 # Exclude a variety of commonly ignored directories.

pyproject.toml

@@ -1,17 +1,25 @@
 from contextlib import asynccontextmanager
 from typing import AsyncGenerator
 
-from fastapi import APIRouter, FastAPI
+from fastapi import APIRouter, FastAPI, HTTPException, Request, status
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
 
 from askui.chat.api.assistants.dependencies import get_assistant_service
 from askui.chat.api.assistants.router import router as assistants_router
 from askui.chat.api.dependencies import SetEnvFromHeadersDep, get_settings
+from askui.chat.api.files.router import router as files_router
 from askui.chat.api.health.router import router as health_router
 from askui.chat.api.mcp_configs.router import router as mcp_configs_router
 from askui.chat.api.messages.router import router as messages_router
 from askui.chat.api.runs.router import router as runs_router
 from askui.chat.api.threads.router import router as threads_router
+from askui.utils.api_utils import (
+    ConflictError,
+    FileTooLargeError,
+    LimitReachedError,
+    NotFoundError,
+)
 
 
 @asynccontextmanager
@@ -38,12 +46,69 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:  # noqa: ARG001
     allow_headers=["*"],
 )
 
+
+@app.exception_handler(NotFoundError)
+def not_found_error_handler(
+    request: Request,  # noqa: ARG001
+    exc: NotFoundError,
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=status.HTTP_404_NOT_FOUND, content={"detail": str(exc)}
+    )
+
+
+@app.exception_handler(ConflictError)
+def conflict_error_handler(
+    request: Request,  # noqa: ARG001
+    exc: ConflictError,
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=status.HTTP_409_CONFLICT, content={"detail": str(exc)}
+    )
+
+
+@app.exception_handler(LimitReachedError)
+def limit_reached_error_handler(
+    request: Request,  # noqa: ARG001
+    exc: LimitReachedError,
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=status.HTTP_400_BAD_REQUEST, content={"detail": str(exc)}
+    )
+
+
+@app.exception_handler(FileTooLargeError)
+def file_too_large_error_handler(
+    request: Request,  # noqa: ARG001
+    exc: FileTooLargeError,
+) -> JSONResponse:
+    return JSONResponse(
+        status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+        content={"detail": str(exc)},
+    )
+
+
+@app.exception_handler(Exception)
+def catch_all_exception_handler(
+    request: Request,  # noqa: ARG001
+    exc: Exception,
+) -> JSONResponse:
+    if isinstance(exc, HTTPException):
+        raise exc
+
+    return JSONResponse(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        content={"detail": "Internal server error"},
+    )
+
+
 # Include routers
 v1_router = APIRouter(prefix="/v1")
 v1_router.include_router(assistants_router)
 v1_router.include_router(threads_router)
 v1_router.include_router(messages_router)
 v1_router.include_router(runs_router)
 v1_router.include_router(mcp_configs_router)
+v1_router.include_router(files_router)
 v1_router.include_router(health_router)
 app.include_router(v1_router)

src/askui/chat/api/app.py

@@ -1,17 +1,53 @@
 from typing import Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
+from askui.chat.api.models import AssistantId
+from askui.utils.api_utils import Resource
 from askui.utils.datetime_utils import UnixDatetime, now
 from askui.utils.id_utils import generate_time_ordered_id
+from askui.utils.not_given import NOT_GIVEN, BaseModelWithNotGiven, NotGiven
 
 
-class Assistant(BaseModel):
-    """An assistant that can be used in a thread."""
+class AssistantBase(BaseModel):
+    """Base assistant model."""
 
-    id: str = Field(default_factory=lambda: generate_time_ordered_id("asst"))
-    created_at: UnixDatetime = Field(default_factory=now)
     name: str | None = None
     description: str | None = None
+    avatar: str | None = None
+
+
+class AssistantCreateParams(AssistantBase):
+    """Parameters for creating an assistant."""
+
+
+class AssistantModifyParams(BaseModelWithNotGiven):
+    """Parameters for modifying an assistant."""
+
+    name: str | NotGiven = NOT_GIVEN
+    description: str | NotGiven = NOT_GIVEN
+    avatar: str | NotGiven = NOT_GIVEN
+
+
+class Assistant(AssistantBase, Resource):
+    """An assistant that can be used in a thread."""
+
+    id: AssistantId
     object: Literal["assistant"] = "assistant"
-    avatar: str | None = Field(default=None, description="URL of the avatar image")
+    created_at: UnixDatetime
+
+    @classmethod
+    def create(cls, params: AssistantCreateParams) -> "Assistant":
+        return cls(
+            id=generate_time_ordered_id("asst"),
+            created_at=now(),
+            **params.model_dump(),
+        )
+
+    def modify(self, params: AssistantModifyParams) -> "Assistant":
+        return Assistant.model_validate(
+            {
+                **self.model_dump(),
+                **params.model_dump(),
+            }
+        )

src/askui/chat/api/assistants/models.py

@@ -1,12 +1,14 @@
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter, status
 
-# from fastapi import status
 from askui.chat.api.assistants.dependencies import AssistantServiceDep
-from askui.chat.api.assistants.models import Assistant
-from askui.chat.api.assistants.service import (
-    AssistantService,  # AssistantModifyRequest, CreateAssistantRequest,
+from askui.chat.api.assistants.models import (
+    Assistant,
+    AssistantCreateParams,
+    AssistantModifyParams,
 )
-from askui.chat.api.models import ListQueryDep
+from askui.chat.api.assistants.service import AssistantService
+from askui.chat.api.dependencies import ListQueryDep
+from askui.chat.api.models import AssistantId
 from askui.utils.api_utils import ListQuery, ListResponse
 
 router = APIRouter(prefix="/assistants", tags=["assistants"])
@@ -17,51 +19,37 @@ def list_assistants(
     query: ListQuery = ListQueryDep,
     assistant_service: AssistantService = AssistantServiceDep,
 ) -> ListResponse[Assistant]:
-    """List all assistants."""
     return assistant_service.list_(query=query)
 
 
-# @router.post("", status_code=status.HTTP_201_CREATED)
-# def create_assistant(
-#     request: CreateAssistantRequest,
-#     assistant_service: AssistantService = AssistantServiceDep,
-# ) -> Assistant:
-#     """Create a new assistant."""
-#     return assistant_service.create(request)
+@router.post("", status_code=status.HTTP_201_CREATED)
+def create_assistant(
+    params: AssistantCreateParams,
+    assistant_service: AssistantService = AssistantServiceDep,
+) -> Assistant:
+    return assistant_service.create(params)
 
 
 @router.get("/{assistant_id}")
 def retrieve_assistant(
-    assistant_id: str,
+    assistant_id: AssistantId,
     assistant_service: AssistantService = AssistantServiceDep,
 ) -> Assistant:
-    """Get an assistant by ID."""
-    try:
-        return assistant_service.retrieve(assistant_id)
-    except FileNotFoundError as e:
-        raise HTTPException(status_code=404, detail=str(e)) from e
+    return assistant_service.retrieve(assistant_id)
 
 
-# @router.post("/{assistant_id}")
-# def modify_assistant(
-#     assistant_id: str,
-#     request: AssistantModifyRequest,
-#     assistant_service: AssistantService = AssistantServiceDep,
-# ) -> Assistant:
-#     """Update an assistant."""
-#     try:
-#         return assistant_service.modify(assistant_id, request)
-#     except FileNotFoundError as e:
-#         raise HTTPException(status_code=404, detail=str(e)) from e
+@router.post("/{assistant_id}")
+def modify_assistant(
+    assistant_id: AssistantId,
+    params: AssistantModifyParams,
+    assistant_service: AssistantService = AssistantServiceDep,
+) -> Assistant:
+    return assistant_service.modify(assistant_id, params)
 
 
-# @router.delete("/{assistant_id}", status_code=status.HTTP_204_NO_CONTENT)
-# def delete_assistant(
-#     assistant_id: str,
-#     assistant_service: AssistantService = AssistantServiceDep,
-# ) -> None:
-#     """Delete an assistant."""
-#     try:
-#         assistant_service.delete(assistant_id)
-#     except FileNotFoundError as e:
-#         raise HTTPException(status_code=404, detail=str(e)) from e
+@router.delete("/{assistant_id}", status_code=status.HTTP_204_NO_CONTENT)
+def delete_assistant(
+    assistant_id: AssistantId,
+    assistant_service: AssistantService = AssistantServiceDep,
+) -> None:
+    assistant_service.delete(assistant_id)