First draft of referenced_documents caching

Signed-off-by: Maysun J Faisal <maysunaneek@gmail.com>

Author: maysunfaisal

src/app/endpoints/conversations_v2.py

@@ -240,13 +240,23 @@ def check_conversation_existence(user_id: str, conversation_id: str) -> None:
 
 def transform_chat_message(entry: CacheEntry) -> dict[str, Any]:
     """Transform the message read from cache into format used by response payload."""
+    user_message = {
+        "content": entry.query, 
+        "type": "user"
+    }
+    assistant_message: dict[str, Any] = {
+        "content": entry.response,
+        "type": "assistant"
+    }
+
+    # Check for additional_kwargs and add it to the assistant message if it exists
+    if entry.additional_kwargs:
+        assistant_message["additional_kwargs"] = entry.additional_kwargs.model_dump()
+
     return {
         "provider": entry.provider,
         "model": entry.model,
-        "messages": [
-            {"content": entry.query, "type": "user"},
-            {"content": entry.response, "type": "assistant"},
-        ],
+        "messages": [user_message, assistant_message],
         "started_at": entry.started_at,
         "completed_at": entry.completed_at,
     }

src/app/endpoints/query.py

@@ -22,6 +22,7 @@
 from llama_stack_client.types.model_list_response import ModelListResponse
 from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
 from llama_stack_client.types.tool_execution_step import ToolExecutionStep
+from pydantic import AnyUrl
 
 import constants
 import metrics
@@ -32,6 +33,7 @@
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
 from metrics.utils import update_llm_token_count_from_turn
+from models.cache_entry import CacheEntry, AdditionalKwargs
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import Attachment, QueryRequest
@@ -331,16 +333,27 @@ async def query_endpoint_handler(  # pylint: disable=R0914
         )
 
         completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+        additional_kwargs_obj = None
+        if referenced_documents:
+            additional_kwargs_obj = AdditionalKwargs(
+                referenced_documents=referenced_documents
+            )
+        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,
+            additional_kwargs=additional_kwargs_obj
+        )
+            
         store_conversation_into_cache(
             configuration,
             user_id,
             conversation_id,
-            provider_id,
-            model_id,
-            query_request.query,
-            summary.llm_response,
-            started_at,
-            completed_at,
+            cache_entry,
             _skip_userid_check,
             topic_summary,
         )

src/app/endpoints/streaming_query.py

@@ -20,6 +20,7 @@
 )
 from llama_stack_client.types.shared import ToolCall
 from llama_stack_client.types.shared.interleaved_content_item import TextContentItem
+from pydantic import AnyUrl
 
 from app.database import get_session
 from app.endpoints.query import (
@@ -42,10 +43,11 @@
 from constants import DEFAULT_RAG_TOOL
 import metrics
 from metrics.utils import update_llm_token_count_from_turn
+from models.cache_entry import CacheEntry, AdditionalKwargs
 from models.config import Action
 from models.database.conversations import UserConversation
 from models.requests import QueryRequest
-from models.responses import ForbiddenResponse, UnauthorizedResponse
+from models.responses import ForbiddenResponse, UnauthorizedResponse, ReferencedDocument
 from utils.endpoints import (
     check_configuration_loaded,
     create_referenced_documents_with_metadata,
@@ -733,16 +735,29 @@ async def response_generator(
                     )
 
             completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+            referenced_documents = create_referenced_documents_with_metadata(summary, metadata_map)
+
+            additional_kwargs_obj = None
+            if referenced_documents:
+                additional_kwargs_obj = AdditionalKwargs(
+                    referenced_documents=referenced_documents
+                )
+            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,
+                additional_kwargs=additional_kwargs_obj
+            )
+            
             store_conversation_into_cache(
                 configuration,
                 user_id,
                 conversation_id,
-                provider_id,
-                model_id,
-                query_request.query,
-                summary.llm_response,
-                started_at,
-                completed_at,
+                cache_entry,
                 _skip_userid_check,
                 topic_summary,
             )

src/cache/cache.py

@@ -2,7 +2,8 @@
 
 from abc import ABC, abstractmethod
 
-from models.cache_entry import CacheEntry, ConversationData
+from models.cache_entry import CacheEntry
+from models.responses import ConversationData
 from utils.suid import check_suid
 
 

src/cache/in_memory_cache.py

@@ -1,8 +1,9 @@
 """In-memory cache implementation."""
 
 from cache.cache import Cache
-from models.cache_entry import CacheEntry, ConversationData
+from models.cache_entry import CacheEntry
 from models.config import InMemoryCacheConfig
+from models.responses import ConversationData
 from log import get_logger
 from utils.connection_decorator import connection
 

src/cache/noop_cache.py

@@ -1,7 +1,8 @@
 """No-operation cache implementation."""
 
 from cache.cache import Cache
-from models.cache_entry import CacheEntry, ConversationData
+from models.cache_entry import CacheEntry
+from models.responses import ConversationData
 from log import get_logger
 from utils.connection_decorator import connection
 

src/cache/postgres_cache.py

@@ -4,8 +4,9 @@
 
 from cache.cache import Cache
 from cache.cache_error import CacheError
-from models.cache_entry import CacheEntry, ConversationData
+from models.cache_entry import CacheEntry, AdditionalKwargs
 from models.config import PostgreSQLDatabaseConfiguration
+from models.responses import ConversationData
 from log import get_logger
 from utils.connection_decorator import connection
 
@@ -37,15 +38,16 @@ class PostgresCache(Cache):
 
     CREATE_CACHE_TABLE = """
         CREATE TABLE IF NOT EXISTS cache (
-            user_id         text NOT NULL,
-            conversation_id text NOT NULL,
-            created_at      timestamp NOT NULL,
-            started_at      text,
-            completed_at    text,
-            query           text,
-            response        text,
-            provider        text,
-            model           text,
+            user_id           text NOT NULL,
+            conversation_id   text NOT NULL,
+            created_at        timestamp NOT NULL,
+            started_at        text,
+            completed_at      text,
+            query             text,
+            response          text,
+            provider          text,
+            model             text,
+            additional_kwargs jsonb,
             PRIMARY KEY(user_id, conversation_id, created_at)
         );
         """
@@ -66,16 +68,16 @@ class PostgresCache(Cache):
         """
 
     SELECT_CONVERSATION_HISTORY_STATEMENT = """
-        SELECT query, response, provider, model, started_at, completed_at
+        SELECT query, response, provider, model, started_at, completed_at, additional_kwargs
           FROM cache
          WHERE user_id=%s AND conversation_id=%s
          ORDER BY created_at
         """
 
     INSERT_CONVERSATION_HISTORY_STATEMENT = """
         INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
-                          query, response, provider, model)
-        VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s)
+                          query, response, provider, model, additional_kwargs)
+        VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s)
         """
 
     QUERY_CACHE_SIZE = """
@@ -211,13 +213,19 @@ def get(
 
             result = []
             for conversation_entry in conversation_entries:
+                # Parse it back into an LLMResponse object
+                additional_kwargs_data = conversation_entry[6]
+                additional_kwargs_obj = None
+                if additional_kwargs_data:
+                    additional_kwargs_obj = AdditionalKwargs.model_validate(additional_kwargs_data)
                 cache_entry = CacheEntry(
                     query=conversation_entry[0],
                     response=conversation_entry[1],
                     provider=conversation_entry[2],
                     model=conversation_entry[3],
                     started_at=conversation_entry[4],
                     completed_at=conversation_entry[5],
+                    additional_kwargs=additional_kwargs_obj,
                 )
                 result.append(cache_entry)
 
@@ -245,6 +253,10 @@ def insert_or_append(
             raise CacheError("insert_or_append: cache is disconnected")
 
         try:
+            additional_kwargs_json = None
+            if cache_entry.additional_kwargs:
+                # Use exclude_none=True to keep JSON clean
+                additional_kwargs_json = cache_entry.additional_kwargs.model_dump_json(exclude_none=True)
             # the whole operation is run in one transaction
             with self.connection.cursor() as cursor:
                 cursor.execute(
@@ -258,6 +270,7 @@ def insert_or_append(
                         cache_entry.response,
                         cache_entry.provider,
                         cache_entry.model,
+                        additional_kwargs_json,
                     ),
                 )
 

src/cache/sqlite_cache.py

@@ -3,11 +3,13 @@
 from time import time
 
 import sqlite3
+import json
 
 from cache.cache import Cache
 from cache.cache_error import CacheError
-from models.cache_entry import CacheEntry, ConversationData
+from models.cache_entry import CacheEntry, AdditionalKwargs
 from models.config import SQLiteDatabaseConfiguration
+from models.responses import ConversationData
 from log import get_logger
 from utils.connection_decorator import connection
 
@@ -41,15 +43,16 @@ class SQLiteCache(Cache):
 
     CREATE_CACHE_TABLE = """
         CREATE TABLE IF NOT EXISTS cache (
-            user_id         text NOT NULL,
-            conversation_id text NOT NULL,
-            created_at      int NOT NULL,
-            started_at      text,
-            completed_at    text,
-            query           text,
-            response        text,
-            provider        text,
-            model           text,
+            user_id           text NOT NULL,
+            conversation_id   text NOT NULL,
+            created_at        int NOT NULL,
+            started_at        text,
+            completed_at      text,
+            query             text,
+            response          text,
+            provider          text,
+            model             text,
+            additional_kwargs text,
             PRIMARY KEY(user_id, conversation_id, created_at)
         );
         """
@@ -70,16 +73,16 @@ class SQLiteCache(Cache):
         """
 
     SELECT_CONVERSATION_HISTORY_STATEMENT = """
-        SELECT query, response, provider, model, started_at, completed_at
+        SELECT query, response, provider, model, started_at, completed_at, additional_kwargs
           FROM cache
          WHERE user_id=? AND conversation_id=?
          ORDER BY created_at
         """
 
     INSERT_CONVERSATION_HISTORY_STATEMENT = """
         INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
-                          query, response, provider, model)
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                          query, response, provider, model, additional_kwargs)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
         """
 
     QUERY_CACHE_SIZE = """
@@ -209,13 +212,18 @@ def get(
 
         result = []
         for conversation_entry in conversation_entries:
+            additional_kwargs_json = conversation_entry[6]
+            additional_kwargs_obj = None
+            if additional_kwargs_json:
+                additional_kwargs_obj = AdditionalKwargs.model_validate_json(additional_kwargs_json)
             cache_entry = CacheEntry(
                 query=conversation_entry[0],
                 response=conversation_entry[1],
                 provider=conversation_entry[2],
                 model=conversation_entry[3],
                 started_at=conversation_entry[4],
                 completed_at=conversation_entry[5],
+                additional_kwargs=additional_kwargs_obj,
             )
             result.append(cache_entry)
 
@@ -244,6 +252,11 @@ def insert_or_append(
 
         cursor = self.connection.cursor()
         current_time = time()
+
+        additional_kwargs_json = None
+        if cache_entry.additional_kwargs:
+            additional_kwargs_json = cache_entry.additional_kwargs.model_dump_json(exclude_none=True)
+
         cursor.execute(
             self.INSERT_CONVERSATION_HISTORY_STATEMENT,
             (
@@ -256,6 +269,7 @@ def insert_or_append(
                 cache_entry.response,
                 cache_entry.provider,
                 cache_entry.model,
+                additional_kwargs_json,
             ),
         )