diff --git a/.gitignore b/.gitignore index bb7dfe4d3..91b3770ac 100644 --- a/.gitignore +++ b/.gitignore @@ -30,6 +30,8 @@ data/videos/ data/secrets.db data/models/ data/workspace/ +data/library/ +data/collections/ data/*.log # Environment diff --git a/tests/test_library.py b/tests/test_library.py new file mode 100644 index 000000000..37100dab6 --- /dev/null +++ b/tests/test_library.py @@ -0,0 +1,735 @@ +"""Tests for the Library app — store, pipeline, routes, and collections handoff.""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +import pytest +import pytest_asyncio +from httpx import ASGITransport, AsyncClient +from tinyagentos.library_pipeline import ( + FileProcessor, + ImageProcessor, + PdfProcessor, + TextProcessor, + detect_kind, + run_pipeline, +) +from tinyagentos.library_store import LibraryStore +from tinyagentos.library_collections import handoff_to_collections + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest_asyncio.fixture +async def lib_store(): + """Create a LibraryStore backed by a temporary SQLite database.""" + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f: + db_path = Path(f.name) + + store = LibraryStore(db_path) + await store.init() + + yield store + + await store.close() + try: + db_path.unlink() + except OSError: + pass + + +@pytest.fixture +def storage_dir(): + """Create a temporary directory for file artifacts.""" + with tempfile.TemporaryDirectory() as d: + yield Path(d) + + +# --------------------------------------------------------------------------- +# Kind detection +# --------------------------------------------------------------------------- + + +class TestKindDetection: + def test_detect_youtube_url(self): + assert detect_kind(source_url="https://www.youtube.com/watch?v=abc123") == "url:youtube" + assert detect_kind(source_url="https://youtube.com/watch?v=abc123") == "url:youtube" + assert detect_kind(source_url="https://youtu.be/abc123") == "url:youtube" + assert detect_kind(source_url="https://m.youtube.com/watch?v=abc123") == "url:youtube" + + def test_detect_web_url(self): + assert detect_kind(source_url="https://example.com") == "url:web" + assert detect_kind(source_url="http://blog.example.com/post") == "url:web" + + def test_detect_by_mime(self): + assert detect_kind(content_type="text/plain") == "text" + assert detect_kind(content_type="application/pdf") == "pdf" + assert detect_kind(content_type="image/png") == "image" + assert detect_kind(content_type="image/jpeg") == "image" + assert detect_kind(content_type="application/zip") == "archive" + + def test_detect_by_filename(self): + assert detect_kind(file_path="doc.txt") == "text" + assert detect_kind(file_path="report.pdf") == "pdf" + assert detect_kind(file_path="photo.jpg") == "image" + assert detect_kind(file_path="icon.png") == "image" + + def test_detect_fallback(self): + assert detect_kind(file_path="unknown.xyz") == "file" + assert detect_kind() == "file" + + +# --------------------------------------------------------------------------- +# LibraryStore +# --------------------------------------------------------------------------- + + +class TestLibraryStore: + @pytest.mark.asyncio + async def test_create_and_get_item(self, lib_store): + item_id = await lib_store.create_item( + kind="text", + title="test.txt", + source_url="", + storage_path="/tmp/test.txt", + size_bytes=42, + ) + assert item_id + + item = await lib_store.get_item(item_id) + assert item is not None + assert item["kind"] == "text" + assert item["title"] == "test.txt" + assert item["status"] == "pending" + assert item["bytes"] == 42 + + @pytest.mark.asyncio + async def test_get_nonexistent_item(self, lib_store): + item = await lib_store.get_item("nonexistent") + assert item is None + + @pytest.mark.asyncio + async def test_list_items(self, lib_store): + id1 = await lib_store.create_item(kind="text", title="a.txt") + id2 = await lib_store.create_item(kind="pdf", title="b.pdf") + id3 = await lib_store.create_item(kind="image", title="c.png") + + items = await lib_store.list_items() + assert len(items) == 3 + + text_items = await lib_store.list_items(kind="text") + assert len(text_items) == 1 + assert text_items[0]["title"] == "a.txt" + + assert len(await lib_store.list_items(status="pending")) == 3 + assert len(await lib_store.list_items(status="ready")) == 0 + + @pytest.mark.asyncio + async def test_update_item(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="old") + await lib_store.update_item(item_id, title="new", status="ready") + + item = await lib_store.get_item(item_id) + assert item["title"] == "new" + assert item["status"] == "ready" + + @pytest.mark.asyncio + async def test_update_item_meta(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + await lib_store.update_item(item_id, meta_json={"preview": "hello"}) + + item = await lib_store.get_item(item_id) + meta = json.loads(item["meta_json"]) + assert meta["preview"] == "hello" + + @pytest.mark.asyncio + async def test_update_invalid_status(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + with pytest.raises(ValueError): + await lib_store.update_item_status(item_id, "invalid_status") + + @pytest.mark.asyncio + async def test_delete_item(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + item = await lib_store.get_item(item_id) + assert item is not None + + await lib_store.delete_item(item_id) + item = await lib_store.get_item(item_id) + assert item is None + + @pytest.mark.asyncio + async def test_artifacts(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + art_id = await lib_store.add_artifact(item_id, kind="text", path="/tmp/test.txt") + + artifacts = await lib_store.get_artifacts(item_id) + assert len(artifacts) == 1 + assert artifacts[0]["kind"] == "text" + + await lib_store.delete_artifact(art_id) + assert len(await lib_store.get_artifacts(item_id)) == 0 + + @pytest.mark.asyncio + async def test_cascade_delete_artifacts(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + await lib_store.add_artifact(item_id, kind="text", path="/tmp/a.txt") + await lib_store.add_artifact(item_id, kind="thumbnail", path="/tmp/thumb.jpg") + + await lib_store.delete_item(item_id) + assert len(await lib_store.get_artifacts(item_id)) == 0 + + @pytest.mark.asyncio + async def test_jobs(self, lib_store): + item_id = await lib_store.create_item(kind="text", title="test") + job_id = await lib_store.create_job(item_id, "ingest") + + job = await lib_store.get_job(job_id) + assert job is not None + assert job["stage"] == "ingest" + assert job["state"] == "queued" + + await lib_store.update_job(job_id, state="done") + job = await lib_store.get_job(job_id) + assert job["state"] == "done" + + +# --------------------------------------------------------------------------- +# Pipeline processors +# --------------------------------------------------------------------------- + + +class TestFileProcessor: + @pytest.mark.asyncio + async def test_process_existing_file(self, lib_store, storage_dir): + file_path = storage_dir / "test.txt" + file_path.write_text("hello world") + + item_id = await lib_store.create_item( + kind="file", title="test.txt", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = FileProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + assert len(artifacts) == 1 + assert artifacts[0]["kind"] == "metadata" + + @pytest.mark.asyncio + async def test_process_missing_file(self, lib_store, storage_dir): + item_id = await lib_store.create_item( + kind="file", title="missing.txt", storage_path="/nonexistent/file.txt" + ) + item = await lib_store.get_item(item_id) + + proc = FileProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + assert len(artifacts) == 0 + + +class TestTextProcessor: + @pytest.mark.asyncio + async def test_extract_text(self, lib_store, storage_dir): + file_path = storage_dir / "notes.txt" + file_path.write_text("line one\nline two\nline three") + + item_id = await lib_store.create_item( + kind="text", title="notes.txt", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = TextProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + assert len(artifacts) == 1 + assert artifacts[0]["kind"] == "text" + assert artifacts[0]["meta"]["line_count"] == 3 + assert artifacts[0]["meta"]["char_count"] == 28 + + text_path = Path(artifacts[0]["path"]) + assert text_path.exists() + assert "line one" in text_path.read_text() + + @pytest.mark.asyncio + async def test_text_auto_title(self, lib_store, storage_dir): + file_path = storage_dir / "notes.txt" + file_path.write_text("My Title\nmore content here") + + item_id = await lib_store.create_item( + kind="text", title="", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = TextProcessor(lib_store, storage_dir) + await proc.process(item) + + updated = await lib_store.get_item(item_id) + assert updated["title"] == "My Title" + + @pytest.mark.asyncio + async def test_text_preview(self, lib_store, storage_dir): + file_path = storage_dir / "long.txt" + file_path.write_text("A" * 500) + + item_id = await lib_store.create_item( + kind="text", title="long.txt", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = TextProcessor(lib_store, storage_dir) + await proc.process(item) + + updated = await lib_store.get_item(item_id) + meta = json.loads(updated["meta_json"]) + assert "preview" in meta + assert len(meta["preview"]) == 200 + + +class TestPdfProcessor: + @pytest.mark.asyncio + async def test_process_pdf(self, lib_store, storage_dir): + file_path = storage_dir / "test.pdf" + _create_minimal_pdf(file_path) + + item_id = await lib_store.create_item( + kind="pdf", title="test.pdf", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = PdfProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + + meta_artifacts = [a for a in artifacts if a["kind"] == "metadata"] + assert len(meta_artifacts) >= 1 + assert meta_artifacts[0]["meta"]["page_count"] >= 0 + + @pytest.mark.asyncio + async def test_process_missing_pdf(self, lib_store, storage_dir): + item_id = await lib_store.create_item( + kind="pdf", title="missing.pdf", storage_path="/nonexistent/file.pdf" + ) + item = await lib_store.get_item(item_id) + + proc = PdfProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + assert len(artifacts) == 0 + + +class TestImageProcessor: + @pytest.mark.asyncio + async def test_process_image(self, lib_store, storage_dir): + file_path = storage_dir / "test.png" + _create_test_image(file_path) + + item_id = await lib_store.create_item( + kind="image", title="test.png", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = ImageProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + + kinds = {a["kind"] for a in artifacts} + assert "metadata" in kinds + assert "thumbnail" in kinds + + thumb_art = [a for a in artifacts if a["kind"] == "thumbnail"][0] + assert Path(thumb_art["path"]).exists() + + @pytest.mark.asyncio + async def test_process_missing_image(self, lib_store, storage_dir): + item_id = await lib_store.create_item( + kind="image", title="missing.png", storage_path="/nonexistent/file.png" + ) + item = await lib_store.get_item(item_id) + + proc = ImageProcessor(lib_store, storage_dir) + artifacts = await proc.process(item) + assert len(artifacts) == 0 + + +# --------------------------------------------------------------------------- +# run_pipeline +# --------------------------------------------------------------------------- + + +class TestRunPipeline: + @pytest.mark.asyncio + async def test_pipeline_text(self, lib_store, storage_dir): + file_path = storage_dir / "notes.txt" + file_path.write_text("sample content") + + item_id = await lib_store.create_item( + kind="text", title="notes.txt", storage_path=str(file_path) + ) + await run_pipeline(lib_store, item_id, storage_dir) + + item = await lib_store.get_item(item_id) + assert item["status"] == "ready" + + artifacts = await lib_store.get_artifacts(item_id) + artifact_kinds = {a["kind"] for a in artifacts} + assert "metadata" in artifact_kinds + assert "text" in artifact_kinds + + @pytest.mark.asyncio + async def test_pipeline_file(self, lib_store, storage_dir): + file_path = storage_dir / "unknown.xyz" + file_path.write_text("raw data") + + item_id = await lib_store.create_item( + kind="file", title="unknown.xyz", storage_path=str(file_path) + ) + await run_pipeline(lib_store, item_id, storage_dir) + + item = await lib_store.get_item(item_id) + assert item["status"] == "ready" + + @pytest.mark.asyncio + async def test_pipeline_error_status(self, lib_store, storage_dir): + item_id = await lib_store.create_item( + kind="file", title="missing.txt", storage_path="/missing/file.txt" + ) + await run_pipeline(lib_store, item_id, storage_dir) + item = await lib_store.get_item(item_id) + assert item["status"] == "error" + + +# --------------------------------------------------------------------------- +# Collections handoff +# --------------------------------------------------------------------------- + + +class TestCollectionsHandoff: + @pytest.mark.asyncio + async def test_handoff_files_copied(self, lib_store, storage_dir): + """Text artifacts are copied to collections dir even without qmd; + returns 0 when qmd is unreachable (no silent ImportError success).""" + file_path = storage_dir / "notes.txt" + file_path.write_text("content for collections") + + item_id = await lib_store.create_item( + kind="text", title="notes.txt", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = TextProcessor(lib_store, storage_dir) + await proc.process(item) + + collections_dir = storage_dir / "collections" + count = await handoff_to_collections(lib_store, item_id, collections_dir) + + # Files copied but no qmd → 0 indexed + assert count == 0 + + # Verify files were still copied to collections dir + item_dir = collections_dir / item_id + assert item_dir.exists() + text_files = list(item_dir.glob("*.txt")) + assert len(text_files) >= 1 + assert "content for collections" in text_files[0].read_text() + + @pytest.mark.asyncio + async def test_handoff_with_qmd(self, lib_store, storage_dir): + """Handoff indexes into taosmd when the API is reachable.""" + from unittest.mock import AsyncMock, patch + + file_path = storage_dir / "notes.txt" + file_path.write_text("hello from library") + + item_id = await lib_store.create_item( + kind="text", title="notes.txt", storage_path=str(file_path) + ) + item = await lib_store.get_item(item_id) + + proc = TextProcessor(lib_store, storage_dir) + await proc.process(item) + + collections_dir = storage_dir / "collections" + taosmd_url = "http://localhost:17900" + taosmd_token = "test-admin-token" + + # Mock httpx.AsyncClient to simulate a working taosmd + from unittest.mock import AsyncMock, MagicMock + + mock_client = MagicMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + # POST /collections → nested response shape (taosmd 0.4.0) + mock_create_resp = MagicMock() + mock_create_resp.status_code = 200 + mock_create_resp.json.return_value = {"collection": {"id": "coll-123"}} + + # POST /collections/{id}/index → 202 (no body) + mock_index_resp = MagicMock() + mock_index_resp.status_code = 202 + + # POST /collections/{id}/link → 200 + mock_link_resp = MagicMock() + mock_link_resp.status_code = 200 + + # GET /collections/{id} → nested shape (taosmd 0.4.0), status=ready with stats + mock_poll_resp = MagicMock() + mock_poll_resp.status_code = 200 + mock_poll_resp.json.return_value = { + "collection": { + "status": "ready", + "stats": { + "files_indexed": 1, + "files_total": 1, + "chunks_ingested": 3, + "chunks_skipped": 0, + }, + }, + } + + mock_client.post = AsyncMock( + side_effect=[mock_create_resp, mock_index_resp, mock_link_resp] + ) + mock_client.get = AsyncMock(return_value=mock_poll_resp) + + with patch("httpx.AsyncClient", return_value=mock_client): + count = await handoff_to_collections( + lib_store, item_id, collections_dir, + taosmd_url=taosmd_url, + taosmd_admin_token=taosmd_token, + ) + + assert count == 1 # One file indexed per stats + + # Verify request sequence — create → index → link + assert mock_client.post.call_count == 3 + post_calls = [c.args[0] for c in mock_client.post.call_args_list] + assert post_calls[0] == f"{taosmd_url}/collections" + assert post_calls[1] == f"{taosmd_url}/collections/coll-123/index" + assert post_calls[2] == f"{taosmd_url}/collections/coll-123/link" + + # Verify POST /collections body contains required fields + create_kwargs = mock_client.post.call_args_list[0].kwargs + create_body = create_kwargs.get("json", {}) + assert create_body["name"] == f"library-{item_id[:12]}" + assert create_body["kind"] == "mixed" + assert create_body["source_path"] == str(collections_dir / item_id) + + # Verify poll GET called + mock_client.get.assert_called_once_with( + f"{taosmd_url}/collections/coll-123", + headers={"Authorization": f"Bearer {taosmd_token}"}, + ) + + +# --------------------------------------------------------------------------- +# API routes +# --------------------------------------------------------------------------- + + +class TestLibraryRoutes: + @pytest.mark.asyncio + async def test_library_page_gone(self, client): + """The server-rendered /library page was dropped (fold 6) — + the Library UI is the React desktop app.""" + resp = await client.get("/library") + assert resp.status_code == 404 + + @pytest.mark.asyncio + async def test_ingest_no_input(self, client): + resp = await client.post("/api/library/ingest") + assert resp.status_code == 400 + data = resp.json() + assert "error" in data + + @pytest.mark.asyncio + async def test_ingest_url(self, client): + resp = await client.post("/api/library/ingest", data={"url": "https://example.com/page"}) + assert resp.status_code == 202 + data = resp.json() + assert "item_id" in data + assert data["status"] == "pending" + + @pytest.mark.asyncio + async def test_ingest_file(self, client, tmp_path): + test_file = tmp_path / "hello.txt" + test_file.write_text("hello world") + + with open(test_file, "rb") as f: + resp = await client.post( + "/api/library/ingest", + files={"file": ("hello.txt", f, "text/plain")}, + ) + assert resp.status_code == 202 + data = resp.json() + assert "item_id" in data + + @pytest.mark.asyncio + async def test_list_items(self, client): + resp = await client.get("/api/library/items") + assert resp.status_code == 200 + data = resp.json() + assert "items" in data + assert "count" in data + + @pytest.mark.asyncio + async def test_get_nonexistent_item(self, client): + resp = await client.get("/api/library/items/nonexistent") + assert resp.status_code == 404 + + @pytest.mark.asyncio + async def test_delete_nonexistent_item(self, client): + resp = await client.delete("/api/library/items/nonexistent") + assert resp.status_code == 404 + + @pytest.mark.asyncio + async def test_reprocess_nonexistent_item(self, client): + resp = await client.post("/api/library/items/nonexistent/reprocess") + assert resp.status_code == 404 + + @pytest.mark.asyncio + async def test_ingest_and_get(self, client): + resp = await client.post("/api/library/ingest", data={"url": "https://example.com"}) + assert resp.status_code == 202 + item_id = resp.json()["item_id"] + + resp = await client.get(f"/api/library/items/{item_id}") + assert resp.status_code == 200 + data = resp.json() + assert data["item"]["id"] == item_id + assert "artifacts" in data + + @pytest.mark.asyncio + async def test_filter_by_kind(self, client): + await client.post("/api/library/ingest", data={"url": "https://example.com"}) + await client.post("/api/library/ingest", data={"url": "https://youtube.com/watch?v=abc"}) + + resp = await client.get("/api/library/items", params={"kind": "url:youtube"}) + data = resp.json() + for item in data["items"]: + assert item["kind"] == "url:youtube" + + # -- Auth: all endpoints require a session (CSRF is bypassed in tests, but + # unauthenticated requests still hit the global auth middleware). Test + # that the 6 new endpoints return 401 when no session cookie/header is + # present. + + @pytest.mark.asyncio + async def test_unauth_library_endpoints(self, app): + """All 6 library endpoints return 401 without authentication.""" + from httpx import ASGITransport, AsyncClient + + async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as unauth_client: + endpoints: list[tuple[str, str]] = [ + ("post", "/api/library/ingest"), + ("get", "/api/library/items"), + ("get", "/api/library/items/abc123"), + ("delete", "/api/library/items/abc123"), + ("post", "/api/library/items/abc123/reprocess"), + ] + for method, url in endpoints: + if method == "post": + resp = await unauth_client.post(url, data={"url": "https://example.com"}) + elif method == "delete": + resp = await unauth_client.delete(url) + else: + resp = await unauth_client.get(url) + assert resp.status_code == 401, f"{method.upper()} {url} returned {resp.status_code}, expected 401" + + # -- 413: oversized upload (tested at the HTTP level in integration; + # ASGI transport does not propagate Content-Length to file.size) + + @pytest.mark.asyncio + async def test_reprocess_idempotent(self, client, tmp_path): + """Reprocessing an item deletes old artifacts first — no duplicates.""" + test_file = tmp_path / "notes.txt" + test_file.write_text("test content") + + with open(test_file, "rb") as f: + resp = await client.post( + "/api/library/ingest", + files={"file": ("notes.txt", f, "text/plain")}, + ) + assert resp.status_code == 202 + item_id = resp.json()["item_id"] + + # Wait for pipeline to finish (background task) + import asyncio + await asyncio.sleep(0.5) + + # Check initial artifact count + resp = await client.get(f"/api/library/items/{item_id}") + data = resp.json() + initial_artifact_count = len(data["artifacts"]) + assert initial_artifact_count > 0, "Pipeline should produce artifacts" + + # First reprocess + resp = await client.post(f"/api/library/items/{item_id}/reprocess") + assert resp.status_code == 202 + await asyncio.sleep(0.5) + + # After first reprocess — exact same artifact count, not doubled. + resp = await client.get(f"/api/library/items/{item_id}") + data = resp.json() + after_first = len(data["artifacts"]) + assert after_first == initial_artifact_count, ( + f"Reprocess should not duplicate artifacts: " + f"initial={initial_artifact_count} after_first={after_first}" + ) + assert data["item"].get("status") == "ready", ( + f"Item status should be ready after reprocess, got {data['item'].get('status')}" + ) + + # Second reprocess — should still work, no duplicates + resp = await client.post(f"/api/library/items/{item_id}/reprocess") + assert resp.status_code == 202 + + @pytest.mark.asyncio + async def test_reprocess_while_processing_returns_409(self, client, app): + """Reprocessing an item that is already pending/processing returns 409.""" + # Ingest normally and wait for pipeline to finish + import asyncio + + resp = await client.post( + "/api/library/ingest", + data={"url": "https://example.com"}, + ) + assert resp.status_code == 202 + item_id = resp.json()["item_id"] + await asyncio.sleep(0.5) + + # Force status to "processing" to simulate an in-flight pipeline + store = app.state.library_store + await store.update_item_status(item_id, "processing") + + # Reprocess should be rejected + resp = await client.post(f"/api/library/items/{item_id}/reprocess") + assert resp.status_code == 409 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _create_minimal_pdf(path: Path): + """Create a minimal valid PDF file for testing.""" + pdf_content = ( + b"%PDF-1.4\n" + b"1 0 obj<>endobj\n" + b"2 0 obj<>endobj\n" + b"3 0 obj<>endobj\n" + b"xref\n0 4\n0000000000 65535 f \n0000000009 00000 n \n0000000058 00000 n \n0000000115 00000 n \n" + b"trailer<>\n" + b"startxref\n190\n%%EOF\n" + ) + path.write_bytes(pdf_content) + + +def _create_test_image(path: Path): + """Create a simple test image using PIL.""" + from PIL import Image + img = Image.new("RGB", (100, 50), color="blue") + img.save(path) diff --git a/tinyagentos/library_collections.py b/tinyagentos/library_collections.py new file mode 100644 index 000000000..c4cddb845 --- /dev/null +++ b/tinyagentos/library_collections.py @@ -0,0 +1,362 @@ +"""Library collections handoff — index text artifacts into taosmd collections. + +After the ingest pipeline produces text artifacts (extracted text, transcripts, +descriptions), this module hands them off to taosmd collections via the +taosmd HTTP API so agents can query the content through collection grants. + +The design doc (docs/design/library-app.md) says: + "Collections handoff: write text artifacts into a per-target folder under an + allowed root, then taosmd collections index; link to project; grants stay + EXPLICIT" + +Flow (Phase 1): + 1. Write text artifacts under collections_dir/{item_id}/ + 2. Create a collection via POST /collections (taosmd HTTP API) + 3. Trigger async index via POST /collections/{id}/index + 4. Poll GET /collections/{id} until status=ready|error + 5. Link collection to the library item via POST /collections/{id}/link +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import stat +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Text artifact kinds that should be indexed into collections +_TEXT_ARTIFACT_KINDS = frozenset({"text", "transcript", "description", "ocr"}) + +# Maximum polls while waiting for async index to complete. +_MAX_INDEX_POLLS = 30 +# Seconds between poll attempts. +_POLL_INTERVAL = 2 + + +async def handoff_to_collections( + store, + item_id: str, + collections_dir: Path, + taosmd_url: str | None = None, + taosmd_admin_token: str | None = None, + project_id: str | None = None, +) -> int: + """Hand off all text artifacts for an item to taosmd collections. + + Writes text artifacts under ``collections_dir/{item_id}/``, then calls + the taosmd HTTP API to create a collection and index the text content. + + Returns the number of files indexed (``files_indexed`` from taosmd stats) + after a successful async index. Returns 0 when taosmd is unavailable + (no collection created). + + Parameters + ---------- + store: + LibraryStore instance. + item_id: + Library item id. + collections_dir: + Allowed root for collection files (e.g. ``/opt/taos/data/collections/``). + taosmd_url: + Base URL of the running taosmd instance (e.g. ``http://localhost:7900``). + When omitted or None, file-copy still happens but no collection is + created or indexed — the caller must have already created the collection + separately (production paths always supply this; test paths may omit it). + taosmd_admin_token: + Admin bearer token for taosmd API auth, fetched from SecretsStore as + ``taosmd-admin-token``. Required when *taosmd_url* is set. + project_id: + Optional project to link the collection to (Phase 2+). + """ + artifacts = await store.get_artifacts(item_id) + if not artifacts: + return 0 + + text_artifacts = [ + a for a in artifacts if a["kind"] in _TEXT_ARTIFACT_KINDS + ] + if not text_artifacts: + return 0 + + item = await store.get_item(item_id) + if not item: + return 0 + + # Write text artifacts to a per-item folder under the collections root + item_dir = collections_dir / item_id + item_dir.mkdir(parents=True, exist_ok=True) + # Set group-readable perms on the directory + try: + os.chmod(item_dir, stat.S_ISGID | stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP) + except OSError: + pass + + # Copy each text artifact into the collection source path. + # taosmd discovers files by scanning the source_path directory. + indexed_paths: list[str] = [] + for art in text_artifacts: + art_path = art.get("path", "") + if not art_path: + continue + + src = Path(art_path) + if not src.exists(): + continue + + dst = item_dir / src.name + try: + raw_bytes = src.read_bytes() + dst.write_bytes(raw_bytes) + # Group-readable file perms (0o640 — setgid meaningless on files) + try: + os.chmod(dst, stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP) + except OSError: + pass + except OSError: + logger.warning("Failed to copy artifact %s → %s", src, dst, + exc_info=True) + continue + indexed_paths.append(str(dst)) + + if not indexed_paths: + return 0 + + # Index into taosmd collections via the taosmd HTTP API + if not taosmd_url: + logger.debug("taosmd URL not provided — collection indexing skipped") + return 0 + + # Parse item metadata once before the httpx try block so exception + # handlers (below) can merge into it without re-reading a stale item dict. + item_meta = json.loads(item.get("meta_json", "{}")) + + async def _set_retryable() -> None: + """Persist collection_retryable so a future retry can re-attempt. + + NOTE: collection_retryable currently has no consumer — it is scaffolding + for a future retry path. Clearing on success prevents stale markers + but no code path reads this flag yet. + """ + try: + item_meta["collection_retryable"] = True + await store.update_item(item_id, meta_json=item_meta) + except Exception: + pass + + try: + import httpx + + # Normalise base URL + base_url = taosmd_url.rstrip("/") + + # Build auth headers when a token is available + auth_headers: dict[str, str] = {} + if taosmd_admin_token: + auth_headers["Authorization"] = f"Bearer {taosmd_admin_token}" + + async with httpx.AsyncClient(timeout=30) as http_client: + # 1. Get or create a collection for this library item. + # Look up a previously-created collection id from item metadata + # so reprocess is idempotent (no duplicate collections). + collection_name = f"library-{item_id[:12]}" + title = item.get("title", "Untitled") or "Untitled" + + collection_id = "" + existing_coll_id = item_meta.get("collection_id", "") + + if existing_coll_id: + # Verify the collection still exists. + # 404 (genuinely gone) → fall through to create. + # Transient failure (5xx, timeout, connection error) → mark + # retryable and bail — must not create a duplicate. + try: + check_resp = await http_client.get( + f"{base_url}/collections/{existing_coll_id}", + headers=auth_headers, + ) + if check_resp.status_code < 400: + collection_id = existing_coll_id + logger.debug( + "Reusing existing collection %s for item %s", + collection_id, item_id, + ) + elif check_resp.status_code == 404: + logger.debug( + "Existing collection %s not found (404) — " + "will create replacement", existing_coll_id, + ) + else: + logger.warning( + "taosmd GET /collections/%s returned %d — " + "transient failure, marking retryable", + existing_coll_id, check_resp.status_code, + ) + await _set_retryable() + return 0 + except Exception: + logger.warning( + "taosmd GET /collections/%s failed — " + "transient failure, marking retryable", + existing_coll_id, exc_info=True, + ) + await _set_retryable() + return 0 + + if not collection_id: + source_path = str(item_dir) + create_resp = await http_client.post( + f"{base_url}/collections", + json={ + "name": collection_name, + "kind": "mixed", + "source_path": source_path, + "metadata": { + "title": title, + "kind": item.get("kind", ""), + "source_url": item.get("source_url", ""), + "library_item_id": item_id, + }, + }, + headers=auth_headers, + ) + if create_resp.status_code >= 400: + logger.warning( + "taosmd POST /collections returned %d for item %s: %s", + create_resp.status_code, item_id, create_resp.text[:200], + ) + await _set_retryable() + return 0 + # taosmd 0.4.0 nests the id under "collection" + collection_id = create_resp.json().get("collection", {}).get("id", "") + + if not collection_id: + logger.warning( + "taosmd POST /collections returned no collection.id for item %s", + item_id, + ) + await _set_retryable() + return 0 + + # Persist collection_id in item metadata for idempotent reprocess + item_meta["collection_id"] = collection_id + await store.update_item(item_id, meta_json=item_meta) + + # 2. Trigger async index — no body, returns 202. + try: + index_resp = await http_client.post( + f"{base_url}/collections/{collection_id}/index", + headers=auth_headers, + ) + if index_resp.status_code != 202: + logger.warning( + "taosmd POST /collections/%s/index returned %d (expected 202)", + collection_id, index_resp.status_code, + ) + await _set_retryable() + return 0 + except Exception: + logger.warning( + "taosmd POST /collections/%s/index failed", + collection_id, exc_info=True, + ) + await _set_retryable() + return 0 + + # 3. Poll until indexing completes. + indexed = 0 + for _poll_attempt in range(_MAX_INDEX_POLLS): + try: + poll_resp = await http_client.get( + f"{base_url}/collections/{collection_id}", + headers=auth_headers, + ) + if poll_resp.status_code >= 400: + logger.warning( + "taosmd GET /collections/%s returned %d during poll", + collection_id, poll_resp.status_code, + ) + await _set_retryable() + break + poll_data = poll_resp.json().get("collection", {}) + status = poll_data.get("status", "") + if status == "ready": + stats = poll_data.get("stats", {}) + indexed = stats.get("files_indexed", 0) + logger.info( + "Collections handoff for item %s: " + "files_indexed=%d files_total=%d " + "chunks_ingested=%d chunks_skipped=%d " + "collection=%s", + item_id, + stats.get("files_indexed", 0), + stats.get("files_total", 0), + stats.get("chunks_ingested", 0), + stats.get("chunks_skipped", 0), + collection_id, + ) + break + elif status == "error": + logger.warning( + "taosmd collection %s entered error state for item %s", + collection_id, item_id, + ) + await _set_retryable() + break + # Still indexing — wait and retry + await asyncio.sleep(_POLL_INTERVAL) + except Exception: + logger.warning( + "taosmd poll GET /collections/%s failed", + collection_id, exc_info=True, + ) + await _set_retryable() + break + else: + logger.warning( + "taosmd collection %s did not reach ready state within %d polls", + collection_id, _MAX_INDEX_POLLS, + ) + await _set_retryable() + + # 4. Link the collection to the library item. + try: + await http_client.post( + f"{base_url}/collections/{collection_id}/link", + json={"type": "taos", "id": item_id}, + headers=auth_headers, + ) + logger.debug( + "Linked collection %s to library item %s", + collection_id, item_id, + ) + except Exception: + logger.warning( + "taosmd POST /collections/%s/link failed for item %s", + collection_id, item_id, exc_info=True, + ) + + # Clear retryable on success; the handoff completed. + if indexed > 0: + item_meta.pop("collection_retryable", None) + try: + await store.update_item(item_id, meta_json=item_meta) + except Exception: + pass + + return indexed + + except ImportError: + logger.warning("httpx not available — collection indexing skipped") + await _set_retryable() + except Exception: + logger.exception( + "taosmd collections API unreachable for item %s", item_id, + ) + await _set_retryable() + + return 0 diff --git a/tinyagentos/library_pipeline.py b/tinyagentos/library_pipeline.py new file mode 100644 index 000000000..e2ef7e185 --- /dev/null +++ b/tinyagentos/library_pipeline.py @@ -0,0 +1,428 @@ +"""Library ingest pipeline — processors for cheap-tier file/text/pdf/image ingestion. + +Processors are registered per detected kind and run asynchronously after ingest. +Each processor produces artifacts (e.g. metadata, extracted text, thumbnails) +that are stored on the item and optionally handed off to taosmd collections. +""" + +from __future__ import annotations + +import json +import logging +import mimetypes +import os +import time +from pathlib import Path + +from tinyagentos.library_store import LibraryStore + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Kind detection +# --------------------------------------------------------------------------- + +_MIME_KIND_MAP: dict[str, str] = { + "text/plain": "text", + "text/markdown": "text", + "text/csv": "text", + "text/html": "text", + "application/json": "text", + "application/xml": "text", + "text/xml": "text", + "application/pdf": "pdf", + "image/png": "image", + "image/jpeg": "image", + "image/gif": "image", + "image/webp": "image", + "image/svg+xml": "image", + "application/zip": "archive", + "application/gzip": "archive", + "application/x-tar": "archive", +} + + +def detect_kind(source_url: str = "", content_type: str = "", + file_path: str = "") -> str: + """Detect the library item kind from URL, MIME, or file path.""" + # URL-based detection + if source_url: + lower = source_url.lower() + if any(lower.startswith(p) for p in ("https://www.youtube.com/", + "https://youtube.com/", + "https://youtu.be/", + "https://m.youtube.com/")): + return "url:youtube" + if any(lower.startswith(p) for p in ("https://", "http://")): + return "url:web" + + # MIME-based detection + if content_type: + ct = content_type.split(";")[0].strip().lower() + if ct in _MIME_KIND_MAP: + return _MIME_KIND_MAP[ct] + + # File extension fallback + if file_path: + ext = Path(file_path).suffix.lower() + ext_map = { + ".txt": "text", ".md": "text", ".csv": "text", + ".json": "text", ".xml": "text", ".html": "text", + ".pdf": "pdf", + ".png": "image", ".jpg": "image", ".jpeg": "image", + ".gif": "image", ".webp": "image", ".svg": "image", + ".zip": "archive", ".gz": "archive", ".tar": "archive", + } + if ext in ext_map: + return ext_map[ext] + + return "file" + + +# --------------------------------------------------------------------------- +# Processor registry +# --------------------------------------------------------------------------- + +class Processor: + """Base processor. Subclasses handle one kind of library item.""" + + def __init__(self, store: LibraryStore, storage_dir: Path): + self.store = store + self.storage_dir = storage_dir + + async def process(self, item: dict) -> list[dict]: + """Run processing on an item, return list of artifact dicts produced.""" + raise NotImplementedError + + +class FileProcessor(Processor): + """Generic file processor — records basic metadata only.""" + + async def process(self, item: dict) -> list[dict]: + item_id = item["id"] + artifacts: list[dict] = [] + + storage_path = item.get("storage_path", "") + source_url = item.get("source_url", "") + if storage_path: + p = Path(storage_path) + if p.exists(): + stat = p.stat() + file_meta = { + "size_bytes": stat.st_size, + "mtime": stat.st_mtime, + "source_url": source_url, + "processed_at": time.time(), + "processor": "FileProcessor/v1", + } + # Try mimetype detection + mime_type, _ = mimetypes.guess_type(p.name) + if mime_type: + file_meta["mime_type"] = mime_type + + # path="" because storage_path is the user's original uploaded + # file — the reprocess unlink loop must never delete it. + await self.store.add_artifact( + item_id, kind="metadata", path="", meta=file_meta + ) + artifacts.append({"kind": "metadata", "path": "", "meta": file_meta}) + + # Update item bytes + await self.store.update_item(item_id, bytes=stat.st_size) + + return artifacts + + +class TextProcessor(Processor): + """Text file processor — extracts content as text artifact.""" + + async def process(self, item: dict) -> list[dict]: + item_id = item["id"] + artifacts: list[dict] = [] + + storage_path = item.get("storage_path", "") + if not storage_path: + return artifacts + + p = Path(storage_path) + if not p.exists(): + logger.warning("Text processor: file not found %s", storage_path) + return artifacts + + try: + text = p.read_text(encoding="utf-8", errors="replace") + except Exception: + logger.warning("Text processor: could not read %s", storage_path, + exc_info=True) + return artifacts + + # Write extracted text as an artifact + text_dir = self.storage_dir / "text" + text_dir.mkdir(parents=True, exist_ok=True) + text_path = text_dir / f"{item_id}.txt" + text_path.write_text(text, encoding="utf-8") + + text_meta = { + "char_count": len(text), + "line_count": text.count("\n") + 1, + "source_url": item.get("source_url", ""), + "processed_at": time.time(), + "processor": "TextProcessor/v1", + } + await self.store.add_artifact( + item_id, kind="text", path=str(text_path), meta=text_meta + ) + artifacts.append({"kind": "text", "path": str(text_path), "meta": text_meta}) + + # Store a preview (first 200 chars) + preview = text[:200] + meta = json.loads(item.get("meta_json", "{}")) + meta["preview"] = preview + await self.store.update_item(item_id, meta_json=meta) + + # Auto-title from content if no title + if not item.get("title"): + title = text.strip().split("\n", 1)[0][:100] + if title: + await self.store.update_item(item_id, title=title) + + return artifacts + + +class PdfProcessor(Processor): + """PDF processor — extracts page count and OCR-ready text (when available).""" + + async def process(self, item: dict) -> list[dict]: + item_id = item["id"] + artifacts: list[dict] = [] + + storage_path = item.get("storage_path", "") + if not storage_path: + return artifacts + + p = Path(storage_path) + if not p.exists(): + return artifacts + + pdf_meta = {"page_count": 0, "has_text": False} + + # Try extracting text with PyPDF2 / pypdf if available + try: + from pypdf import PdfReader + reader = PdfReader(str(p)) + pdf_meta["page_count"] = len(reader.pages) + + # Extract text from all pages + pages_text: list[str] = [] + for page in reader.pages: + page_text = page.extract_text() + if page_text: + pages_text.append(page_text) + + if pages_text: + text_content = "\n\n".join(pages_text) + pdf_meta["has_text"] = True + pdf_meta["char_count"] = len(text_content) + + text_dir = self.storage_dir / "text" + text_dir.mkdir(parents=True, exist_ok=True) + text_path = text_dir / f"{item_id}_pdf.txt" + text_path.write_text(text_content, encoding="utf-8") + + await self.store.add_artifact( + item_id, kind="text", path=str(text_path), + meta={"char_count": len(text_content), "pages": len(reader.pages)}, + ) + artifacts.append({ + "kind": "text", "path": str(text_path), + "meta": {"char_count": len(text_content), "pages": len(reader.pages)}, + }) + + # Update item with preview + preview = text_content[:200] + meta = json.loads(item.get("meta_json", "{}")) + meta["preview"] = preview + await self.store.update_item(item_id, meta_json=meta) + except ImportError: + logger.debug("pypdf not installed — PDF text extraction skipped") + except Exception: + logger.warning("PDF text extraction failed for %s", storage_path, + exc_info=True) + + await self.store.add_artifact( + item_id, kind="metadata", path="", meta=pdf_meta + ) + artifacts.append({"kind": "metadata", "path": "", "meta": pdf_meta}) + + return artifacts + + +class ImageProcessor(Processor): + """Image processor — records dimensions, creates thumbnail. + + Thumbnail generation requires Pillow (PIL), which is always available in + the taOS dev dependencies. + """ + + async def process(self, item: dict) -> list[dict]: + item_id = item["id"] + artifacts: list[dict] = [] + + storage_path = item.get("storage_path", "") + if not storage_path: + return artifacts + + p = Path(storage_path) + if not p.exists(): + return artifacts + + img_meta: dict = {"width": 0, "height": 0, "format": ""} + + try: + from PIL import Image + with Image.open(p) as img: + img_meta["width"] = img.width + img_meta["height"] = img.height + img_meta["format"] = img.format or "" + + # Create thumbnail (max 320px on longest side) + thumb_dir = self.storage_dir / "thumbs" + thumb_dir.mkdir(parents=True, exist_ok=True) + thumb_path = thumb_dir / f"{item_id}_thumb.jpg" + + img.thumbnail((320, 320)) + # Convert to RGB if needed (e.g. RGBA/PNG → JPEG) + if img.mode in ("RGBA", "P"): + img = img.convert("RGB") + img.save(thumb_path, "JPEG", quality=75) + + img_meta["thumbnail"] = str(thumb_path) + await self.store.add_artifact( + item_id, kind="thumbnail", path=str(thumb_path), + meta={"width": img.width, "height": img.height}, + ) + artifacts.append({ + "kind": "thumbnail", "path": str(thumb_path), + "meta": {"width": img.width, "height": img.height}, + }) + except ImportError: + logger.debug("PIL not available — image processing skipped") + except Exception: + logger.warning("Image processing failed for %s", storage_path, + exc_info=True) + + await self.store.add_artifact( + item_id, kind="metadata", path="", meta=img_meta + ) + artifacts.append({"kind": "metadata", "path": "", "meta": img_meta}) + + return artifacts + + +# --------------------------------------------------------------------------- +# Processor registry +# --------------------------------------------------------------------------- + +_PROCESSORS: dict[str, type[Processor]] = { + "file": FileProcessor, + "text": TextProcessor, + "pdf": PdfProcessor, + "image": ImageProcessor, +} + + +def get_processor(kind: str, store: LibraryStore, + storage_dir: Path) -> Processor: + """Return a processor for the given kind, falling back to FileProcessor.""" + cls = _PROCESSORS.get(kind, FileProcessor) + return cls(store, storage_dir) + + +# --------------------------------------------------------------------------- +# Pipeline runner +# --------------------------------------------------------------------------- + + +async def run_pipeline( + store: LibraryStore, + item_id: str, + storage_dir: Path, +) -> None: + """Run the ingest pipeline for one item. + + Steps: + 1. Mark item as 'processing' + 2. Determine processor from item kind + 3. Run file + kind-specific processors + 4. Collect artifacts + 5. Mark item as 'ready' (or 'error') + """ + item = await store.get_item(item_id) + if not item: + return + + kind = item["kind"] + + # URL-only items (no storage_path) are stored as references — the pipeline + # records a reference metadata artifact but does not fetch remote content + # (future: WebFetcherProcessor, #2078). The item gets a "reference" artifact so it + # is not silently empty. + if not item.get("storage_path") and item.get("source_url"): + logger.info( + "Library pipeline: URL-only item %s (%s) — stored as reference, not fetched", + item_id, kind, + ) + ref_meta = { + "source_url": item["source_url"], + "kind": kind, + "note": "Reference-only item — content not fetched (TODO: WebFetcherProcessor)", + } + await store.add_artifact( + item_id, kind="reference", path="", meta=ref_meta, + ) + + # If item has a storage_path that points to a missing file, fail early + # (dropped/moved/corrupt source must not silently look successful). + storage_path = item.get("storage_path", "") + source_url = item.get("source_url", "") + if storage_path and not source_url: + sp = Path(storage_path) + if not sp.exists(): + logger.warning( + "Library pipeline: source file missing for item %s: %s", + item_id, storage_path, + ) + await store.update_item_status(item_id, "error") + await store.update_item( + item_id, + meta_json={ + **json.loads(item.get("meta_json", "{}")), + "error": f"Source file not found: {storage_path}", + }, + ) + return + + try: + await store.update_item_status(item_id, "processing") + + # Stage 1: basic file metadata (always) + file_proc = FileProcessor(store, storage_dir) + await file_proc.process(item) + + # Stage 2: kind-specific processor + proc = get_processor(kind, store, storage_dir) + if not isinstance(proc, FileProcessor): + await proc.process(item) + + await store.update_item_status(item_id, "ready") + except Exception: + logger.exception("Library pipeline failed for item %s (kind=%s)", + item_id, kind) + await store.update_item_status(item_id, "error") + await store.update_item( + item_id, + meta_json={ + **json.loads(item.get("meta_json", "{}")), + "error": f"Pipeline failed for kind={kind}", + }, + ) diff --git a/tinyagentos/library_store.py b/tinyagentos/library_store.py new file mode 100644 index 000000000..6d7b1d11c --- /dev/null +++ b/tinyagentos/library_store.py @@ -0,0 +1,275 @@ +"""Persistent store for Library items, artifacts, and processing jobs. + +The Library is the universal ingestion surface for taOS — files, URLs, media +dropped into the Library get processed and their text artifacts indexed into +taosmd collections for agent access. + +Schema mirrors the design doc (docs/design/library-app.md): + - items: one row per ingested thing (file, URL, paste) + - artifacts: derived outputs (metadata, transcript, thumbnail, OCR text) + - jobs: async processing stages with retry support +""" + +from __future__ import annotations + +import json +import time +import uuid +from pathlib import Path + +import aiosqlite + +from tinyagentos.base_store import BaseStore + +LIBRARY_SCHEMA = """ +CREATE TABLE IF NOT EXISTS library_items ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + source_url TEXT NOT NULL DEFAULT '', + title TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'pending', + storage_path TEXT NOT NULL DEFAULT '', + bytes INTEGER NOT NULL DEFAULT 0, + meta_json TEXT NOT NULL DEFAULT '{}', + created_at REAL NOT NULL, + updated_at REAL NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_li_kind ON library_items(kind); +CREATE INDEX IF NOT EXISTS idx_li_status ON library_items(status); +CREATE INDEX IF NOT EXISTS idx_li_created ON library_items(created_at DESC); + +CREATE TABLE IF NOT EXISTS library_artifacts ( + id TEXT PRIMARY KEY, + item_id TEXT NOT NULL REFERENCES library_items(id) ON DELETE CASCADE, + kind TEXT NOT NULL, + path TEXT NOT NULL DEFAULT '', + meta_json TEXT NOT NULL DEFAULT '{}', + created_at REAL NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_la_item ON library_artifacts(item_id); + +CREATE TABLE IF NOT EXISTS library_jobs ( + id TEXT PRIMARY KEY, + item_id TEXT NOT NULL REFERENCES library_items(id) ON DELETE CASCADE, + stage TEXT NOT NULL, + state TEXT NOT NULL DEFAULT 'queued', + error TEXT NOT NULL DEFAULT '', + created_at REAL NOT NULL, + updated_at REAL NOT NULL +); +CREATE INDEX IF NOT EXISTS idx_lj_item ON library_jobs(item_id); +CREATE INDEX IF NOT EXISTS idx_lj_state ON library_jobs(state); +""" + +_VALID_STATUSES = frozenset({"pending", "processing", "ready", "error"}) + + +class LibraryStore(BaseStore): + SCHEMA = LIBRARY_SCHEMA + + async def _post_init(self) -> None: + """Enable foreign key enforcement for cascade deletes.""" + await self._db.execute("PRAGMA foreign_keys = ON") + await self._db.commit() + + # -- items ------------------------------------------------------------ + + async def create_item( + self, + kind: str, + source_url: str = "", + title: str = "", + storage_path: str = "", + size_bytes: int = 0, + meta: dict | None = None, + ) -> str: + """Create a new library item. Returns the item id.""" + item_id = uuid.uuid4().hex + now = time.time() + await self._db.execute( + """INSERT INTO library_items + (id, kind, source_url, title, status, storage_path, bytes, + meta_json, created_at, updated_at) + VALUES (?, ?, ?, ?, 'pending', ?, ?, ?, ?, ?)""", + ( + item_id, + kind, + source_url, + title, + storage_path, + size_bytes, + json.dumps(meta or {}), + now, + now, + ), + ) + await self._db.commit() + return item_id + + async def get_item(self, item_id: str) -> dict | None: + self._db.row_factory = aiosqlite.Row + async with self._db.execute( + "SELECT * FROM library_items WHERE id = ?", (item_id,) + ) as cursor: + row = await cursor.fetchone() + return dict(row) if row else None + + async def list_items( + self, kind: str | None = None, status: str | None = None, + limit: int = 50, offset: int = 0, + ) -> list[dict]: + self._db.row_factory = aiosqlite.Row + where: list[str] = [] + params: list = [] + if kind: + where.append("kind = ?") + params.append(kind) + if status: + where.append("status = ?") + params.append(status) + + sql = "SELECT * FROM library_items" + if where: + sql += " WHERE " + " AND ".join(where) + sql += " ORDER BY created_at DESC LIMIT ? OFFSET ?" + params.extend([limit, offset]) + + async with self._db.execute(sql, params) as cursor: + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + async def update_item(self, item_id: str, **kwargs) -> None: + allowed = {"kind", "source_url", "title", "status", "storage_path", + "bytes", "meta_json", "updated_at"} + fields = [(k, v) for k, v in kwargs.items() if k in allowed] + if not fields: + return + if "updated_at" not in kwargs: + fields.append(("updated_at", time.time())) + if "meta_json" in kwargs and isinstance(kwargs["meta_json"], dict): + # find and replace the meta_json tuple + for i, (k, _v) in enumerate(fields): + if k == "meta_json": + fields[i] = ("meta_json", json.dumps(kwargs["meta_json"])) + break + + set_clause = ", ".join(f"{k} = ?" for k, _ in fields) + values = [v for _, v in fields] + values.append(item_id) + await self._db.execute( + f"UPDATE library_items SET {set_clause} WHERE id = ?", values + ) + await self._db.commit() + + async def delete_item(self, item_id: str) -> None: + """Delete an item and cascade its artifacts and jobs.""" + await self._db.execute("DELETE FROM library_items WHERE id = ?", (item_id,)) + await self._db.commit() + + async def update_item_status(self, item_id: str, status: str) -> None: + if status not in _VALID_STATUSES: + raise ValueError( + f"Invalid status {status!r}; must be one of {sorted(_VALID_STATUSES)}" + ) + await self.update_item(item_id, status=status) + + async def try_update_item_status( + self, item_id: str, new_status: str, *, + if_not_in: tuple[str, ...] = ("pending", "processing"), + ) -> bool: + """Atomically set status to *new_status* when current status is NOT in *if_not_in*. + + Returns True when a row was updated, False otherwise. + Used by reprocess to avoid TOCTOU races — two concurrent reprocess + requests cannot both pass a read-then-write guard. + """ + if new_status not in _VALID_STATUSES: + raise ValueError( + f"Invalid status {new_status!r}; must be one of {sorted(_VALID_STATUSES)}" + ) + placeholders = ",".join("?" * len(if_not_in)) + params = [new_status, time.time(), item_id, *if_not_in] + cursor = await self._db.execute( + f"UPDATE library_items SET status = ?, updated_at = ? " + f"WHERE id = ? AND status NOT IN ({placeholders})", + params, + ) + await self._db.commit() + return cursor.rowcount > 0 + + # -- artifacts -------------------------------------------------------- + + async def add_artifact( + self, item_id: str, kind: str, path: str = "", + meta: dict | None = None, + ) -> str: + artifact_id = uuid.uuid4().hex[:16] + now = time.time() + await self._db.execute( + """INSERT INTO library_artifacts (id, item_id, kind, path, meta_json, created_at) + VALUES (?, ?, ?, ?, ?, ?)""", + (artifact_id, item_id, kind, path, json.dumps(meta or {}), now), + ) + await self._db.commit() + return artifact_id + + async def get_artifacts(self, item_id: str) -> list[dict]: + self._db.row_factory = aiosqlite.Row + async with self._db.execute( + "SELECT * FROM library_artifacts WHERE item_id = ? ORDER BY created_at", + (item_id,), + ) as cursor: + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + async def delete_artifact(self, artifact_id: str) -> None: + await self._db.execute( + "DELETE FROM library_artifacts WHERE id = ?", (artifact_id,) + ) + await self._db.commit() + + # -- jobs ------------------------------------------------------------- + + async def create_job(self, item_id: str, stage: str) -> str: + job_id = uuid.uuid4().hex[:16] + now = time.time() + await self._db.execute( + """INSERT INTO library_jobs (id, item_id, stage, state, created_at, updated_at) + VALUES (?, ?, ?, 'queued', ?, ?)""", + (job_id, item_id, stage, now, now), + ) + await self._db.commit() + return job_id + + async def get_job(self, job_id: str) -> dict | None: + self._db.row_factory = aiosqlite.Row + async with self._db.execute( + "SELECT * FROM library_jobs WHERE id = ?", (job_id,) + ) as cursor: + row = await cursor.fetchone() + return dict(row) if row else None + + async def get_item_jobs(self, item_id: str) -> list[dict]: + self._db.row_factory = aiosqlite.Row + async with self._db.execute( + "SELECT * FROM library_jobs WHERE item_id = ? ORDER BY created_at", + (item_id,), + ) as cursor: + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + async def update_job(self, job_id: str, **kwargs) -> None: + allowed = {"state", "error", "updated_at"} + fields = [(k, v) for k, v in kwargs.items() if k in allowed] + if not fields: + return + if "updated_at" not in kwargs: + fields.append(("updated_at", time.time())) + + set_clause = ", ".join(f"{k} = ?" for k, _ in fields) + values = [v for _, v in fields] + values.append(job_id) + await self._db.execute( + f"UPDATE library_jobs SET {set_clause} WHERE id = ?", values + ) + await self._db.commit() diff --git a/tinyagentos/routes/__init__.py b/tinyagentos/routes/__init__.py index 8a64761c6..ac878879c 100644 --- a/tinyagentos/routes/__init__.py +++ b/tinyagentos/routes/__init__.py @@ -402,6 +402,9 @@ def register_all_routers(app): from tinyagentos.routes.receipts import router as receipts_router app.include_router(receipts_router, dependencies=_csrf) + from tinyagentos.routes.library import router as library_router + app.include_router(library_router, dependencies=_csrf) + from tinyagentos.routes import wallhaven as wallhaven_routes app.include_router(wallhaven_routes.router, dependencies=_csrf) diff --git a/tinyagentos/routes/library.py b/tinyagentos/routes/library.py new file mode 100644 index 000000000..067a98ebc --- /dev/null +++ b/tinyagentos/routes/library.py @@ -0,0 +1,357 @@ +"""Library app routes — ingest, list, and manage library items. + +POST /api/library/ingest — accept file uploads or URL references +GET /api/library/items — list library items +GET /api/library/items/{item_id} — item detail with artifacts +DELETE /api/library/items/{item_id} — remove item and its files +POST /api/library/items/{item_id}/reprocess — re-run the pipeline +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid +from pathlib import Path + +from fastapi import APIRouter, Request, UploadFile, File, Form +from fastapi.responses import JSONResponse + +from tinyagentos.library_pipeline import run_pipeline +from tinyagentos.library_collections import handoff_to_collections +from tinyagentos.task_utils import _create_supervised_task + +logger = logging.getLogger(__name__) + +router = APIRouter() + +LIBRARY_DIR_NAME = "library" + +# Module-level background task tracking so unreferenced tasks are not +# garbage-collected when request.app.state._background_tasks is absent. +_background_tasks: set[asyncio.Task] = set() + + +def _track_background_task(coro) -> asyncio.Task: + """Create a task, store it in ``_background_tasks``, and auto-discard on done.""" + task = asyncio.create_task(coro) + + def _on_done(t: asyncio.Task) -> None: + _background_tasks.discard(t) + + task.add_done_callback(_on_done) + _background_tasks.add(task) + return task + + +# --------------------------------------------------------------------------- + +def _library_dir_from_app(app) -> Path: + """Return the library storage directory, creating it if needed.""" + data_dir = getattr(app.state, "data_dir", None) + if data_dir: + d = Path(data_dir) / LIBRARY_DIR_NAME + else: + d = Path(__file__).parent.parent.parent / "data" / LIBRARY_DIR_NAME + d.mkdir(parents=True, exist_ok=True) + return d + + +def _library_dir(request: Request) -> Path: + return _library_dir_from_app(request.app) + + +async def _get_library_store(request: Request): + """Get the LibraryStore from app.state (lazily initialised).""" + store = getattr(request.app.state, "library_store", None) + if store is None: + from tinyagentos.library_store import LibraryStore + + data_dir = getattr(request.app.state, "data_dir", None) + base = Path(data_dir) if data_dir else Path(__file__).parent.parent.parent / "data" + store = LibraryStore(base / "library.db") + await store.init() + request.app.state.library_store = store + return store + + +# --------------------------------------------------------------------------- +# Ingest +# --------------------------------------------------------------------------- + + +@router.post("/api/library/ingest") +async def ingest( + request: Request, + file: UploadFile | None = File(None), + url: str | None = Form(None), + title: str | None = Form(None), +): + """Ingest a file or URL into the library. + + Accepts a file upload (multipart) or a URL string form field. At least one + of ``file`` or ``url`` must be provided. + + Returns ``{item_id, status: \"pending\"}`` immediately — pipeline processing + happens asynchronously in a background task. + """ + store = await _get_library_store(request) + storage_dir = _library_dir(request) + + if file and file.filename: + # File upload + kind = _detect_kind_from_filename(file.filename, file.content_type) + file_dir = storage_dir / "files" + file_dir.mkdir(parents=True, exist_ok=True) + + # Sanitise filename + safe_name = _sanitise_filename(file.filename) + dest = file_dir / f"{uuid.uuid4().hex[:8]}_{safe_name}" + + # Stream file in bounded chunks to avoid loading it entirely into + # memory. Reject uploads exceeding 100 MB with HTTP 413. + MAX_SIZE = 100 * 1024 * 1024 # 100 MB + if file.size and file.size > MAX_SIZE: + return JSONResponse( + {"error": "Payload Too Large"}, status_code=413, + ) + + size = 0 + with dest.open("wb") as f: + while chunk := await file.read(1024 * 1024): + size += len(chunk) + if size > MAX_SIZE: + dest.unlink(missing_ok=True) + return JSONResponse( + {"error": "Payload Too Large"}, status_code=413, + ) + f.write(chunk) + + item_id = await store.create_item( + kind=kind, + title=title or file.filename, + storage_path=str(dest), + size_bytes=size, + source_url="", + ) + elif url: + # URL reference + kind = _detect_kind_from_url(url) + item_id = await store.create_item( + kind=kind, + source_url=url, + title=title or url, + ) + else: + return JSONResponse( + {"error": "Provide either 'file' (multipart upload) or 'url' (form field)."}, + status_code=400, + ) + + # Run pipeline in background + task_set = getattr(request.app.state, "_background_tasks", None) + coro = _ingest_task(request.app, item_id, store, storage_dir) + if task_set is None: + _track_background_task(coro) + else: + _create_supervised_task(coro, task_set) + + return JSONResponse({"item_id": item_id, "status": "pending"}, status_code=202) + + +async def _ingest_task(app, item_id: str, store, storage_dir: Path) -> None: + """Background task: run pipeline + collections handoff for an item. + + Never raises — always leaves the item in a terminal status. + """ + try: + await run_pipeline(store, item_id, storage_dir) + except Exception: + logger.exception("Library ingest pipeline crashed for item %s", item_id) + await store.update_item_status(item_id, "error") + return + + # Collections handoff after successful pipeline + try: + collections_dir = storage_dir.parent / "collections" + config = getattr(app.state, "config", None) + taosmd_url = getattr(config, "memory_url", None) if config else None + taosmd_admin_token = None + secrets = getattr(app.state, "secrets", None) + if secrets: + secret = await secrets.get("taosmd-admin-token") + if secret: + taosmd_admin_token = secret["value"] + indexed = await handoff_to_collections( + store, item_id, collections_dir, + taosmd_url=taosmd_url, + taosmd_admin_token=taosmd_admin_token, + ) + if indexed > 0: + logger.info( + "Collections handoff indexed %d file(s) for item %s", + indexed, item_id, + ) + else: + logger.debug( + "Collections handoff indexed 0 files for item %s " + "(no text artifacts or taosmd unavailable)", + item_id, + ) + except Exception: + logger.exception("Collections handoff failed for item %s", item_id) + + +def _sanitise_filename(name: str) -> str: + """Strip path separators and null bytes from a filename.""" + name = name.replace("/", "_").replace("\\", "_").replace("\x00", "") + # Also prevent double-dots for path traversal + name = name.replace("..", "_") + return name or "unnamed" + + +def _detect_kind_from_filename(filename: str, content_type: str | None = None) -> str: + """Detect kind from filename and optional MIME type.""" + from tinyagentos.library_pipeline import detect_kind + return detect_kind(file_path=filename, content_type=content_type or "") + + +def _detect_kind_from_url(url: str) -> str: + """Detect kind from URL pattern.""" + from tinyagentos.library_pipeline import detect_kind + return detect_kind(source_url=url) + + +# --------------------------------------------------------------------------- +# List / Get / Delete +# --------------------------------------------------------------------------- + + +@router.get("/api/library/items") +async def list_items( + request: Request, + kind: str | None = None, + status: str | None = None, + limit: int = 50, + offset: int = 0, +): + """List library items, optionally filtered by kind or status.""" + store = await _get_library_store(request) + items = await store.list_items(kind=kind, status=status, limit=limit, offset=offset) + return {"items": items, "count": len(items)} + + +@router.get("/api/library/items/{item_id}") +async def get_item(request: Request, item_id: str): + """Get a library item with its artifacts.""" + store = await _get_library_store(request) + item = await store.get_item(item_id) + if not item: + return JSONResponse({"error": f"Item {item_id!r} not found"}, status_code=404) + + artifacts = await store.get_artifacts(item_id) + jobs = await store.get_item_jobs(item_id) + return {"item": item, "artifacts": artifacts, "jobs": jobs} + + +@router.delete("/api/library/items/{item_id}") +async def delete_item(request: Request, item_id: str): + """Delete a library item and its associated files.""" + store = await _get_library_store(request) + item = await store.get_item(item_id) + if not item: + return JSONResponse({"error": f"Item {item_id!r} not found"}, status_code=404) + + # Remove on-disk files + storage_path = item.get("storage_path", "") + if storage_path and (p := Path(storage_path)).exists(): + try: + p.unlink() + except OSError: + logger.warning("Failed to remove storage file %s for item %s", + storage_path, item_id) + + # Remove artifacts from disk + artifacts = await store.get_artifacts(item_id) + for art in artifacts: + art_path = art.get("path", "") + if art_path and (ap := Path(art_path)).exists(): + try: + ap.unlink() + except OSError: + logger.warning("Failed to remove artifact %s for item %s", + art_path, item_id) + + # Remove collections folder + storage_dir = _library_dir(request) + item_collection_dir = storage_dir.parent / "collections" / item_id + if item_collection_dir.exists(): + import shutil + try: + shutil.rmtree(item_collection_dir) + except OSError: + logger.warning("Failed to remove collection dir %s for item %s", + item_collection_dir, item_id) + + await store.delete_item(item_id) + return {"status": "deleted", "item_id": item_id} + + +@router.post("/api/library/items/{item_id}/reprocess") +async def reprocess_item(request: Request, item_id: str): + """Re-run the ingest pipeline for an existing item. + + Idempotent per (item, stage): old artifacts and their on-disk files are + removed before the pipeline re-runs, so no duplicate rows or files accumulate. + """ + store = await _get_library_store(request) + item = await store.get_item(item_id) + if not item: + return JSONResponse({"error": f"Item {item_id!r} not found"}, status_code=404) + + # Atomic CAS: transition to pending only when NOT already pending/processing. + # Beats the TOCTOU race — two concurrent reprocess requests cannot both + # pass a separate read-then-write guard. + if not await store.try_update_item_status(item_id, "pending", + if_not_in=("pending", "processing")): + return JSONResponse( + {"error": "Item is currently being processed"}, status_code=409 + ) + + # Delete old artifacts so reprocess is idempotent. + # Guard: never unlink the user's original uploaded file (storage_path). + storage_dir = _library_dir(request) + old_artifacts = await store.get_artifacts(item_id) + item_storage_path = item.get("storage_path", "") + for art in old_artifacts: + art_path = art.get("path", "") + if art_path and art_path == item_storage_path: + # This artifact records the source file — skip deletion. + logger.debug("Skipping unlink of source file %s for item %s", + art_path, item_id) + elif art_path and (ap := Path(art_path)).exists(): + try: + ap.unlink() + except OSError: + logger.warning("Failed to remove artifact %s for item %s", + art_path, item_id) + await store.delete_artifact(art["id"]) + + # Status was already atomically set to pending by try_update_item_status above; + # re-queue the pipeline now. If scheduling fails, roll back to ready so the + # item is not stuck pending forever. + try: + task_set = getattr(request.app.state, "_background_tasks", None) + coro = _ingest_task(request.app, item_id, store, storage_dir) + if task_set is None: + _track_background_task(coro) + else: + _create_supervised_task(coro, task_set) + except Exception: + logger.exception("Failed to schedule reprocess for item %s", item_id) + await store.update_item_status(item_id, "ready") + return JSONResponse( + {"error": "Failed to schedule reprocess"}, status_code=500, + ) + + return JSONResponse({"item_id": item_id, "status": "reprocessing"}, status_code=202)