Skip to content

Commit 2ded492

Browse files
tisniktisnik
authored
Merge pull request lightspeed-core#618 from tisnik/lcore-776
LCORE-776: Allow started_at and completed_at timestamps to be store in database
2 parents f72b573 + 7b71c8b commit 2ded492

11 files changed

Lines changed: 89 additions & 13 deletions

File tree

‎src/app/endpoints/conversations_v2.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,4 +247,6 @@ def transform_chat_message(entry: CacheEntry) -> dict[str, Any]:
247247
{"content": entry.query, "type": "user"},
248248
{"content": entry.response, "type": "assistant"},
249249
],
250+
"started_at": entry.started_at,
251+
"completed_at": entry.completed_at,
250252
}

‎src/app/endpoints/query.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ async def query_endpoint_handler( # pylint: disable=R0914
241241

242242
user_id, _, _skip_userid_check, token = auth
243243

244+
started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
244245
user_conversation: UserConversation | None = None
245246
if query_request.conversation_id:
246247
logger.debug(
@@ -330,6 +331,7 @@ async def query_endpoint_handler( # pylint: disable=R0914
330331
topic_summary=topic_summary,
331332
)
332333

334+
completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
333335
store_conversation_into_cache(
334336
configuration,
335337
user_id,
@@ -338,6 +340,8 @@ async def query_endpoint_handler( # pylint: disable=R0914
338340
model_id,
339341
query_request.query,
340342
summary.llm_response,
343+
started_at,
344+
completed_at,
341345
_skip_userid_check,
342346
topic_summary,
343347
)

‎src/app/endpoints/streaming_query.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import json
55
import logging
66
import re
7+
from datetime import UTC, datetime
78
from typing import Annotated, Any, AsyncIterator, Iterator, cast
89

910
from fastapi import APIRouter, Depends, HTTPException, Request, status
@@ -596,6 +597,7 @@ async def streaming_query_endpoint_handler( # pylint: disable=R0915,R0914
596597
_ = request
597598

598599
check_configuration_loaded(configuration)
600+
started_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
599601

600602
# Enforce RBAC: optionally disallow overriding model/provider in requests
601603
validate_model_provider_override(query_request, request.state.authorized_actions)
@@ -719,6 +721,7 @@ async def response_generator(
719721
query_request.query, client, model_id
720722
)
721723

724+
completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
722725
store_conversation_into_cache(
723726
configuration,
724727
user_id,
@@ -727,6 +730,8 @@ async def response_generator(
727730
model_id,
728731
query_request.query,
729732
summary.llm_response,
733+
started_at,
734+
completed_at,
730735
_skip_userid_check,
731736
topic_summary,
732737
)

‎src/cache/postgres_cache.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ class PostgresCache(Cache):
2323
user_id | text | not null |
2424
conversation_id | text | not null |
2525
created_at | timestamp without time zone | not null |
26+
started_at | text | |
27+
completed_at | text | |
2628
query | text | |
2729
response | text | |
2830
provider | text | |
@@ -38,6 +40,8 @@ class PostgresCache(Cache):
3840
user_id text NOT NULL,
3941
conversation_id text NOT NULL,
4042
created_at timestamp NOT NULL,
43+
started_at text,
44+
completed_at text,
4145
query text,
4246
response text,
4347
provider text,
@@ -62,15 +66,16 @@ class PostgresCache(Cache):
6266
"""
6367

6468
SELECT_CONVERSATION_HISTORY_STATEMENT = """
65-
SELECT query, response, provider, model
69+
SELECT query, response, provider, model, started_at, completed_at
6670
FROM cache
6771
WHERE user_id=%s AND conversation_id=%s
6872
ORDER BY created_at
6973
"""
7074

7175
INSERT_CONVERSATION_HISTORY_STATEMENT = """
72-
INSERT INTO cache(user_id, conversation_id, created_at, query, response, provider, model)
73-
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s)
76+
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
77+
query, response, provider, model)
78+
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s)
7479
"""
7580

7681
QUERY_CACHE_SIZE = """
@@ -211,6 +216,8 @@ def get(
211216
response=conversation_entry[1],
212217
provider=conversation_entry[2],
213218
model=conversation_entry[3],
219+
started_at=conversation_entry[4],
220+
completed_at=conversation_entry[5],
214221
)
215222
result.append(cache_entry)
216223

@@ -245,6 +252,8 @@ def insert_or_append(
245252
(
246253
user_id,
247254
conversation_id,
255+
cache_entry.started_at,
256+
cache_entry.completed_at,
248257
cache_entry.query,
249258
cache_entry.response,
250259
cache_entry.provider,

‎src/cache/sqlite_cache.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ class SQLiteCache(Cache):
2525
user_id | text | not null |
2626
conversation_id | text | not null |
2727
created_at | int | not null |
28+
started_at | text | |
29+
completed_at | text | |
2830
query | text | |
2931
response | text | |
3032
provider | text | |
@@ -42,6 +44,8 @@ class SQLiteCache(Cache):
4244
user_id text NOT NULL,
4345
conversation_id text NOT NULL,
4446
created_at int NOT NULL,
47+
started_at text,
48+
completed_at text,
4549
query text,
4650
response text,
4751
provider text,
@@ -66,15 +70,16 @@ class SQLiteCache(Cache):
6670
"""
6771

6872
SELECT_CONVERSATION_HISTORY_STATEMENT = """
69-
SELECT query, response, provider, model
73+
SELECT query, response, provider, model, started_at, completed_at
7074
FROM cache
7175
WHERE user_id=? AND conversation_id=?
7276
ORDER BY created_at
7377
"""
7478

7579
INSERT_CONVERSATION_HISTORY_STATEMENT = """
76-
INSERT INTO cache(user_id, conversation_id, created_at, query, response, provider, model)
77-
VALUES (?, ?, ?, ?, ?, ?, ?)
80+
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
81+
query, response, provider, model)
82+
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
7883
"""
7984

8085
QUERY_CACHE_SIZE = """
@@ -209,6 +214,8 @@ def get(
209214
response=conversation_entry[1],
210215
provider=conversation_entry[2],
211216
model=conversation_entry[3],
217+
started_at=conversation_entry[4],
218+
completed_at=conversation_entry[5],
212219
)
213220
result.append(cache_entry)
214221

@@ -243,6 +250,8 @@ def insert_or_append(
243250
user_id,
244251
conversation_id,
245252
current_time,
253+
cache_entry.started_at,
254+
cache_entry.completed_at,
246255
cache_entry.query,
247256
cache_entry.response,
248257
cache_entry.provider,

‎src/models/cache_entry.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ class CacheEntry(BaseModel):
1717
response: str
1818
provider: str
1919
model: str
20+
started_at: str
21+
completed_at: str
2022

2123

2224
class ConversationData(BaseModel):

‎src/utils/endpoints.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,8 @@ def store_conversation_into_cache(
189189
model_id: str,
190190
query: str,
191191
response: str,
192+
started_at: str,
193+
completed_at: str,
192194
_skip_userid_check: bool,
193195
topic_summary: str | None,
194196
) -> None:
@@ -203,6 +205,8 @@ def store_conversation_into_cache(
203205
response=response,
204206
provider=provider_id,
205207
model=model_id,
208+
started_at=started_at,
209+
completed_at=completed_at,
206210
)
207211
cache.insert_or_append(
208212
user_id, conversation_id, cache_entry, _skip_userid_check

‎tests/unit/app/endpoints/test_conversations_v2.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,12 @@
1010
def test_transform_message() -> None:
1111
"""Test the transform_chat_message transformation function."""
1212
entry = CacheEntry(
13-
query="query", response="response", provider="provider", model="model"
13+
query="query",
14+
response="response",
15+
provider="provider",
16+
model="model",
17+
started_at="2024-01-01T00:00:00Z",
18+
completed_at="2024-01-01T00:00:05Z",
1419
)
1520
transformed = transform_chat_message(entry)
1621
assert transformed is not None
@@ -21,6 +26,12 @@ def test_transform_message() -> None:
2126
assert "model" in transformed
2227
assert transformed["model"] == "model"
2328

29+
assert "started_at" in transformed
30+
assert transformed["started_at"] == "2024-01-01T00:00:00Z"
31+
32+
assert "completed_at" in transformed
33+
assert transformed["completed_at"] == "2024-01-01T00:00:05Z"
34+
2435
assert "messages" in transformed
2536
assert len(transformed["messages"]) == 2
2637

‎tests/unit/cache/test_noop_cache.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,20 @@
1010
CONVERSATION_ID = suid.get_suid()
1111
USER_PROVIDED_USER_ID = "test-user1"
1212
cache_entry_1 = CacheEntry(
13-
query="user message1", response="AI message1", provider="foo", model="bar"
13+
query="user message1",
14+
response="AI message1",
15+
provider="foo",
16+
model="bar",
17+
started_at="2025-10-03T09:31:25Z",
18+
completed_at="2025-10-03T09:31:29Z",
1419
)
1520
cache_entry_2 = CacheEntry(
16-
query="user message2", response="AI message2", provider="foo", model="bar"
21+
query="user message2",
22+
response="AI message2",
23+
provider="foo",
24+
model="bar",
25+
started_at="2025-10-03T09:31:25Z",
26+
completed_at="2025-10-03T09:31:29Z",
1727
)
1828

1929

‎tests/unit/cache/test_postgres_cache.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,20 @@
1616
CONVERSATION_ID_1 = suid.get_suid()
1717
CONVERSATION_ID_2 = suid.get_suid()
1818
cache_entry_1 = CacheEntry(
19-
query="user message1", response="AI message1", provider="foo", model="bar"
19+
query="user message1",
20+
response="AI message1",
21+
provider="foo",
22+
model="bar",
23+
started_at="2025-10-03T09:31:25Z",
24+
completed_at="2025-10-03T09:31:29Z",
2025
)
2126
cache_entry_2 = CacheEntry(
22-
query="user message2", response="AI message2", provider="foo", model="bar"
27+
query="user message2",
28+
response="AI message2",
29+
provider="foo",
30+
model="bar",
31+
started_at="2025-10-03T09:31:25Z",
32+
completed_at="2025-10-03T09:31:29Z",
2333
)
2434

2535
# pylint: disable=fixme

0 commit comments

Comments
 (0)