From 54d4fd0e084db3de9392e157671512d61b01b843 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Fri, 10 Nov 2023 16:46:34 -0800
Subject: [PATCH] Add chat_model data for logging selected models to telemetry

---
 src/khoj/routers/api.py     | 25 ++++++++++++++-----------
 src/khoj/routers/helpers.py |  6 +++++-
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py
index 32793d7c..ddfe9bc1 100644
--- a/src/khoj/routers/api.py
+++ b/src/khoj/routers/api.py
@@ -606,7 +606,7 @@ async def chat(
         return StreamingResponse(iter([formatted_help]), media_type="text/event-stream", status_code=200)
 
     # Get the (streamed) chat response from the LLM of choice.
-    llm_response = await agenerate_chat_response(
+    llm_response, chat_metadata = await agenerate_chat_response(
         defiltered_query,
         meta_log,
         compiled_references,
@@ -615,6 +615,19 @@ async def chat(
         user,
     )
 
+    chat_metadata.update({"conversation_command": conversation_command.value})
+
+    update_telemetry_state(
+        request=request,
+        telemetry_type="api",
+        api="chat",
+        client=client,
+        user_agent=user_agent,
+        referer=referer,
+        host=host,
+        metadata=chat_metadata,
+    )
+
     if llm_response is None:
         return Response(content=llm_response, media_type="text/plain", status_code=500)
 
@@ -634,16 +647,6 @@ async def chat(
 
     response_obj = {"response": actual_response, "context": compiled_references}
 
-    update_telemetry_state(
-        request=request,
-        telemetry_type="api",
-        api="chat",
-        client=client,
-        user_agent=user_agent,
-        referer=referer,
-        host=host,
-    )
-
     return Response(content=json.dumps(response_obj), media_type="application/json", status_code=200)
 
 
diff --git a/src/khoj/routers/helpers.py b/src/khoj/routers/helpers.py
index ebb661ef..99c1b281 100644
--- a/src/khoj/routers/helpers.py
+++ b/src/khoj/routers/helpers.py
@@ -132,6 +132,8 @@ def generate_chat_response(
     chat_response = None
     logger.debug(f"Conversation Type: {conversation_command.name}")
 
+    metadata = {}
+
     try:
         partial_completion = partial(
             _save_to_conversation_log,
@@ -179,8 +181,10 @@ def generate_chat_response(
                 tokenizer_name=conversation_config.tokenizer,
             )
 
+        metadata.update({"chat_model": conversation_config.chat_model})
+
     except Exception as e:
         logger.error(e, exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))
 
-    return chat_response
+    return chat_response, metadata