From 8e9f661efc6e43d24f293289f53a917e3ad36476 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 28 Apr 2026 00:38:27 +0100 Subject: [PATCH] fix: attachments --- application/llm/base.py | 32 +++++ application/llm/google_ai.py | 89 ++++++++++--- .../storage/db/repositories/attachments.py | 23 +++- tests/llm/test_base.py | 87 +++++++++++++ tests/llm/test_google_ai.py | 121 +++++++++++++++++- tests/llm/test_google_llm.py | 8 +- .../db/repositories/test_attachments.py | 17 +++ 7 files changed, 351 insertions(+), 26 deletions(-) diff --git a/application/llm/base.py b/application/llm/base.py index 58bee9b6..0971a886 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -123,6 +123,26 @@ class BaseLLM(ABC): return args_dict return {k: v for k, v in args_dict.items() if v is not None} + @staticmethod + def _is_non_retriable_client_error(exc: BaseException) -> bool: + """4xx errors mean the request itself is malformed — retrying with + a different model fails identically and doubles the work. Only + transient/5xx/connection errors should trigger fallback.""" + try: + from google.genai.errors import ClientError as _GenaiClientError + + if isinstance(exc, _GenaiClientError): + return True + except ImportError: + pass + for attr in ("status_code", "code", "http_status"): + v = getattr(exc, attr, None) + if isinstance(v, int) and 400 <= v < 500: + return True + resp = getattr(exc, "response", None) + v = getattr(resp, "status_code", None) + return isinstance(v, int) and 400 <= v < 500 + def _execute_with_fallback( self, method_name: str, decorators: list, *args, **kwargs ): @@ -152,6 +172,12 @@ class BaseLLM(ABC): try: return decorated_method() except Exception as e: + if self._is_non_retriable_client_error(e): + logger.error( + f"Primary LLM failed with non-retriable client error; " + f"skipping fallback: {str(e)}" + ) + raise if not self.fallback_llm: logger.error(f"Primary LLM failed and no fallback configured: {str(e)}") raise @@ -181,6 +207,12 @@ class BaseLLM(ABC): try: yield from decorated_method() except Exception as e: + if self._is_non_retriable_client_error(e): + logger.error( + f"Primary LLM failed mid-stream with non-retriable client " + f"error; skipping fallback: {str(e)}" + ) + raise if not self.fallback_llm: logger.error( f"Primary LLM failed and no fallback configured: {str(e)}" diff --git a/application/llm/google_ai.py b/application/llm/google_ai.py index c73077e9..0fe1ff8e 100644 --- a/application/llm/google_ai.py +++ b/application/llm/google_ai.py @@ -81,24 +81,39 @@ class GoogleLLM(BaseLLM): for attachment in attachments: mime_type = attachment.get("mime_type") - if mime_type in self.get_supported_attachment_types(): - try: + if mime_type not in self.get_supported_attachment_types(): + continue + try: + # Images go inline as bytes per Google's guidance for + # requests under 20MB; the Files API can return before + # the upload reaches ACTIVE state and yield an empty URI. + if mime_type.startswith("image/"): + file_bytes = self._read_attachment_bytes(attachment) + files.append( + {"file_bytes": file_bytes, "mime_type": mime_type} + ) + else: file_uri = self._upload_file_to_google(attachment) + if not file_uri: + raise ValueError( + f"Google Files API returned empty URI for " + f"{attachment.get('path', 'unknown')}" + ) logging.info( f"GoogleLLM: Successfully uploaded file, got URI: {file_uri}" ) files.append({"file_uri": file_uri, "mime_type": mime_type}) - except Exception as e: - logging.error( - f"GoogleLLM: Error uploading file: {e}", exc_info=True + except Exception as e: + logging.error( + f"GoogleLLM: Error processing attachment: {e}", exc_info=True + ) + if "content" in attachment: + prepared_messages[user_message_index]["content"].append( + { + "type": "text", + "text": f"[File could not be processed: {attachment.get('path', 'unknown')}]", + } ) - if "content" in attachment: - prepared_messages[user_message_index]["content"].append( - { - "type": "text", - "text": f"[File could not be processed: {attachment.get('path', 'unknown')}]", - } - ) if files: logging.info(f"GoogleLLM: Adding {len(files)} files to message") prepared_messages[user_message_index]["content"].append({"files": files}) @@ -114,7 +129,9 @@ class GoogleLLM(BaseLLM): Returns: str: Google AI file URI for the uploaded file. """ - if "google_file_uri" in attachment: + # Truthy check, not membership: a poisoned cache row of "" or + # None must be treated as a miss and trigger a fresh upload. + if attachment.get("google_file_uri"): return attachment["google_file_uri"] file_path = attachment.get("path") if not file_path: @@ -128,6 +145,10 @@ class GoogleLLM(BaseLLM): file=local_path ).uri, ) + if not file_uri: + raise ValueError( + f"Google Files API upload returned empty URI for {file_path}" + ) # Cache the Google file URI on the attachment row so we don't # re-upload on the next LLM call. Accept either a PG UUID @@ -161,6 +182,26 @@ class GoogleLLM(BaseLLM): logging.error(f"Error uploading file to Google AI: {e}", exc_info=True) raise + def _read_attachment_bytes(self, attachment): + """ + Read attachment bytes from storage for inline transmission. + + Args: + attachment (dict): Attachment dictionary with path and metadata. + + Returns: + bytes: Raw file bytes. + """ + file_path = attachment.get("path") + if not file_path: + raise ValueError("No file path provided in attachment") + if not self.storage.file_exists(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + return self.storage.process_file( + file_path, + lambda local_path, **kwargs: open(local_path, "rb").read(), + ) + def _clean_messages_google(self, messages): """ Convert OpenAI format messages to Google AI format and collect system prompts. @@ -300,12 +341,24 @@ class GoogleLLM(BaseLLM): ) elif "files" in item: for file_data in item["files"]: - parts.append( - types.Part.from_uri( - file_uri=file_data["file_uri"], - mime_type=file_data["mime_type"], + if "file_bytes" in file_data: + parts.append( + types.Part.from_bytes( + data=file_data["file_bytes"], + mime_type=file_data["mime_type"], + ) + ) + elif file_data.get("file_uri"): + parts.append( + types.Part.from_uri( + file_uri=file_data["file_uri"], + mime_type=file_data["mime_type"], + ) + ) + else: + logging.warning( + "GoogleLLM: dropping file part with empty URI and no bytes" ) - ) else: raise ValueError( f"Unexpected content dictionary format:{item}" diff --git a/application/storage/db/repositories/attachments.py b/application/storage/db/repositories/attachments.py index f250e8cc..08662920 100644 --- a/application/storage/db/repositories/attachments.py +++ b/application/storage/db/repositories/attachments.py @@ -17,6 +17,21 @@ _UPDATABLE_SCALARS = { _UPDATABLE_JSONB = {"metadata"} +def _attachment_to_dict(row: Any) -> dict: + """row_to_dict + ``upload_path``→``path`` alias. + + Pre-Postgres, the Mongo attachment shape used ``path``. The PG column + is ``upload_path``; LLM provider code (google_ai/openai/anthropic and + handlers/base) still reads ``attachment.get("path")``. Mirroring the + ``id``/``_id`` dual-emit in row_to_dict so consumers don't need to + know which storage backend produced the dict. + """ + out = row_to_dict(row) + if "upload_path" in out and out.get("path") is None: + out["path"] = out["upload_path"] + return out + + class AttachmentsRepository: def __init__(self, conn: Connection) -> None: self._conn = conn @@ -66,7 +81,7 @@ class AttachmentsRepository: "legacy_mongo_id": legacy_mongo_id, }, ) - return row_to_dict(result.fetchone()) + return _attachment_to_dict(result.fetchone()) def get(self, attachment_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( @@ -76,7 +91,7 @@ class AttachmentsRepository: {"id": attachment_id, "user_id": user_id}, ) row = result.fetchone() - return row_to_dict(row) if row is not None else None + return _attachment_to_dict(row) if row is not None else None def get_any(self, attachment_id: str, user_id: str) -> Optional[dict]: """Resolve an attachment by either PG UUID or legacy Mongo ObjectId string.""" @@ -155,14 +170,14 @@ class AttachmentsRepository: params["user_id"] = user_id result = self._conn.execute(text(sql), params) row = result.fetchone() - return row_to_dict(row) if row is not None else None + return _attachment_to_dict(row) if row is not None else None def list_for_user(self, user_id: str) -> list[dict]: result = self._conn.execute( text("SELECT * FROM attachments WHERE user_id = :user_id ORDER BY created_at DESC"), {"user_id": user_id}, ) - return [row_to_dict(r) for r in result.fetchall()] + return [_attachment_to_dict(r) for r in result.fetchall()] def update(self, attachment_id: str, user_id: str, fields: dict) -> bool: """Partial update. Used by the LLM providers to cache their diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index c67de984..56a3caa4 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -261,6 +261,93 @@ class TestStreamWithFallback: assert fallback.gen_stream_called +# --------------------------------------------------------------------------- +# Non-retriable client error guard +# --------------------------------------------------------------------------- + + +class _StatusError(Exception): + """Mimics openai/anthropic-shaped client errors with a status_code.""" + + def __init__(self, status_code, message="bad request"): + super().__init__(message) + self.status_code = status_code + + +class _ClientErrorLLM(BaseLLM): + def __init__(self, status_code, **kwargs): + super().__init__(**kwargs) + self._status = status_code + + def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kw): + raise _StatusError(self._status) + + def _raw_gen_stream(self, baseself, model, messages, stream=True, tools=None, **kw): + raise _StatusError(self._status) + + +@pytest.mark.unit +class TestNonRetriableClientError: + + def test_helper_detects_4xx_status_code(self): + assert BaseLLM._is_non_retriable_client_error(_StatusError(400)) + assert BaseLLM._is_non_retriable_client_error(_StatusError(404)) + assert BaseLLM._is_non_retriable_client_error(_StatusError(429)) + + def test_helper_passes_5xx_through(self): + # 5xx and connection errors should still trigger fallback. + assert not BaseLLM._is_non_retriable_client_error(_StatusError(500)) + assert not BaseLLM._is_non_retriable_client_error(_StatusError(503)) + assert not BaseLLM._is_non_retriable_client_error(RuntimeError("oops")) + + def test_helper_detects_genai_client_error(self): + try: + from google.genai.errors import ClientError + except ImportError: + pytest.skip("google-genai not installed") + # ClientError(code, response_json, response=None) + exc = ClientError(400, {"error": {"message": "bad", "code": 400}}, None) + assert BaseLLM._is_non_retriable_client_error(exc) + + def test_helper_detects_response_status_code(self): + exc = RuntimeError("wrapped") + exc.response = type("R", (), {"status_code": 401})() + assert BaseLLM._is_non_retriable_client_error(exc) + + @patch("application.llm.base.gen_cache", lambda f: f) + @patch("application.llm.base.gen_token_usage", lambda f: f) + def test_4xx_skips_fallback(self): + fallback = FallbackLLM(model_id="fallback-model") + llm = _ClientErrorLLM(status_code=400) + llm._fallback_llm = fallback + + with pytest.raises(_StatusError): + llm.gen(model="m", messages=[]) + assert not fallback.gen_called + + @patch("application.llm.base.stream_cache", lambda f: f) + @patch("application.llm.base.stream_token_usage", lambda f: f) + def test_4xx_skips_stream_fallback(self): + fallback = FallbackLLM(model_id="fallback-model") + llm = _ClientErrorLLM(status_code=400) + llm._fallback_llm = fallback + + with pytest.raises(_StatusError): + list(llm.gen_stream(model="m", messages=[])) + assert not fallback.gen_stream_called + + @patch("application.llm.base.gen_cache", lambda f: f) + @patch("application.llm.base.gen_token_usage", lambda f: f) + def test_5xx_still_falls_back(self): + fallback = FallbackLLM(model_id="fallback-model") + llm = _ClientErrorLLM(status_code=503) + llm._fallback_llm = fallback + + result = llm.gen(model="m", messages=[]) + assert result == "fallback_gen_result" + assert fallback.gen_called + + # --------------------------------------------------------------------------- # fallback_llm property: backup model resolution # --------------------------------------------------------------------------- diff --git a/tests/llm/test_google_ai.py b/tests/llm/test_google_ai.py index 8f213c43..7d91069c 100644 --- a/tests/llm/test_google_ai.py +++ b/tests/llm/test_google_ai.py @@ -28,10 +28,11 @@ from application.llm.google_ai import GoogleLLM class _FakePart: - def __init__(self, text=None, function_call=None, file_data=None, thought=False, **kwargs): + def __init__(self, text=None, function_call=None, file_data=None, inline_data=None, thought=False, **kwargs): self.text = text self.function_call = function_call or kwargs.get("functionCall") self.file_data = file_data + self.inline_data = inline_data self.thought = thought self.thoughtSignature = kwargs.get("thoughtSignature") @@ -53,6 +54,12 @@ class _FakePart: file_data=types.SimpleNamespace(file_uri=file_uri, mime_type=mime_type) ) + @staticmethod + def from_bytes(data, mime_type): + return _FakePart( + inline_data=types.SimpleNamespace(data=data, mime_type=mime_type) + ) + class _FakeContent: def __init__(self, role, parts): @@ -223,6 +230,43 @@ class TestCleanMessagesGoogle: for p in cleaned[0].parts ) + def test_files_with_inline_bytes(self, llm): + msgs = [ + { + "role": "user", + "content": [ + { + "files": [ + {"file_bytes": b"\x89PNG", "mime_type": "image/png"} + ] + }, + ], + } + ] + cleaned, _ = llm._clean_messages_google(msgs) + assert len(cleaned) == 1 + inline_parts = [ + p for p in cleaned[0].parts + if getattr(p, "inline_data", None) is not None + ] + assert len(inline_parts) == 1 + assert inline_parts[0].inline_data.data == b"\x89PNG" + assert inline_parts[0].inline_data.mime_type == "image/png" + + def test_files_with_empty_uri_dropped(self, llm): + msgs = [ + { + "role": "user", + "content": [ + {"files": [{"file_uri": "", "mime_type": "image/png"}]}, + ], + } + ] + cleaned, _ = llm._clean_messages_google(msgs) + # Empty URI part is dropped; no other parts means the whole + # content is empty and the message itself is not appended. + assert cleaned == [] + def test_unexpected_list_item_raises(self, llm): msgs = [{"role": "user", "content": [{"unknown_key": "val"}]}] with pytest.raises(ValueError, match="Unexpected content dictionary"): @@ -710,7 +754,9 @@ class TestPrepareMessagesWithAttachments: def test_upload_error_adds_text_fallback(self, llm, monkeypatch): monkeypatch.setattr( - llm, "_upload_file_to_google", lambda a: (_ for _ in ()).throw(Exception("fail")) + llm, + "_read_attachment_bytes", + lambda a: (_ for _ in ()).throw(Exception("fail")), ) msgs = [{"role": "user", "content": "hi"}] attachments = [ @@ -724,8 +770,57 @@ class TestPrepareMessagesWithAttachments: ] assert len(text_parts) == 1 + def test_pdf_upload_error_adds_text_fallback(self, llm, monkeypatch): + monkeypatch.setattr( + llm, + "_upload_file_to_google", + lambda a: (_ for _ in ()).throw(Exception("fail")), + ) + msgs = [{"role": "user", "content": "hi"}] + attachments = [ + {"mime_type": "application/pdf", "path": "/tmp/doc.pdf", "content": "x"}, + ] + result = llm.prepare_messages_with_attachments(msgs, attachments) + user_msg = next(m for m in result if m["role"] == "user") + text_parts = [ + p for p in user_msg["content"] + if isinstance(p, dict) and p.get("type") == "text" and "could not" in p.get("text", "").lower() + ] + assert len(text_parts) == 1 + + def test_pdf_empty_uri_adds_text_fallback(self, llm, monkeypatch): + monkeypatch.setattr(llm, "_upload_file_to_google", lambda a: "") + msgs = [{"role": "user", "content": "hi"}] + attachments = [ + {"mime_type": "application/pdf", "path": "/tmp/doc.pdf", "content": "x"}, + ] + result = llm.prepare_messages_with_attachments(msgs, attachments) + user_msg = next(m for m in result if m["role"] == "user") + files_entries = [ + p for p in user_msg["content"] if isinstance(p, dict) and "files" in p + ] + assert files_entries == [] + text_parts = [ + p for p in user_msg["content"] + if isinstance(p, dict) and p.get("type") == "text" and "could not" in p.get("text", "").lower() + ] + assert len(text_parts) == 1 + + def test_image_uses_inline_bytes(self, llm, monkeypatch): + monkeypatch.setattr(llm, "_read_attachment_bytes", lambda a: b"\x89PNG-bytes") + msgs = [{"role": "user", "content": "hi"}] + attachments = [{"mime_type": "image/png", "path": "/img.png"}] + result = llm.prepare_messages_with_attachments(msgs, attachments) + user_msg = next(m for m in result if m["role"] == "user") + files_entry = next( + p for p in user_msg["content"] if isinstance(p, dict) and "files" in p + ) + assert files_entry["files"] == [ + {"file_bytes": b"\x89PNG-bytes", "mime_type": "image/png"} + ] + def test_no_user_message_creates_one(self, llm, monkeypatch): - monkeypatch.setattr(llm, "_upload_file_to_google", lambda a: "gs://uri") + monkeypatch.setattr(llm, "_read_attachment_bytes", lambda a: b"png") msgs = [{"role": "system", "content": "sys"}] attachments = [{"mime_type": "image/png", "path": "/img.png"}] result = llm.prepare_messages_with_attachments(msgs, attachments) @@ -746,6 +841,26 @@ class TestUploadFileToGoogle: result = llm._upload_file_to_google(attachment) assert result == "gs://cached" + def test_empty_cached_uri_triggers_reupload(self, llm, monkeypatch): + # Poisoned-cache repro: an empty-string google_file_uri must be + # treated as a miss and re-upload, not returned as-is. + monkeypatch.setattr( + "application.llm.google_ai.settings", + types.SimpleNamespace(GOOGLE_API_KEY="k", API_KEY="k"), + ) + result = llm._upload_file_to_google( + {"google_file_uri": "", "path": "/tmp/file.pdf"} + ) + assert result == "gs://fake-uri" + + def test_empty_upload_uri_raises(self, llm): + llm.storage = types.SimpleNamespace( + file_exists=lambda p: True, + process_file=lambda path, fn, **kw: "", + ) + with pytest.raises(ValueError, match="empty URI"): + llm._upload_file_to_google({"path": "/tmp/file.pdf"}) + def test_raises_for_no_path(self, llm): with pytest.raises(ValueError, match="No file path"): llm._upload_file_to_google({}) diff --git a/tests/llm/test_google_llm.py b/tests/llm/test_google_llm.py index 54d86bcf..b5555047 100644 --- a/tests/llm/test_google_llm.py +++ b/tests/llm/test_google_llm.py @@ -228,6 +228,7 @@ def test_prepare_messages_with_attachments_appends_files(monkeypatch): process_file=lambda path, processor_func, **kwargs: "gs://file_uri" ) monkeypatch.setattr(llm, "_upload_file_to_google", lambda att: "gs://file_uri") + monkeypatch.setattr(llm, "_read_attachment_bytes", lambda att: b"png-bytes") messages = [{"role": "user", "content": "Hi"}] attachments = [ @@ -240,4 +241,9 @@ def test_prepare_messages_with_attachments_appends_files(monkeypatch): assert isinstance(user_msg["content"], list) files_entry = next((p for p in user_msg["content"] if isinstance(p, dict) and "files" in p), None) assert files_entry is not None - assert isinstance(files_entry["files"], list) and len(files_entry["files"]) == 2 + files = files_entry["files"] + assert len(files) == 2 + image_part = next(f for f in files if f["mime_type"] == "image/png") + pdf_part = next(f for f in files if f["mime_type"] == "application/pdf") + assert image_part == {"file_bytes": b"png-bytes", "mime_type": "image/png"} + assert pdf_part == {"file_uri": "gs://file_uri", "mime_type": "application/pdf"} diff --git a/tests/storage/db/repositories/test_attachments.py b/tests/storage/db/repositories/test_attachments.py index 00cc636c..6420fb26 100644 --- a/tests/storage/db/repositories/test_attachments.py +++ b/tests/storage/db/repositories/test_attachments.py @@ -31,6 +31,23 @@ class TestCreate: doc = repo.create("u", "f", "/p") assert doc["_id"] == doc["id"] + def test_create_aliases_upload_path_as_path(self, pg_conn): + # LLM provider code (google_ai/openai/anthropic and handlers/base) + # reads attachment.get("path") — preserved from the legacy Mongo + # shape. Repo emits both keys so consumers don't need to know + # which storage backend produced the dict. + repo = _repo(pg_conn) + doc = repo.create("u", "f", "/uploads/x.png") + assert doc["path"] == "/uploads/x.png" + assert doc["upload_path"] == "/uploads/x.png" + + def test_get_aliases_upload_path_as_path(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("u", "f", "/uploads/y.pdf") + fetched = repo.get(created["id"], "u") + assert fetched is not None + assert fetched["path"] == "/uploads/y.pdf" + def test_create_with_legacy_mongo_id(self, pg_conn): repo = _repo(pg_conn) doc = repo.create(