Merge pull request #52 from YangZhiBoGreenHand/feat/yzb/update-1

feat: use async session

Author: liuooo

app/api/deps.py

@@ -1,19 +1,21 @@
-from sqlmodel import Session
+from typing import AsyncGenerator
 
 from fastapi import Depends, Request
-
 from fastapi.security import APIKeyHeader
-from config.config import settings
+from sqlalchemy.ext.asyncio import AsyncSession
+
 from app.exceptions.exception import AuthenticationError, AuthorizationError, ResourceNotFoundError
-from app.providers import database
-from app.models.token_relation import RelationType, TokenRelationQuery
 from app.models.token import Token
-from app.services.token.token_relation import TokenRelationService
+from app.models.token_relation import RelationType, TokenRelationQuery
+from app.providers import database
 from app.services.token.token import TokenService
+from app.services.token.token_relation import TokenRelationService
+from config.config import settings
 
 
-def get_session():
-    with Session(database.engine) as session:
+async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+    """session生成器 作为fast api的Depends选项"""
+    async with database.async_session_local() as session:
         yield session
 
 
@@ -49,13 +51,13 @@ async def verify_admin_token(token=Depends(oauth_token)) -> Token:
         raise AuthorizationError()
 
 
-async def get_token(session=Depends(get_session), token=Depends(oauth_token)) -> Token:
+async def get_token(session=Depends(get_async_session), token=Depends(oauth_token)) -> Token:
     """
     get token info
     """
     if token and token != "":
         try:
-            return TokenService.get_token(session=session, token=token)
+            return await TokenService.get_token(session=session, token=token)
         except ResourceNotFoundError:
             pass
     return None
@@ -90,13 +92,21 @@ async def get_param_from_request(request: Request):
     return get_param_from_request
 
 
-def verify_token_relation(relation_type: RelationType, name: str):
+def verify_token_relation(relation_type: RelationType, name: str, ignore_none_relation_id: bool = False):
+    """
+    param relation_type: relation type
+    param name: param name
+    param ignore_none_relation_id: if ignore_none_relation_id is set, return where relation_id is None, use for copy thread api
+    """
+
     async def verify_authorization(
-        session=Depends(get_session), token_id=Depends(get_token_id), relation_id=Depends(get_param(name))
+        session=Depends(get_async_session), token_id=Depends(get_token_id), relation_id=Depends(get_param(name))
     ):
+        if token_id and ignore_none_relation_id:
+            return
         if token_id and relation_id:
             verify = TokenRelationQuery(token_id=token_id, relation_type=relation_type, relation_id=relation_id)
-            if TokenRelationService.verify_relation(session=session, verify=verify):
+            if await TokenRelationService.verify_relation(session=session, verify=verify):
                 return
         raise AuthorizationError()
 

app/api/routes.py

@@ -13,4 +13,3 @@ def router_init():
     api_router.include_router(files.router, prefix="/files", tags=["files"])
     api_router.include_router(token.router, prefix="/tokens", tags=["tokens"])
     api_router.include_router(action.router, prefix="/actions", tags=["actions"])
-

app/api/v1/action.py

@@ -1,73 +1,75 @@
 from typing import Dict, List
 
 from fastapi import APIRouter, Depends
-from sqlmodel import Session, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import select
 
-from app.api.deps import get_session, get_token_id
+from app.api.deps import get_async_session, get_token_id
+from app.libs.paginate import cursor_page, CommonPage
 from app.models.action import Action
 from app.models.token_relation import RelationType
 from app.providers.auth_provider import auth_policy
 from app.schemas.common import DeleteResponse, BaseSuccessDataResponse
 from app.schemas.tool.action import ActionBulkCreateRequest, ActionUpdateRequest, ActionRunRequest
 from app.services.tool.action import ActionService
-from app.libs.paginate import cursor_page, CommonPage
-
 
 router = APIRouter()
 
 
 @router.get("", response_model=CommonPage[Action])
-def list_actions(*, session: Session = Depends(get_session), token_id=Depends(get_token_id)):
+async def list_actions(*, session: AsyncSession = Depends(get_async_session), token_id=Depends(get_token_id)):
     """
     Returns a list of Actions.
     """
     statement = auth_policy.token_filter(
         select(Action), field=Action.id, relation_type=RelationType.Action, token_id=token_id
     )
-    return cursor_page(statement, session)
+    return await cursor_page(statement, session)
 
 
 @router.post("", response_model=List[Action])
-def create_actions(
-    *, session: Session = Depends(get_session), body: ActionBulkCreateRequest, token_id=Depends(get_token_id)
+async def create_actions(
+    *, session: AsyncSession = Depends(get_async_session), body: ActionBulkCreateRequest, token_id=Depends(get_token_id)
 ) -> List[Action]:
     """
     Create an action with openapi schema.
     """
 
-    return ActionService.create_actions(session=session, body=body, token_id=token_id)
+    return await ActionService.create_actions(session=session, body=body, token_id=token_id)
 
 
 @router.get("/{action_id}", response_model=Action)
-def get_action(*, session: Session = Depends(get_session), action_id: str) -> Action:
+async def get_action(*, session: AsyncSession = Depends(get_async_session), action_id: str) -> Action:
     """
     Retrieves an action.
     """
-    return ActionService.get_action(session=session, action_id=action_id)
+    return await ActionService.get_action(session=session, action_id=action_id)
 
 
 @router.post("/{action_id}", response_model=Action)
-def modify_action(*, session: Session = Depends(get_session), action_id: str, body: ActionUpdateRequest) -> Action:
+async def modify_action(
+    *, session: AsyncSession = Depends(get_async_session), action_id: str, body: ActionUpdateRequest
+) -> Action:
     """
     Modifies an action.
     """
-    return ActionService.modify_action(session=session, action_id=action_id, body=body)
+    return await ActionService.modify_action(session=session, action_id=action_id, body=body)
 
 
 @router.delete("/{action_id}", response_model=DeleteResponse)
-def delete_action(*, session: Session = Depends(get_session), action_id: str) -> DeleteResponse:
+async def delete_action(*, session: AsyncSession = Depends(get_async_session), action_id: str) -> DeleteResponse:
     """
     Delete an action.
     """
-    return ActionService.delete_action(session=session, action_id=action_id)
+    return await ActionService.delete_action(session=session, action_id=action_id)
 
 
 @router.post(
     "/{action_id}/run",
     response_model=BaseSuccessDataResponse,
 )
-def api_run_action(*, session: Session = Depends(get_session), action_id: str, body: ActionRunRequest):
-    response: Dict = ActionService.run_action(
+async def api_run_action(*, session: AsyncSession = Depends(get_async_session), action_id: str, body: ActionRunRequest):
+    response: Dict = await ActionService.run_action(
         session=session,
         action_id=action_id,
         parameters=body.parameters,

app/api/v1/assistant.py

@@ -1,7 +1,8 @@
 from fastapi import APIRouter, Depends
-from sqlmodel import Session, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import select
 
-from app.api.deps import get_token_id, get_session
+from app.api.deps import get_token_id, get_async_session
 from app.models.assistant import Assistant, AssistantUpdate, AssistantCreate
 from app.libs.paginate import cursor_page, CommonPage
 from app.models.token_relation import RelationType
@@ -13,45 +14,47 @@
 
 
 @router.get("", response_model=CommonPage[Assistant])
-def list_assistants(*, session: Session = Depends(get_session), token_id=Depends(get_token_id)):
+async def list_assistants(*, session: AsyncSession = Depends(get_async_session), token_id=Depends(get_token_id)):
     """
     Returns a list of assistants.
     """
     statement = auth_policy.token_filter(
         select(Assistant), field=Assistant.id, relation_type=RelationType.Assistant, token_id=token_id
     )
-    return cursor_page(statement, session)
+    return await cursor_page(statement, session)
 
 
 @router.post("", response_model=Assistant)
-def create_assistant(
-    *, session: Session = Depends(get_session), body: AssistantCreate, token_id=Depends(get_token_id)
+async def create_assistant(
+    *, session: AsyncSession = Depends(get_async_session), body: AssistantCreate, token_id=Depends(get_token_id)
 ) -> Assistant:
     """
     Create an assistant with a model and instructions.
     """
-    return AssistantService.create_assistant(session=session, body=body, token_id=token_id)
+    return await AssistantService.create_assistant(session=session, body=body, token_id=token_id)
 
 
 @router.get("/{assistant_id}", response_model=Assistant)
-def get_assistant(*, session: Session = Depends(get_session), assistant_id: str) -> Assistant:
+async def get_assistant(*, session: AsyncSession = Depends(get_async_session), assistant_id: str) -> Assistant:
     """
     Retrieves an assistant.
     """
-    return AssistantService.get_assistant(session=session, assistant_id=assistant_id)
+    return await AssistantService.get_assistant(session=session, assistant_id=assistant_id)
 
 
 @router.post("/{assistant_id}", response_model=Assistant)
-def modify_assistant(*, session: Session = Depends(get_session), assistant_id: str, body: AssistantUpdate) -> Assistant:
+async def modify_assistant(
+    *, session: AsyncSession = Depends(get_async_session), assistant_id: str, body: AssistantUpdate
+) -> Assistant:
     """
     Modifies an assistant.
     """
-    return AssistantService.modify_assistant(session=session, assistant_id=assistant_id, body=body)
+    return await AssistantService.modify_assistant(session=session, assistant_id=assistant_id, body=body)
 
 
 @router.delete("/{assistant_id}", response_model=DeleteResponse)
-def delete_assistant(*, session: Session = Depends(get_session), assistant_id: str) -> DeleteResponse:
+async def delete_assistant(*, session: AsyncSession = Depends(get_async_session), assistant_id: str) -> DeleteResponse:
     """
     Delete an assistant.
     """
-    return AssistantService.delete_assistant(session=session, assistant_id=assistant_id)
+    return await AssistantService.delete_assistant(session=session, assistant_id=assistant_id)

app/api/v1/assistant_file.py

@@ -1,57 +1,60 @@
 from fastapi import APIRouter, Depends
-from sqlmodel import Session, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import select
 
-from app.api.deps import get_session
-from app.models.assistant_file import AssistantFileCreate, AssistantFile
+from app.api.deps import get_async_session
 from app.libs.paginate import cursor_page, CommonPage
+from app.models.assistant_file import AssistantFileCreate, AssistantFile
 from app.schemas.common import DeleteResponse
 from app.services.assistant.assistant_file import AssistantFileService
 
 router = APIRouter()
 
 
 @router.get("/{assistant_id}/files", response_model=CommonPage[AssistantFile])
-def list_assistant_files(
+async def list_assistant_files(
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_async_session),
     assistant_id: str,
 ):
     """
     Returns a list of assistant files.
     """
-    return cursor_page(select(AssistantFile).where(AssistantFile.assistant_id == assistant_id), db=session)
+    return await cursor_page(select(AssistantFile).where(AssistantFile.assistant_id == assistant_id), db=session)
 
 
 @router.post("/{assistant_id}/files", response_model=AssistantFile)
-def create_assistant_file(
+async def create_assistant_file(
     *,
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_async_session),
     assistant_id: str,
     body: AssistantFileCreate,
 ) -> AssistantFile:
     """
     Create an assistant file by attaching a [File](/docs/api-reference/files)
     to an [assistant](/docs/api-reference/assistants).
     """
-    return AssistantFileService.create_assistant_file(session=session, assistant_id=assistant_id, body=body)
+    return await AssistantFileService.create_assistant_file(session=session, assistant_id=assistant_id, body=body)
 
 
 @router.get("/{assistant_id}/files/{file_id}", response_model=AssistantFile)
-def get_assistant_file(*, session: Session = Depends(get_session), assistant_id: str, file_id: str) -> AssistantFile:
+async def get_assistant_file(
+    *, session: AsyncSession = Depends(get_async_session), assistant_id: str, file_id: str
+) -> AssistantFile:
     """
     Retrieves an AssistantFile.
     """
-    return AssistantFileService.get_assistant_file(session=session, assistant_id=assistant_id, file_id=file_id)
+    return await AssistantFileService.get_assistant_file(session=session, assistant_id=assistant_id, file_id=file_id)
 
 
 @router.delete(
     "/{assistant_id}/files/{file_id}",
     response_model=DeleteResponse,
 )
-def delete_assistant_file(
-    *, session: Session = Depends(get_session), assistant_id: str, file_id: str
+async def delete_assistant_file(
+    *, session: AsyncSession = Depends(get_async_session), assistant_id: str, file_id: str
 ) -> DeleteResponse:
     """
     Delete an assistant file.
     """
-    return AssistantFileService.delete_assistant_file(session=session, assistant_id=assistant_id, file_id=file_id)
+    return await AssistantFileService.delete_assistant_file(session=session, assistant_id=assistant_id, file_id=file_id)

app/api/v1/files.py

@@ -4,10 +4,10 @@
 from typing import Optional, List
 
 from fastapi import APIRouter, Depends, UploadFile, Form, HTTPException, Query
-from sqlmodel import Session
 from starlette.responses import StreamingResponse
+from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.api.deps import get_session
+from app.api.deps import get_async_session
 from app.models import File
 from app.schemas.common import DeleteResponse
 from app.schemas.files import ListFilesResponse
@@ -22,22 +22,22 @@
 
 
 @router.get("", response_model=ListFilesResponse)
-def list_files(
+async def list_files(
     *,
     purpose: Optional[str] = None,
     file_ids: Optional[List[str]] = Query(None, alias="ids[]"),
-    session: Session = Depends(get_session),
+    session: AsyncSession = Depends(get_async_session),
 ) -> ListFilesResponse:
     """
     Returns a list of files that belong to the user's organization.
     """
-    files = FileService.get_file_list(session=session, purpose=purpose, file_ids=file_ids)
+    files = await FileService.get_file_list(session=session, purpose=purpose, file_ids=file_ids)
     return ListFilesResponse(data=files)
 
 
 @router.post("", response_model=File)
-def create_file(
-    *, session: Session = Depends(get_session), purpose: str = Form(default="assistants"), file: UploadFile
+async def create_file(
+    *, session: AsyncSession = Depends(get_async_session), purpose: str = Form(default="assistants"), file: UploadFile
 ) -> File:
     """
     The size of individual files can be a maximum of 512 MB. See the [Assistants Tools guide]
@@ -51,31 +51,31 @@ def create_file(
     if file.size == 0 or file.size > max_size:
         raise HTTPException(status_code=413, detail="File too large")
 
-    return FileService.create_file(session=session, purpose=purpose, file=file)
+    return await FileService.create_file(session=session, purpose=purpose, file=file)
 
 
 @router.delete("/{file_id}", response_model=DeleteResponse)
-def delete_file(*, session: Session = Depends(get_session), file_id: str) -> DeleteResponse:
+async def delete_file(*, session: AsyncSession = Depends(get_async_session), file_id: str) -> DeleteResponse:
     """
     Delete a file.
     """
-    return FileService.delete_file(session=session, file_id=file_id)
+    return await FileService.delete_file(session=session, file_id=file_id)
 
 
 @router.get("/{file_id}", response_model=File)
-def retrieve_file(*, session: Session = Depends(get_session), file_id: str) -> File:
+async def retrieve_file(*, session: AsyncSession = Depends(get_async_session), file_id: str) -> File:
     """
     Returns information about a specific file.
     """
-    return FileService.get_file(session=session, file_id=file_id)
+    return await FileService.get_file(session=session, file_id=file_id)
 
 
 @router.get("/{file_id}/content", response_class=StreamingResponse)
-def download_file(*, file_id: str, session: Session = Depends(get_session)):
+async def download_file(*, file_id: str, session: AsyncSession = Depends(get_async_session)):
     """
     Returns the contents of the specified file.
     """
-    file_data, filename = FileService.get_file_content(session=session, file_id=file_id)
+    file_data, filename = await FileService.get_file_content(session=session, file_id=file_id)
 
     response = StreamingResponse(io.BytesIO(file_data), media_type="application/octet-stream")
     response.headers["Content-Disposition"] = f"attachment; filename*=UTF-8''{urllib.parse.quote(filename)}"