Context Managers - Flexible Tracing

Understanding context managers and how they provide flexible, granular control over tracing in your AI applications

Context managers are Python's powerful with statement syntax that provides the most flexible way to trace specific parts of your code. They give you granular control over what to trace, when to trace it, and how to add metadata: all without requiring decorators or modifying function signatures.

🎯 What are Context Managers?

Context managers in Noveum provide:

Granular control - Trace only specific parts of functions
No decorators required - Add tracing without modifying function signatures
Flexible metadata - Add attributes and events dynamically
Automatic cleanup - Spans are properly closed even if errors occur
Nested tracing - Easily create parent-child span relationships

🏗️ Basic Syntax

The basic pattern for using context managers:

from noveum_trace import trace_operation

with trace_operation("operation-name") as span:
    # Your code here
    result = do_something()
    
    # Add attributes to the span
    span.set_attributes({
        "result.length": len(result),
        "success": True
    })

🔧 Core Context Managers

trace_llm_call

Trace LLM API calls with automatic token and cost tracking.

from noveum_trace import trace_llm_call
from openai import OpenAI

client = OpenAI()

with trace_llm_call(model="gpt-4", provider="openai") as span:
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "What is AI?"}
        ]
    )
    
    # Add token usage attributes
    span.set_attributes({
        "llm.input_tokens": response.usage.prompt_tokens,
        "llm.output_tokens": response.usage.completion_tokens,
        "llm.total_tokens": response.usage.total_tokens,
        "llm.response": response.choices[0].message.content
    })

For more structured tracing with function calling and tools, use set_input_attributes and set_output_attributes:

from noveum_trace import trace_llm_call
from openai import OpenAI

client = OpenAI()

with trace_llm_call(model="gpt-4o", provider="openai", operation="chat_completion_with_tools") as span:
    # Capture input: the entire messages dict and tools
    span.set_input_attributes(messages=messages, tools_available=tools_available)
    
    # Make the first API call with function calling enabled
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=messages,
        tools=tools_available,
        tool_choice="auto",  # Let the LLM decide whether to call a tool
        max_tokens=150,
        temperature=0.7,
    )
    
    # Automatically extract tokens, costs, and metadata from response
    span.capture_response(response)
    span.set_output_attributes(response=response)

Parameters:

model - The AI model name (e.g., "gpt-4", "claude-3-opus")
provider - The AI provider (e.g., "openai", "anthropic")
operation - Optional operation name for context

Span Methods:

set_input_attributes(**kwargs) - Set input attributes for the LLM call (e.g., messages, tools_available)
set_output_attributes(**kwargs) - Set output attributes from the LLM response
capture_response(response) - Automatically extract tokens, costs, and metadata from the response object

trace_operation

Trace any generic operation or business logic.

from noveum_trace import trace_operation

def process_customer_data(customer_id: str):
    with trace_operation("process-customer-data") as span:
        # Add context attributes
        span.set_attributes({
            "customer.id": customer_id,
            "operation.type": "data_processing"
        })
        
        # Your business logic
        data = fetch_customer_data(customer_id)
        processed = transform_data(data)
        
        # Add result attributes
        span.set_attributes({
            "records.processed": len(processed),
            "processing.success": True
        })
        
        return processed

Parameters:

name - Operation name (required)
attributes - Optional initial attributes dictionary

trace_agent_operation

Trace agent operations in multi-agent workflows.

from noveum_trace import trace_agent_operation

with trace_agent_operation(agent_type="researcher", agent_id="researcher_001") as span:
    span.set_attributes({
        "agent.capabilities": ["web_search", "analysis"],
        "agent.task": "research_topic",
        "agent.input": topic
    })
    
    # Agent performs research
    results = research_agent.analyze(topic)
    
    span.set_attributes({
        "agent.output": results,
        "agent.confidence": results.confidence,
        "agent.sources_used": len(results.sources)
    })

Parameters:

agent_type - Type of agent (e.g., "researcher", "writer")
agent_id - Unique identifier for the agent instance
operation - Optional operation name

trace_batch_operation

Trace batch processing operations.

from noveum_trace import trace_batch_operation

with trace_batch_operation("process-documents-batch", len(documents)) as span:
    span.set_attributes({
        "batch.size": len(documents),
        "batch.type": "document_processing"
    })
    
    results = []
    for doc in documents:
        result = process_document(doc)
        results.append(result)
    
    span.set_attributes({
        "batch.processed": len(results),
        "batch.success_rate": sum(r.success for r in results) / len(results)
    })

Parameters:

name - Batch operation name (required)
batch_size - Size of the batch (required)

trace_pipeline_stage

Trace individual stages in data pipelines.

from noveum_trace import trace_pipeline_stage

def rag_pipeline(query: str):
    # Stage 1: Generate embeddings
    with trace_pipeline_stage("generate-embeddings", stage_index=0) as span:
        embeddings = embed_query(query)
        span.set_attributes({
            "embeddings.dimensions": len(embeddings),
            "embeddings.model": "text-embedding-ada-002"
        })
    
    # Stage 2: Retrieve documents
    with trace_pipeline_stage("retrieve-documents", stage_index=1) as span:
        documents = vector_search(embeddings)
        span.set_attributes({
            "documents.retrieved": len(documents),
            "documents.top_score": documents[0].score
        })
    
    # Stage 3: Generate answer
    with trace_pipeline_stage("generate-answer", stage_index=2) as span:
        answer = generate_answer(query, documents)
        span.set_attributes({
            "answer.length": len(answer),
            "answer.confidence": 0.95
        })
        
    return answer

Parameters:

name - Stage name (required)
stage_index - Optional stage position in pipeline (0-based)

trace_operation with initial attributes

For tracing generic function calls or operations with initial attribute context:

from noveum_trace import trace_operation

# Trace any operation with initial attributes
with trace_operation("calculate-total") as span:
    span.set_attributes({"items.count": len(items), "discount": 0.1})
    result = calculate_total(items, discount=0.1)
    span.set_attribute("result", result)

create_child_span

Explicitly create child spans for nested operations.

from noveum_trace import trace_operation, create_child_span

with trace_operation("parent-operation") as parent_span:
    parent_span.set_attributes({"operation.type": "complex"})
    
    # Create explicit child span
    with create_child_span("child-operation", parent_span) as child_span:
        child_span.set_attributes({
            "child.operation": "sub_task",
            "child.input": data
        })
        
        result = perform_sub_task(data)
        
        child_span.set_attributes({
            "child.result": result
        })

Parameters:

name - Child span name (required)
parent_span - Parent span object (required)

🌊 Streaming Context Managers

streaming_llm

Handle streaming LLM responses with automatic token tracking.

from noveum_trace import streaming_llm
from openai import OpenAI

client = OpenAI()

with streaming_llm(model="gpt-4", provider="openai") as span:
    stream = client.chat.completions.create(
        model="gpt-4",
        messages=[{"role": "user", "content": "Tell me a story"}],
        stream=True
    )
    
    full_response = ""
    for chunk in stream:
        if chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            full_response += content
            print(content, end="")
    
    # Add final attributes
    span.set_attributes({
        "llm.response": full_response,
        "llm.response_length": len(full_response),
        "llm.streaming": True
    })

trace_streaming

Wrap any stream iterator for tracing.

from noveum_trace import trace_streaming

def process_stream(data_stream):
    with trace_streaming("process-data-stream") as span:
        items_processed = 0
        
        for item in data_stream:
            process_item(item)
            items_processed += 1
        
        span.set_attributes({
            "stream.items_processed": items_processed,
            "stream.success": True
        })

💡 Practical Examples

Multi-Step Operation with Context Managers

import os
from openai import OpenAI
from noveum_trace import trace_operation, trace_llm_call

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def process_user_query(user_query: str) -> dict:
    """Process a user query through multiple steps"""
    results = {}
    
    # Step 1: Preprocess (not traced)
    cleaned_query = user_query.strip().lower()
    
    # Step 2: Enhance query with LLM (traced)
    with trace_llm_call(model="gpt-3.5-turbo", provider="openai") as span:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a query enhancement assistant."},
                {"role": "user", "content": f"Enhance this search query: {cleaned_query}"}
            ]
        )
        
        enhanced_query = response.choices[0].message.content
        results["enhanced_query"] = enhanced_query
        
        span.set_attributes({
            "llm.input_tokens": response.usage.prompt_tokens,
            "llm.output_tokens": response.usage.completion_tokens,
            "llm.enhanced_query": enhanced_query
        })
    
    # Step 3: Database lookup (traced as operation)
    with trace_operation("database-lookup") as span:
        search_results = perform_database_search(enhanced_query)
        results["search_results"] = search_results
        
        span.set_attributes({
            "db.query": enhanced_query,
            "db.results_count": len(search_results)
        })
    
    # Step 4: Generate final response (traced)
    with trace_llm_call(model="gpt-4", provider="openai") as span:
        context = str(search_results[:2])
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": f"Use this context: {context}"},
                {"role": "user", "content": cleaned_query}
            ]
        )
        
        final_response = response.choices[0].message.content
        results["final_response"] = final_response
        
        span.set_attributes({
            "llm.input_tokens": response.usage.prompt_tokens,
            "llm.output_tokens": response.usage.completion_tokens,
            "llm.response_length": len(final_response)
        })
    
    return results

Error Handling with Context Managers

from noveum_trace import trace_operation

def safe_operation(data: str):
    with trace_operation("safe-operation") as span:
        span.set_attributes({
            "operation.input": data,
            "operation.start_time": time.time()
        })
        
        try:
            # Perform operation
            result = risky_operation(data)
            
            # Set success status
            span.set_status("ok")
            span.set_attributes({
                "operation.result": result,
                "operation.success": True
            })
            
            return result
            
        except ValueError as e:
            # Handle specific error
            span.set_status("error", f"ValueError: {str(e)}")
            span.set_attributes({
                "error.type": "ValueError",
                "error.message": str(e),
                "operation.success": False
            })
            return None
            
        except Exception as e:
            # Handle general errors
            span.set_status("error", str(e))
            span.set_attributes({
                "error.type": type(e).__name__,
                "error.message": str(e),
                "operation.success": False
            })
            raise

Nested Spans for Complex Workflows

from noveum_trace import trace_operation

def complex_workflow(user_input: str):
    """Demonstrate nested tracing for complex workflows"""
    
    with trace_operation("complex-workflow") as workflow_span:
        workflow_span.set_attributes({
            "workflow.input": user_input,
            "workflow.start_time": time.time()
        })
        
        results = {}
        
        # Step 1: Custom processing (nested)
        with trace_operation("custom-processing") as process_span:
            processed_input = f"Processed: {user_input}"
            
            process_span.set_attributes({
                "process.input_length": len(user_input),
                "process.output_length": len(processed_input)
            })
            
            results["processed_input"] = processed_input
        
        # Step 2: Data transformation (nested)
        with trace_operation("data-transformation") as transform_span:
            transformed_data = processed_input.replace(" ", "_")
            
            transform_span.set_attributes({
                "transform.method": "replace_spaces",
                "transform.result": transformed_data
            })
            
            results["transformed_data"] = transformed_data
        
        # Step 3: Final LLM call (nested)
        with trace_llm_call(model="gpt-3.5-turbo", provider="openai") as llm_span:
            response = generate_summary(results)
            results["final_result"] = response
            
            llm_span.set_attributes({
                "llm.operation": "summarization",
                "llm.input_size": len(str(results))
            })
        
        # Update workflow span with final results
        workflow_span.set_attributes({
            "workflow.end_time": time.time(),
            "workflow.steps_completed": 3,
            "workflow.success": True
        })
        
        return results

Customer Support Bot Example

import os
import time
from noveum_trace import trace_llm_call, trace_operation
from openai import OpenAI
import noveum_trace

# Initialize Noveum
noveum_trace.init(
    api_key=os.getenv("NOVEUM_API_KEY"),
    project="customer-support-bot",
    environment="development"
)

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def customer_support_bot(user_question: str, customer_id: str = None):
    """A customer support chatbot with comprehensive tracing"""
    
    # Create trace for the entire interaction
    with trace_operation("customer-support-query") as main_span:
        # Add customer context
        main_span.set_attributes({
            "customer.id": customer_id or "anonymous",
            "query.length": len(user_question),
            "query.type": "customer_support",
            "bot.version": "1.0.0"
        })
        
        # Add start event
        main_span.add_event("customer.query.received", {
            "timestamp": time.time(),
            "query.preview": user_question[:50]
        })
        
        try:
            # Trace the LLM call
            with trace_llm_call(model="gpt-4", provider="openai") as llm_span:
                # Add LLM-specific attributes
                llm_span.set_attributes({
                    "ai.model": "gpt-4",
                    "ai.provider": "openai",
                    "ai.temperature": 0.7,
                    "ai.max_tokens": 1000
                })
                
                # Make the LLM call
                response = client.chat.completions.create(
                    model="gpt-4",
                    messages=[
                        {
                            "role": "system",
                            "content": "You are a helpful customer support assistant."
                        },
                        {"role": "user", "content": user_question}
                    ],
                    temperature=0.7,
                    max_tokens=1000
                )
                
                ai_response = response.choices[0].message.content
                
                # Set usage attributes for cost tracking
                llm_span.set_attributes({
                    "llm.input_tokens": response.usage.prompt_tokens,
                    "llm.output_tokens": response.usage.completion_tokens,
                    "llm.total_tokens": response.usage.total_tokens,
                    "response.length": len(ai_response)
                })
            
            # Add success event
            main_span.add_event("customer.query.answered", {
                "timestamp": time.time(),
                "response.length": len(ai_response),
                "success": True
            })
            
            main_span.set_status("ok")
            return ai_response
            
        except Exception as e:
            # Add error event
            main_span.add_event("customer.query.failed", {
                "timestamp": time.time(),
                "error.type": type(e).__name__,
                "error.message": str(e)
            })
            
            main_span.set_status("error", str(e))
            return "I'm sorry, I'm having trouble right now. Please try again later."

✨ Best Practices

Why Use Context Managers

Context Managers provide:

Ability to trace only specific parts of a function
Dynamic attributes based on runtime conditions
Easy integration with legacy code without modifications
Fine-grained control over span boundaries
Automatic resource cleanup even when errors occur

Structure Nested Traces Properly

# Good: Clear hierarchy
with trace_operation("main-workflow") as main:
    with trace_operation("step-1") as step1:
        # Step 1 work
        pass
    
    with trace_operation("step-2") as step2:
        # Step 2 work
        pass

# Avoid: Too deep nesting
with trace_operation("main") as main:
    with trace_operation("sub") as sub:
        with trace_operation("subsub") as subsub:
            with trace_operation("subsubsub") as subsubsub:
                # Too deep - hard to follow
                pass

Set Attributes at the Right Time

with trace_operation("process-data") as span:
    # Set input attributes early
    span.set_attributes({
        "input.size": len(data),
        "input.type": type(data).__name__
    })
    
    # Perform operation
    result = process(data)
    
    # Set output attributes after completion
    span.set_attributes({
        "output.size": len(result),
        "processing.success": True
    })

Handle Errors Gracefully

with trace_operation("risky-operation") as span:
    try:
        result = risky_function()
        span.set_status("ok")
        return result
    except SpecificError as e:
        span.set_status("error", f"SpecificError: {str(e)}")
        span.set_attributes({
            "error.handled": True,
            "error.recovery_attempted": True
        })
        return fallback_value
    except Exception as e:
        span.set_status("error", str(e))
        raise  # Re-raise unexpected errors

Add Meaningful Events

with trace_operation("long-running-task") as span:
    span.add_event("task.started", {"timestamp": time.time()})
    
    # Phase 1
    phase1_result = do_phase1()
    span.add_event("task.phase1.completed", {
        "timestamp": time.time(),
        "phase1.result": phase1_result
    })
    
    # Phase 2
    phase2_result = do_phase2()
    span.add_event("task.phase2.completed", {
        "timestamp": time.time(),
        "phase2.result": phase2_result
    })
    
    span.add_event("task.completed", {"timestamp": time.time()})

🔗 Context Manager vs Manual Spans

Context managers provide automatic span management, but you can also create spans manually for more control:

# With Context Manager (Recommended)
with trace_operation("my-operation") as span:
    # Span automatically started
    result = do_work()
    span.set_attributes({"result": result})
    # Span automatically closed

# Manual Span Creation (Advanced)
client = noveum_trace.get_client()
span = client.start_span("my-operation")
try:
    result = do_work()
    span.set_attributes({"result": result})
    span.set_status("ok")
finally:
    client.finish_span(span)

Recommendation: Use context managers unless you need explicit control over span lifecycle.

🚀 Next Steps

Now that you understand context managers, explore these related concepts:

Traces - Complete request journeys
Spans - Individual operations within traces
Attributes - Metadata and context
Events - Point-in-time occurrences

Integration Examples

Simple LLM Integration - Basic LLM tracing with context managers
SDK Integration Guide - Complete integration patterns

Context managers are the most flexible way to add tracing to your AI applications. They provide granular control while maintaining clean, readable code with automatic resource management.

Context Managers - Flexible Tracing

On this page