fix: attachments

This commit is contained in:
Alex
2026-04-28 00:38:27 +01:00
parent 82c71be819
commit 8e9f661efc
7 changed files with 351 additions and 26 deletions
+32
View File
@@ -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)}"
+71 -18
View File
@@ -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}"
@@ -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
+87
View File
@@ -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
# ---------------------------------------------------------------------------
+118 -3
View File
@@ -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({})
+7 -1
View File
@@ -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"}
@@ -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(