From 330623246b680d71cc31f86c9deb2e80d67b7bd2 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Fri, 3 Oct 2025 11:47:54 +0200 Subject: [PATCH 1/2] Allow started_at and completed_at timestamps to be store in database --- src/app/endpoints/conversations_v2.py | 2 ++ src/app/endpoints/query.py | 4 ++++ src/app/endpoints/streaming_query.py | 5 +++++ src/cache/postgres_cache.py | 15 ++++++++++++--- src/cache/sqlite_cache.py | 15 ++++++++++++--- src/models/cache_entry.py | 2 ++ src/utils/endpoints.py | 4 ++++ 7 files changed, 41 insertions(+), 6 deletions(-) diff --git a/src/app/endpoints/conversations_v2.py b/src/app/endpoints/conversations_v2.py index caf5b590..ad3529c8 100644 --- a/src/app/endpoints/conversations_v2.py +++ b/src/app/endpoints/conversations_v2.py @@ -247,4 +247,6 @@ def transform_chat_message(entry: CacheEntry) -> dict[str, Any]: {"content": entry.query, "type": "user"}, {"content": entry.response, "type": "assistant"}, ], + "started_at": entry.started_at, + "completed_at": entry.completed_at, } diff --git a/src/app/endpoints/query.py b/src/app/endpoints/query.py index 1af26f4d..6cb93b5a 100644 --- a/src/app/endpoints/query.py +++ b/src/app/endpoints/query.py @@ -241,6 +241,7 @@ async def query_endpoint_handler( # pylint: disable=R0914 user_id, _, _skip_userid_check, token = auth + started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") user_conversation: UserConversation | None = None if query_request.conversation_id: logger.debug( @@ -330,6 +331,7 @@ async def query_endpoint_handler( # pylint: disable=R0914 topic_summary=topic_summary, ) + completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") store_conversation_into_cache( configuration, user_id, @@ -338,6 +340,8 @@ async def query_endpoint_handler( # pylint: disable=R0914 model_id, query_request.query, summary.llm_response, + started_at, + completed_at, _skip_userid_check, topic_summary, ) diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index b4139ae9..33e4e644 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -4,6 +4,7 @@ import json import logging import re +from datetime import UTC, datetime from typing import Annotated, Any, AsyncIterator, Iterator, cast from fastapi import APIRouter, Depends, HTTPException, Request, status @@ -596,6 +597,7 @@ async def streaming_query_endpoint_handler( # pylint: disable=R0915,R0914 _ = request check_configuration_loaded(configuration) + started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") # Enforce RBAC: optionally disallow overriding model/provider in requests validate_model_provider_override(query_request, request.state.authorized_actions) @@ -719,6 +721,7 @@ async def response_generator( query_request.query, client, model_id ) + completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") store_conversation_into_cache( configuration, user_id, @@ -727,6 +730,8 @@ async def response_generator( model_id, query_request.query, summary.llm_response, + started_at, + completed_at, _skip_userid_check, topic_summary, ) diff --git a/src/cache/postgres_cache.py b/src/cache/postgres_cache.py index a8e5d033..0881b66c 100644 --- a/src/cache/postgres_cache.py +++ b/src/cache/postgres_cache.py @@ -23,6 +23,8 @@ class PostgresCache(Cache): user_id | text | not null | conversation_id | text | not null | created_at | timestamp without time zone | not null | + started_at | text | | + completed_at | text | | query | text | | response | text | | provider | text | | @@ -38,6 +40,8 @@ class PostgresCache(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, @@ -62,15 +66,16 @@ class PostgresCache(Cache): """ SELECT_CONVERSATION_HISTORY_STATEMENT = """ - SELECT query, response, provider, model + SELECT query, response, provider, model, started_at, completed_at 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, query, response, provider, model) - VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s) + 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_CACHE_SIZE = """ @@ -211,6 +216,8 @@ def get( response=conversation_entry[1], provider=conversation_entry[2], model=conversation_entry[3], + started_at=conversation_entry[4], + completed_at=conversation_entry[5], ) result.append(cache_entry) @@ -245,6 +252,8 @@ def insert_or_append( ( user_id, conversation_id, + cache_entry.started_at, + cache_entry.completed_at, cache_entry.query, cache_entry.response, cache_entry.provider, diff --git a/src/cache/sqlite_cache.py b/src/cache/sqlite_cache.py index a39f8ade..b8a91fac 100644 --- a/src/cache/sqlite_cache.py +++ b/src/cache/sqlite_cache.py @@ -25,6 +25,8 @@ class SQLiteCache(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 | | @@ -42,6 +44,8 @@ class SQLiteCache(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, @@ -66,15 +70,16 @@ class SQLiteCache(Cache): """ SELECT_CONVERSATION_HISTORY_STATEMENT = """ - SELECT query, response, provider, model + SELECT query, response, provider, model, started_at, completed_at 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, query, response, provider, model) - VALUES (?, ?, ?, ?, ?, ?, ?) + INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at, + query, response, provider, model) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """ QUERY_CACHE_SIZE = """ @@ -209,6 +214,8 @@ def get( response=conversation_entry[1], provider=conversation_entry[2], model=conversation_entry[3], + started_at=conversation_entry[4], + completed_at=conversation_entry[5], ) result.append(cache_entry) @@ -243,6 +250,8 @@ def insert_or_append( user_id, conversation_id, current_time, + cache_entry.started_at, + cache_entry.completed_at, cache_entry.query, cache_entry.response, cache_entry.provider, diff --git a/src/models/cache_entry.py b/src/models/cache_entry.py index f87445be..9f3119f1 100644 --- a/src/models/cache_entry.py +++ b/src/models/cache_entry.py @@ -17,6 +17,8 @@ class CacheEntry(BaseModel): response: str provider: str model: str + started_at: str + completed_at: str class ConversationData(BaseModel): diff --git a/src/utils/endpoints.py b/src/utils/endpoints.py index 0b0c1510..092c0d73 100644 --- a/src/utils/endpoints.py +++ b/src/utils/endpoints.py @@ -189,6 +189,8 @@ def store_conversation_into_cache( model_id: str, query: str, response: str, + started_at: str, + completed_at: str, _skip_userid_check: bool, topic_summary: str | None, ) -> None: @@ -203,6 +205,8 @@ def store_conversation_into_cache( response=response, provider=provider_id, model=model_id, + started_at=started_at, + completed_at=completed_at, ) cache.insert_or_append( user_id, conversation_id, cache_entry, _skip_userid_check From 7b71c8b5eb9dd5638063b63c2b2e6c834899b2ec Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Fri, 3 Oct 2025 11:50:06 +0200 Subject: [PATCH 2/2] Updated unit tests accordingly --- tests/unit/app/endpoints/test_conversations_v2.py | 13 ++++++++++++- tests/unit/cache/test_noop_cache.py | 14 ++++++++++++-- tests/unit/cache/test_postgres_cache.py | 14 ++++++++++++-- tests/unit/cache/test_sqlite_cache.py | 14 ++++++++++++-- 4 files changed, 48 insertions(+), 7 deletions(-) diff --git a/tests/unit/app/endpoints/test_conversations_v2.py b/tests/unit/app/endpoints/test_conversations_v2.py index d24cabd0..21cb71b5 100644 --- a/tests/unit/app/endpoints/test_conversations_v2.py +++ b/tests/unit/app/endpoints/test_conversations_v2.py @@ -10,7 +10,12 @@ def test_transform_message() -> None: """Test the transform_chat_message transformation function.""" entry = CacheEntry( - query="query", response="response", provider="provider", model="model" + query="query", + response="response", + provider="provider", + model="model", + started_at="2024-01-01T00:00:00Z", + completed_at="2024-01-01T00:00:05Z", ) transformed = transform_chat_message(entry) assert transformed is not None @@ -21,6 +26,12 @@ def test_transform_message() -> None: assert "model" in transformed assert transformed["model"] == "model" + assert "started_at" in transformed + assert transformed["started_at"] == "2024-01-01T00:00:00Z" + + assert "completed_at" in transformed + assert transformed["completed_at"] == "2024-01-01T00:00:05Z" + assert "messages" in transformed assert len(transformed["messages"]) == 2 diff --git a/tests/unit/cache/test_noop_cache.py b/tests/unit/cache/test_noop_cache.py index 55dd2cc0..aad543d2 100644 --- a/tests/unit/cache/test_noop_cache.py +++ b/tests/unit/cache/test_noop_cache.py @@ -10,10 +10,20 @@ CONVERSATION_ID = suid.get_suid() USER_PROVIDED_USER_ID = "test-user1" cache_entry_1 = CacheEntry( - query="user message1", response="AI message1", provider="foo", model="bar" + query="user message1", + response="AI message1", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", ) cache_entry_2 = CacheEntry( - query="user message2", response="AI message2", provider="foo", model="bar" + query="user message2", + response="AI message2", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", ) diff --git a/tests/unit/cache/test_postgres_cache.py b/tests/unit/cache/test_postgres_cache.py index 51e06448..32646997 100644 --- a/tests/unit/cache/test_postgres_cache.py +++ b/tests/unit/cache/test_postgres_cache.py @@ -16,10 +16,20 @@ CONVERSATION_ID_1 = suid.get_suid() CONVERSATION_ID_2 = suid.get_suid() cache_entry_1 = CacheEntry( - query="user message1", response="AI message1", provider="foo", model="bar" + query="user message1", + response="AI message1", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", ) cache_entry_2 = CacheEntry( - query="user message2", response="AI message2", provider="foo", model="bar" + query="user message2", + response="AI message2", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", ) # pylint: disable=fixme diff --git a/tests/unit/cache/test_sqlite_cache.py b/tests/unit/cache/test_sqlite_cache.py index 32cc9a46..6c9fd0fa 100644 --- a/tests/unit/cache/test_sqlite_cache.py +++ b/tests/unit/cache/test_sqlite_cache.py @@ -18,10 +18,20 @@ CONVERSATION_ID_1 = suid.get_suid() CONVERSATION_ID_2 = suid.get_suid() cache_entry_1 = CacheEntry( - query="user message1", response="AI message1", provider="foo", model="bar" + query="user message1", + response="AI message1", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", ) cache_entry_2 = CacheEntry( - query="user message2", response="AI message2", provider="foo", model="bar" + query="user message2", + response="AI message2", + provider="foo", + model="bar", + started_at="2025-10-03T09:31:25Z", + completed_at="2025-10-03T09:31:29Z", )