feat: Context management using with keyword for a more Pythonic way

[dot-agi]

📥 Pull Request

📘 Description

Implements Python context manager protocol for AgentOps, enabling automatic trace lifecycle management with the with statement. This provides guaranteed cleanup, exception safety, and a more Pythonic API while maintaining full backward compatibility.

Key Changes:

• Added AgentOpsContextManager class with __enter__ and __exit__ methods for automatic trace management
• Created InitializationProxy that supports both regular and context manager usage patterns
• Modified agentops.init() to return the proxy, enabling with agentops.init(...) as session: syntax
• Automatic trace status setting based on execution flow ("Success" vs "Error")
• Exception safety with guaranteed cleanup even when errors occur
• Zero breaking changes - all existing code continues to work unchanged

Examples Added:

• 01_basic_usage.py - Demonstrates basic context manager usage vs traditional approach with agent operations
• 02_exception_handling.py - Shows exception handling patterns and guaranteed cleanup in various error scenarios
• 03_backward_compatibility.py - Proves existing code works unchanged and demonstrates mixed usage patterns
• 04_retry_logic.py - Implements retry patterns with exponential backoff and batch processing using individual traces

Documentation Added:

• docs/v2/usage/context-manager.mdx - Comprehensive context manager documentation

Usage Example:

# New context manager approach
with agentops.init(api_key="key", trace_name="my_workflow") as session:
    # Your agent code here
    pass  # Trace automatically ends with "Success"

# Traditional approach still works
session = agentops.init(api_key="key", auto_start_session=True)
# ... existing code ...
agentops.end_session()

🧪 Testing

• TestAgentOpsContextManager - Tests context manager lifecycle, exception handling, and session management
• TestInitializationProxy - Tests dual-purpose proxy behavior for both regular and context manager usage
• TestIntegration - Tests complete flow integration between proxy and context manager components
• Exception propagation tests ensuring exceptions are never suppressed

[Copilot]

Pull Request Overview

This PR introduces a new context manager implementation to simplify AgentOps lifecycle management while retaining backward compatibility. Key changes include:

• Implementation of AgentOpsContextManager with enter and exit methods, and the InitializationProxy to support both context manager and traditional usage.
• Updates to examples and tests to demonstrate context manager usage with improved session lifecycle handling.
• Documentation updates to include context manager usage in the docs and examples.

Reviewed Changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit/test_context_manager.py	Added tests covering various context manager usage scenarios.
examples/context_manager/README.md	New documentation and usage examples for the context manager.
examples/context_manager/04_retry_logic.py	Added examples demonstrating retry logic with context management.
examples/context_manager/03_backward_compatibility.py	Ensured backward compatibility with legacy APIs using both methods.
examples/context_manager/02_exception_handling.py	Added examples for exception handling with context manager.
examples/context_manager/01_basic_usage.py	Provided basic usage examples contrasting traditional and context methods.
docs/mint.json	Updated documentation index to include context manager usage page.
agentops/context_manager.py	Introduced new context manager and proxy classes for AgentOps.
agentops/init.py	Modified init() to return InitializationProxy supporting new usage.

Comments suppressed due to low confidence (1)

agentops/context_manager.py:63

• Consider explicitly checking the auto_start_session parameter rather than relying on a None check. Clarifying this conditional logic may improve code readability and maintainability.

if not tracing_core.initialized:

[codecov]

Codecov Report

Attention: Patch coverage is 66.03774% with 54 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
agentops/context_manager.py	65.16%	54 Missing ⚠️

📢 Thoughts on this report? Let us know!

[Copilot]

Pull Request Overview

This PR introduces a Python context manager for AgentOps to improve trace lifecycle management by supporting the “with” statement and maintaining backward compatibility through the new InitializationProxy.

• Adds AgentOpsContextManager with proper enter and exit methods for automatic trace cleanup.
• Introduces InitializationProxy to support both the traditional and context manager usage patterns.
• Updates tests, examples, and documentation to cover the new context manager functionality.

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/unit/test_session_legacy.py	Adjusted type assertions and added check for wrapped Session.
tests/unit/test_session.py	Updated tests to verify that the proxy returns a wrapped Session.
tests/unit/test_context_manager.py	Comprehensive tests for context manager operations and exception handling.
examples/context_manager/README.md	Added documentation and examples for context manager usage.
examples/context_manager/*	Multiple usage examples illustrating basic, exception, retry, and legacy behaviors.
agentops/context_manager.py	Introduces context manager implementation and proxy for dual usage.
agentops/init.py	Adjusted init() to return the InitializationProxy and added "tool" decorator.

Comments suppressed due to low confidence (1)

agentops/context_manager.py:37

• [nitpick] Consider renaming 'init_result' to 'initialization_result' to improve clarity in what the variable represents.

self.init_result = None

[Dwij1704]

Using with agentops.init(...) merely wraps client initialization in a context manager without adding any real value. What we really need is a dedicated tracing context, like with agentops.start_trace(...), that clearly scopes exactly the code we want to monitor and then automatically closes the trace when the block exits. Just as with open('file.txt') as f: implicitly calls f.close() at the end, with agentops.start_trace() should implicitly call end_trace(), ensuring that spans start and end cleanly without any extra effort of end_trace.

[Dwij1704]

Maybe a better title that clearly points to the feature?