Tracing and Observability

def configure(disable_propagation=None, sampling_rate=None):
    """
    Configure global tracing settings.
    
    Parameters:
    - disable_propagation: bool, optional - Disable trace context propagation
    - sampling_rate: float, optional - Trace sampling rate (0.0 to 1.0)
    """

def enable(experiment_id=None):
    """
    Enable MLflow tracing for current session.
    
    Parameters:
    - experiment_id: str, optional - Experiment ID for traces
    """

def disable():
    """
    Disable MLflow tracing for current session.
    """

def reset():
    """
    Reset tracing configuration to defaults.
    """

def set_destination(destination):
    """
    Set tracing destination for spans and traces.
    
    Parameters:
    - destination: str or TraceDestination - Destination for trace data
    """

def start_span(name, span_type=None, inputs=None, attributes=None):
    """
    Start a new span in current trace context.
    
    Parameters:
    - name: str - Span name for identification
    - span_type: str, optional - Type of span (LLM, CHAIN, AGENT, etc.)
    - inputs: Any, optional - Input data for the span
    - attributes: dict, optional - Additional span attributes
    
    Returns:
    Span object representing the active span
    """

@contextmanager
def trace(name, span_type=None, inputs=None, attributes=None):
    """
    Context manager for automatic span lifecycle management.
    
    Parameters:
    - name: str - Span name for identification
    - span_type: str, optional - Type of span
    - inputs: Any, optional - Input data for the span
    - attributes: dict, optional - Additional span attributes
    
    Yields:
    Span object for the duration of the context
    """

def get_current_active_span():
    """
    Get currently active span in execution context.
    
    Returns:
    Span object or None if no active span
    """

def end_span():
    """
    End the currently active span.
    """

def set_span_attribute(key, value):
    """
    Set attribute on current active span.
    
    Parameters:
    - key: str - Attribute key
    - value: Any - Attribute value (will be serialized)
    """

def set_span_attributes(attributes):
    """
    Set multiple attributes on current active span.
    
    Parameters:
    - attributes: dict - Dictionary of attributes to set
    """

def set_span_chat_tools(tools):
    """
    Set chat tools metadata for current span.
    
    Parameters:
    - tools: list - List of tool specifications for chat models
    """

def enable_notebook_display():
    """
    Enable automatic trace display in Jupyter notebooks.
    """

def disable_notebook_display():
    """
    Disable automatic trace display in Jupyter notebooks.
    """

class SpanType:
    """Enumeration of span types for categorizing operations."""
    UNKNOWN = "UNKNOWN"
    LLM = "LLM"
    CHAT_MODEL = "CHAT_MODEL"
    CHAIN = "CHAIN"
    AGENT = "AGENT"
    TOOL = "TOOL"
    RETRIEVER = "RETRIEVER"
    EMBEDDING = "EMBEDDING"
    RERANKER = "RERANKER"
    PARSER = "PARSER"
    
class SpanStatus:
    """Enumeration of span status codes."""
    OK = "OK"
    ERROR = "ERROR"
    CANCELLED = "CANCELLED"
    TIMEOUT = "TIMEOUT"

class SpanStatusCode:
    """Standard status codes for spans."""
    UNSET = "UNSET"
    OK = "OK"  
    ERROR = "ERROR"

@mlflow_trace
def instrumented_function(param1, param2):
    """
    Decorator for automatic function instrumentation.
    
    The decorator automatically creates spans for function execution,
    capturing inputs, outputs, and execution metadata.
    """

def instrument_langchain():
    """
    Enable automatic instrumentation for LangChain components.
    """

def instrument_openai():
    """
    Enable automatic instrumentation for OpenAI API calls.
    """

def instrument_anthropic():
    """
    Enable automatic instrumentation for Anthropic API calls.
    """

def instrument_transformers():
    """
    Enable automatic instrumentation for Transformers library.
    """

def get_trace(trace_id):
    """
    Retrieve trace by ID.
    
    Parameters:
    - trace_id: str - Unique trace identifier
    
    Returns:
    Trace object with all spans and metadata
    """

def search_traces(experiment_ids=None, filter_string=None, max_results=None):
    """
    Search traces matching criteria.
    
    Parameters:
    - experiment_ids: list, optional - List of experiment IDs to search
    - filter_string: str, optional - Filter expression for traces
    - max_results: int, optional - Maximum number of traces to return
    
    Returns:
    List of Trace objects matching criteria
    """

def get_last_active_trace():
    """
    Get the most recently active trace.
    
    Returns:
    Trace object or None if no recent traces
    """

def log_span_event(name, attributes=None, timestamp=None):
    """
    Log event to current active span.
    
    Parameters:
    - name: str - Event name
    - attributes: dict, optional - Event attributes
    - timestamp: int, optional - Event timestamp (defaults to current time)
    """

def set_span_output(output):
    """
    Set output data for current active span.
    
    Parameters:
    - output: Any - Output data (will be serialized)
    """

def set_span_inputs(inputs):
    """
    Set input data for current active span.
    
    Parameters:
    - inputs: Any - Input data (will be serialized)
    """

def add_span_exception(exception):
    """
    Record exception in current active span.
    
    Parameters:
    - exception: Exception - Exception object to record
    """

import mlflow
import mlflow.tracing

# Configure tracing
mlflow.tracing.configure(
    disable_propagation=False,
    sampling_rate=1.0  # Trace all requests
)

# Enable tracing with specific experiment
mlflow.tracing.enable(experiment_id="123")

# Enable notebook display for interactive development
mlflow.tracing.enable_notebook_display()

# Simple function tracing
@mlflow.tracing.trace
def process_text(text):
    # Simulate text processing
    processed = text.upper().strip()
    return processed

# Use traced function
result = process_text("hello world")
print(f"Processed: {result}")

# Disable tracing when done
mlflow.tracing.disable()

import mlflow.tracing
from mlflow.tracing import SpanType, SpanStatus

# Start experiment and enable tracing
mlflow.set_experiment("llm-application")
mlflow.tracing.enable()

# Create manual spans with hierarchy
with mlflow.tracing.trace("user_query_processing", SpanType.CHAIN) as parent_span:
    # Set input data
    mlflow.tracing.set_span_inputs({"query": "What is machine learning?"})
    
    # Child span for retrieval
    with mlflow.tracing.trace("document_retrieval", SpanType.RETRIEVER) as retrieval_span:
        # Simulate document retrieval
        documents = ["Doc1", "Doc2", "Doc3"]
        mlflow.tracing.set_span_output({"retrieved_docs": documents})
        mlflow.tracing.set_span_attribute("num_documents", len(documents))
    
    # Child span for LLM processing
    with mlflow.tracing.trace("llm_generation", SpanType.LLM) as llm_span:
        mlflow.tracing.set_span_inputs({
            "prompt": "Answer based on context",
            "context": documents
        })
        
        # Simulate LLM call
        response = "Machine learning is a subset of AI..."
        mlflow.tracing.set_span_output({"response": response})
        mlflow.tracing.set_span_attribute("model_name", "gpt-4")
        mlflow.tracing.set_span_attribute("tokens_used", 150)
    
    # Set final output
    mlflow.tracing.set_span_output({"final_answer": response})

import mlflow.tracing
from mlflow.tracing import mlflow_trace

# Decorate functions for automatic tracing
@mlflow_trace(span_type=SpanType.EMBEDDING)
def get_embeddings(texts):
    """Get embeddings for list of texts."""
    # Simulate embedding generation
    embeddings = [[0.1, 0.2, 0.3] for _ in texts]
    return embeddings

@mlflow_trace(span_type=SpanType.RETRIEVER)
def retrieve_documents(query_embedding, top_k=5):
    """Retrieve relevant documents."""
    # Simulate document retrieval
    documents = [f"Document {i}" for i in range(top_k)]
    return documents

@mlflow_trace(span_type=SpanType.LLM)
def generate_response(query, context_docs):
    """Generate response using LLM."""
    # Simulate LLM generation
    response = f"Based on {len(context_docs)} documents: Answer here"
    return response

# Enable tracing
mlflow.tracing.enable()

# Use instrumented functions - traces created automatically
query = "What is MLflow?"
query_embedding = get_embeddings([query])
relevant_docs = retrieve_documents(query_embedding[0])
response = generate_response(query, relevant_docs)

print(f"Response: {response}")

import mlflow.tracing
from mlflow.tracing import SpanStatus

mlflow.tracing.enable()

def risky_operation(data):
    """Function that might fail."""
    if len(data) == 0:
        raise ValueError("Empty data provided")
    return sum(data)

# Handle errors in traces
with mlflow.tracing.trace("data_processing") as span:
    try:
        mlflow.tracing.set_span_inputs({"data": [1, 2, 3]})
        result = risky_operation([1, 2, 3])
        mlflow.tracing.set_span_output({"result": result})
        # Span automatically gets OK status
        
    except Exception as e:
        # Record exception in span
        mlflow.tracing.add_span_exception(e)
        # Span automatically gets ERROR status
        raise

# Example with empty data (will fail)
try:
    with mlflow.tracing.trace("failing_operation") as span:
        mlflow.tracing.set_span_inputs({"data": []})
        result = risky_operation([])
except ValueError as e:
    print(f"Operation failed: {e}")

import mlflow.tracing
from mlflow.tracing import SpanType

mlflow.tracing.enable()

# Simulate chat model with tools
def chat_with_tools(messages, tools=None):
    """Simulate chat model with tool calling capability."""
    
    with mlflow.tracing.trace("chat_completion", SpanType.CHAT_MODEL) as span:
        # Set chat tools if provided
        if tools:
            mlflow.tracing.set_span_chat_tools(tools)
        
        # Set input messages
        mlflow.tracing.set_span_inputs({
            "messages": messages,
            "tools": tools or []
        })
        
        # Set model attributes
        mlflow.tracing.set_span_attributes({
            "model_name": "gpt-4-turbo",
            "temperature": 0.7,
            "max_tokens": 1000
        })
        
        # Simulate tool calling
        if tools and "calculator" in [t["name"] for t in tools]:
            with mlflow.tracing.trace("tool_call", SpanType.TOOL) as tool_span:
                mlflow.tracing.set_span_inputs({
                    "tool_name": "calculator",
                    "arguments": {"expression": "2 + 2"}
                })
                tool_result = 4
                mlflow.tracing.set_span_output({"result": tool_result})
        
        # Generate final response
        response = {
            "role": "assistant",
            "content": "The calculation result is 4."
        }
        
        mlflow.tracing.set_span_output({"response": response})
        return response

# Define available tools
tools = [
    {
        "name": "calculator",
        "description": "Perform mathematical calculations",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {"type": "string"}
            }
        }
    }
]

# Use chat model with tools
messages = [{"role": "user", "content": "What is 2 + 2?"}]
response = chat_with_tools(messages, tools=tools)

import mlflow
import mlflow.tracing

# Set up experiment
mlflow.set_experiment("trace_analysis")
mlflow.tracing.enable()

# Generate some traces
for i in range(5):
    with mlflow.tracing.trace(f"operation_{i}") as span:
        mlflow.tracing.set_span_inputs({"iteration": i})
        # Simulate varying execution times
        import time
        time.sleep(0.1 * i)
        mlflow.tracing.set_span_output({"processed": True})

# Retrieve and analyze traces
experiment = mlflow.get_experiment_by_name("trace_analysis")
traces = mlflow.tracing.search_traces(
    experiment_ids=[experiment.experiment_id],
    max_results=10
)

print(f"Found {len(traces)} traces")

for trace in traces:
    print(f"Trace ID: {trace.request_id}")
    print(f"  Start time: {trace.timestamp_ms}")
    print(f"  Duration: {trace.execution_time_ms}ms")
    print(f"  Status: {trace.status}")
    
    # Analyze spans
    for span in trace.spans:
        print(f"    Span: {span.name} ({span.span_type})")
        print(f"      Duration: {span.end_time_ns - span.start_time_ns}ns")
        if span.inputs:
            print(f"      Inputs: {span.inputs}")
        if span.outputs:
            print(f"      Outputs: {span.outputs}")

# Get specific trace
if traces:
    last_trace = mlflow.tracing.get_trace(traces[0].request_id)
    print(f"\nDetailed trace: {last_trace.request_id}")
    print(f"Execution time: {last_trace.execution_time_ms}ms")

import mlflow.tracing

# Enable automatic LangChain instrumentation
mlflow.tracing.instrument_langchain()
mlflow.tracing.enable()

# LangChain components will be automatically traced
from langchain.llms import OpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# Create LangChain components (automatically instrumented)
prompt_template = PromptTemplate(
    input_variables=["topic"],
    template="Write a short summary about {topic}"
)

llm = OpenAI(temperature=0.7)
chain = LLMChain(llm=llm, prompt=prompt_template)

# Run chain - automatically traced
with mlflow.tracing.trace("langchain_summary_generation"):
    result = chain.run(topic="machine learning")
    print(f"Summary: {result}")

# Traces will show LangChain component hierarchy
# with automatic span creation for each component

from mlflow.entities import Trace, Span, SpanEvent
from mlflow.tracing.provider import SpanType, SpanStatus, SpanStatusCode

class Trace:
    request_id: str
    experiment_id: str
    timestamp_ms: int
    execution_time_ms: int
    status: str
    request_metadata: Dict[str, Any]
    tags: Dict[str, str]
    spans: List[Span]

class Span:
    name: str
    context: SpanContext
    parent_id: Optional[str]
    start_time_ns: int
    end_time_ns: int
    status_code: str
    status_message: str
    attributes: Dict[str, Any]
    events: List[SpanEvent]
    span_type: str
    inputs: Optional[Any]
    outputs: Optional[Any]

class SpanContext:
    span_id: str
    trace_id: str

class SpanEvent:
    name: str
    timestamp_ns: int
    attributes: Dict[str, Any]

class SpanType:
    UNKNOWN: str
    LLM: str
    CHAT_MODEL: str
    CHAIN: str
    AGENT: str
    TOOL: str
    RETRIEVER: str
    EMBEDDING: str
    RERANKER: str
    PARSER: str

class SpanStatus:
    OK: str
    ERROR: str
    CANCELLED: str
    TIMEOUT: str

class SpanStatusCode:
    UNSET: str
    OK: str
    ERROR: str

# Context manager and decorator types
from typing import ContextManager, Callable, Any

def trace(
    name: str,
    span_type: Optional[str] = None,
    inputs: Optional[Any] = None,
    attributes: Optional[Dict[str, Any]] = None
) -> ContextManager[Span]: ...

def mlflow_trace(
    span_type: Optional[str] = None,
    name: Optional[str] = None
) -> Callable[[Callable], Callable]: ...