add tracking

app.py

@@ -1,9 +1,15 @@
 import os
 import time
+import datetime
+
 import gradio as gr
+
 import google.generativeai as genai
 from src.llamaindex_palm import LlamaIndexPaLM
 
+import wandb
+from wandb.sdk.data_types.trace_tree import Trace
+
 import logging
 logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%Y-%m-%d %I:%M:%S %p', level=logging.INFO)
 logger = logging.getLogger('llm')
@@ -15,6 +21,9 @@ llm.set_index_from_pinecone()
 # Credentials
 genai.configure(api_key=os.getenv('PALM_API_KEY'))
 
+# W&B
+wandb.init(project="ChatExp")
+
 # Gradio
 chat_history = []
 
@@ -23,31 +32,71 @@ def clear_chat() -> None:
     chat_history = []
     return None
 
+def get_chat_history(chat_history) -> str:
+    ind = 0
+    formatted_chat_history = ""
+    for message in chat_history:
+        formatted_chat_history += f"User: \n{message}\n" if ind % 2 == 0 else f"Bot: \n{message}\n"
+        ind += 1
+    return formatted_chat_history
+
 def generate_chat(prompt: str, llamaindex_llm: LlamaIndexPaLM):
     global chat_history
     # get chat history
-    context_chat_history = "\n".join(chat_history)
+    context_chat_history = "\n".join(list(filter(None, chat_history)))
 
     logger.info("Generating Message...")
     logger.info(f"User Message:\n{prompt}\n")
     chat_history.append(prompt)
 
+    # w&b trace start
+    start_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+
+    root_span = Trace(
+        name="LLMChain",
+        kind="chain",
+        start_time_ms=start_time_ms,
+        metadata={"user": "Gradio"},
+    )
+
     # get context
     context_from_index = llamaindex_llm.generate_response(prompt)
     logger.info(f"Context from Llama-Index:\n{context_from_index}\n")
 
+    # w&b trace agent
+    agent_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+    agent_span = Trace(
+        name="Agent", 
+        kind="agent", 
+        status_code="success",
+        metadata={
+            "framework": "Llama-Index",
+            "index_type": "VectorStoreIndex",
+            "vector_store": "Pinecone",
+            "model_name": "models/text-bison-001",
+            "temperture": 0.7,
+            "top_k": 40,
+            "top_p": 0.95,
+        },
+        start_time_ms=start_time_ms,
+        end_time_ms=agent_end_time_ms,
+        inputs={"query": prompt},
+        outputs={"response": context_from_index},
+    )
+    root_span.add_child(agent_span)
+
     prompt_with_context = f"""
-    Rule: 
+    [System]
         You are in a role play of Gerard Lee and you need to pretend to be him to answer questions from people who interested in Gerard's background.
-        Respond in 2 to 3 complete sentences, unless specifically asked by the user to elaborate on something. Use only the History and Context to inform your answers.
-    --
-    History:
+        Respond the User Query below in no more than 5 complete sentences, unless specifically asked by the user to elaborate on something. Use only the History and Context to inform your answers.
+    
+    [History]
         {context_chat_history}
-    --
-    Context:
+    
+    [Context]
         {context_from_index}
-    --
-    User Query:
+    
+    [User Query]
         {prompt}
     """
 
@@ -62,13 +111,40 @@ def generate_chat(prompt: str, llamaindex_llm: LlamaIndexPaLM):
             ]
         )
         result = response.result
+        success_flag = "success"
+        if result is None:
+            result = "Seems something went wrong. Please try again later."
+            logger.error(f"Result with 'None' received\n")
+            success_flag = "fail"
 
     except Exception as e:
         result = "Seems something went wrong. Please try again later."
         logger.error(f"Exception {e} occured\n")
+        success_flag = "fail"
 
     chat_history.append(result)
     logger.info(f"Bot Message:\n{result}\n")
+
+    # w&b trace llm
+    llm_end_time_ms = round(datetime.datetime.now().timestamp() * 1000)
+    llm_span = Trace(
+        name="LLM",
+        kind="llm",
+        status_code=success_flag,
+        start_time_ms=agent_end_time_ms,
+        end_time_ms=llm_end_time_ms,
+        inputs={"input": prompt_with_context},
+        outputs={"result": result},
+    )
+    root_span.add_child(llm_span)
+
+    # w&b finalize trace
+    root_span.add_inputs_and_outputs(
+        inputs={"query": prompt}, outputs={"result": result}
+    )
+    root_span._span.end_time_ms = llm_end_time_ms
+    root_span.log(name="llm_app_trace")
+
     return result
 
 with gr.Blocks() as app:
@@ -76,7 +152,7 @@ with gr.Blocks() as app:
         bubble_full_width=False, 
         container=False, 
         show_share_button=False, 
-        avatar_images=[None, './akag-g-only.png']
+        avatar_images=[None, './asset/akag-g-only.png']
     )
     with gr.Row():
         msg = gr.Textbox(
@@ -90,7 +166,7 @@ with gr.Blocks() as app:
          send = gr.Button(
              value="", 
              variant="primary",
-             icon="./send-message.png", 
+             icon="./asset/send-message.png", 
              scale=1
          )