Removing attachment metadata and streamlining conversation attachments.

Author: sriroopar

src/app/endpoints/conversations_v2.py

@@ -257,7 +257,7 @@ def build_conversation_turn_from_cache_entry(entry: CacheEntry) -> ConversationT
     """
     # Create Message objects for user and assistant
     messages = [
-        Message(content=entry.query, type="user"),
+        Message(content=entry.query, type="user", attachments=entry.attachments),
         Message(content=entry.response, type="assistant"),
     ]
 

src/cache/postgres_cache.py

@@ -36,6 +36,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 +61,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 +83,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 +92,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 +323,24 @@ def get(  # pylint: disable=R0914
                             e,
                         )
 
+                # Parse attachments back into Attachment objects
+                attachments_data = conversation_entry[9]
+                attachments_obj = None
+                if attachments_data:
+                    try:
+                        from models.requests import Attachment
+
+                        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 +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)
 
@@ -405,6 +426,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 +456,7 @@ def insert_or_append(
                         referenced_documents_json,
                         tool_calls_json,
                         tool_results_json,
+                        attachments_json,
                     ),
                 )
 

src/cache/sqlite_cache.py

@@ -37,6 +37,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)
@@ -59,6 +60,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 +82,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 +91,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 +303,24 @@ 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:
+                    from models.requests import Attachment
+
+                    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 +331,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)
 
@@ -386,6 +407,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,
+                )
+
         cursor.execute(
             self.INSERT_CONVERSATION_HISTORY_STATEMENT,
             (
@@ -401,6 +436,7 @@ def insert_or_append(
                 referenced_documents_json,
                 tool_calls_json,
                 tool_results_json,
+                attachments_json,
             ),
         )
 

src/models/cache_entry.py

@@ -1,9 +1,12 @@
 """Model for conversation history cache entry."""
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 from pydantic import BaseModel
 from utils.types import ReferencedDocument, ToolCallSummary, ToolResultSummary
 
+if TYPE_CHECKING:
+    from models.requests import Attachment
+
 
 class CacheEntry(BaseModel):
     """Model representing a cache entry.
@@ -16,6 +19,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 +31,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

@@ -2,7 +2,7 @@
 
 """Models for REST API responses."""
 
-from typing import Any, ClassVar, Literal, Optional, Union
+from typing import Any, ClassVar, Literal, Optional, Union, TYPE_CHECKING
 
 from fastapi import status
 from pydantic import BaseModel, Field
@@ -13,6 +13,9 @@
 from quota.quota_exceed_error import QuotaExceedError
 from utils.types import RAGChunk, ReferencedDocument, ToolCallSummary, ToolResultSummary
 
+if TYPE_CHECKING:
+    from models.requests import Attachment
+
 SUCCESSFUL_RESPONSE_DESCRIPTION = "Successful response"
 BAD_REQUEST_DESCRIPTION = "Invalid request format"
 UNAUTHORIZED_DESCRIPTION = "Unauthorized"
@@ -810,6 +813,7 @@ class Message(BaseModel):
     Attributes:
         content: The message content.
         type: The type of message.
+        attachments: Optional list of attachments included with the message.
     """
 
     content: str = Field(
@@ -822,6 +826,10 @@ class Message(BaseModel):
         description="The type of message",
         examples=["user", "assistant", "system", "developer"],
     )
+    attachments: Optional[list["Attachment"]] = Field(
+        default=None,
+        description="Optional attachments included with this message",
+    )
 
 
 class ConversationTurn(BaseModel):

src/utils/query.py

@@ -420,6 +420,7 @@ def store_query_results(  # pylint: disable=too-many-arguments,too-many-locals
             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")