Removing attachment metadata and streamlining conversation attachments.

Author: sriroopar

src/cache/postgres_cache.py

@@ -8,6 +8,7 @@
 from cache.cache_error import CacheError
 from models.cache_entry import CacheEntry
 from models.config import PostgreSQLDatabaseConfiguration
+from models.requests import Attachment
 from models.responses import ConversationData
 from utils.connection_decorator import connection
 from utils.types import ReferencedDocument, ToolCallSummary, ToolResultSummary
@@ -36,6 +37,7 @@ class PostgresCache(Cache):
      referenced_documents  | jsonb                          |          |
      tool_calls            | jsonb                          |          |
      tool_results          | jsonb                          |          |
+     attachments           | jsonb                          |          |
     Indexes:
         "cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
         "timestamps" btree (created_at)
@@ -60,6 +62,7 @@ class PostgresCache(Cache):
             referenced_documents jsonb,
             tool_calls           jsonb,
             tool_results         jsonb,
+            attachments          jsonb,
             PRIMARY KEY(user_id, conversation_id, created_at)
         );
         """
@@ -81,7 +84,7 @@ class PostgresCache(Cache):
 
     SELECT_CONVERSATION_HISTORY_STATEMENT = """
         SELECT query, response, provider, model, started_at, completed_at,
-               referenced_documents, tool_calls, tool_results
+               referenced_documents, tool_calls, tool_results, attachments
           FROM cache
          WHERE user_id=%s AND conversation_id=%s
          ORDER BY created_at
@@ -90,8 +93,8 @@ class PostgresCache(Cache):
     INSERT_CONVERSATION_HISTORY_STATEMENT = """
         INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
                           query, response, provider, model, referenced_documents,
-                          tool_calls, tool_results)
-        VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                          tool_calls, tool_results, attachments)
+        VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
         """
 
     QUERY_CACHE_SIZE = """
@@ -321,6 +324,22 @@ def get(  # pylint: disable=R0914
                             e,
                         )
 
+                # Parse attachments back into Attachment objects
+                attachments_data = conversation_entry[9]
+                attachments_obj = None
+                if attachments_data:
+                    try:
+                        attachments_obj = [
+                            Attachment.model_validate(att) for att in attachments_data
+                        ]
+                    except (ValueError, TypeError) as e:
+                        logger.warning(
+                            "Failed to deserialize attachments for "
+                            "conversation %s: %s",
+                            conversation_id,
+                            e,
+                        )
+
                 cache_entry = CacheEntry(
                     query=conversation_entry[0],
                     response=conversation_entry[1],
@@ -331,6 +350,7 @@ def get(  # pylint: disable=R0914
                     referenced_documents=docs_obj,
                     tool_calls=tool_calls_obj,
                     tool_results=tool_results_obj,
+                    attachments=attachments_obj,
                 )
                 result.append(cache_entry)
 
@@ -405,6 +425,20 @@ def insert_or_append(
                         e,
                     )
 
+            attachments_json = None
+            if cache_entry.attachments:
+                try:
+                    attachments_as_dicts = [
+                        att.model_dump(mode="json") for att in cache_entry.attachments
+                    ]
+                    attachments_json = json.dumps(attachments_as_dicts)
+                except (TypeError, ValueError) as e:
+                    logger.warning(
+                        "Failed to serialize attachments for conversation %s: %s",
+                        conversation_id,
+                        e,
+                    )
+
             # the whole operation is run in one transaction
             with self.connection.cursor() as cursor:
                 cursor.execute(
@@ -421,6 +455,7 @@ def insert_or_append(
                         referenced_documents_json,
                         tool_calls_json,
                         tool_results_json,
+                        attachments_json,
                     ),
                 )
 

src/cache/sqlite_cache.py

@@ -1,6 +1,7 @@
 """Cache that uses SQLite to store cached values."""
 
 from time import time
+from typing import Any
 
 import sqlite3
 import json
@@ -9,6 +10,7 @@
 from cache.cache_error import CacheError
 from models.cache_entry import CacheEntry
 from models.config import SQLiteDatabaseConfiguration
+from models.requests import Attachment
 from models.responses import ConversationData
 from utils.connection_decorator import connection
 from utils.types import ReferencedDocument, ToolCallSummary, ToolResultSummary
@@ -37,6 +39,7 @@ class SQLiteCache(Cache):
      referenced_documents  | text                        |          |
      tool_calls            | text                        |          |
      tool_results          | text                        |          |
+     attachments           | text                        |          |
     Indexes:
         "cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
         "cache_key_key" UNIQUE CONSTRAINT, btree (key)
@@ -45,6 +48,26 @@ class SQLiteCache(Cache):
     ```
     """
 
+    @staticmethod
+    def _safe_json_dumps_models(
+        items: list[Any] | None, conversation_id: str, field_name: str
+    ) -> str | None:
+        """Serialize a list of Pydantic models to JSON, returning None on failure."""
+        if not items:
+            return None
+
+        try:
+            as_dicts = [item.model_dump(mode="json") for item in items]
+            return json.dumps(as_dicts)
+        except (TypeError, ValueError) as e:
+            logger.warning(
+                "Failed to serialize %s for conversation %s: %s",
+                field_name,
+                conversation_id,
+                e,
+            )
+            return None
+
     CREATE_CACHE_TABLE = """
         CREATE TABLE IF NOT EXISTS cache (
             user_id              text NOT NULL,
@@ -59,6 +82,7 @@ class SQLiteCache(Cache):
             referenced_documents text,
             tool_calls           text,
             tool_results         text,
+            attachments          text,
             PRIMARY KEY(user_id, conversation_id, created_at)
         );
         """
@@ -80,7 +104,7 @@ class SQLiteCache(Cache):
 
     SELECT_CONVERSATION_HISTORY_STATEMENT = """
         SELECT query, response, provider, model, started_at, completed_at,
-               referenced_documents, tool_calls, tool_results
+               referenced_documents, tool_calls, tool_results, attachments
           FROM cache
          WHERE user_id=? AND conversation_id=?
          ORDER BY created_at
@@ -89,8 +113,8 @@ class SQLiteCache(Cache):
     INSERT_CONVERSATION_HISTORY_STATEMENT = """
         INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
                           query, response, provider, model, referenced_documents,
-                          tool_calls, tool_results)
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                          tool_calls, tool_results, attachments)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
         """
 
     QUERY_CACHE_SIZE = """
@@ -301,6 +325,22 @@ def get(  # pylint: disable=R0914
                         e,
                     )
 
+            # Parse attachments back into Attachment objects
+            attachments_json_str = conversation_entry[9]
+            attachments_obj = None
+            if attachments_json_str:
+                try:
+                    attachments_data = json.loads(attachments_json_str)
+                    attachments_obj = [
+                        Attachment.model_validate(att) for att in attachments_data
+                    ]
+                except (json.JSONDecodeError, ValueError) as e:
+                    logger.warning(
+                        "Failed to deserialize attachments for conversation %s: %s",
+                        conversation_id,
+                        e,
+                    )
+
             cache_entry = CacheEntry(
                 query=conversation_entry[0],
                 response=conversation_entry[1],
@@ -311,6 +351,7 @@ def get(  # pylint: disable=R0914
                 referenced_documents=docs_obj,
                 tool_calls=tool_calls_obj,
                 tool_results=tool_results_obj,
+                attachments=attachments_obj,
             )
             result.append(cache_entry)
 
@@ -342,49 +383,20 @@ def insert_or_append(
         cursor = self.connection.cursor()
         current_time = time()
 
-        referenced_documents_json = None
-        if cache_entry.referenced_documents:
-            try:
-                docs_as_dicts = [
-                    doc.model_dump(mode="json")
-                    for doc in cache_entry.referenced_documents
-                ]
-                referenced_documents_json = json.dumps(docs_as_dicts)
-            except (TypeError, ValueError) as e:
-                logger.warning(
-                    "Failed to serialize referenced_documents for "
-                    "conversation %s: %s",
-                    conversation_id,
-                    e,
-                )
-
-        tool_calls_json = None
-        if cache_entry.tool_calls:
-            try:
-                tool_calls_as_dicts = [
-                    tc.model_dump(mode="json") for tc in cache_entry.tool_calls
-                ]
-                tool_calls_json = json.dumps(tool_calls_as_dicts)
-            except (TypeError, ValueError) as e:
-                logger.warning(
-                    "Failed to serialize tool_calls for conversation %s: %s",
-                    conversation_id,
-                    e,
-                )
-
-        tool_results_json = None
-        if cache_entry.tool_results:
-            try:
-                tool_results_as_dicts = [
-                    tr.model_dump(mode="json") for tr in cache_entry.tool_results
-                ]
-                tool_results_json = json.dumps(tool_results_as_dicts)
-            except (TypeError, ValueError) as e:
-                logger.warning(
-                    "Failed to serialize tool_results for conversation %s: %s",
-                    conversation_id,
-                    e,
-                )
+        referenced_documents_json = self._safe_json_dumps_models(
+            cache_entry.referenced_documents,
+            conversation_id,
+            "referenced_documents",
+        )
+        tool_calls_json = self._safe_json_dumps_models(
+            cache_entry.tool_calls, conversation_id, "tool_calls"
+        )
+        tool_results_json = self._safe_json_dumps_models(
+            cache_entry.tool_results, conversation_id, "tool_results"
+        )
+        attachments_json = self._safe_json_dumps_models(
+            cache_entry.attachments, conversation_id, "attachments"
+        )
 
         cursor.execute(
             self.INSERT_CONVERSATION_HISTORY_STATEMENT,
@@ -401,6 +413,7 @@ def insert_or_append(
                 referenced_documents_json,
                 tool_calls_json,
                 tool_results_json,
+                attachments_json,
             ),
         )
 

src/models/cache_entry.py

@@ -2,6 +2,8 @@
 
 from typing import Optional
 from pydantic import BaseModel
+
+from models.requests import Attachment
 from utils.types import ReferencedDocument, ToolCallSummary, ToolResultSummary
 
 
@@ -16,6 +18,7 @@ class CacheEntry(BaseModel):
         referenced_documents: List of documents referenced by the response
         tool_calls: List of tool calls made during response generation
         tool_results: List of tool results from tool calls
+        attachments: Optional list of attachments included with the query
     """
 
     query: str
@@ -27,3 +30,4 @@ class CacheEntry(BaseModel):
     referenced_documents: Optional[list[ReferencedDocument]] = None
     tool_calls: Optional[list[ToolCallSummary]] = None
     tool_results: Optional[list[ToolResultSummary]] = None
+    attachments: Optional[list[Attachment]] = None

src/models/responses.py

@@ -19,6 +19,7 @@
 
 from constants import MEDIA_TYPE_EVENT_STREAM
 from models.config import Action, Configuration
+from models.requests import Attachment
 from quota.quota_exceed_error import QuotaExceedError
 from utils.types import RAGChunk, ReferencedDocument, ToolCallSummary, ToolResultSummary
 
@@ -867,6 +868,7 @@ class Message(BaseModel):
         content: The message content.
         type: The type of message.
         referenced_documents: Optional list of documents referenced in an assistant response.
+        attachments: Optional list of attachments included with the message.
     """
 
     content: str = Field(
@@ -879,10 +881,6 @@ class Message(BaseModel):
         description="The type of message",
         examples=["user", "assistant", "system", "developer"],
     )
-    referenced_documents: Optional[list[ReferencedDocument]] = Field(
-        None,
-        description="List of documents referenced in the response (assistant messages only)",
-    )
 
 
 class ConversationTurn(BaseModel):

src/utils/query.py

@@ -309,6 +309,19 @@ def store_query_results(  # pylint: disable=too-many-arguments
         tool_results=summary.tool_results,
     )
     try:
+        cache_entry = CacheEntry(
+            query=query_request.query,
+            response=summary.llm_response,
+            provider=provider_id,
+            model=model_id,
+            started_at=started_at,
+            completed_at=completed_at,
+            referenced_documents=summary.referenced_documents,
+            tool_calls=summary.tool_calls,
+            tool_results=summary.tool_results,
+            attachments=query_request.attachments,
+        )
+
         logger.info("Storing conversation in cache")
         store_conversation_into_cache(
             user_id=user_id,