Documentation
Core Concepts/Context Managers - Flexible Tracing

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

Trace LLM API calls with automatic token and cost tracking.

from noveum_trace import trace_llm
from openai import OpenAI
 
client = OpenAI()
 
with trace_llm(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
    })

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

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

Trace agent operations in multi-agent workflows.

from noveum_trace import trace_agent
 
with trace_agent(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") 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 - Optional size of the batch

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_number=1) 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_number=2) 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_number=3) 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_number - Optional stage position in pipeline

trace_function_calls

Trace function calls with granular control.

from noveum_trace import trace_function_calls
 
with trace_function_calls("calculate-total") as span:
    total = sum(items)
    span.set_attributes({
        "function.input_count": len(items),
        "function.result": total
    })

Parameters:

  • name - Function name (required)

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
 
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(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(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("success")
            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(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, 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(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("success")
            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("success")
        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

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.

Exclusive Early Access

Get Early Access to Noveum.ai Platform

Be the first one to get notified when we open Noveum Platform to more users. All users get access to Observability suite for free, early users get free eval jobs and premium support for the first year.

Sign up now. We send access to new batch every week.

Early access members receive premium onboarding support and influence our product roadmap. Limited spots available.

Previous

Events - Point-in-Time Occurrences

Next

Observability Best Practices

On this page