Update app.py

app.py

@@ -121,6 +121,7 @@ sqlite_agent = create_sql_agent(
     verbose=True
 )
 #### Let's convert the sql agent into a tool that our fin agent can use.
+@weave.op()
 @tool
 def sql_tool(user_input):
     """
@@ -210,7 +211,21 @@ client = AzureOpenAI(
     api_version=os.environ["AZURE_OPENAI_APIVERSION_mini"]
 )
 
+@weave.op()
+def get_rag_response(prompt):
 
+
+    try:
+        response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=prompt
+        )
+
+        prediction = response.choices[0].message.content
+    except Exception as e:
+        prediction = f'Sorry, I encountered the following error: \n {e}'
+
+    
 @tool
 def rag(user_input: str) -> str:
   
@@ -231,25 +246,16 @@ def rag(user_input: str) -> str:
     context_for_query = ". ".join(context_list)
 
     prompt = [
-        {'role':'system', 'content': qna_system_message},
-        {'role': 'user', 'content': qna_user_message_template.format(
-            context=context_for_query,
-            question=user_input
-            )
-        }
-    ]
-
-    try:
-        response = client.chat.completions.create(
-        model="gpt-4o-mini",
-        messages=prompt
+    {'role':'system', 'content': qna_system_message},
+    {'role': 'user', 'content': qna_user_message_template.format(
+        context=context_for_query,
+        question=user_input
         )
+    }
+    ]
 
-        prediction = response.choices[0].message.content
-    except Exception as e:
-        prediction = f'Sorry, I encountered the following error: \n {e}'
-
-    return prediction
+    
+    return get_rag_response(prompt)
 
 
 #=================================== Other TOOLS======================================#
@@ -489,6 +495,11 @@ def fetch_details(email):
             connection.close()      
 
 # Function to process user input and generate a chatbot response
+@weave.op()
+def get_response(conversation_input):
+    return agent_executor.invoke({"input": conversation_input})
+    
+
 def chatbot_interface():
     st.title("E-Commerce Chatbot")
     
@@ -532,7 +543,7 @@ def chatbot_interface():
 
             try:
                 # Pass the history to the agent
-                response = agent_executor.invoke({"input": conversation_input})
+                response = get_response(conversation_input)
     
                 # Add the chatbot's response to the history
                 chatbot_response = response['output']