OpenTelemetry Cohere Instrumentation

from opentelemetry.instrumentation.cohere import CohereInstrumentor

from opentelemetry.instrumentation.cohere import CohereInstrumentor
import cohere

# Instrument Cohere calls with default settings
CohereInstrumentor().instrument()

# Create Cohere client (instrumentation applies automatically)
client = cohere.Client("your-api-key")

# All Cohere API calls are now automatically traced
response = client.chat(message="Hello, world!")
completion = client.generate(prompt="Once upon a time...")
rerank_result = client.rerank(
    query="What is machine learning?",
    documents=["AI overview", "ML tutorial", "Deep learning guide"]
)

class CohereInstrumentor(BaseInstrumentor):
    """An instrumentor for Cohere's client library."""
    
    def __init__(self, exception_logger=None, use_legacy_attributes=True):
        """
        Initialize the Cohere instrumentor.
        
        Parameters:
        - exception_logger: Optional custom exception logger function
        - use_legacy_attributes: Whether to use legacy span attributes (default: True)
        """
    
    def instrument(self, **kwargs):
        """
        Enable instrumentation of Cohere API calls.
        
        Parameters:
        - tracer_provider: Optional OpenTelemetry TracerProvider
        - event_logger_provider: Optional EventLoggerProvider for non-legacy mode
        """
    
    def uninstrument(self, **kwargs):
        """Disable instrumentation and restore original Cohere client methods."""
    
    def instrumentation_dependencies(self) -> Collection[str]:
        """
        Return collection of required package dependencies.
        
        Returns:
        Collection[str]: Tuple containing package dependency specifications
        """

class Config:
    """Configuration class for instrumentation settings."""
    
    exception_logger = None  # Custom exception logger function
    use_legacy_attributes: bool = True  # Boolean flag for legacy attributes

# From opentelemetry.instrumentation.cohere.utils
def should_send_prompts() -> bool:
    """
    Determine if prompts should be traced based on TRACELOOP_TRACE_CONTENT environment variable.
    
    Returns:
    bool: True if content should be traced, False otherwise
    """

def should_emit_events() -> bool:
    """
    Check if instrumentation should emit events instead of using legacy attributes.
    
    Returns:
    bool: True if events should be emitted, False for legacy attributes
    """

def dont_throw(func):
    """
    Decorator that wraps functions to log exceptions instead of throwing them.
    Uses the Config.exception_logger if available for custom exception handling.
    
    Parameters:
    - func: The function to wrap
    
    Returns:
    Wrapper function that logs exceptions instead of raising them
    """

# From opentelemetry.instrumentation.cohere.span_utils
def _set_span_chat_response(span, response):
    """
    Set span attributes specific to chat responses.
    
    Parameters:
    - span: OpenTelemetry span object
    - response: Cohere chat response object
    """

def _set_span_generations_response(span, response):
    """
    Set span attributes specific to completion generation responses.
    
    Parameters:
    - span: OpenTelemetry span object
    - response: Cohere generation response object
    """

def _set_span_rerank_response(span, response):
    """
    Set span attributes specific to rerank responses.
    
    Parameters:
    - span: OpenTelemetry span object
    - response: Cohere rerank response object
    """

def set_input_attributes(span, llm_request_type, kwargs):
    """
    Set input attributes on span for LLM requests.
    
    Parameters:
    - span: OpenTelemetry span object
    - llm_request_type: Type of LLM request (completion, chat, rerank)
    - kwargs: Request parameters from Cohere API call
    """

def set_response_attributes(span, llm_request_type, response):
    """
    Set response attributes on span for LLM responses.
    
    Parameters:
    - span: OpenTelemetry span object
    - llm_request_type: Type of LLM request (completion, chat, rerank)
    - response: Response object from Cohere API
    """

def set_span_request_attributes(span, kwargs):
    """
    Set general request attributes on span.
    
    Parameters:
    - span: OpenTelemetry span object
    - kwargs: Request parameters including model, temperature, tokens, etc.
    """

def _set_span_attribute(span, name, value):
    """
    Helper function to safely set span attributes, only if value is not None or empty.
    
    Parameters:
    - span: OpenTelemetry span object
    - name: Attribute name
    - value: Attribute value
    """

# From opentelemetry.instrumentation.cohere.event_emitter
def _parse_response_event(index: int, llm_request_type: str, response) -> ChoiceEvent:
    """
    Parse Cohere API response into a ChoiceEvent for event emission.
    
    Parameters:
    - index: Index of the response choice
    - llm_request_type: Type of LLM request
    - response: Cohere API response object
    
    Returns:
    ChoiceEvent: Structured event object
    """

def _emit_message_event(event: MessageEvent, event_logger: EventLogger) -> None:
    """
    Emit a message event to OpenTelemetry event logger.
    
    Parameters:
    - event: MessageEvent to emit
    - event_logger: OpenTelemetry EventLogger instance
    """

def _emit_choice_event(event: ChoiceEvent, event_logger: EventLogger) -> None:
    """
    Emit a choice event to OpenTelemetry event logger.
    
    Parameters:
    - event: ChoiceEvent to emit
    - event_logger: OpenTelemetry EventLogger instance
    """

def emit_input_event(event_logger, llm_request_type: str, kwargs):
    """
    Emit input events for LLM requests.
    
    Parameters:
    - event_logger: OpenTelemetry EventLogger instance
    - llm_request_type: Type of LLM request
    - kwargs: Request parameters
    """

def emit_response_events(event_logger, llm_request_type: str, response):
    """
    Emit response events for LLM responses.
    
    Parameters:
    - event_logger: OpenTelemetry EventLogger instance
    - llm_request_type: Type of LLM request
    - response: Response object from Cohere API
    """

def emit_event(event: Union[MessageEvent, ChoiceEvent], event_logger: Union[EventLogger, None]):
    """
    General event emitter for MessageEvent or ChoiceEvent.
    
    Parameters:
    - event: Event object to emit
    - event_logger: OpenTelemetry EventLogger instance
    """

# From opentelemetry.instrumentation.cohere.event_models
@dataclass
class MessageEvent:
    """Represents an input event for the AI model."""
    
    content: Any  # Message content
    role: str = "user"  # Message role
    tool_calls: Optional[List[ToolCall]] = None  # Tool calls

@dataclass
class ChoiceEvent:
    """Represents a completion event for the AI model."""
    
    index: int  # Choice index
    message: CompletionMessage  # Completion message
    finish_reason: str = "unknown"  # Completion finish reason
    tool_calls: Optional[List[ToolCall]] = None  # Tool calls

class ToolCall(TypedDict):
    """Represents a tool call in the AI model."""
    
    id: str  # Tool call ID
    function: _FunctionToolCall  # Function call details
    type: Literal["function"]  # Tool call type

class CompletionMessage(TypedDict):
    """Represents a message in the AI model."""
    
    content: Any  # Message content
    role: str = "assistant"  # Message role

class _FunctionToolCall(TypedDict):
    function_name: str  # Function name
    arguments: Optional[dict[str, Any]]  # Function arguments

WRAPPED_METHODS = [
    {
        "object": "Client",
        "method": "generate", 
        "span_name": "cohere.completion",
    },
    {
        "object": "Client",
        "method": "chat",
        "span_name": "cohere.chat", 
    },
    {
        "object": "Client",
        "method": "rerank",
        "span_name": "cohere.rerank",
    },
]

_instruments = ("cohere >=4.2.7, <6",)

# Request type values for different Cohere operations
LLMRequestTypeValues.CHAT = "llm.chat"
LLMRequestTypeValues.COMPLETION = "llm.completion" 
LLMRequestTypeValues.RERANK = "llm.rerank"
LLMRequestTypeValues.UNKNOWN = "llm.unknown"

TRACELOOP_TRACE_CONTENT = "TRACELOOP_TRACE_CONTENT"

# From opentelemetry.instrumentation.cohere.event_emitter
EVENT_ATTRIBUTES = {
    "gen_ai.system": "cohere"
}

VALID_MESSAGE_ROLES = {"user", "assistant", "system", "tool"}

class Roles(Enum):
    USER = "user"
    ASSISTANT = "assistant" 
    SYSTEM = "system"
    TOOL = "tool"

# From opentelemetry.instrumentation.cohere.__init__
def _with_tracer_wrapper(func):
    """
    Helper decorator for providing tracer context to wrapper functions.
    
    Parameters:
    - func: Function to wrap with tracer context
    
    Returns:
    Decorated function with tracer context
    """

def _llm_request_type_by_method(method_name: str):
    """
    Map Cohere client method names to LLM request type values.
    
    Parameters:
    - method_name: Name of the Cohere client method
    
    Returns:
    LLMRequestTypeValues: Corresponding request type enum value
    """

def _wrap(tracer, event_logger, to_wrap, wrapped, instance, args, kwargs):
    """
    Core wrapper function that instruments Cohere API calls with OpenTelemetry spans.
    
    Parameters:
    - tracer: OpenTelemetry Tracer instance
    - event_logger: EventLogger for modern telemetry mode
    - to_wrap: Method configuration from WRAPPED_METHODS
    - wrapped: Original method being wrapped
    - instance: Instance of the Cohere client
    - args: Method arguments
    - kwargs: Method keyword arguments
    
    Returns:
    Result from the original Cohere API call
    """

export TRACELOOP_TRACE_CONTENT=false

def custom_exception_logger(exception):
    print(f"Instrumentation error: {exception}")

instrumentor = CohereInstrumentor(exception_logger=custom_exception_logger)
instrumentor.instrument()

from opentelemetry.sdk._events import EventLoggerProvider
from opentelemetry.sdk._logs import LoggerProvider

# Setup for modern event-based telemetry
log_provider = LoggerProvider()
event_logger_provider = EventLoggerProvider(log_provider)

instrumentor = CohereInstrumentor(use_legacy_attributes=False)
instrumentor.instrument(event_logger_provider=event_logger_provider)

from opentelemetry.sdk.trace import TracerProvider

tracer_provider = TracerProvider()
instrumentor = CohereInstrumentor()
instrumentor.instrument(tracer_provider=tracer_provider)