From 6f9904895642c7b6a11114bb3b203907850a4602 Mon Sep 17 00:00:00 2001
From: Rathijit Paul <30369246+rathijitpapon@users.noreply.github.com>
Date: Tue, 2 Jul 2024 19:27:24 +0600
Subject: [PATCH] :construction: added tracing in prompt

---
 server/src/llms/query_rephraser.rs | 5 +++--
 server/src/llms/summarizer.rs      | 2 ++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/server/src/llms/query_rephraser.rs b/server/src/llms/query_rephraser.rs
index cdc43be2..a5e0168c 100644
--- a/server/src/llms/query_rephraser.rs
+++ b/server/src/llms/query_rephraser.rs
@@ -52,7 +52,8 @@ pub struct QueryRephraserOutput {
     pub rephrased_query: String,
 }
 
-fn prepare_prompt(query_rephraser_input: &QueryRephraserInput) -> String {
+#[tracing::instrument(level = "info", ret)]
+fn prepare_rephrase_query_prompt(query_rephraser_input: &QueryRephraserInput) -> String {
     "[INST] Rephrase the input text based on the context and the final sentence. So that it can be understood without the context. Return the rephrased question only\n\n---\n\nFollow the following format.\n\nContext: contains the chat history\n\nQuestion: ${question}\n\nReasoning: Let's think step by step in order to ${produce the answer}. We ...\n\nAnswer: Given a chat history and the latest user question, which might reference the context from the chat history, formulate a standalone question that can be understood from the history without needing the chat history. DO NOT ANSWER THE QUESTION - just reformulate it and return the rephrased question only \n\n---\n\nContext: ".to_string()
         + query_rephraser_input.previous_context.iter().map(|x| format!("{}: {}", x.query, x.result)).collect::<Vec<String>>().join("\n").as_str()
         + "\n\nQuestion: "
@@ -72,7 +73,7 @@ pub async fn rephrase_query(
         HeaderValue::from_str(&settings.api_key.expose())?,
     );
 
-    let prompt = prepare_prompt(query_rephraser_input);
+    let prompt = prepare_rephrase_query_prompt(query_rephraser_input);
 
     let response = client
         .post(&settings.api_url)
diff --git a/server/src/llms/summarizer.rs b/server/src/llms/summarizer.rs
index dac9ba41..3baafb0f 100644
--- a/server/src/llms/summarizer.rs
+++ b/server/src/llms/summarizer.rs
@@ -48,6 +48,7 @@ pub struct SummarizerStreamOutput {
     pub generated_text: Option<String>,
 }
 
+#[tracing::instrument(level = "info", ret)]
 fn prepare_llm_context_string(
     settings: &SummarizerSettings,
     summarizer_input: SummarizerInput,
@@ -121,6 +122,7 @@ pub async fn generate_text_with_llm(
     Ok(())
 }
 
+#[tracing::instrument(level = "info", ret)]
 fn prepare_openai_input(
     settings: &OpenAISettings,
     summarizer_input: SummarizerInput,