Merge pull request #28 from AstraBert/feat/long-polling-and-docs

feat: wait-task command and docs

Author: AstraBert

packages/lobsterx/README.md

@@ -41,6 +41,8 @@ You then need to set three required env variables:
 - `TELEGRAM_BOT_TOKEN`: token for the Telegram bot
 - `LLAMA_CLOUD_API_KEY`: API key for LlamaCloud
 
+If you wish to setup LobsterX as an API server, you will need to set an API key that only you can use to interact with it, set in the environment as `LOBSTERX_SERVER_KEY`. The key has to be at least 32 charachters long and contain only lowercase and uppercase alphanumeric characters, `-` and `_`.
+
 You can use the setup wizard to configure LobsterX interactively on the terminal:
 
 ```bash
@@ -54,7 +56,8 @@ lobsterx setup --provider google \
     --model gemini-3-flash-preview \
     --api-key $GOOGLE_API_KEY \
     --llama-cloud-key $LLAMA_CLOUD_API_KEY \
-    --telegram-token $TELEGRAM_BOT_TOKEN
+    --telegram-token $TELEGRAM_BOT_TOKEN \
+    --server-key $SERVER_KEY
 ```
 
 This will create a `.env` file with the necessary variables, which will be loaded by LobsterX at runtime (make sure not to share it with anyone).
@@ -63,7 +66,9 @@ If you wish to further customize the instructions that LobsterX has access to, y
 
 ## Run
 
-Run LobsterX as a CLI app:
+### As a Telegram Bot
+
+Run LobsterX as a Telegram Bot:
 
 ```bash
 lobsterx run 
@@ -88,15 +93,85 @@ docker run ghcr.io/astrabert/lobsterx:main \
     --env="TELEGRAM_BOT_TOKEN=tok-xxx"
 ```
 
-## Use as a Telegram Bot
-
 When on Telegram, you can perform two actions:
 
 - Sending PDF files, which will be downloaded by the bot
 - Sending text messages, which will work as prompts for the bot to start a new task
 
 > _With `/start` command, you will have a welcome message explaining how to use the bot_
 
+### As an API server
+
+To run as an API server, you need to specify a series of options that are necessary for authentication, rate limiting and CORS.
+
+- For **authentication**, you need to set the `LOBSTERX_SERVER_KEY` within the environment or in a `.env` file in the same working directory as the agent
+- For **CORS**, you can set a list of allowed origins
+- For **rate limiting**, you can set the maximum limits of file uploads, task creations, task polling and task deletion per minute
+
+In addition to these, you will also need to provide the host (`0.0.0.0` e.g.), port (`8000` e.g.) and protocol (`http` or `https`) on which the server will run.
+
+You can provide all of these details directly from the CLI:
+
+```bash
+lobsterx serve \
+    --file-downloads-per-minute 300 \
+    --create-tasks-per-minute 60 \
+    --delete-tasks-per-minute 60 \
+    --poll-tasks-per-minute 300 \
+    --bind 0.0.0.0 \
+    --port 8000 \
+    --protocol http \
+    --allow https://example.com \
+    --allow https://anotherexample.com
+```
+
+> All of these options have sensible defaults, but personalization is always recommended
+
+Or create a JSON configuration ([as in thie example](config.api.json)) following this specification:
+
+```json
+{
+  "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": 8000,
+  "protocol": "http"
+}
+```
+
+And provide it to the CLI:
+
+```bash
+lobsterx serve --config config.api.json
+```
+
+> The configuration approach is recommended, as it can be re-use through different API-related commands.
+
+Once you are serving your API through `lobsterx serve`, you can:
+
+- Upload files, by sending a POST request to `/files`
+- Create tasks, by sending a POST request to `/task`
+- Get the status of a task, by sending a GET request to `/task/{task_id}`
+- Cancel a task, by sending a DELETE request to `/task/{task_id}`
+
+You don't have to do this through raw API calls, the LobsterX CLI provides several commands to perform these operations on your behalf:
+
+```bash
+# upload a file
+lobsterx upload-file path/to/file.pdf --config config.api.json # pass the server configuration
+# start a task
+lobsterx create-task "Your prompt" --config config.api.json # this will return a task ID
+# check the status of a task
+lobsterx get-task some-task-id --config config.api.json
+# cancel a task 
+lobsterx cancel-task some-task-id --config config.api.json
+# wait until a task is complete
+lobsterx wait-task some-task-id --config config.api.json --polling-interval 2.0 --max-attempts 900 --verbose
+```
+
 ## How LobsterX Works
 
 LobsterX is a generalist AI agent based on three main principles:
@@ -111,6 +186,16 @@ Here is what happens when you send a prompt to LobsterX:
 
 Along with the final response, the agent will also send you a report of everything it did during its session as a markdown file (namedd `session-<random-id>-report.md`).
 
+### The API server
+
+While sharing the core desing principles outlined above, the API server has some more features related to the data flow:
+
+- When a POST request to the `/tasks` endpoint (task creation) is made, a new `asyncio.Task` is spawned and stored within a in-memory task manager, using a locked dictionary to associate a task ID with an async Task.
+- When a GET request is sent to `/task/{task_id}`, the task manager provides details on the status of the task (`success`, `failed`, `cancelled`, `pending`). If the task was succesfull, failed or was cancelled, it is removed from the dictionary.
+- When a DELETE request is sent to `/task/{task_id}`, the async Task is cancelled and removed from the dictionary
+
+Besides the Task Manager, the API server uses an in-memory rate limiter ([`fastapi-throttle`](https://github.com/AliYmn/fastapi-throttle)) and Starlette CORS and Auth middleawares to provide authentication (through a `Bearer` token provided with an `Authorization` header) and CORS servicres.
+
 ## License
 
 This package is provided under [MIT License](./LICENSE)

packages/lobsterx/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "lobsterx"
-version = "0.2.0-beta"
+version = "0.2.1-beta"
 description = "Background AI assistant working as a Telegram bot, built specifically for document-related use cases"
 readme = "README.md"
 requires-python = ">=3.11"

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

@@ -1,4 +1,6 @@
+import asyncio
 import os
+import sys
 from mimetypes import guess_type
 from typing import Literal
 
@@ -81,3 +83,38 @@ async def cancel_task(self, task_id: str) -> None:
         ) as client:
             response = await client.delete(f"/tasks/{task_id}")
             response.raise_for_status()
+
+    async def poll_for_task(
+        self,
+        task_id: str,
+        polling_interval: float = 2.0,
+        max_attempts: int = 900,  # 30 minutes
+        verbose: bool = True,
+    ) -> GetTaskResponse | None:
+        attempts = 0
+        async with AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            timeout=600,
+        ) as client:
+            while True:
+                attempts += 1
+                response = await client.get(f"/tasks/{task_id}")
+                response.raise_for_status()
+                json_response = response.json()
+                validated = GetTaskResponse.model_validate(json_response)
+                if validated.status.value == "pending" and attempts < max_attempts:
+                    if verbose:
+                        print(
+                            f"Attempt {attempts}: Task still pending...",
+                            file=sys.stderr,
+                        )
+                    await asyncio.sleep(polling_interval)
+                elif validated.status.value == "pending" and attempts >= max_attempts:
+                    print(
+                        "Maximum number of attempts reached, exiting...",
+                        file=sys.stderr,
+                    )
+                    return
+                else:
+                    return validated

packages/lobsterx/src/lobsterx/cli.py

@@ -15,6 +15,7 @@
 from .api.shared import LobsterXApiConfig
 from .bot import run_bot
 from .constants import LOG_LEVELS
+from .utils import _setup_agentfs
 
 app = Typer()
 
@@ -161,7 +162,6 @@ def serve(
         int | None,
         Option(
             "--file-downloads-per-minute",
-            "-a",
             help="Rate limit (per minute) on file downloads. Defaults to 300.",
         ),
     ] = None,
@@ -216,6 +216,7 @@ def serve(
             file_downloads_per_minute=file_downloads_per_minute,
             server_api_key=server_api_key,
         )
+    asyncio.run(_setup_agentfs(with_print=True))
     uvicorn.run(app, host=host, port=port)
 
 
@@ -407,7 +408,114 @@ def get_task(
         )
         rprint(
             Markdown(
-                f"## Final Outout\n\n{final_output}\n\n## Activity Report\n\n{report}"
+                f"## Final Output\n\n{final_output}\n\n## Activity Report\n\n{report}"
+            )
+        )
+
+
+@app.command(
+    name="wait-task",
+    help="Poll for a task until it is completed.",
+)
+def wait_task(
+    task_id: str,
+    protocol: Annotated[
+        Literal["http", "https"],
+        Option(
+            "--protocol",
+            "-t",
+            help="Protocol for the connection. Defaults to 'http'.",
+        ),
+    ] = "http",
+    host: Annotated[
+        str,
+        Option(
+            "--bind",
+            "-b",
+            help="Host to bind the server to. Defaults to 0.0.0.0",
+        ),
+    ] = "0.0.0.0",
+    port: Annotated[
+        int,
+        Option(
+            "--port",
+            "-p",
+            help="Port to bind the server to. Defaults to 8000",
+        ),
+    ] = 8000,
+    polling_interval: Annotated[
+        float,
+        Option(
+            "--polling-interval",
+            "-i",
+            help="Interval (in seconds) between a polling request and the following one. Defaults to 2 seconds.",
+        ),
+    ] = 2.0,
+    max_attempts: Annotated[
+        int,
+        Option(
+            "--max-attempts",
+            "-m",
+            help="Maximum number of polling attempts. Defaults to 900 (for a total of 30 minutes with the default polling interval).",
+        ),
+    ] = 900,
+    server_api_key: Annotated[
+        str | None,
+        Option(
+            "--server-key",
+            help="API key to be used within the server to authorize requests. Reads from LOBSTERX_SERVER_KEY env variable if not provided.",
+        ),
+    ] = None,
+    config_file: Annotated[
+        str | None,
+        Option(
+            "--config",
+            "-c",
+            help="Config file from which to read the LobsterX server configuration. Configured options have precedence over CLI.",
+        ),
+    ] = None,
+    verbose: Annotated[
+        bool,
+        Option(
+            "--verbose/--no-verbose",
+            help="Whether or not to enable verbose logging.",
+        ),
+    ] = True,
+) -> None:
+    if config_file is not None:
+        args = LobsterXApiConfig.load_from_config(config_file)
+        port = args.port or port
+        host = args.host or host
+        protocol = args.protocol or protocol
+    client = LobsterXClient(
+        api_key=server_api_key, host=host, port=port, protocol=protocol
+    )
+    response = asyncio.run(
+        client.poll_for_task(
+            task_id,
+            polling_interval=polling_interval,
+            max_attempts=max_attempts,
+            verbose=verbose,
+        )
+    )
+    if response is None:
+        return
+    if response.status.value in ("cancelled", "failed"):
+        rprint(f"[bold red]Task {task_id} was cancelled or produced an error[/]")
+        if response.error is not None:
+            rprint(f"[bold red]Error: {response.error}[/]")
+    elif response.status.value == "pending":
+        rprint(f"[bold cyan]Task {task_id} is still being executed[/]")
+    else:
+        final_output = (
+            response.output[1] if response.output is not None else "No final output"
+        )
+        report = (
+            response.output[0] if response.output is not None else "No activity report"
+        )
+        rprint(
+            Markdown(
+                f"## Final Output\n\n{final_output}\n\n## Activity Report\n\n{report}"
             )
         )
 

packages/lobsterx/src/lobsterx/utils.py

@@ -3,6 +3,7 @@
 import logging
 import mimetypes
 import os
+import sys
 from typing import cast
 
 import aiofiles
@@ -168,17 +169,37 @@ async def _remove_temporary_report_file(path: str) -> None:
         pass
 
 
-async def _setup_agentfs() -> None:
+async def _setup_agentfs(with_print: bool = False) -> None:
     if not AGENTFS_FILE.exists():
-        logging.info("Loading all files in the current working directory to AgentFS")
+        if not with_print:
+            logging.info(
+                "Loading all files in the current working directory to AgentFS"
+            )
+        else:
+            print(
+                "Loading all files in the current working directory to AgentFS",
+                file=sys.stderr,
+            )
         await load_all_files(DEFAULT_TO_AVOID, DEFAULT_TO_AVOID_FILES, progress=True)
-        logging.info(
-            "Finished loading all files in the current working directory to AgentFS"
-        )
+        if not with_print:
+            logging.info(
+                "Finished loading all files in the current working directory to AgentFS"
+            )
+        else:
+            print(
+                "Finished loading all files in the current working directory to AgentFS",
+                file=sys.stderr,
+            )
     else:
-        logging.info(
-            f"Detected {str(AGENTFS_FILE)} in current working directory, will not load files."
-        )
+        if not with_print:
+            logging.info(
+                f"Detected {str(AGENTFS_FILE)} in current working directory, will not load files."
+            )
+        else:
+            print(
+                f"Detected {str(AGENTFS_FILE)} in current working directory, will not load files.",
+                file=sys.stderr,
+            )
 
 
 def _escape_markdow_for_tg(markdown: str) -> str: