Merge pull request #24 from AstraBert/feat/more-documents-formats

feat: support for more document formats for telegram upload

Author: AstraBert

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"

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:

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:

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:

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)