fix: upload Thread attachments when sending messages (#1394)

## Threads
### Resolved Issues
- Fixed: After the first message, sending a thread message with file
search attachments may fail to send.

Author: ekassos

pingpong/ai.py

@@ -3356,6 +3356,39 @@ async def on_response_canceled(self) -> None:
         )
 
 
+async def poll_vector_store_files(
+    cli: openai.AsyncClient,
+    *,
+    vector_store_id: str,
+    file_ids: list[str],
+) -> None:
+    async def poll_single_file(file_id: str) -> bool:
+        try:
+            await cli.vector_stores.files.poll(
+                file_id=file_id,
+                vector_store_id=vector_store_id,
+            )
+            return True
+        except openai.NotFoundError:
+            return False
+
+    poll_results = await asyncio.gather(
+        *[poll_single_file(file_id) for file_id in file_ids]
+    )
+    missing_file_ids = [
+        file_id
+        for file_id, was_found in zip(file_ids, poll_results, strict=False)
+        if not was_found
+    ]
+    if missing_file_ids:
+        logger.warning(
+            "Skipping %s missing file(s) during vector store poll for vector store %s: %s",
+            len(missing_file_ids),
+            vector_store_id,
+            ", ".join(missing_file_ids),
+        )
+
+
 async def run_response(
     cli: openai.AsyncClient,
     *,
@@ -3463,14 +3496,10 @@ async def run_response(
                 if attached_file_search_file_ids:
                     if not thread_vector_store_id:
                         raise ValueError("Vector store ID is required for file search")
-                    await asyncio.gather(
-                        *[
-                            cli.vector_stores.files.poll(
-                                file_id=file_id,
-                                vector_store_id=thread_vector_store_id,
-                            )
-                            for file_id in attached_file_search_file_ids
-                        ]
+                    await poll_vector_store_files(
+                        cli,
+                        vector_store_id=thread_vector_store_id,
+                        file_ids=attached_file_search_file_ids,
                     )
                 if vector_store_ids:
                     tools.append(

pingpong/server.py

@@ -145,6 +145,7 @@
 from .state_types import AppState, StateRequest, StateWebSocket
 from .vector_stores import (
     add_vector_store_files_to_db,
+    append_vector_store_files,
     create_vector_store,
     delete_vector_store,
     delete_vector_store_db_returning_file_ids,
@@ -6267,11 +6268,19 @@ async def send_message(
         if data.file_search_file_ids:
             if thread.vector_store_id:
                 # Vector store already exists, update
-                await add_vector_store_files_to_db(
-                    request.state["db"],
-                    thread.vector_store_id,
-                    data.file_search_file_ids,
-                )
+                if thread.version == 3:
+                    await append_vector_store_files(
+                        request.state["db"],
+                        openai_client,
+                        thread.vector_store_id,
+                        data.file_search_file_ids,
+                    )
+                else:
+                    await add_vector_store_files_to_db(
+                        request.state["db"],
+                        thread.vector_store_id,
+                        data.file_search_file_ids,
+                    )
             else:
                 # Store doesn't exist, create a new one
                 # (empty, since we're adding files as attachments)

pingpong/test_ai_vector_store_polling.py

@@ -0,0 +1,52 @@
+import logging
+from unittest.mock import AsyncMock
+
+import httpx
+import openai
+import pytest
+
+from pingpong.ai import poll_vector_store_files
+
+pytestmark = pytest.mark.asyncio
+
+
+def _not_found_error(file_id: str, vector_store_id: str) -> openai.NotFoundError:
+    request = httpx.Request(
+        "GET",
+        f"https://api.openai.com/v1/vector_stores/{vector_store_id}/files/{file_id}",
+    )
+    response = httpx.Response(404, request=request)
+    return openai.NotFoundError(
+        f"No file found with id '{file_id}' in vector store '{vector_store_id}'.",
+        response=response,
+        body={"error": {"message": "not found"}},
+    )
+
+
+async def test_poll_vector_store_files_skips_missing_files(caplog):
+    cli = AsyncMock()
+
+    async def fake_poll(*, file_id: str, vector_store_id: str):
+        if file_id == "file-missing":
+            raise _not_found_error(file_id, vector_store_id)
+        return None
+
+    cli.vector_stores.files.poll = AsyncMock(side_effect=fake_poll)
+
+    with caplog.at_level(logging.WARNING):
+        await poll_vector_store_files(
+            cli, vector_store_id="vs-test", file_ids=["file-ok", "file-missing"]
+        )
+
+    assert cli.vector_stores.files.poll.await_count == 2
+    assert "file-missing" in caplog.text
+    assert "vs-test" in caplog.text
+
+
+async def test_poll_vector_store_files_noop_for_empty_file_list():
+    cli = AsyncMock()
+    cli.vector_stores.files.poll = AsyncMock()
+
+    await poll_vector_store_files(cli, vector_store_id="vs-test", file_ids=[])
+
+    assert cli.vector_stores.files.poll.await_count == 0