Merge pull request #380 from keitwb/model-provider-transcripts

Add resolved model/provider to transcripts

Author: tisnik

src/app/endpoints/query.py

@@ -205,6 +205,8 @@ async def query_endpoint_handler(
             store_transcript(
                 user_id=user_id,
                 conversation_id=conversation_id,
+                model_id=model_id,
+                provider_id=provider_id,
                 query_is_valid=True,  # TODO(lucasagomes): implement as part of query validation
                 query=query_request.query,
                 query_request=query_request,
@@ -456,6 +458,8 @@ def construct_transcripts_path(user_id: str, conversation_id: str) -> Path:
 def store_transcript(  # pylint: disable=too-many-arguments,too-many-positional-arguments
     user_id: str,
     conversation_id: str,
+    model_id: str,
+    provider_id: str | None,
     query_is_valid: bool,
     query: str,
     query_request: QueryRequest,
@@ -482,8 +486,10 @@ def store_transcript(  # pylint: disable=too-many-arguments,too-many-positional-
 
     data_to_store = {
         "metadata": {
-            "provider": query_request.provider,
-            "model": query_request.model,
+            "provider": provider_id,
+            "model": model_id,
+            "query_provider": query_request.provider,
+            "query_model": query_request.model,
             "user_id": user_id,
             "conversation_id": conversation_id,
             "timestamp": datetime.now(UTC).isoformat(),

src/app/endpoints/streaming_query.py

@@ -464,6 +464,8 @@ async def response_generator(turn_response: Any) -> AsyncIterator[str]:
                 store_transcript(
                     user_id=user_id,
                     conversation_id=conversation_id,
+                    model_id=model_id,
+                    provider_id=provider_id,
                     query_is_valid=True,  # TODO(lucasagomes): implement as part of query validation
                     query=query_request.query,
                     query_request=query_request,

tests/unit/app/endpoints/test_query.py

@@ -159,6 +159,8 @@ async def _test_query_endpoint_handler(mocker, store_transcript_to_file=False):
         mock_transcript.assert_called_once_with(
             user_id="mock_user_id",
             conversation_id=conversation_id,
+            model_id="fake_model_id",
+            provider_id="fake_provider_id",
             query_is_valid=True,
             query=query,
             query_request=query_request,
@@ -1048,6 +1050,8 @@ def test_store_transcript(mocker):
     store_transcript(
         user_id,
         conversation_id,
+        model,
+        provider,
         query_is_valid,
         query,
         query_request,
@@ -1061,8 +1065,10 @@ def test_store_transcript(mocker):
     mock_json.dump.assert_called_once_with(
         {
             "metadata": {
-                "provider": query_request.provider,
-                "model": query_request.model,
+                "provider": "fake-provider",
+                "model": "fake-model",
+                "query_provider": query_request.provider,
+                "query_model": query_request.model,
                 "user_id": user_id,
                 "conversation_id": conversation_id,
                 "timestamp": mocker.ANY,

tests/unit/app/endpoints/test_streaming_query.py

@@ -318,6 +318,8 @@ async def _test_streaming_query_endpoint_handler(mocker, store_transcript=False)
         mock_transcript.assert_called_once_with(
             user_id="mock_user_id",
             conversation_id="test_conversation_id",
+            model_id="fake_model_id",
+            provider_id="fake_provider_id",
             query_is_valid=True,
             query=query,
             query_request=query_request,