fix: improve session management and notebook execution

- Add proper AgentOps session creation and cleanup - Fix async handling in notebook execution - Update documentation with session management examples - Add named sessions for better monitoring Fixes #543 Co-Authored-By: Alex Reibman <[email protected]>
AgentOps-AI · Dec 12, 2024 · f4fa3b8 · f4fa3b8
1 parent 3a0db02
commit f4fa3b8
Show file tree

Hide file tree

Showing 3 changed files with 90 additions and 77 deletions.
diff --git a/cookbook/AgentOps_Mistral/create_notebook.py b/cookbook/AgentOps_Mistral/create_notebook.py
@@ -293,40 +293,39 @@ def analyze_costs(prompts):
 @agentops.track_agent(name="mistral-analyzer")
 def comprehensive_analysis():
     """Run a comprehensive analysis of Mistral model behavior."""
-    # Test different scenarios
-    scenarios = [
-        ("basic", "What is artificial intelligence?"),
-        ("technical", "Explain how transformers work in deep learning."),
-        ("creative", "Write a short story about a robot learning to paint."),
-        ("complex", "Compare and contrast different machine learning algorithms.")
-    ]
-
-    results = {}
-    for scenario_type, prompt in scenarios:
-        try:
-            print(f"\\nTesting {scenario_type} scenario...")
+    try:
+        # Test different scenarios
+        prompts = [
+            "What is AI?",  # Short prompt
+            "Explain the concept of machine learning.",  # Medium prompt
+            "Write a detailed analysis of artificial intelligence.",  # Long prompt
+        ]
+
+        results = []
+        for prompt in prompts:
             response = client.chat.complete(
                 model="mistral-small-latest",
                 messages=[{"role": "user", "content": prompt}]
             )
-            results[scenario_type] = response.choices[0].message.content
-        except Exception as e:
-            print(f"Error in {scenario_type} scenario: {str(e)}")
-            results[scenario_type] = None
+            results.append(response.choices[0].message.content)
 
-    return results
+        # Analyze results
+        for i, (prompt, result) in enumerate(zip(prompts, results)):
+            print(f"Analysis {i+1}:")
+            print(f"Prompt length: {len(prompt)} chars")
+            print(f"Response length: {len(result)} chars")
 
-# Run analysis
-results = comprehensive_analysis()
+        return "Analysis completed successfully"
+    except Exception as e:
+        print(f"Error in analysis: {str(e)}")
+        return str(e)
 
-# End session with summary
-agentops.end_session("Analysis complete")
+# Run the analysis
+result = comprehensive_analysis()
+print(f"Analysis result: {result}")
 
-print("\\nView the session replay in the AgentOps dashboard to analyze:")
-print("- Response timing patterns")
-print("- Token usage distribution")
-print("- Error patterns")
-print("- Cost optimization opportunities")'''
+# End the session with status
+agentops.end_session("Analysis completed")'''
         ),
     ]
 )

diff --git a/cookbook/AgentOps_Mistral/monitoring_mistral.ipynb b/cookbook/AgentOps_Mistral/monitoring_mistral.ipynb
@@ -30,32 +30,32 @@
    "id": "83cdfb5f",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:26.699562Z",
-     "iopub.status.busy": "2024-12-12T05:53:26.699303Z",
-     "iopub.status.idle": "2024-12-12T05:53:27.744496Z",
-     "shell.execute_reply": "2024-12-12T05:53:27.743624Z"
+     "iopub.execute_input": "2024-12-12T21:48:41.386023Z",
+     "iopub.status.busy": "2024-12-12T21:48:41.385807Z",
+     "iopub.status.idle": "2024-12-12T21:48:43.143389Z",
+     "shell.execute_reply": "2024-12-12T21:48:43.142549Z"
     }
    },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Requirement already satisfied: mistralai in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (1.2.5)\r\n",
-      "Requirement already satisfied: agentops in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (0.3.21)\r\n"
+      "Requirement already satisfied: mistralai in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (1.2.5)\r\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Requirement already satisfied: python-dotenv in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (1.0.1)\r\n"
+      "Requirement already satisfied: agentops in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (0.3.21)\r\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Requirement already satisfied: python-dotenv in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (1.0.1)\r\n",
       "Requirement already satisfied: eval-type-backport<0.3.0,>=0.2.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from mistralai) (0.2.0)\r\n",
       "Requirement already satisfied: httpx<0.28.0,>=0.27.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from mistralai) (0.27.2)\r\n",
       "Requirement already satisfied: jsonpath-python<2.0.0,>=1.0.6 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from mistralai) (1.0.6)\r\n",
@@ -69,7 +69,13 @@
       "Requirement already satisfied: opentelemetry-api<2.0.0,>=1.22.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from agentops) (1.28.2)\r\n",
       "Requirement already satisfied: opentelemetry-sdk<2.0.0,>=1.22.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from agentops) (1.28.2)\r\n",
       "Requirement already satisfied: opentelemetry-exporter-otlp-proto-http<2.0.0,>=1.22.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from agentops) (1.28.2)\r\n",
-      "Requirement already satisfied: anyio in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from httpx<0.28.0,>=0.27.0->mistralai) (4.7.0)\r\n",
+      "Requirement already satisfied: anyio in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from httpx<0.28.0,>=0.27.0->mistralai) (4.7.0)\r\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "Requirement already satisfied: certifi in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from httpx<0.28.0,>=0.27.0->mistralai) (2024.8.30)\r\n",
       "Requirement already satisfied: httpcore==1.* in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from httpx<0.28.0,>=0.27.0->mistralai) (1.0.7)\r\n",
       "Requirement already satisfied: idna in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from httpx<0.28.0,>=0.27.0->mistralai) (3.10)\r\n",
@@ -88,14 +94,14 @@
       "Requirement already satisfied: six>=1.5 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from python-dateutil<3.0.0,>=2.8.2->mistralai) (1.17.0)\r\n",
       "Requirement already satisfied: charset-normalizer<4,>=2 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from requests<3.0.0,>=2.0.0->agentops) (3.4.0)\r\n",
       "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from requests<3.0.0,>=2.0.0->agentops) (2.2.3)\r\n",
-      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from typing-inspect<0.10.0,>=0.9.0->mistralai) (1.0.0)\r\n"
+      "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from typing-inspect<0.10.0,>=0.9.0->mistralai) (1.0.0)\r\n",
+      "Requirement already satisfied: wrapt<2,>=1.10 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from deprecated>=1.2.6->opentelemetry-api<2.0.0,>=1.22.0->agentops) (1.17.0)\r\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Requirement already satisfied: wrapt<2,>=1.10 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from deprecated>=1.2.6->opentelemetry-api<2.0.0,>=1.22.0->agentops) (1.17.0)\r\n",
       "Requirement already satisfied: zipp>=3.20 in /home/ubuntu/.pyenv/versions/3.12.7/lib/python3.12/site-packages (from importlib-metadata<=8.5.0,>=6.0->opentelemetry-api<2.0.0,>=1.22.0->agentops) (3.21.0)\r\n"
      ]
     },
@@ -134,10 +140,10 @@
    "id": "afe63572",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:27.747242Z",
-     "iopub.status.busy": "2024-12-12T05:53:27.747003Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.381901Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.381001Z"
+     "iopub.execute_input": "2024-12-12T21:48:43.146524Z",
+     "iopub.status.busy": "2024-12-12T21:48:43.146189Z",
+     "iopub.status.idle": "2024-12-12T21:48:43.930912Z",
+     "shell.execute_reply": "2024-12-12T21:48:43.930130Z"
     }
    },
    "outputs": [
@@ -202,10 +208,10 @@
    "id": "8eb7d2ea",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.384469Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.384240Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.389834Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.389063Z"
+     "iopub.execute_input": "2024-12-12T21:48:43.933528Z",
+     "iopub.status.busy": "2024-12-12T21:48:43.933303Z",
+     "iopub.status.idle": "2024-12-12T21:48:43.938566Z",
+     "shell.execute_reply": "2024-12-12T21:48:43.937895Z"
     }
    },
    "outputs": [
@@ -255,10 +261,10 @@
    "id": "7a43ee78",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.392228Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.392017Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.397684Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.396936Z"
+     "iopub.execute_input": "2024-12-12T21:48:43.940865Z",
+     "iopub.status.busy": "2024-12-12T21:48:43.940655Z",
+     "iopub.status.idle": "2024-12-12T21:48:43.946230Z",
+     "shell.execute_reply": "2024-12-12T21:48:43.945587Z"
     }
    },
    "outputs": [
@@ -315,10 +321,10 @@
    "id": "a877f95c",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.400497Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.399844Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.404899Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.404167Z"
+     "iopub.execute_input": "2024-12-12T21:48:43.948659Z",
+     "iopub.status.busy": "2024-12-12T21:48:43.948405Z",
+     "iopub.status.idle": "2024-12-12T21:48:43.953295Z",
+     "shell.execute_reply": "2024-12-12T21:48:43.952641Z"
     }
    },
    "outputs": [
@@ -373,10 +379,10 @@
    "id": "1bf51283",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.407159Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.406951Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.540402Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.539552Z"
+     "iopub.execute_input": "2024-12-12T21:48:43.955646Z",
+     "iopub.status.busy": "2024-12-12T21:48:43.955393Z",
+     "iopub.status.idle": "2024-12-12T21:48:44.092667Z",
+     "shell.execute_reply": "2024-12-12T21:48:44.091833Z"
     }
    },
    "outputs": [
@@ -424,10 +430,10 @@
    "id": "51928376",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.543005Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.542769Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.547772Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.547023Z"
+     "iopub.execute_input": "2024-12-12T21:48:44.094954Z",
+     "iopub.status.busy": "2024-12-12T21:48:44.094746Z",
+     "iopub.status.idle": "2024-12-12T21:48:44.099509Z",
+     "shell.execute_reply": "2024-12-12T21:48:44.098826Z"
     }
    },
    "outputs": [
@@ -480,10 +486,10 @@
    "id": "493db2c3",
    "metadata": {
     "execution": {
-     "iopub.execute_input": "2024-12-12T05:53:28.550119Z",
-     "iopub.status.busy": "2024-12-12T05:53:28.549906Z",
-     "iopub.status.idle": "2024-12-12T05:53:28.554697Z",
-     "shell.execute_reply": "2024-12-12T05:53:28.553817Z"
+     "iopub.execute_input": "2024-12-12T21:48:44.102047Z",
+     "iopub.status.busy": "2024-12-12T21:48:44.101577Z",
+     "iopub.status.idle": "2024-12-12T21:48:44.105813Z",
+     "shell.execute_reply": "2024-12-12T21:48:44.105165Z"
     }
    },
    "outputs": [

diff --git a/cookbook/AgentOps_Mistral/monitoring_mistral.py b/cookbook/AgentOps_Mistral/monitoring_mistral.py
@@ -40,7 +40,8 @@ def validate_api_keys():
 
 try:
     if has_valid_keys:
-        agentops.init(os.getenv("AGENTOPS_API_KEY"))
+        # Create a new session for monitoring Mistral interactions
+        agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-monitoring")
         client = Mistral(api_key=os.getenv("MISTRAL_API_KEY"))
         print("Successfully initialized AgentOps and Mistral clients")
     else:
@@ -69,9 +70,11 @@ def get_completion(prompt):
         return f"Error: {str(e)}"
 
 
-# Example usage
+# Example usage with session management
+agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-basic")
 response = get_completion("Explain quantum computing in simple terms")
 print(response)
+agentops.end_session("Basic completion completed")
 
 
 # ## Streaming Responses
@@ -101,8 +104,10 @@ def get_streaming_completion(prompt):
         return f"Error: {str(e)}"
 
 
-# Example usage
+# Example usage with session management
+agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-stream")
 response = get_streaming_completion("What is machine learning?")
+agentops.end_session("Streaming completion completed")
 
 
 # ## Async Operations
@@ -126,17 +131,20 @@ async def get_async_completion(prompt):
         return f"Error: {str(e)}"
 
 
-# Example usage
+# Example usage with session management
 async def main():
+    agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-async")
     response = await get_async_completion("What are the benefits of async programming?")
     print(response)
+    agentops.end_session("Async completion completed")
 
 
 if __name__ == "__main__":
     asyncio.run(main())
 else:
-    # For notebook execution
-    loop = asyncio.get_event_loop()
+    # For notebook execution, create a new event loop
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
     loop.run_until_complete(main())
 
 
@@ -158,16 +166,20 @@ def process_response(response):
         return 0
 
 
-# Example usage combining completion and processing
+# Example usage with session management
+agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-processing")
 response = get_completion("Explain the theory of relativity")
 word_count = process_response(response)
 print(f"Response word count: {word_count}")
+agentops.end_session("Processing completed")
 
 
 # ## Multiple Prompts
 #
 # Track multiple interactions in a single session:
 
+# Start a new session for multiple prompts
+agentops.init(os.getenv("AGENTOPS_API_KEY"), session_name="mistral-multi")
 prompts = ["What is artificial intelligence?", "How does natural language processing work?", "Explain neural networks"]
 
 responses = []
@@ -176,9 +188,5 @@ def process_response(response):
     responses.append(response)
     print(f"\nPrompt: {prompt}\nResponse: {response}\n")
 
-
-# ## End Session
-#
-# Always remember to end your AgentOps session:
-
-agentops.end_session("Success")
+# End the final session
+agentops.end_session("Multiple prompts completed")