From c9f3e38a85f35a56dbb611a53906f03a99e4ff2d Mon Sep 17 00:00:00 2001
From: sarath2496 <sarathc225@gmail.com>
Date: Tue, 23 Jul 2024 21:07:56 -0500
Subject: [PATCH] AgentOps integration

---
 .../code_interpreter_tool.py                  | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/examples/langchain-python/langchain_e2b_python/code_interpreter_tool.py b/examples/langchain-python/langchain_e2b_python/code_interpreter_tool.py
index bec6eaeb..843472e2 100644
--- a/examples/langchain-python/langchain_e2b_python/code_interpreter_tool.py
+++ b/examples/langchain-python/langchain_e2b_python/code_interpreter_tool.py
@@ -1,5 +1,7 @@
 import os
 import json
+import agentops
+from dotenv import load_dotenv
 
 from typing import Any, List
 from langchain_core.tools import Tool
@@ -10,11 +12,15 @@
     ToolAgentAction,
 )
 
+# Load environment variables
+load_dotenv()
+
+# Initialize AgentOps
+agentops.init(os.getenv('AGENTOPS_API_KEY'))
 
 class LangchainCodeInterpreterToolInput(BaseModel):
     code: str = Field(description="Python code to execute.")
 
-
 class CodeInterpreterFunctionTool:
     """
     This class calls arbitrary code against a Python Jupyter notebook.
@@ -23,6 +29,7 @@ class CodeInterpreterFunctionTool:
 
     tool_name: str = "code_interpreter"
 
+    @agentops.record_function('CodeInterpreterFunctionTool_init')
     def __init__(self):
         # Instantiate the E2B sandbox - this is a long lived object
         # that's pinging E2B cloud to keep the sandbox alive.
@@ -32,13 +39,22 @@ def __init__(self):
             )
         self.code_interpreter = CodeInterpreter()
 
+    @agentops.record_function('CodeInterpreterFunctionTool_close')
     def close(self):
         self.code_interpreter.close()
 
+    @agentops.record_function('CodeInterpreterFunctionTool_call')
     def call(self, parameters: dict, **kwargs: Any):
         code = parameters.get("code", "")
         print(f"***Code Interpreting...\n{code}\n====")
-        execution = self.code_interpreter.notebook.exec_cell(code)
+        
+        with agentops.record_span('code_execution'):
+            execution = self.code_interpreter.notebook.exec_cell(code)
+        
+        # Record metrics
+        agentops.record_metric('code_execution_success', 1 if not execution.error else 0)
+        agentops.record_metric('code_execution_time', execution.duration)
+        
         return {
             "results": execution.results,
             "stdout": execution.logs.stdout,
@@ -46,10 +62,11 @@ def call(self, parameters: dict, **kwargs: Any):
             "error": execution.error,
         }
 
-    # langchain does not return a dict as a parameter, only a code string
+    @agentops.record_function('CodeInterpreterFunctionTool_langchain_call')
     def langchain_call(self, code: str):
         return self.call({"code": code})
 
+    @agentops.record_function('CodeInterpreterFunctionTool_to_langchain_tool')
     def to_langchain_tool(self) -> Tool:
         tool = Tool(
             name=self.tool_name,
@@ -60,6 +77,7 @@ def to_langchain_tool(self) -> Tool:
         return tool
 
     @staticmethod
+    @agentops.record_function('CodeInterpreterFunctionTool_format_to_tool_message')
     def format_to_tool_message(
         agent_action: ToolAgentAction,
         observation: dict,
@@ -77,4 +95,10 @@ def format_to_tool_message(
             ToolMessage(content=content, tool_call_id=agent_action.tool_call_id)
         )
 
+        # Record metric for message count
+        agentops.record_metric('tool_message_count', len(new_messages))
+
         return new_messages
+
+# End of program
+agentops.end_session('Success')
\ No newline at end of file