Merge pull request #27 from AstraBert/feat/lobsterx-api

wip: lobsterx api

Author: AstraBert

.github/workflows/docker-image.yaml

@@ -1,8 +1,6 @@
 name: Build and Publish Docker Image
 
 on:
-  push:
-    branches: [main]
   workflow_dispatch:
 
 env:

packages/lobsterx/config.api.json

@@ -0,0 +1,9 @@
+{
+  "allow_origins": [],
+  "file_downloads_per_minute": 300,
+  "create_tasks_per_minute": 60,
+  "delete_tasks_per_minute": 60,
+  "poll_tasks_per_minute": 300,
+  "host": "0.0.0.0",
+  "port": 9000
+}

packages/lobsterx/pyproject.toml

@@ -4,15 +4,19 @@ build-backend = "uv_build"
 
 [project]
 name = "lobsterx"
-version = "0.1.1-beta"
+version = "0.2.0-beta"
 description = "Background AI assistant working as a Telegram bot, built specifically for document-related use cases"
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "aiofiles>=25.1.0",
     "diskcache>=5.6.3",
+    "fastapi>=0.129.0",
+    "fastapi-throttle>=0.1.8",
+    "httpx>=0.28.1",
     "llama-cloud>=1.2.0,<1.3",
     "python-dotenv>=1.2.1",
+    "python-multipart>=0.0.21",
     "python-telegram-bot>=22.6",
     "random-name>=0.1.1",
     "workflows-acp",

packages/lobsterx/src/lobsterx/api/__init__.py

@@ -0,0 +1,132 @@
+import asyncio
+import mimetypes
+import os
+
+from fastapi import FastAPI, HTTPException
+from fastapi.datastructures import UploadFile
+from fastapi.param_functions import Depends, File
+from fastapi_throttle import RateLimiter
+from random_name import generate_name
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse
+
+from ..constants import DATA_DIR
+from ..utils import _download_file_to_agentfs, handle_prompt
+from .auth import LobsterXAuthentication, on_auth_error
+from .shared import (
+    GetTaskResponse,
+    TaskRequest,
+    TaskResponse,
+    UploadFileResponse,
+    get_server_key_from_env,
+    validate_api_key,
+)
+from .task_manager import get_task_manager
+
+DEFAULT_FILE_DOWNLOADS_PER_MINUTE = 300
+DEFAULT_TASKS_PER_MINUTE = 60
+DEFAULT_DELETE_TASKS_PER_MINUTE = 60
+DEFAULT_POLL_TASKS_PER_MINUTE = 300
+
+
+def _get_file_name(document: UploadFile) -> str:
+    extension = (
+        mimetypes.guess_extension(document.content_type or "application/pdf") or ".pdf"
+    )
+    if document.filename is None:
+        return generate_name() + extension
+    else:
+        if document.filename.endswith(extension):
+            return document.filename
+        return document.filename + extension
+
+
+def create_api_app(
+    allow_origins: list[str],
+    file_downloads_per_minute: int | None,
+    create_tasks_per_minute: int | None,
+    delete_tasks_per_minute: int | None,
+    poll_tasks_per_minute: int | None,
+    server_api_key: str | None,
+) -> FastAPI:
+    app = FastAPI()
+
+    file_downloads_per_minute = (
+        file_downloads_per_minute or DEFAULT_FILE_DOWNLOADS_PER_MINUTE
+    )
+    tasks_per_minute = create_tasks_per_minute or DEFAULT_TASKS_PER_MINUTE
+    delete_tasks_per_minute = delete_tasks_per_minute or DEFAULT_DELETE_TASKS_PER_MINUTE
+    poll_tasks_per_minute = poll_tasks_per_minute or DEFAULT_POLL_TASKS_PER_MINUTE
+    api_key = server_api_key or get_server_key_from_env()
+    if api_key is None:
+        raise ValueError(
+            "API key not provided and `LOBSTERX_SERVER_KEY` not found within the current environment"
+        )
+
+    if not validate_api_key(api_key):
+        raise ValueError(
+            "API key should be an a string of letters, numbers, hyphens and underscores, with a minimum length of 32."
+        )
+
+    app.add_middleware(
+        CORSMiddleware,  # type: ignore[invalid-argument-type]
+        allow_origins=allow_origins,
+        allow_methods=["POST", "GET", "DELETE"],
+        allow_headers=["Content-Type", "Authorization"],
+    )
+
+    app.add_middleware(
+        AuthenticationMiddleware,  # type: ignore[invalid-argument-type]
+        backend=LobsterXAuthentication(api_key=api_key),
+        on_error=on_auth_error,
+    )
+
+    @app.post(
+        "/files",
+        dependencies=[
+            Depends(RateLimiter(times=file_downloads_per_minute, seconds=60))
+        ],
+    )
+    async def download_file(
+        file: UploadFile = File(...),
+    ) -> UploadFileResponse:
+        file_name = _get_file_name(file)
+        file_content = await file.read()
+        path = os.path.join(DATA_DIR, file_name)
+        await _download_file_to_agentfs(path, file_content)
+        return UploadFileResponse(new_file_path=path)
+
+    @app.post(
+        "/tasks",
+        dependencies=[Depends(RateLimiter(times=tasks_per_minute, seconds=60))],
+    )
+    async def create_task(request: TaskRequest) -> TaskResponse:
+        task_manager = get_task_manager()
+        task = asyncio.create_task(handle_prompt(request.prompt))
+        task_id = await task_manager.add_task(task)
+        return TaskResponse(task_id=task_id)
+
+    @app.delete(
+        "/tasks/{task_id}",
+        dependencies=[Depends(RateLimiter(times=delete_tasks_per_minute, seconds=60))],
+    )
+    async def cancel_task(task_id: str) -> JSONResponse:
+        task_manager = get_task_manager()
+        await task_manager.cancel_task(task_id)
+        return JSONResponse(status_code=204, content={})
+
+    @app.get(
+        "/tasks/{task_id}",
+        dependencies=[Depends(RateLimiter(times=poll_tasks_per_minute, seconds=60))],
+    )
+    async def get_task(task_id: str) -> GetTaskResponse:
+        task_manager = get_task_manager()
+        task = await task_manager.check_task(task_id)
+        if task is None:
+            raise HTTPException(
+                status_code=404, detail=f"Task {task_id} does not exist"
+            )
+        return GetTaskResponse.from_dataclass(task)
+
+    return app

packages/lobsterx/src/lobsterx/api/api.py

@@ -0,0 +1,55 @@
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    AuthenticationError,
+    BaseUser,
+)
+from starlette.requests import HTTPConnection
+from starlette.responses import PlainTextResponse, Response
+
+from .shared import get_auth_header_pattern
+
+DEFAULT_DISPLAY_NAME = "lobsterx"
+DEFAULT_IDENTITY = "user"
+
+
+class LobsterXUser(BaseUser):
+    def __init__(self, authenticated: bool) -> None:
+        self._authenticated = authenticated
+
+    @property
+    def is_authenticated(self) -> bool:
+        return self._authenticated
+
+    @property
+    def identity(self) -> str:
+        return DEFAULT_IDENTITY
+
+    @property
+    def display_name(self) -> str:
+        return DEFAULT_DISPLAY_NAME
+
+
+class LobsterXAuthentication(AuthenticationBackend):
+    def __init__(self, api_key: str) -> None:
+        self.api_key = api_key
+
+    async def authenticate(
+        self, conn: HTTPConnection
+    ) -> tuple[AuthCredentials, BaseUser] | None:
+        auth_header = conn.headers.get("Authorization", None)
+        if auth_header is None:
+            raise AuthenticationError("No authorization header in request")
+        matches = get_auth_header_pattern().findall(auth_header)
+        try:
+            assert len(matches) == 1, "Should only provide one bearer token"
+        except AssertionError as e:
+            raise AuthenticationError("Should only provide one bearer token") from e
+        api_key = matches[0]
+        if api_key == self.api_key:
+            return AuthCredentials(scopes=["http"]), LobsterXUser(authenticated=True)
+        raise AuthenticationError("API key not authorized")
+
+
+def on_auth_error(conn: HTTPConnection, exc: AuthenticationError) -> Response:
+    return PlainTextResponse(str(exc), status_code=401)

packages/lobsterx/src/lobsterx/api/auth.py

@@ -0,0 +1,83 @@
+import os
+from mimetypes import guess_type
+from typing import Literal
+
+from httpx import AsyncClient
+
+from .shared import (
+    GetTaskResponse,
+    TaskRequest,
+    TaskResponse,
+    UploadFileResponse,
+    get_server_key_from_env,
+    validate_api_key,
+)
+
+
+class LobsterXClient:
+    def __init__(
+        self,
+        api_key: str | None,
+        host: str,
+        port: int,
+        protocol: Literal["http", "https"],
+    ) -> None:
+        self.base_url = f"{protocol}://{host}:{port}"
+        self.api_key = api_key or get_server_key_from_env()
+        if self.api_key is None:
+            raise ValueError(
+                "API key not provided and `LOBSTERX_SERVER_KEY` not found within the current environment"
+            )
+        if not validate_api_key(self.api_key):
+            raise ValueError(
+                "API key should be an a string of letters, numbers, hyphens and underscores, with a minimum length of 32."
+            )
+
+    async def upload_file(self, file_path: str) -> str:
+        async with AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            timeout=600,
+        ) as client:
+            with open(file_path, "rb") as f:
+                mimetype, _ = guess_type(file_path)
+                file_type = mimetype or "application/pdf"
+                file = (os.path.basename(file_path), f, file_type)
+                response = await client.post("/files", files={"file": file})
+                response.raise_for_status()
+                payload = response.json()
+                validated = UploadFileResponse.model_validate(payload)
+                return validated.new_file_path
+
+    async def create_task(self, prompt: str) -> str:
+        async with AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            timeout=600,
+        ) as client:
+            payload = TaskRequest(prompt=prompt).model_dump()
+            response = await client.post("/tasks", json=payload)
+            response.raise_for_status()
+            json_response = response.json()
+            validated = TaskResponse.model_validate(json_response)
+            return validated.task_id
+
+    async def get_task(self, task_id: str) -> GetTaskResponse:
+        async with AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            timeout=600,
+        ) as client:
+            response = await client.get(f"/tasks/{task_id}")
+            response.raise_for_status()
+            json_response = response.json()
+            return GetTaskResponse.model_validate(json_response)
+
+    async def cancel_task(self, task_id: str) -> None:
+        async with AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            timeout=600,
+        ) as client:
+            response = await client.delete(f"/tasks/{task_id}")
+            response.raise_for_status()

packages/lobsterx/src/lobsterx/api/client.py

@@ -0,0 +1,85 @@
+import functools
+import json
+import os
+import re
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from dotenv import load_dotenv
+from pydantic import BaseModel
+
+from .task_manager import StatusEnum, TaskRepr
+
+
+@dataclass
+class LobsterXApiConfig:
+    allow_origins: list[str]
+    file_downloads_per_minute: int | None = None
+    create_tasks_per_minute: int | None = None
+    delete_tasks_per_minute: int | None = None
+    poll_tasks_per_minute: int | None = None
+    server_api_key: str | None = None
+    host: str | None = None
+    port: int | None = None
+    protocol: Literal["http", "https"] = "http"
+
+    @classmethod
+    def load_from_config(cls, config_file: str) -> "LobsterXApiConfig":
+        with open(config_file, "r") as f:
+            config = json.load(f)
+        return cls(**config)
+
+    def to_args(self) -> dict[str, Any]:
+        return {
+            "allow_origins": self.allow_origins,
+            "file_downloads_per_minute": self.file_downloads_per_minute,
+            "create_tasks_per_minute": self.create_tasks_per_minute,
+            "delete_tasks_per_minute": self.delete_tasks_per_minute,
+            "poll_tasks_per_minute": self.poll_tasks_per_minute,
+            "server_api_key": self.server_api_key,
+        }
+
+
+@functools.lru_cache(maxsize=1)
+def get_api_key_pattern() -> re.Pattern:
+    return re.compile(r"[a-zA-Z0-9_-]{32,}")
+
+
+@functools.lru_cache(maxsize=1)
+def get_auth_header_pattern() -> re.Pattern:
+    return re.compile(r"Bearer\s([a-zA-Z0-9_-]{32,})")
+
+
+def validate_api_key(api_key: str) -> bool:
+    pattern = get_api_key_pattern()
+    return pattern.match(api_key) is not None
+
+
+@functools.lru_cache(maxsize=1)
+def get_server_key_from_env() -> str | None:
+    load_dotenv(".env")
+    return os.getenv("LOBSTERX_SERVER_KEY")
+
+
+class TaskRequest(BaseModel):
+    prompt: str
+
+
+class TaskResponse(BaseModel):
+    task_id: str
+
+
+class GetTaskResponse(BaseModel):
+    status: StatusEnum
+    output: tuple[str, str] | None = None
+    error: str | None = None
+
+    @classmethod
+    def from_dataclass(cls, task_repr: TaskRepr) -> "GetTaskResponse":
+        return cls(
+            status=task_repr.status, output=task_repr.output, error=task_repr.error
+        )
+
+
+class UploadFileResponse(BaseModel):
+    new_file_path: str