From 9f34ba5d87718acbfe712020746481cee1e259eb Mon Sep 17 00:00:00 2001 From: "Clelia (Astra) Bertelli" Date: Mon, 9 Feb 2026 14:37:42 +0100 Subject: [PATCH] feat: support for more document formats for telegram upload --- packages/lobsterx/pyproject.toml | 2 +- packages/lobsterx/src/lobsterx/bot.py | 6 ++ packages/lobsterx/src/lobsterx/utils.py | 10 ++- packages/lobsterx/tests/conftest.py | 1 + packages/lobsterx/tests/test_utils.py | 110 ++++++++++++++++++++++-- uv.lock | 2 +- 6 files changed, 118 insertions(+), 13 deletions(-) diff --git a/packages/lobsterx/pyproject.toml b/packages/lobsterx/pyproject.toml index f17adb7..8920934 100644 --- a/packages/lobsterx/pyproject.toml +++ b/packages/lobsterx/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "uv_build" [project] name = "lobsterx" -version = "0.1.0-beta" +version = "0.1.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" diff --git a/packages/lobsterx/src/lobsterx/bot.py b/packages/lobsterx/src/lobsterx/bot.py index a68a946..d00454b 100644 --- a/packages/lobsterx/src/lobsterx/bot.py +++ b/packages/lobsterx/src/lobsterx/bot.py @@ -91,6 +91,12 @@ async def run_bot(log_level: str) -> None: application.add_handler( MessageHandler(filters.Document.PDF, handle_documents_tg) ) + application.add_handler( + MessageHandler(filters.Document.DOCX, handle_documents_tg) + ) + application.add_handler( + MessageHandler(filters.Document.TXT, handle_documents_tg) + ) application.add_handler(MessageHandler(filters.TEXT, handle_prompt_tg)) application.add_error_handler(cast(HandlerCallback, error_handler)) if application.updater is not None: diff --git a/packages/lobsterx/src/lobsterx/utils.py b/packages/lobsterx/src/lobsterx/utils.py index 7952213..0b0ab3a 100644 --- a/packages/lobsterx/src/lobsterx/utils.py +++ b/packages/lobsterx/src/lobsterx/utils.py @@ -1,6 +1,7 @@ import functools import json import logging +import mimetypes import os from typing import cast @@ -85,12 +86,15 @@ def get_workflow() -> AgentWorkflow: def _get_file_name(document: Document) -> str: + extension = ( + mimetypes.guess_extension(document.mime_type or "application/pdf") or ".pdf" + ) if document.file_name is None: - return generate_name() + ".pdf" + return generate_name() + extension else: - if document.file_name.endswith(".pdf"): + if document.file_name.endswith(extension): return document.file_name - return document.file_name + ".pdf" + return document.file_name + extension async def handle_documents(document: Document, context: CallbackContext) -> str: diff --git a/packages/lobsterx/tests/conftest.py b/packages/lobsterx/tests/conftest.py index e25f29a..4f61e7d 100644 --- a/packages/lobsterx/tests/conftest.py +++ b/packages/lobsterx/tests/conftest.py @@ -59,6 +59,7 @@ class TelegramUserMock: class TelegramDocumentMock: file_name: str | None file_id: str + mime_type: str | None class ParsingMock: diff --git a/packages/lobsterx/tests/test_utils.py b/packages/lobsterx/tests/test_utils.py index 8df5b88..478fbed 100644 --- a/packages/lobsterx/tests/test_utils.py +++ b/packages/lobsterx/tests/test_utils.py @@ -4,6 +4,13 @@ from unittest.mock import Mock, patch import pytest +from telegram import Document, User +from telegram.ext import CallbackContext +from workflows_acp.llm_wrapper import LLMWrapper +from workflows_acp.llms.openai_llm import OpenAILLM +from workflows_acp.tools.agentfs import load_all_files +from workflows_acp.workflow import AgentWorkflow + from lobsterx.constants import ( DATA_DIR, DEFAULT_TO_AVOID, @@ -23,12 +30,6 @@ handle_prompt, start, ) -from telegram import Document, User -from telegram.ext import CallbackContext -from workflows_acp.llm_wrapper import LLMWrapper -from workflows_acp.llms.openai_llm import OpenAILLM -from workflows_acp.tools.agentfs import load_all_files -from workflows_acp.workflow import AgentWorkflow from .conftest import ( AgentWorkflowMock, @@ -79,7 +80,10 @@ async def test_handle_documents_success( to_avoid_files=DEFAULT_TO_AVOID_FILES, ) document = cast( - Document, TelegramDocumentMock(file_name="hello.pdf", file_id="123") + Document, + TelegramDocumentMock( + file_name="hello.pdf", file_id="123", mime_type="application/pdf" + ), ) callback_context = cast(CallbackContext, TelegramCallBackContextMock()) result = await handle_documents(document, callback_context) @@ -90,6 +94,93 @@ async def test_handle_documents_success( ) +@pytest.mark.asyncio +async def test_handle_documents_no_name_mimetype( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + (tmp_path / "test.txt").write_text("Hello world") + await load_all_files( + to_avoid_dirs=DEFAULT_TO_AVOID, + to_avoid_files=DEFAULT_TO_AVOID_FILES, + ) + document = cast( + Document, + TelegramDocumentMock( + file_name=None, + file_id="123", + mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document", + ), + ) + callback_context = cast(CallbackContext, TelegramCallBackContextMock()) + result = await handle_documents(document, callback_context) + assert ( + result.startswith("Your file has been successfully downloaded at: ") + and result.endswith( + ". Use this path to reference the file in your follow-up requests to the agent" + ) + and ".docx" in result + ) + + +@pytest.mark.asyncio +async def test_handle_documents_no_name_no_mimetype( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + (tmp_path / "test.txt").write_text("Hello world") + await load_all_files( + to_avoid_dirs=DEFAULT_TO_AVOID, + to_avoid_files=DEFAULT_TO_AVOID_FILES, + ) + document = cast( + Document, + TelegramDocumentMock( + file_name=None, + file_id="123", + mime_type=None, + ), + ) + callback_context = cast(CallbackContext, TelegramCallBackContextMock()) + result = await handle_documents(document, callback_context) + assert ( + result.startswith("Your file has been successfully downloaded at: ") + and result.endswith( + ". Use this path to reference the file in your follow-up requests to the agent" + ) + and ".pdf" in result + ) + + +@pytest.mark.asyncio +async def test_handle_documents_no_extension( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + (tmp_path / "test.txt").write_text("Hello world") + await load_all_files( + to_avoid_dirs=DEFAULT_TO_AVOID, + to_avoid_files=DEFAULT_TO_AVOID_FILES, + ) + document = cast( + Document, + TelegramDocumentMock( + file_name="file", + file_id="123", + mime_type="text/plain", + ), + ) + callback_context = cast(CallbackContext, TelegramCallBackContextMock()) + result = await handle_documents(document, callback_context) + assert ( + result.startswith("Your file has been successfully downloaded at: ") + and result.endswith( + ". Use this path to reference the file in your follow-up requests to the agent" + ) + and "file.txt" in result + ) + + @pytest.mark.asyncio async def test_handle_documents_fail( tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -101,7 +192,10 @@ async def test_handle_documents_fail( to_avoid_files=DEFAULT_TO_AVOID_FILES, ) document = cast( - Document, TelegramDocumentMock(file_name="hello.pdf", file_id="123") + Document, + TelegramDocumentMock( + file_name="hello.pdf", file_id="123", mime_type="application/pdf" + ), ) callback_context = cast( CallbackContext, TelegramCallBackContextMock(should_fail=True) diff --git a/uv.lock b/uv.lock index e997382..35b34c6 100644 --- a/uv.lock +++ b/uv.lock @@ -838,7 +838,7 @@ wheels = [ [[package]] name = "lobsterx" -version = "0.1.0b0" +version = "0.1.1b0" source = { editable = "packages/lobsterx" } dependencies = [ { name = "aiofiles" },