wehub-resource-sync
fed8b2eed7
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Backend release / release (push) Has been cancelled
Bandit Security Scan / bandit_scan (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push multi-arch DocsGPT Docker image / manifest (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Has been cancelled
Build and push DocsGPT FE Docker image for development / manifest (push) Has been cancelled
Python linting / ruff (push) Has been cancelled
Run python tests with pytest / Run tests and count coverage (3.12) (push) Has been cancelled
React Widget Build / build (push) Has been cancelled
1065 lines
42 KiB
Python
1065 lines
42 KiB
Python
"""Repository for the ``conversations`` and ``conversation_messages`` tables.
|
|
|
|
Covers every operation the legacy Mongo code performs on
|
|
``conversations_collection``:
|
|
|
|
- create / get / list / delete conversations
|
|
- append message (transactional position allocation)
|
|
- update message at index (overwrite + optional truncation)
|
|
- set / unset feedback on a message
|
|
- rename conversation
|
|
- update compression metadata
|
|
- shared_with access checks
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from enum import Enum
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import Connection, text
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
|
|
from application.storage.db.base_repository import looks_like_uuid, row_to_dict
|
|
from application.storage.db.models import conversations_table, conversation_messages_table
|
|
from application.storage.db.serialization import PGNativeJSONEncoder
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MessageUpdateOutcome(str, Enum):
|
|
"""Discriminated result of ``update_message_by_id``.
|
|
|
|
Distinguishes the row-actually-updated case from the row-already-at-
|
|
the-requested-terminal-state case so an abort handler can journal
|
|
``end`` instead of ``error`` when the normal-path finalize already
|
|
flipped the row to ``complete``.
|
|
"""
|
|
|
|
UPDATED = "updated"
|
|
ALREADY_COMPLETE = "already_complete"
|
|
ALREADY_FAILED = "already_failed"
|
|
NOT_FOUND = "not_found"
|
|
INVALID = "invalid"
|
|
|
|
|
|
def _message_row_to_dict(row) -> dict:
|
|
"""Like ``row_to_dict`` but renames the DB column ``message_metadata``
|
|
back to the public API key ``metadata`` so callers keep the Mongo-era
|
|
shape. See migration 0001_initial for the column rename rationale."""
|
|
out = row_to_dict(row)
|
|
if "message_metadata" in out:
|
|
out["metadata"] = out.pop("message_metadata")
|
|
return out
|
|
|
|
|
|
class ConversationsRepository:
|
|
def __init__(self, conn: Connection) -> None:
|
|
self._conn = conn
|
|
|
|
# ------------------------------------------------------------------
|
|
# Reference translation helpers
|
|
# ------------------------------------------------------------------
|
|
#
|
|
# During the Mongo→Postgres dual-write window, callers routinely
|
|
# hand us Mongo ObjectId strings (24-char hex) for fields that are
|
|
# UUID FKs in Postgres (``agent_id``, ``attachments`` entries, ...).
|
|
# Casting those straight to ``uuid`` raises and the outer dual-write
|
|
# shim swallows the exception, so the write silently drops. These
|
|
# helpers translate via the ``legacy_mongo_id`` columns we added
|
|
# precisely for this purpose.
|
|
|
|
def _resolve_agent_ref(self, agent_id_raw: str | None) -> str | None:
|
|
"""Translate ``agent_id_raw`` to a Postgres UUID string.
|
|
|
|
- ``None``/empty → ``None`` (no agent).
|
|
- Already-UUID-shaped → returned as-is.
|
|
- Otherwise treated as a Mongo ObjectId and looked up via
|
|
``agents.legacy_mongo_id``. Returns ``None`` if no PG row
|
|
exists yet (e.g. the agent was created before the backfill
|
|
ran).
|
|
"""
|
|
if not agent_id_raw:
|
|
return None
|
|
value = str(agent_id_raw)
|
|
if looks_like_uuid(value):
|
|
return value
|
|
result = self._conn.execute(
|
|
text("SELECT id FROM agents WHERE legacy_mongo_id = :lid LIMIT 1"),
|
|
{"lid": value},
|
|
)
|
|
row = result.fetchone()
|
|
return str(row[0]) if row is not None else None
|
|
|
|
def _resolve_attachment_refs(
|
|
self, ids: list[str] | None,
|
|
) -> list[str]:
|
|
"""Translate a list of attachment ids to canonical PG
|
|
``attachments.id`` UUIDs.
|
|
|
|
Inputs may be:
|
|
|
|
- A Mongo ObjectId string (24-hex), legacy dual-write era —
|
|
must be looked up via ``attachments.legacy_mongo_id``.
|
|
- A UUID string that is a real ``attachments.id`` PK.
|
|
- A UUID string that is *only* present as
|
|
``attachments.legacy_mongo_id`` — this is the post-cutover
|
|
shape: ``/store_attachment`` mints a UUID, hands it to the
|
|
worker, and the worker stashes it in ``legacy_mongo_id``
|
|
while the row gets a freshly-generated PK. Trusting the
|
|
input UUID as a PK here orphans the array entry: the column
|
|
is ``uuid[]`` (no FK), so PG accepts the bad value and all
|
|
downstream reads via ``AttachmentsRepository.get_any`` miss.
|
|
|
|
Resolution therefore tries ``legacy_mongo_id`` first for every
|
|
id (UUID-shaped or not), then falls back to the direct PK
|
|
match. Unknown ids are dropped — they'd have failed the
|
|
``uuid[]`` cast otherwise and the whole row would have vanished
|
|
via dual-write's exception swallow.
|
|
"""
|
|
if not ids:
|
|
return []
|
|
# Defer to AttachmentsRepository for the batched lookup so the
|
|
# legacy-first semantics live in one place.
|
|
from application.storage.db.repositories.attachments import (
|
|
AttachmentsRepository,
|
|
)
|
|
|
|
clean: list[str] = [str(raw) for raw in ids if raw is not None]
|
|
if not clean:
|
|
return []
|
|
repo = AttachmentsRepository(self._conn)
|
|
mapping = repo.resolve_ids(clean)
|
|
out: list[str] = []
|
|
for value in clean:
|
|
mapped = mapping.get(value)
|
|
if mapped is not None:
|
|
out.append(mapped)
|
|
return out
|
|
|
|
# ------------------------------------------------------------------
|
|
# Conversation CRUD
|
|
# ------------------------------------------------------------------
|
|
|
|
def create(
|
|
self,
|
|
user_id: str,
|
|
name: str | None = None,
|
|
*,
|
|
agent_id: str | None = None,
|
|
api_key: str | None = None,
|
|
is_shared_usage: bool = False,
|
|
shared_token: str | None = None,
|
|
visibility: str = "listed",
|
|
legacy_mongo_id: str | None = None,
|
|
) -> dict:
|
|
"""Create a new conversation.
|
|
|
|
``legacy_mongo_id`` is used by the dual-write shim so that a
|
|
Postgres row inserted *after* a successful Mongo insert carries
|
|
the Mongo ``_id`` as a lookup key. Subsequent appends/updates
|
|
can then resolve the PG row by that id via
|
|
:meth:`get_by_legacy_id`.
|
|
"""
|
|
values: dict = {
|
|
"user_id": user_id,
|
|
"name": name,
|
|
"visibility": visibility,
|
|
}
|
|
# ``agent_id`` may arrive as a Mongo ObjectId during the dual-write
|
|
# window; resolve to a UUID (or drop silently if not yet backfilled).
|
|
resolved_agent_id = self._resolve_agent_ref(agent_id)
|
|
if resolved_agent_id:
|
|
values["agent_id"] = resolved_agent_id
|
|
if api_key:
|
|
values["api_key"] = api_key
|
|
if is_shared_usage:
|
|
values["is_shared_usage"] = True
|
|
if shared_token:
|
|
values["shared_token"] = shared_token
|
|
if legacy_mongo_id:
|
|
values["legacy_mongo_id"] = legacy_mongo_id
|
|
|
|
stmt = pg_insert(conversations_table).values(**values).returning(conversations_table)
|
|
result = self._conn.execute(stmt)
|
|
return row_to_dict(result.fetchone())
|
|
|
|
def get_by_legacy_id(
|
|
self, legacy_mongo_id: str, user_id: str | None = None,
|
|
) -> Optional[dict]:
|
|
"""Look up a conversation by the original Mongo ObjectId string.
|
|
|
|
Used by the dual-write helpers to translate a Mongo ``_id`` into
|
|
the Postgres UUID for follow-up writes. When ``user_id`` is
|
|
provided, the lookup is scoped to rows owned by that user so
|
|
callers can't accidentally resolve another user's conversation.
|
|
"""
|
|
legacy_mongo_id = str(legacy_mongo_id) if legacy_mongo_id is not None else None
|
|
if user_id is not None:
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversations "
|
|
"WHERE legacy_mongo_id = :legacy_id AND user_id = :user_id"
|
|
),
|
|
{"legacy_id": legacy_mongo_id, "user_id": user_id},
|
|
)
|
|
else:
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversations WHERE legacy_mongo_id = :legacy_id"
|
|
),
|
|
{"legacy_id": legacy_mongo_id},
|
|
)
|
|
row = result.fetchone()
|
|
return row_to_dict(row) if row is not None else None
|
|
|
|
def get(self, conversation_id: str, user_id: str) -> Optional[dict]:
|
|
"""Fetch a conversation the user owns or has shared access to."""
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversations "
|
|
"WHERE id = CAST(:id AS uuid) "
|
|
"AND (user_id = :user_id OR :user_id = ANY(shared_with))"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id},
|
|
)
|
|
row = result.fetchone()
|
|
return row_to_dict(row) if row is not None else None
|
|
|
|
def get_any(self, conversation_id: str, user_id: str) -> Optional[dict]:
|
|
"""Resolve a conversation by either PG UUID or legacy Mongo ObjectId string.
|
|
|
|
Returns a conversation the user owns or has shared access to.
|
|
"""
|
|
if looks_like_uuid(conversation_id):
|
|
row = self.get(conversation_id, user_id)
|
|
if row is not None:
|
|
return row
|
|
return self.get_by_legacy_id(conversation_id, user_id)
|
|
|
|
def get_owned(self, conversation_id: str, user_id: str) -> Optional[dict]:
|
|
"""Fetch a conversation owned by the user (no shared access)."""
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversations "
|
|
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id},
|
|
)
|
|
row = result.fetchone()
|
|
return row_to_dict(row) if row is not None else None
|
|
|
|
def list_for_user(self, user_id: str, limit: int = 30) -> list[dict]:
|
|
"""List a user's sidebar conversations, most recent first.
|
|
|
|
Only ``visibility = 'listed'`` rows surface; agent/API/OpenAI-compat
|
|
traffic persists as ``'hidden'`` and is excluded.
|
|
"""
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversations "
|
|
"WHERE user_id = :user_id "
|
|
"AND visibility = 'listed' "
|
|
"ORDER BY date DESC LIMIT :limit"
|
|
),
|
|
{"user_id": user_id, "limit": limit},
|
|
)
|
|
return [row_to_dict(r) for r in result.fetchall()]
|
|
|
|
def search_for_user(
|
|
self, user_id: str, query: str, limit: int = 30,
|
|
) -> list[dict]:
|
|
"""Search a user's conversations by name or message content.
|
|
|
|
Same visibility filter as :meth:`list_for_user`. Matches against
|
|
``conversations.name`` or any of the conversation's messages'
|
|
``prompt`` / ``response`` columns (case-insensitive substring).
|
|
|
|
Each returned row includes ``match_field`` (one of ``name``,
|
|
``prompt``, ``response``) and ``match_text`` (the full text of the
|
|
first matching field, ``name`` taking precedence over messages,
|
|
``prompt`` over ``response``) so callers can render a snippet.
|
|
"""
|
|
if not query:
|
|
return []
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT c.*, mt.match_field, mt.match_text "
|
|
"FROM conversations c "
|
|
"JOIN LATERAL ( "
|
|
" SELECT field AS match_field, txt AS match_text "
|
|
" FROM ( "
|
|
" SELECT 'name'::text AS field, c.name AS txt, 0 AS prio "
|
|
" WHERE c.name ILIKE :pattern "
|
|
" UNION ALL "
|
|
" SELECT 'prompt'::text, m.prompt, 1 "
|
|
" FROM conversation_messages m "
|
|
" WHERE m.conversation_id = c.id "
|
|
" AND m.prompt ILIKE :pattern "
|
|
" UNION ALL "
|
|
" SELECT 'response'::text, m.response, 2 "
|
|
" FROM conversation_messages m "
|
|
" WHERE m.conversation_id = c.id "
|
|
" AND m.response ILIKE :pattern "
|
|
" ) s "
|
|
" ORDER BY prio "
|
|
" LIMIT 1 "
|
|
") mt ON TRUE "
|
|
"WHERE c.user_id = :user_id "
|
|
"AND c.visibility = 'listed' "
|
|
"ORDER BY c.date DESC LIMIT :limit"
|
|
),
|
|
{
|
|
"user_id": user_id,
|
|
"pattern": f"%{query}%",
|
|
"limit": limit,
|
|
},
|
|
)
|
|
return [row_to_dict(r) for r in result.fetchall()]
|
|
|
|
def rename(self, conversation_id: str, user_id: str, name: str) -> bool:
|
|
# Shape-gate so a non-UUID id (legacy Mongo ObjectId still floating
|
|
# around in client-side state during the cutover) never reaches the
|
|
# ``CAST(:id AS uuid)`` — that cast raises on the server and poisons
|
|
# the enclosing transaction, making every subsequent query on the
|
|
# same connection fail.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE conversations SET name = :name, updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id, "name": name},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def add_shared_user(self, conversation_id: str, user_to_add: str) -> bool:
|
|
"""Idempotently append ``user_to_add`` to ``shared_with``.
|
|
|
|
Accepts either a PG UUID or a legacy Mongo ObjectId as the
|
|
conversation id. Mirrors Mongo ``$addToSet`` semantics via the
|
|
``NOT (:user = ANY(shared_with))`` guard.
|
|
"""
|
|
if not user_to_add:
|
|
return False
|
|
if looks_like_uuid(conversation_id):
|
|
sql = (
|
|
"UPDATE conversations "
|
|
"SET shared_with = array_append(shared_with, :user), "
|
|
" updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) "
|
|
"AND NOT (:user = ANY(shared_with))"
|
|
)
|
|
else:
|
|
sql = (
|
|
"UPDATE conversations "
|
|
"SET shared_with = array_append(shared_with, :user), "
|
|
" updated_at = now() "
|
|
"WHERE legacy_mongo_id = :id "
|
|
"AND NOT (:user = ANY(shared_with))"
|
|
)
|
|
result = self._conn.execute(
|
|
text(sql), {"id": conversation_id, "user": user_to_add},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def remove_shared_user(self, conversation_id: str, user_to_remove: str) -> bool:
|
|
"""Remove ``user_to_remove`` from ``shared_with``. Mirror of Mongo ``$pull``."""
|
|
if not user_to_remove:
|
|
return False
|
|
if looks_like_uuid(conversation_id):
|
|
sql = (
|
|
"UPDATE conversations "
|
|
"SET shared_with = array_remove(shared_with, :user), "
|
|
" updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) "
|
|
"AND :user = ANY(shared_with)"
|
|
)
|
|
else:
|
|
sql = (
|
|
"UPDATE conversations "
|
|
"SET shared_with = array_remove(shared_with, :user), "
|
|
" updated_at = now() "
|
|
"WHERE legacy_mongo_id = :id "
|
|
"AND :user = ANY(shared_with)"
|
|
)
|
|
result = self._conn.execute(
|
|
text(sql), {"id": conversation_id, "user": user_to_remove},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def set_shared_token(self, conversation_id: str, user_id: str, token: str) -> bool:
|
|
# Shape-gate: see ``rename`` — prevents transaction poisoning when
|
|
# a non-UUID id reaches this code path.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE conversations SET shared_token = :token, updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id, "token": token},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def update_compression_metadata(
|
|
self, conversation_id: str, user_id: str, metadata: dict,
|
|
) -> bool:
|
|
"""Replace the entire ``compression_metadata`` JSONB blob.
|
|
|
|
Prefer :meth:`append_compression_point` + :meth:`set_compression_flags`
|
|
to match the Mongo service semantics exactly (those two mirror
|
|
``$set`` + ``$push $slice``). This method is retained for callers
|
|
that already compute the full merged blob client-side.
|
|
"""
|
|
# Shape-gate: see ``rename`` — prevents transaction poisoning when
|
|
# a non-UUID id reaches this code path.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE conversations "
|
|
"SET compression_metadata = CAST(:meta AS jsonb), updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id, "meta": json.dumps(metadata)},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def set_compression_flags(
|
|
self,
|
|
conversation_id: str,
|
|
*,
|
|
is_compressed: bool,
|
|
last_compression_at,
|
|
) -> bool:
|
|
"""Update ``compression_metadata.is_compressed`` and
|
|
``compression_metadata.last_compression_at`` without touching
|
|
``compression_points``.
|
|
|
|
Mirrors the Mongo ``$set`` on those two subfields in
|
|
``ConversationService.update_compression_metadata``. Initialises
|
|
the surrounding object when the row has no ``compression_metadata``
|
|
yet.
|
|
"""
|
|
# Shape-gate: the streaming pipeline may pass through a legacy id
|
|
# that ``get_by_legacy_id`` couldn't resolve; in that case the id
|
|
# remains a non-UUID string and the CAST would poison the txn.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"""
|
|
UPDATE conversations SET
|
|
compression_metadata = jsonb_set(
|
|
jsonb_set(
|
|
COALESCE(compression_metadata, '{}'::jsonb),
|
|
'{is_compressed}',
|
|
to_jsonb(CAST(:is_compressed AS boolean)), true
|
|
),
|
|
'{last_compression_at}',
|
|
to_jsonb(CAST(:last_compression_at AS text)), true
|
|
),
|
|
updated_at = now()
|
|
WHERE id = CAST(:id AS uuid)
|
|
"""
|
|
),
|
|
{
|
|
"id": conversation_id,
|
|
"is_compressed": bool(is_compressed),
|
|
"last_compression_at": (
|
|
str(last_compression_at) if last_compression_at is not None else None
|
|
),
|
|
},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def append_compression_point(
|
|
self,
|
|
conversation_id: str,
|
|
point: dict,
|
|
*,
|
|
max_points: int,
|
|
) -> bool:
|
|
"""Append one compression point, keeping at most ``max_points``.
|
|
|
|
Mirrors Mongo's ``$push {"$each": [point], "$slice": -max_points}``
|
|
on ``compression_metadata.compression_points``. Preserves the
|
|
other top-level keys in ``compression_metadata``.
|
|
"""
|
|
# Shape-gate: see ``set_compression_flags``.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"""
|
|
UPDATE conversations SET
|
|
compression_metadata = jsonb_set(
|
|
COALESCE(compression_metadata, '{}'::jsonb),
|
|
'{compression_points}',
|
|
COALESCE(
|
|
(
|
|
SELECT jsonb_agg(elem ORDER BY rn)
|
|
FROM (
|
|
SELECT
|
|
elem,
|
|
row_number() OVER () AS rn,
|
|
count(*) OVER () AS cnt
|
|
FROM jsonb_array_elements(
|
|
COALESCE(
|
|
compression_metadata -> 'compression_points',
|
|
'[]'::jsonb
|
|
) || jsonb_build_array(CAST(:point AS jsonb))
|
|
) AS elem
|
|
) ranked
|
|
WHERE rn > cnt - :max_points
|
|
),
|
|
'[]'::jsonb
|
|
),
|
|
true
|
|
),
|
|
updated_at = now()
|
|
WHERE id = CAST(:id AS uuid)
|
|
"""
|
|
),
|
|
{
|
|
"id": conversation_id,
|
|
"point": json.dumps(point, cls=PGNativeJSONEncoder),
|
|
"max_points": int(max_points),
|
|
},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def delete(self, conversation_id: str, user_id: str) -> bool:
|
|
# Shape-gate: see ``rename`` — prevents transaction poisoning when
|
|
# a non-UUID id reaches this code path.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
# Artifacts are parent-derived: deleting a conversation must take its
|
|
# artifacts with it (rows + bytes + the quota they consume). conversation_id
|
|
# is a bare uuid (no FK cascade), so delete the rows explicitly and reap the
|
|
# stored bytes best-effort.
|
|
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
|
|
|
artifacts = ArtifactsRepository(self._conn)
|
|
paths = artifacts.storage_paths_for_conversation(conversation_id)
|
|
result = self._conn.execute(
|
|
text(
|
|
"DELETE FROM conversations "
|
|
"WHERE id = CAST(:id AS uuid) AND user_id = :user_id"
|
|
),
|
|
{"id": conversation_id, "user_id": user_id},
|
|
)
|
|
deleted = result.rowcount > 0
|
|
if deleted:
|
|
artifacts.delete_for_conversation(conversation_id)
|
|
self._reap_artifact_storage(paths)
|
|
return deleted
|
|
|
|
def delete_all_for_user(self, user_id: str) -> int:
|
|
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
|
|
|
# Reap artifacts BEFORE the conversations: delete_for_user_conversations
|
|
# resolves them via a subquery over the still-present conversation rows.
|
|
artifacts = ArtifactsRepository(self._conn)
|
|
paths = artifacts.storage_paths_for_user_conversations(user_id)
|
|
artifacts.delete_for_user_conversations(user_id)
|
|
result = self._conn.execute(
|
|
text("DELETE FROM conversations WHERE user_id = :user_id"),
|
|
{"user_id": user_id},
|
|
)
|
|
self._reap_artifact_storage(paths)
|
|
return result.rowcount
|
|
|
|
@staticmethod
|
|
def _reap_artifact_storage(paths: list) -> None:
|
|
"""Best-effort delete artifact bytes whose rows were cascade-removed; never raises."""
|
|
if not paths:
|
|
return
|
|
try:
|
|
from application.storage.storage_creator import StorageCreator
|
|
|
|
storage = StorageCreator.get_storage()
|
|
except Exception:
|
|
logger.warning("conversation delete: storage unavailable for artifact cleanup", exc_info=True)
|
|
return
|
|
for path in paths:
|
|
try:
|
|
storage.delete_file(path)
|
|
except Exception:
|
|
logger.warning("conversation delete: failed to delete artifact bytes %s", path, exc_info=True)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Messages
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_messages(self, conversation_id: str) -> list[dict]:
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid) "
|
|
"ORDER BY position ASC"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
)
|
|
return [_message_row_to_dict(r) for r in result.fetchall()]
|
|
|
|
def first_n_message_ids(self, conversation_id: str, first_n: int) -> set[str]:
|
|
"""Return the ids of the first ``first_n`` messages (by position) of the conversation.
|
|
|
|
Same ordering a share snapshot uses (``get_messages``' ``position ASC``
|
|
then ``[:first_n]``); empty set for ``first_n <= 0`` or a non-UUID id.
|
|
"""
|
|
if first_n <= 0 or not looks_like_uuid(conversation_id):
|
|
return set()
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT id FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:cid AS uuid) "
|
|
"ORDER BY position ASC LIMIT :n"
|
|
),
|
|
{"cid": conversation_id, "n": int(first_n)},
|
|
)
|
|
return {str(r[0]) for r in result.fetchall()}
|
|
|
|
def message_in_first_n(self, conversation_id: str, message_id: str, first_n: int) -> bool:
|
|
"""True if message_id is among the first ``first_n`` messages (by position) of the conversation."""
|
|
if not message_id or first_n <= 0:
|
|
return False
|
|
# Shape-gate both ids: a non-UUID reaching CAST(... AS uuid) would raise
|
|
# and poison the enclosing transaction (see ``rename``).
|
|
if not looks_like_uuid(conversation_id) or not looks_like_uuid(message_id):
|
|
return False
|
|
row = self._conn.execute(
|
|
text(
|
|
"SELECT 1 FROM ( "
|
|
" SELECT id FROM conversation_messages "
|
|
" WHERE conversation_id = CAST(:cid AS uuid) "
|
|
" ORDER BY position ASC LIMIT :n "
|
|
") t WHERE t.id = CAST(:mid AS uuid)"
|
|
),
|
|
{"cid": conversation_id, "mid": message_id, "n": int(first_n)},
|
|
).fetchone()
|
|
return row is not None
|
|
|
|
def get_message_at(self, conversation_id: str, position: int) -> Optional[dict]:
|
|
# Shape-gate: see ``rename``. Callers today always pass a resolved
|
|
# UUID (via ``get_any`` first), but the guard costs nothing and
|
|
# keeps future callers safe from txn-poisoning.
|
|
if not looks_like_uuid(conversation_id):
|
|
return None
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT * FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid) "
|
|
"AND position = :pos"
|
|
),
|
|
{"conv_id": conversation_id, "pos": position},
|
|
)
|
|
row = result.fetchone()
|
|
return _message_row_to_dict(row) if row is not None else None
|
|
|
|
def append_message(self, conversation_id: str, message: dict) -> dict:
|
|
"""Append a message to a conversation.
|
|
|
|
Uses ``SELECT ... FOR UPDATE`` to allocate the next position
|
|
atomically. The caller must be inside a transaction.
|
|
|
|
Mirrors Mongo's ``$push`` on the ``queries`` array.
|
|
"""
|
|
# Lock the parent conversation row to serialize concurrent appends.
|
|
self._conn.execute(
|
|
text(
|
|
"SELECT id FROM conversations "
|
|
"WHERE id = CAST(:conv_id AS uuid) FOR UPDATE"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
)
|
|
next_pos_result = self._conn.execute(
|
|
text(
|
|
"SELECT COALESCE(MAX(position), -1) + 1 AS next_pos "
|
|
"FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid)"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
)
|
|
next_pos = next_pos_result.scalar()
|
|
|
|
values = {
|
|
"conversation_id": conversation_id,
|
|
"position": next_pos,
|
|
"prompt": message.get("prompt"),
|
|
"response": message.get("response"),
|
|
"thought": message.get("thought"),
|
|
"sources": message.get("sources") or [],
|
|
"tool_calls": message.get("tool_calls") or [],
|
|
"model_id": message.get("model_id"),
|
|
"message_metadata": message.get("metadata") or {},
|
|
}
|
|
if message.get("timestamp") is not None:
|
|
values["timestamp"] = message["timestamp"]
|
|
|
|
attachments = message.get("attachments")
|
|
if attachments:
|
|
# Attachment ids may arrive as Mongo ObjectIds during the
|
|
# dual-write window — resolve each to a PG UUID or drop it.
|
|
resolved = self._resolve_attachment_refs(
|
|
[str(a) for a in attachments],
|
|
)
|
|
if resolved:
|
|
values["attachments"] = resolved
|
|
|
|
stmt = (
|
|
pg_insert(conversation_messages_table)
|
|
.values(**values)
|
|
.returning(conversation_messages_table)
|
|
)
|
|
result = self._conn.execute(stmt)
|
|
# Touch the parent conversation's updated_at.
|
|
self._conn.execute(
|
|
text(
|
|
"UPDATE conversations SET updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": conversation_id},
|
|
)
|
|
return _message_row_to_dict(result.fetchone())
|
|
|
|
def update_message_at(
|
|
self, conversation_id: str, position: int, fields: dict,
|
|
) -> bool:
|
|
"""Update specific fields on a message at a given position.
|
|
|
|
Mirrors Mongo's ``$set`` on ``queries.{index}.*``.
|
|
"""
|
|
allowed = {
|
|
"prompt", "response", "thought", "sources", "tool_calls",
|
|
"attachments", "model_id", "metadata", "timestamp",
|
|
# Feedback can be re-set in rare continuation flows; without
|
|
# it in the whitelist an upstream re-append that happens to
|
|
# carry feedback would silently lose it. Mirrors
|
|
# ``set_feedback`` — column is JSONB.
|
|
"feedback", "feedback_timestamp",
|
|
}
|
|
filtered = {k: v for k, v in fields.items() if k in allowed}
|
|
if not filtered:
|
|
return False
|
|
|
|
# Map public API key ``metadata`` → DB column ``message_metadata``.
|
|
api_to_col = {"metadata": "message_metadata"}
|
|
|
|
set_parts = []
|
|
params: dict = {"conv_id": conversation_id, "pos": position}
|
|
for key, val in filtered.items():
|
|
col = api_to_col.get(key, key)
|
|
if key in ("sources", "tool_calls", "metadata", "feedback"):
|
|
set_parts.append(f"{col} = CAST(:{col} AS jsonb)")
|
|
if val is None:
|
|
params[col] = None
|
|
else:
|
|
params[col] = (
|
|
json.dumps(val) if not isinstance(val, str) else val
|
|
)
|
|
elif key == "attachments":
|
|
# Attachment ids may be Mongo ObjectIds during the
|
|
# dual-write window; translate via attachments.legacy_mongo_id.
|
|
set_parts.append(f"{col} = CAST(:{col} AS uuid[])")
|
|
params[col] = self._resolve_attachment_refs(
|
|
[str(a) for a in val] if val else [],
|
|
)
|
|
else:
|
|
set_parts.append(f"{col} = :{col}")
|
|
params[col] = val
|
|
|
|
if "timestamp" not in filtered:
|
|
set_parts.append("timestamp = now()")
|
|
sql = (
|
|
f"UPDATE conversation_messages SET {', '.join(set_parts)} "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid) AND position = :pos"
|
|
)
|
|
result = self._conn.execute(text(sql), params)
|
|
return result.rowcount > 0
|
|
|
|
def reserve_message(
|
|
self,
|
|
conversation_id: str,
|
|
*,
|
|
prompt: str,
|
|
placeholder_response: str,
|
|
request_id: str | None = None,
|
|
status: str = "pending",
|
|
attachments: list[str] | None = None,
|
|
model_id: str | None = None,
|
|
metadata: dict | None = None,
|
|
) -> dict:
|
|
"""Pre-persist a placeholder assistant message before the LLM call."""
|
|
self._conn.execute(
|
|
text(
|
|
"SELECT id FROM conversations "
|
|
"WHERE id = CAST(:conv_id AS uuid) FOR UPDATE"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
)
|
|
next_pos = self._conn.execute(
|
|
text(
|
|
"SELECT COALESCE(MAX(position), -1) + 1 AS next_pos "
|
|
"FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid)"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
).scalar()
|
|
|
|
values = {
|
|
"conversation_id": conversation_id,
|
|
"position": next_pos,
|
|
"prompt": prompt,
|
|
"response": placeholder_response,
|
|
"status": status,
|
|
"request_id": request_id,
|
|
"model_id": model_id,
|
|
"message_metadata": metadata or {},
|
|
}
|
|
if attachments:
|
|
resolved = self._resolve_attachment_refs(
|
|
[str(a) for a in attachments],
|
|
)
|
|
if resolved:
|
|
values["attachments"] = resolved
|
|
|
|
stmt = (
|
|
pg_insert(conversation_messages_table)
|
|
.values(**values)
|
|
.returning(conversation_messages_table)
|
|
)
|
|
result = self._conn.execute(stmt)
|
|
self._conn.execute(
|
|
text(
|
|
"UPDATE conversations SET updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid)"
|
|
),
|
|
{"id": conversation_id},
|
|
)
|
|
return _message_row_to_dict(result.fetchone())
|
|
|
|
def update_message_by_id(
|
|
self, message_id: str, fields: dict,
|
|
*, only_if_non_terminal: bool = False,
|
|
) -> MessageUpdateOutcome:
|
|
"""Update specific fields on a message identified by its UUID.
|
|
|
|
``metadata`` is merged into the existing JSONB rather than
|
|
overwritten, so a reconciler-set ``reconcile_attempts`` survives
|
|
a successful late finalize. When ``only_if_non_terminal`` is
|
|
True, the update is gated so a late finalize cannot retract a
|
|
reconciler-set ``failed`` (or a prior ``complete``).
|
|
|
|
The return value discriminates "I updated the row" from "the
|
|
row was already at a terminal state" so the abort handler can
|
|
journal ``end`` when the normal-path finalize already ran.
|
|
"""
|
|
if not looks_like_uuid(message_id):
|
|
return MessageUpdateOutcome.INVALID
|
|
allowed = {
|
|
"prompt", "response", "thought", "sources", "tool_calls",
|
|
"attachments", "model_id", "metadata", "timestamp", "status",
|
|
"request_id", "feedback", "feedback_timestamp",
|
|
}
|
|
filtered = {k: v for k, v in fields.items() if k in allowed}
|
|
if not filtered:
|
|
return MessageUpdateOutcome.INVALID
|
|
|
|
api_to_col = {"metadata": "message_metadata"}
|
|
|
|
set_parts = []
|
|
params: dict = {"id": message_id}
|
|
for key, val in filtered.items():
|
|
col = api_to_col.get(key, key)
|
|
if key == "metadata":
|
|
if val is None:
|
|
set_parts.append(f"{col} = NULL")
|
|
else:
|
|
set_parts.append(
|
|
f"{col} = COALESCE({col}, '{{}}'::jsonb) "
|
|
f"|| CAST(:{col} AS jsonb)"
|
|
)
|
|
params[col] = (
|
|
json.dumps(val) if not isinstance(val, str) else val
|
|
)
|
|
elif key in ("sources", "tool_calls", "feedback"):
|
|
set_parts.append(f"{col} = CAST(:{col} AS jsonb)")
|
|
if val is None:
|
|
params[col] = None
|
|
else:
|
|
params[col] = (
|
|
json.dumps(val) if not isinstance(val, str) else val
|
|
)
|
|
elif key == "attachments":
|
|
set_parts.append(f"{col} = CAST(:{col} AS uuid[])")
|
|
params[col] = self._resolve_attachment_refs(
|
|
[str(a) for a in val] if val else [],
|
|
)
|
|
else:
|
|
set_parts.append(f"{col} = :{col}")
|
|
params[col] = val
|
|
|
|
set_parts.append("updated_at = now()")
|
|
update_where = ["id = CAST(:id AS uuid)"]
|
|
if only_if_non_terminal:
|
|
update_where.append("status NOT IN ('complete', 'failed')")
|
|
# Single-statement attempt + prior-status probe. Both CTEs see
|
|
# the same MVCC snapshot, so ``prior.status`` reflects the row
|
|
# state before the UPDATE — exactly what we need to tell
|
|
# ``ALREADY_COMPLETE`` apart from ``ALREADY_FAILED`` apart from
|
|
# ``NOT_FOUND`` without a follow-up SELECT.
|
|
sql = (
|
|
"WITH attempted AS ("
|
|
f" UPDATE conversation_messages SET {', '.join(set_parts)} "
|
|
f" WHERE {' AND '.join(update_where)} "
|
|
" RETURNING 1 AS updated"
|
|
"), "
|
|
"prior AS ("
|
|
" SELECT status FROM conversation_messages "
|
|
" WHERE id = CAST(:id AS uuid)"
|
|
") "
|
|
"SELECT (SELECT updated FROM attempted) AS updated, "
|
|
" (SELECT status FROM prior) AS prior_status"
|
|
)
|
|
row = self._conn.execute(text(sql), params).fetchone()
|
|
if row is None:
|
|
return MessageUpdateOutcome.NOT_FOUND
|
|
updated, prior_status = row[0], row[1]
|
|
if updated:
|
|
return MessageUpdateOutcome.UPDATED
|
|
if prior_status is None:
|
|
return MessageUpdateOutcome.NOT_FOUND
|
|
if prior_status == "complete":
|
|
return MessageUpdateOutcome.ALREADY_COMPLETE
|
|
if prior_status == "failed":
|
|
return MessageUpdateOutcome.ALREADY_FAILED
|
|
# ``only_if_non_terminal=False`` always updates an existing row,
|
|
# so reaching here means the gate excluded it for some status
|
|
# the terminal set doesn't cover — treat as "not found" rather
|
|
# than inventing a new variant.
|
|
return MessageUpdateOutcome.NOT_FOUND
|
|
|
|
def update_message_status(
|
|
self, message_id: str, status: str,
|
|
) -> bool:
|
|
"""Cheap status-only transition (e.g. pending → streaming).
|
|
|
|
Only flips non-terminal rows: a reconciler-set ``failed`` row
|
|
stays put so the late streaming chunk doesn't silently retract
|
|
the alert.
|
|
"""
|
|
if not looks_like_uuid(message_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE conversation_messages SET status = :status, "
|
|
"updated_at = now() "
|
|
"WHERE id = CAST(:id AS uuid) "
|
|
"AND status NOT IN ('complete', 'failed')"
|
|
),
|
|
{"id": message_id, "status": status},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def heartbeat_message(self, message_id: str) -> bool:
|
|
"""Stamp ``message_metadata.last_heartbeat_at`` with ``clock_timestamp()``.
|
|
|
|
The reconciler's staleness check uses ``GREATEST(timestamp,
|
|
last_heartbeat_at)``, so this call extends a long-running
|
|
stream's effective freshness without touching ``timestamp`` (the
|
|
creation time, used for history sort) or ``status`` (the WAL
|
|
marker). Skips terminal rows so a late heartbeat can't silently
|
|
retract a reconciler-set ``failed``.
|
|
"""
|
|
if not looks_like_uuid(message_id):
|
|
return False
|
|
result = self._conn.execute(
|
|
text(
|
|
"""
|
|
UPDATE conversation_messages
|
|
SET message_metadata = jsonb_set(
|
|
COALESCE(message_metadata, '{}'::jsonb),
|
|
'{last_heartbeat_at}',
|
|
to_jsonb(clock_timestamp())
|
|
)
|
|
WHERE id = CAST(:id AS uuid)
|
|
AND status NOT IN ('complete', 'failed')
|
|
"""
|
|
),
|
|
{"id": message_id},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def confirm_executed_tool_calls(self, message_id: str) -> int:
|
|
"""Flip ``tool_call_attempts.status='executed' → 'confirmed'`` for the message."""
|
|
if not looks_like_uuid(message_id):
|
|
return 0
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE tool_call_attempts SET status = 'confirmed', "
|
|
"updated_at = now() "
|
|
"WHERE message_id = CAST(:mid AS uuid) AND status = 'executed'"
|
|
),
|
|
{"mid": message_id},
|
|
)
|
|
return result.rowcount or 0
|
|
|
|
def truncate_after(self, conversation_id: str, keep_up_to: int) -> int:
|
|
"""Delete messages with position > keep_up_to.
|
|
|
|
Mirrors Mongo's ``$push`` + ``$slice`` that trims queries after an
|
|
index-based update.
|
|
"""
|
|
# Shape-gate: see ``rename`` — prevents transaction poisoning when
|
|
# a non-UUID id reaches this code path.
|
|
if not looks_like_uuid(conversation_id):
|
|
return 0
|
|
result = self._conn.execute(
|
|
text(
|
|
"DELETE FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid) "
|
|
"AND position > :pos"
|
|
),
|
|
{"conv_id": conversation_id, "pos": keep_up_to},
|
|
)
|
|
return result.rowcount
|
|
|
|
def set_feedback(
|
|
self, conversation_id: str, position: int, feedback: dict | None,
|
|
) -> bool:
|
|
"""Set or unset feedback on a message.
|
|
|
|
``feedback`` is a JSONB value, e.g. ``{"text": "thumbs_up",
|
|
"timestamp": "..."}`` or ``None`` to unset.
|
|
"""
|
|
# Shape-gate: see ``rename`` — prevents transaction poisoning when
|
|
# a non-UUID id reaches this code path.
|
|
if not looks_like_uuid(conversation_id):
|
|
return False
|
|
fb_json = json.dumps(feedback) if feedback is not None else None
|
|
result = self._conn.execute(
|
|
text(
|
|
"UPDATE conversation_messages "
|
|
"SET feedback = CAST(:fb AS jsonb) "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid) AND position = :pos"
|
|
),
|
|
{"conv_id": conversation_id, "pos": position, "fb": fb_json},
|
|
)
|
|
return result.rowcount > 0
|
|
|
|
def message_count(self, conversation_id: str) -> int:
|
|
result = self._conn.execute(
|
|
text(
|
|
"SELECT COUNT(*) FROM conversation_messages "
|
|
"WHERE conversation_id = CAST(:conv_id AS uuid)"
|
|
),
|
|
{"conv_id": conversation_id},
|
|
)
|
|
return result.scalar() or 0
|