chainlit 0.7.604rc2__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

chainlit/server.py

@@ -1,6 +1,7 @@
 import glob
 import json
 import mimetypes
+import shutil
 import urllib.parse
 from typing import Optional, Union
 
@@ -17,30 +18,31 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 
 from chainlit.auth import create_jwt, get_configuration, get_current_user
-from chainlit.client.cloud import AppUser, PersistedAppUser
 from chainlit.config import (
     APP_ROOT,
     BACKEND_ROOT,
     DEFAULT_HOST,
+    FILES_DIRECTORY,
     PACKAGE_ROOT,
     config,
     load_module,
     reload_config,
 )
-from chainlit.data import chainlit_client
-from chainlit.data.acl import is_conversation_author
+from chainlit.data import get_data_layer
+from chainlit.data.acl import is_thread_author
 from chainlit.logger import logger
 from chainlit.markdown import get_markdown_str
 from chainlit.playground.config import get_llm_providers
 from chainlit.telemetry import trace_event
 from chainlit.types import (
-    CompletionRequest,
-    DeleteConversationRequest,
-    GetConversationsRequest,
+    DeleteThreadRequest,
+    GenerationRequest,
+    GetThreadsRequest,
     Theme,
     UpdateFeedbackRequest,
 )
-from fastapi import Depends, FastAPI, HTTPException, Query, Request, status
+from chainlit.user import PersistedUser, User
+from fastapi import Depends, FastAPI, HTTPException, Query, Request, UploadFile, status
 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi.staticfiles import StaticFiles
@@ -117,6 +119,9 @@ async def lifespan(app: FastAPI):
             except asyncio.exceptions.CancelledError:
                 pass
 
+        if FILES_DIRECTORY.is_dir():
+            shutil.rmtree(FILES_DIRECTORY)
+
         # Force exit the process to avoid potential AnyIO threads still running
         os._exit(0)
 
@@ -134,6 +139,7 @@ def get_build_dir():
 
 build_dir = get_build_dir()
 
+
 app = FastAPI(lifespan=lifespan)
 
 app.mount("/public", StaticFiles(directory="public", check_dir=False), name="public")
@@ -156,14 +162,10 @@ app.add_middleware(
 )
 
 
-# Define max HTTP data size to 100 MB
-max_message_size = 100 * 1024 * 1024
-
 socket = SocketManager(
     app,
     cors_allowed_origins=[],
     async_mode="asgi",
-    max_http_buffer_size=max_message_size,
 )
 
 
@@ -244,18 +246,18 @@ async def login(form_data: OAuth2PasswordRequestForm = Depends()):
             status_code=status.HTTP_400_BAD_REQUEST, detail="No auth_callback defined"
         )
 
-    app_user = await config.code.password_auth_callback(
+    user = await config.code.password_auth_callback(
         form_data.username, form_data.password
     )
 
-    if not app_user:
+    if not user:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="credentialssignin",
         )
-    access_token = create_jwt(app_user)
-    if chainlit_client:
-        await chainlit_client.create_app_user(app_user=app_user)
+    access_token = create_jwt(user)
+    if data_layer := get_data_layer():
+        await data_layer.create_user(user)
     return {
         "access_token": access_token,
         "token_type": "bearer",
@@ -270,17 +272,17 @@ async def header_auth(request: Request):
             detail="No header_auth_callback defined",
         )
 
-    app_user = await config.code.header_auth_callback(request.headers)
+    user = await config.code.header_auth_callback(request.headers)
 
-    if not app_user:
+    if not user:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Unauthorized",
         )
 
-    access_token = create_jwt(app_user)
-    if chainlit_client:
-        await chainlit_client.create_app_user(app_user=app_user)
+    access_token = create_jwt(user)
+    if data_layer := get_data_layer():
+        await data_layer.create_user(user)
     return {
         "access_token": access_token,
         "token_type": "bearer",
@@ -369,21 +371,22 @@ async def oauth_callback(
     url = get_user_facing_url(request.url)
     token = await provider.get_token(code, url)
 
-    (raw_user_data, default_app_user) = await provider.get_user_info(token)
+    (raw_user_data, default_user) = await provider.get_user_info(token)
 
-    app_user = await config.code.oauth_callback(
-        provider_id, token, raw_user_data, default_app_user
+    user = await config.code.oauth_callback(
+        provider_id, token, raw_user_data, default_user
     )
 
-    if not app_user:
+    if not user:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Unauthorized",
         )
 
-    access_token = create_jwt(app_user)
-    if chainlit_client:
-        await chainlit_client.create_app_user(app_user=app_user)
+    access_token = create_jwt(user)
+
+    if data_layer := get_data_layer():
+        await data_layer.create_user(user)
 
     params = urllib.parse.urlencode(
         {
@@ -399,23 +402,21 @@ async def oauth_callback(
     return response
 
 
-@app.post("/completion")
-async def completion(
-    request: CompletionRequest,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+@app.post("/generation")
+async def generation(
+    request: GenerationRequest,
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
     """Handle a completion request from the prompt playground."""
 
     providers = get_llm_providers()
 
     try:
-        provider = [p for p in providers if p.id == request.prompt.provider][0]
+        provider = [p for p in providers if p.id == request.generation.provider][0]
     except IndexError:
         raise HTTPException(
             status_code=404,
-            detail=f"LLM provider '{request.prompt.provider}' not found",
+            detail=f"LLM provider '{request.generation.provider}' not found",
         )
 
     trace_event("pp_create_completion")
@@ -426,7 +427,7 @@ async def completion(
 
 @app.get("/project/llm-providers")
 async def get_providers(
-    current_user: Annotated[Union[AppUser, PersistedAppUser], Depends(get_current_user)]
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)]
 ):
     """List the providers."""
     trace_event("pp_get_llm_providers")
@@ -437,7 +438,7 @@ async def get_providers(
 
 @app.get("/project/settings")
 async def project_settings(
-    current_user: Annotated[Union[AppUser, PersistedAppUser], Depends(get_current_user)]
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)]
 ):
     """Return project settings. This is called by the UI before the establishing the websocket connection."""
     profiles = []
@@ -450,126 +451,181 @@ async def project_settings(
             "ui": config.ui.to_dict(),
             "features": config.features.to_dict(),
             "userEnv": config.project.user_env,
-            "dataPersistence": config.data_persistence,
-            "conversationResumable": bool(config.code.on_chat_resume),
+            "dataPersistence": get_data_layer() is not None,
+            "threadResumable": bool(config.code.on_chat_resume),
             "markdown": get_markdown_str(config.root),
             "chatProfiles": profiles,
         }
     )
 
 
-@app.put("/message/feedback")
+@app.put("/feedback")
 async def update_feedback(
     request: Request,
     update: UpdateFeedbackRequest,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
     """Update the human feedback for a particular message."""
-
-    # TODO: check that message belong to a user's conversation
-
-    if not chainlit_client:
+    data_layer = get_data_layer()
+    if not data_layer:
         raise HTTPException(status_code=400, detail="Data persistence is not enabled")
 
     try:
-        await chainlit_client.set_human_feedback(
-            message_id=update.messageId,
-            feedback=update.feedback,
-            feedbackComment=update.feedbackComment,
-        )
+        feedback_id = await data_layer.upsert_feedback(feedback=update.feedback)
     except Exception as e:
         raise HTTPException(detail=str(e), status_code=401)
 
-    return JSONResponse(content={"success": True})
+    return JSONResponse(content={"success": True, "feedbackId": feedback_id})
 
 
-@app.post("/project/conversations")
-async def get_user_conversations(
+@app.post("/project/threads")
+async def get_user_threads(
     request: Request,
-    payload: GetConversationsRequest,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    payload: GetThreadsRequest,
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
-    """Get the conversations page by page."""
-    # Only show the current user conversations
+    """Get the threads page by page."""
+    # Only show the current user threads
 
-    if not chainlit_client:
+    data_layer = get_data_layer()
+
+    if not data_layer:
         raise HTTPException(status_code=400, detail="Data persistence is not enabled")
 
-    payload.filter.username = current_user.username
-    res = await chainlit_client.get_conversations(payload.pagination, payload.filter)
+    payload.filter.userIdentifier = current_user.identifier
+
+    res = await data_layer.list_threads(payload.pagination, payload.filter)
     return JSONResponse(content=res.to_dict())
 
 
-@app.get("/project/conversation/{conversation_id}")
-async def get_conversation(
+@app.get("/project/thread/{thread_id}")
+async def get_thread(
     request: Request,
-    conversation_id: str,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    thread_id: str,
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
-    """Get a specific conversation."""
+    """Get a specific thread."""
+    data_layer = get_data_layer()
 
-    if not chainlit_client:
+    if not data_layer:
         raise HTTPException(status_code=400, detail="Data persistence is not enabled")
 
-    await is_conversation_author(current_user.username, conversation_id)
+    await is_thread_author(current_user.identifier, thread_id)
 
-    res = await chainlit_client.get_conversation(conversation_id)
+    res = await data_layer.get_thread(thread_id)
     return JSONResponse(content=res)
 
 
-@app.get("/project/conversation/{conversation_id}/element/{element_id}")
-async def get_conversation_element(
+@app.get("/project/thread/{thread_id}/element/{element_id}")
+async def get_thread_element(
     request: Request,
-    conversation_id: str,
+    thread_id: str,
     element_id: str,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
-    """Get a specific conversation element."""
+    """Get a specific thread element."""
+    data_layer = get_data_layer()
 
-    if not chainlit_client:
+    if not data_layer:
         raise HTTPException(status_code=400, detail="Data persistence is not enabled")
 
-    await is_conversation_author(current_user.username, conversation_id)
+    await is_thread_author(current_user.identifier, thread_id)
 
-    res = await chainlit_client.get_element(conversation_id, element_id)
+    res = await data_layer.get_element(thread_id, element_id)
     return JSONResponse(content=res)
 
 
-@app.delete("/project/conversation")
-async def delete_conversation(
+@app.delete("/project/thread")
+async def delete_thread(
     request: Request,
-    payload: DeleteConversationRequest,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    payload: DeleteThreadRequest,
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
-    """Delete a conversation."""
+    """Delete a thread."""
+
+    data_layer = get_data_layer()
 
-    if not chainlit_client:
+    if not data_layer:
         raise HTTPException(status_code=400, detail="Data persistence is not enabled")
 
-    conversation_id = payload.conversationId
+    thread_id = payload.threadId
 
-    await is_conversation_author(current_user.username, conversation_id)
+    await is_thread_author(current_user.identifier, thread_id)
 
-    await chainlit_client.delete_conversation(conversation_id)
+    await data_layer.delete_thread(thread_id)
     return JSONResponse(content={"success": True})
 
 
+@app.post("/project/file")
+async def upload_file(
+    session_id: str,
+    file: UploadFile,
+    current_user: Annotated[
+        Union[None, User, PersistedUser], Depends(get_current_user)
+    ],
+):
+    from chainlit.session import WebsocketSession
+
+    session = WebsocketSession.get_by_id(session_id)
+
+    if not session:
+        raise HTTPException(
+            status_code=404,
+            detail="Session not found",
+        )
+
+    if current_user:
+        if not session.user or session.user.identifier != current_user.identifier:
+            raise HTTPException(
+                status_code=401,
+                detail="You are not authorized to upload files for this session",
+            )
+
+    session.files_dir.mkdir(exist_ok=True)
+
+    content = await file.read()
+
+    file_response = await session.persist_file(
+        name=file.filename, content=content, mime=file.content_type
+    )
+
+    return JSONResponse(file_response)
+
+
+@app.get("/project/file/{file_id}")
+async def get_file(
+    file_id: str,
+    session_id: Optional[str] = None,
+    token: Optional[str] = None,
+):
+    from chainlit.session import WebsocketSession
+
+    session = WebsocketSession.get_by_id(session_id) if session_id else None
+
+    if not session:
+        raise HTTPException(
+            status_code=404,
+            detail="Session not found",
+        )
+
+    if current_user := await get_current_user(token or ""):
+        if not session.user or session.user.identifier != current_user.identifier:
+            raise HTTPException(
+                status_code=401,
+                detail="You are not authorized to upload files for this session",
+            )
+
+    if file_id in session.files:
+        file = session.files[file_id]
+        return FileResponse(file["path"], media_type=file["type"])
+    else:
+        raise HTTPException(status_code=404, detail="File not found")
+
+
 @app.get("/files/{filename:path}")
 async def serve_file(
     filename: str,
-    current_user: Annotated[
-        Union[AppUser, PersistedAppUser], Depends(get_current_user)
-    ],
+    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
 ):
     base_path = Path(config.project.local_fs_path).resolve()
     file_path = (base_path / filename).resolve()

chainlit/session.py

@@ -1,13 +1,16 @@
-import asyncio
 import json
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+import mimetypes
+import shutil
+import uuid
+from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional, Union
+
+import aiofiles
 
 if TYPE_CHECKING:
     from chainlit.message import Message
-    from chainlit.types import AskResponse
-
-from chainlit.client.cloud import AppUser, PersistedAppUser
-from chainlit.data import chainlit_client
+    from chainlit.step import Step
+    from chainlit.types import FileDict, FileReference
+    from chainlit.user import PersistedUser, User
 
 
 class JSONEncoderIgnoreNonSerializable(json.JSONEncoder):
@@ -27,12 +30,17 @@ def clean_metadata(metadata: Dict):
 class BaseSession:
     """Base object."""
 
+    active_steps: List["Step"]
+    thread_id_to_resume: Optional[str] = None
+
     def __init__(
         self,
         # Id of the session
         id: str,
+        # Thread id
+        thread_id: Optional[str],
         # Logged-in user informations
-        user: Optional[Union["AppUser", "PersistedAppUser"]],
+        user: Optional[Union["User", "PersistedUser"]],
         # Logged-in user token
         token: Optional[str],
         # User specific environment variables. Empty if no user environment variables are required.
@@ -41,41 +49,30 @@ class BaseSession:
         root_message: Optional["Message"] = None,
         # Chat profile selected before the session was created
         chat_profile: Optional[str] = None,
-        # Conversation id to resume
-        conversation_id: Optional[str] = None,
     ):
+        if thread_id:
+            self.thread_id_to_resume = thread_id
+        self.thread_id = thread_id or str(uuid.uuid4())
         self.user = user
         self.token = token
         self.root_message = root_message
         self.has_user_message = False
         self.user_env = user_env or {}
         self.chat_profile = chat_profile
+        self.active_steps = []
 
         self.id = id
-        self.conversation_id = conversation_id
 
         self.chat_settings: Dict[str, Any] = {}
 
-        self.lock = asyncio.Lock()
-
-    async def get_conversation_id(self) -> Optional[str]:
-        if not chainlit_client:
-            return None
-
-        if isinstance(self, HTTPSession):
-            tags = ["api"]
-        else:
-            tags = ["chat"]
-
-        async with self.lock:
-            if not self.conversation_id:
-                app_user_id = (
-                    self.user.id if isinstance(self.user, PersistedAppUser) else None
-                )
-                self.conversation_id = await chainlit_client.create_conversation(
-                    app_user_id=app_user_id, tags=tags
-                )
-        return self.conversation_id
+    async def persist_file(
+        self,
+        name: str,
+        mime: str,
+        path: Optional[str] = None,
+        content: Optional[Union[bytes, str]] = None,
+    ):
+        return None
 
     def to_persistable(self) -> Dict:
         from chainlit.user_session import user_sessions
@@ -94,17 +91,24 @@ class HTTPSession(BaseSession):
         self,
         # Id of the session
         id: str,
+        # Thread id
+        thread_id: Optional[str] = None,
         # Logged-in user informations
-        user: Optional[Union["AppUser", "PersistedAppUser"]],
+        user: Optional[Union["User", "PersistedUser"]] = None,
         # Logged-in user token
-        token: Optional[str],
-        user_env: Optional[Dict[str, str]],
+        token: Optional[str] = None,
+        user_env: Optional[Dict[str, str]] = None,
         # Last message at the root of the chat
         root_message: Optional["Message"] = None,
         # User specific environment variables. Empty if no user environment variables are required.
     ):
         super().__init__(
-            id=id, user=user, token=token, user_env=user_env, root_message=root_message
+            id=id,
+            thread_id=thread_id,
+            user=user,
+            token=token,
+            user_env=user_env,
+            root_message=root_message,
         )
 
 
@@ -129,28 +133,28 @@ class WebsocketSession(BaseSession):
         # Function to emit a message to the user
         emit: Callable[[str, Any], None],
         # Function to ask the user a question
-        ask_user: Callable[[Any, Optional[int]], Union["AskResponse", None]],
+        ask_user: Callable[[Any, Optional[int]], Any],
         # User specific environment variables. Empty if no user environment variables are required.
         user_env: Dict[str, str],
+        # Thread id
+        thread_id: Optional[str] = None,
         # Logged-in user informations
-        user: Optional[Union["AppUser", "PersistedAppUser"]],
+        user: Optional[Union["User", "PersistedUser"]] = None,
         # Logged-in user token
-        token: Optional[str],
+        token: Optional[str] = None,
         # Last message at the root of the chat
         root_message: Optional["Message"] = None,
         # Chat profile selected before the session was created
         chat_profile: Optional[str] = None,
-        # Conversation id to resume
-        conversation_id: Optional[str] = None,
     ):
         super().__init__(
             id=id,
+            thread_id=thread_id,
             user=user,
             token=token,
             user_env=user_env,
             root_message=root_message,
             chat_profile=chat_profile,
-            conversation_id=conversation_id,
         )
 
         self.socket_id = socket_id
@@ -160,9 +164,66 @@ class WebsocketSession(BaseSession):
         self.should_stop = False
         self.restored = False
 
+        self.thread_queues = {}  # type: Dict[str, Deque[Callable]]
+        self.files = {}  # type: Dict[str, "FileDict"]
+
         ws_sessions_id[self.id] = self
         ws_sessions_sid[socket_id] = self
 
+    @property
+    def files_dir(self):
+        from chainlit.config import FILES_DIRECTORY
+
+        return FILES_DIRECTORY / self.id
+
+    async def persist_file(
+        self,
+        name: str,
+        mime: str,
+        path: Optional[str] = None,
+        content: Optional[Union[bytes, str]] = None,
+    ) -> "FileReference":
+        if not path and not content:
+            raise ValueError(
+                "Either path or content must be provided to persist a file"
+            )
+
+        self.files_dir.mkdir(exist_ok=True)
+
+        file_id = str(uuid.uuid4())
+
+        file_path = self.files_dir / file_id
+
+        file_extension = mimetypes.guess_extension(mime)
+        if file_extension:
+            file_path = file_path.with_suffix(file_extension)
+
+        if path:
+            # Copy the file from the given path
+            async with aiofiles.open(path, "rb") as src, aiofiles.open(
+                file_path, "wb"
+            ) as dst:
+                await dst.write(await src.read())
+        elif content:
+            # Write the provided content to the file
+            async with aiofiles.open(file_path, "wb") as buffer:
+                if isinstance(content, str):
+                    content = content.encode("utf-8")
+                await buffer.write(content)
+
+        # Get the file size
+        file_size = file_path.stat().st_size
+        # Store the file content in memory
+        self.files[file_id] = {
+            "id": file_id,
+            "path": file_path,
+            "name": name,
+            "type": mime,
+            "size": file_size,
+        }
+
+        return {"id": file_id}
+
     def restore(self, new_socket_id: str):
         """Associate a new socket id to the session."""
         ws_sessions_sid.pop(self.socket_id, None)
@@ -172,9 +233,17 @@ class WebsocketSession(BaseSession):
 
     def delete(self):
         """Delete the session."""
+        if self.files_dir.is_dir():
+            shutil.rmtree(self.files_dir)
         ws_sessions_sid.pop(self.socket_id, None)
         ws_sessions_id.pop(self.id, None)
 
+    async def flush_method_queue(self):
+        for method_name, queue in self.thread_queues.items():
+            while queue:
+                method, self, args, kwargs = queue.popleft()
+                await method(self, *args, **kwargs)
+
     @classmethod
     def get(cls, socket_id: str):
         """Get session by socket id."""