Skip to content

Commit f7beb32

Browse files
committed
Add resolved model/provider to transcripts
This makes it easier to analyze differences between models. Before it would only add the model/provider if it was customized in the query.
1 parent 93f1079 commit f7beb32

File tree

4 files changed

+20
-4
lines changed

4 files changed

+20
-4
lines changed

src/app/endpoints/query.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,8 @@ async def query_endpoint_handler(
205205
store_transcript(
206206
user_id=user_id,
207207
conversation_id=conversation_id,
208+
model_id=model_id,
209+
provider_id=provider_id,
208210
query_is_valid=True, # TODO(lucasagomes): implement as part of query validation
209211
query=query_request.query,
210212
query_request=query_request,
@@ -456,6 +458,8 @@ def construct_transcripts_path(user_id: str, conversation_id: str) -> Path:
456458
def store_transcript( # pylint: disable=too-many-arguments,too-many-positional-arguments
457459
user_id: str,
458460
conversation_id: str,
461+
model_id: str,
462+
provider_id: str | None,
459463
query_is_valid: bool,
460464
query: str,
461465
query_request: QueryRequest,
@@ -482,8 +486,10 @@ def store_transcript( # pylint: disable=too-many-arguments,too-many-positional-
482486

483487
data_to_store = {
484488
"metadata": {
485-
"provider": query_request.provider,
486-
"model": query_request.model,
489+
"provider": provider_id,
490+
"model": model_id,
491+
"query_provider": query_request.provider,
492+
"query_model": query_request.model,
487493
"user_id": user_id,
488494
"conversation_id": conversation_id,
489495
"timestamp": datetime.now(UTC).isoformat(),

src/app/endpoints/streaming_query.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,8 @@ async def response_generator(turn_response: Any) -> AsyncIterator[str]:
464464
store_transcript(
465465
user_id=user_id,
466466
conversation_id=conversation_id,
467+
model_id=model_id,
468+
provider_id=provider_id,
467469
query_is_valid=True, # TODO(lucasagomes): implement as part of query validation
468470
query=query_request.query,
469471
query_request=query_request,

tests/unit/app/endpoints/test_query.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,8 @@ async def _test_query_endpoint_handler(mocker, store_transcript_to_file=False):
159159
mock_transcript.assert_called_once_with(
160160
user_id="mock_user_id",
161161
conversation_id=conversation_id,
162+
model_id="fake_model_id",
163+
provider_id="fake_provider_id",
162164
query_is_valid=True,
163165
query=query,
164166
query_request=query_request,
@@ -1045,6 +1047,8 @@ def test_store_transcript(mocker):
10451047
store_transcript(
10461048
user_id,
10471049
conversation_id,
1050+
model,
1051+
provider,
10481052
query_is_valid,
10491053
query,
10501054
query_request,
@@ -1058,8 +1062,10 @@ def test_store_transcript(mocker):
10581062
mock_json.dump.assert_called_once_with(
10591063
{
10601064
"metadata": {
1061-
"provider": query_request.provider,
1062-
"model": query_request.model,
1065+
"provider": "fake-provider",
1066+
"model": "fake-model",
1067+
"query_provider": query_request.provider,
1068+
"query_model": query_request.model,
10631069
"user_id": user_id,
10641070
"conversation_id": conversation_id,
10651071
"timestamp": mocker.ANY,

tests/unit/app/endpoints/test_streaming_query.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,8 @@ async def _test_streaming_query_endpoint_handler(mocker, store_transcript=False)
318318
mock_transcript.assert_called_once_with(
319319
user_id="mock_user_id",
320320
conversation_id="test_conversation_id",
321+
model_id="fake_model_id",
322+
provider_id="fake_provider_id",
321323
query_is_valid=True,
322324
query=query,
323325
query_request=query_request,

0 commit comments

Comments
 (0)