Skip to content

Commit 539c630

Browse files
yangcao77yangcao77
committed
store tool calls and tool summary in cache entry
Signed-off-by: Stephanie <yangcao@redhat.com>
1 parent 2bc4bdb commit 539c630

7 files changed

Lines changed: 358 additions & 11 deletions

File tree

‎src/app/endpoints/query.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ async def query_endpoint_handler_base( # pylint: disable=R0914
363363
truncated=False, # TODO(lucasagomes): implement truncation as part of quota work
364364
attachments=query_request.attachments or [],
365365
)
366-
366+
367367
logger.info("Persisting conversation details...")
368368
persist_user_conversation_details(
369369
user_id=user_id,
@@ -374,7 +374,6 @@ async def query_endpoint_handler_base( # pylint: disable=R0914
374374
)
375375

376376
completed_at = datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
377-
378377
cache_entry = CacheEntry(
379378
query=query_request.query,
380379
response=summary.llm_response,
@@ -383,6 +382,8 @@ async def query_endpoint_handler_base( # pylint: disable=R0914
383382
started_at=started_at,
384383
completed_at=completed_at,
385384
referenced_documents=referenced_documents if referenced_documents else None,
385+
tool_calls=summary.tool_calls if summary.tool_calls else None,
386+
tool_results=summary.tool_results if summary.tool_results else None,
386387
)
387388

388389
consume_tokens(

‎src/cache/postgres_cache.py‎

Lines changed: 79 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from models.cache_entry import CacheEntry
1010
from models.config import PostgreSQLDatabaseConfiguration
1111
from models.responses import ConversationData, ReferencedDocument
12+
from utils.types import ToolCallSummary, ToolResultSummary
1213
from log import get_logger
1314
from utils.connection_decorator import connection
1415

@@ -32,7 +33,9 @@ class PostgresCache(Cache):
3233
response | text | |
3334
provider | text | |
3435
model | text | |
35-
referenced_documents | jsonb | |
36+
referenced_documents | jsonb | |
37+
tool_calls | jsonb | |
38+
tool_results | jsonb | |
3639
Indexes:
3740
"cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
3841
"timestamps" btree (created_at)
@@ -55,6 +58,8 @@ class PostgresCache(Cache):
5558
provider text,
5659
model text,
5760
referenced_documents jsonb,
61+
tool_calls jsonb,
62+
tool_results jsonb,
5863
PRIMARY KEY(user_id, conversation_id, created_at)
5964
);
6065
"""
@@ -75,16 +80,18 @@ class PostgresCache(Cache):
7580
"""
7681

7782
SELECT_CONVERSATION_HISTORY_STATEMENT = """
78-
SELECT query, response, provider, model, started_at, completed_at, referenced_documents
83+
SELECT query, response, provider, model, started_at, completed_at,
84+
referenced_documents, tool_calls, tool_results
7985
FROM cache
8086
WHERE user_id=%s AND conversation_id=%s
8187
ORDER BY created_at
8288
"""
8389

8490
INSERT_CONVERSATION_HISTORY_STATEMENT = """
8591
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
86-
query, response, provider, model, referenced_documents)
87-
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s)
92+
query, response, provider, model, referenced_documents,
93+
tool_calls, tool_results)
94+
VALUES (%s, %s, CURRENT_TIMESTAMP, %s, %s, %s, %s, %s, %s, %s, %s, %s)
8895
"""
8996

9097
QUERY_CACHE_SIZE = """
@@ -260,6 +267,40 @@ def get(
260267
conversation_id,
261268
e,
262269
)
270+
271+
# Parse tool_calls back into ToolCallSummary objects
272+
tool_calls_data = conversation_entry[7]
273+
tool_calls_obj = None
274+
if tool_calls_data:
275+
try:
276+
tool_calls_obj = [
277+
ToolCallSummary.model_validate(tc) for tc in tool_calls_data
278+
]
279+
except (ValueError, TypeError) as e:
280+
logger.warning(
281+
"Failed to deserialize tool_calls for "
282+
"conversation %s: %s",
283+
conversation_id,
284+
e,
285+
)
286+
287+
# Parse tool_results back into ToolResultSummary objects
288+
tool_results_data = conversation_entry[8]
289+
tool_results_obj = None
290+
if tool_results_data:
291+
try:
292+
tool_results_obj = [
293+
ToolResultSummary.model_validate(tr)
294+
for tr in tool_results_data
295+
]
296+
except (ValueError, TypeError) as e:
297+
logger.warning(
298+
"Failed to deserialize tool_results for "
299+
"conversation %s: %s",
300+
conversation_id,
301+
e,
302+
)
303+
263304
cache_entry = CacheEntry(
264305
query=conversation_entry[0],
265306
response=conversation_entry[1],
@@ -268,6 +309,8 @@ def get(
268309
started_at=conversation_entry[4],
269310
completed_at=conversation_entry[5],
270311
referenced_documents=docs_obj,
312+
tool_calls=tool_calls_obj,
313+
tool_results=tool_results_obj,
271314
)
272315
result.append(cache_entry)
273316

@@ -311,6 +354,36 @@ def insert_or_append(
311354
e,
312355
)
313356

357+
tool_calls_json = None
358+
if cache_entry.tool_calls:
359+
try:
360+
tool_calls_as_dicts = [
361+
tc.model_dump(mode="json") for tc in cache_entry.tool_calls
362+
]
363+
tool_calls_json = json.dumps(tool_calls_as_dicts)
364+
except (TypeError, ValueError) as e:
365+
logger.warning(
366+
"Failed to serialize tool_calls for "
367+
"conversation %s: %s",
368+
conversation_id,
369+
e,
370+
)
371+
372+
tool_results_json = None
373+
if cache_entry.tool_results:
374+
try:
375+
tool_results_as_dicts = [
376+
tr.model_dump(mode="json") for tr in cache_entry.tool_results
377+
]
378+
tool_results_json = json.dumps(tool_results_as_dicts)
379+
except (TypeError, ValueError) as e:
380+
logger.warning(
381+
"Failed to serialize tool_results for "
382+
"conversation %s: %s",
383+
conversation_id,
384+
e,
385+
)
386+
314387
# the whole operation is run in one transaction
315388
with self.connection.cursor() as cursor:
316389
cursor.execute(
@@ -325,6 +398,8 @@ def insert_or_append(
325398
cache_entry.provider,
326399
cache_entry.model,
327400
referenced_documents_json,
401+
tool_calls_json,
402+
tool_results_json,
328403
),
329404
)
330405

‎src/cache/sqlite_cache.py‎

Lines changed: 79 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from models.cache_entry import CacheEntry
1111
from models.config import SQLiteDatabaseConfiguration
1212
from models.responses import ConversationData, ReferencedDocument
13+
from utils.types import ToolCallSummary, ToolResultSummary
1314
from log import get_logger
1415
from utils.connection_decorator import connection
1516

@@ -34,6 +35,8 @@ class SQLiteCache(Cache):
3435
provider | text | |
3536
model | text | |
3637
referenced_documents | text | |
38+
tool_calls | text | |
39+
tool_results | text | |
3740
Indexes:
3841
"cache_pkey" PRIMARY KEY, btree (user_id, conversation_id, created_at)
3942
"cache_key_key" UNIQUE CONSTRAINT, btree (key)
@@ -54,6 +57,8 @@ class SQLiteCache(Cache):
5457
provider text,
5558
model text,
5659
referenced_documents text,
60+
tool_calls text,
61+
tool_results text,
5762
PRIMARY KEY(user_id, conversation_id, created_at)
5863
);
5964
"""
@@ -74,16 +79,18 @@ class SQLiteCache(Cache):
7479
"""
7580

7681
SELECT_CONVERSATION_HISTORY_STATEMENT = """
77-
SELECT query, response, provider, model, started_at, completed_at, referenced_documents
82+
SELECT query, response, provider, model, started_at, completed_at,
83+
referenced_documents, tool_calls, tool_results
7884
FROM cache
7985
WHERE user_id=? AND conversation_id=?
8086
ORDER BY created_at
8187
"""
8288

8389
INSERT_CONVERSATION_HISTORY_STATEMENT = """
8490
INSERT INTO cache(user_id, conversation_id, created_at, started_at, completed_at,
85-
query, response, provider, model, referenced_documents)
86-
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
91+
query, response, provider, model, referenced_documents,
92+
tool_calls, tool_results)
93+
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
8794
"""
8895

8996
QUERY_CACHE_SIZE = """
@@ -228,6 +235,41 @@ def get(
228235
conversation_id,
229236
e,
230237
)
238+
239+
# Parse tool_calls back into ToolCallSummary objects
240+
tool_calls_json_str = conversation_entry[7]
241+
tool_calls_obj = None
242+
if tool_calls_json_str:
243+
try:
244+
tool_calls_data = json.loads(tool_calls_json_str)
245+
tool_calls_obj = [
246+
ToolCallSummary.model_validate(tc) for tc in tool_calls_data
247+
]
248+
except (json.JSONDecodeError, ValueError) as e:
249+
logger.warning(
250+
"Failed to deserialize tool_calls for "
251+
"conversation %s: %s",
252+
conversation_id,
253+
e,
254+
)
255+
256+
# Parse tool_results back into ToolResultSummary objects
257+
tool_results_json_str = conversation_entry[8]
258+
tool_results_obj = None
259+
if tool_results_json_str:
260+
try:
261+
tool_results_data = json.loads(tool_results_json_str)
262+
tool_results_obj = [
263+
ToolResultSummary.model_validate(tr) for tr in tool_results_data
264+
]
265+
except (json.JSONDecodeError, ValueError) as e:
266+
logger.warning(
267+
"Failed to deserialize tool_results for "
268+
"conversation %s: %s",
269+
conversation_id,
270+
e,
271+
)
272+
231273
cache_entry = CacheEntry(
232274
query=conversation_entry[0],
233275
response=conversation_entry[1],
@@ -236,6 +278,8 @@ def get(
236278
started_at=conversation_entry[4],
237279
completed_at=conversation_entry[5],
238280
referenced_documents=docs_obj,
281+
tool_calls=tool_calls_obj,
282+
tool_results=tool_results_obj,
239283
)
240284
result.append(cache_entry)
241285

@@ -281,6 +325,36 @@ def insert_or_append(
281325
e,
282326
)
283327

328+
tool_calls_json = None
329+
if cache_entry.tool_calls:
330+
try:
331+
tool_calls_as_dicts = [
332+
tc.model_dump(mode="json") for tc in cache_entry.tool_calls
333+
]
334+
tool_calls_json = json.dumps(tool_calls_as_dicts)
335+
except (TypeError, ValueError) as e:
336+
logger.warning(
337+
"Failed to serialize tool_calls for "
338+
"conversation %s: %s",
339+
conversation_id,
340+
e,
341+
)
342+
343+
tool_results_json = None
344+
if cache_entry.tool_results:
345+
try:
346+
tool_results_as_dicts = [
347+
tr.model_dump(mode="json") for tr in cache_entry.tool_results
348+
]
349+
tool_results_json = json.dumps(tool_results_as_dicts)
350+
except (TypeError, ValueError) as e:
351+
logger.warning(
352+
"Failed to serialize tool_results for "
353+
"conversation %s: %s",
354+
conversation_id,
355+
e,
356+
)
357+
284358
cursor.execute(
285359
self.INSERT_CONVERSATION_HISTORY_STATEMENT,
286360
(
@@ -294,6 +368,8 @@ def insert_or_append(
294368
cache_entry.provider,
295369
cache_entry.model,
296370
referenced_documents_json,
371+
tool_calls_json,
372+
tool_results_json,
297373
),
298374
)
299375

‎src/models/cache_entry.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from typing import Optional
44
from pydantic import BaseModel
55
from models.responses import ReferencedDocument
6+
from utils.types import ToolCallSummary, ToolResultSummary
67

78

89
class CacheEntry(BaseModel):
@@ -14,6 +15,8 @@ class CacheEntry(BaseModel):
1415
provider: Provider identification
1516
model: Model identification
1617
referenced_documents: List of documents referenced by the response
18+
tool_calls: List of tool calls made during response generation
19+
tool_results: List of tool results from tool calls
1720
"""
1821

1922
query: str
@@ -23,3 +26,5 @@ class CacheEntry(BaseModel):
2326
started_at: str
2427
completed_at: str
2528
referenced_documents: Optional[list[ReferencedDocument]] = None
29+
tool_calls: Optional[list[ToolCallSummary]] = None
30+
tool_results: Optional[list[ToolResultSummary]] = None

‎src/utils/endpoints.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -806,6 +806,8 @@ async def cleanup_after_streaming(
806806
started_at=started_at,
807807
completed_at=completed_at,
808808
referenced_documents=referenced_documents if referenced_documents else None,
809+
tool_calls=summary.tool_calls if summary.tool_calls else None,
810+
tool_results=summary.tool_results if summary.tool_results else None,
809811
)
810812

811813
store_conversation_into_cache(

0 commit comments

Comments
 (0)