Vector Database Instrumentation Guide

Vector Database Instrumentation Guide

This guide provides step-by-step instructions for adding a new vector database instrumentation to OpenLIT, following the established patterns and conventions learned from Pinecone and ChromaDB implementations.

Table of Contents

1. File Structure
2. Code Structure
3. Naming Conventions
4. Code Style Guidelines
5. Implementation Steps
6. Testing
7. Integration

File Structure

Every vector database instrumentation follows a consistent 3-4 file structure:

sdk/python/src/openlit/instrumentation/{database_name}/
├── __init__.py          # Instrumentor class with dependency requirements
├── {database_name}.py   # Sync wrapper functions
├── async_{database_name}.py  # Async wrapper functions (only if database supports async)
└── utils.py            # Telemetry processing logic

Example Structure (using "pinecone" as example):

sdk/python/src/openlit/instrumentation/pinecone/
├── __init__.py
├── pinecone.py
├── async_pinecone.py    # Only if async operations exist
└── utils.py

Note: If the database doesn't support async operations (like ChromaDB), omit the async_{database_name}.py file.

Code Structure

1. DB_OPERATION_MAP Pattern

All vector database instrumentations use a centralized operation mapping in utils.py:

# Operation mapping for simple span naming
DB_OPERATION_MAP = {
    "{database}.create_collection": SemanticConvention.DB_OPERATION_CREATE_COLLECTION,
    "{database}.add": SemanticConvention.DB_OPERATION_INSERT,
    "{database}.get": SemanticConvention.DB_OPERATION_GET,
    "{database}.query": SemanticConvention.DB_OPERATION_GET,  # Both get() and query() are GET operations
    "{database}.update": SemanticConvention.DB_OPERATION_UPDATE,
    "{database}.upsert": SemanticConvention.DB_OPERATION_UPSERT,
    "{database}.delete": SemanticConvention.DB_OPERATION_DELETE,
}

Key Learning: Multiple endpoints can map to the same semantic operation. Use the endpoint parameter to differentiate implementation details.

2. General Wrap Pattern

Use the general_wrap pattern for efficiency:

def general_wrap(gen_ai_endpoint, version, environment, application_name, tracer, 
                pricing_info, capture_message_content, metrics, disable_metrics):
    """
    Create a wrapper for {database} operations using the general wrap pattern.
    """
    def wrapper(wrapped, instance, args, kwargs):
        # CRITICAL: Suppression check
        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
            return wrapped(*args, **kwargs)
        
        # Get server address and port using the standard helper
        server_address, server_port = set_server_address_and_port(instance)
        
        db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "unknown")
        if db_operation == "create_collection":
            namespace = kwargs.get("name") or (args[0] if args else "unknown")
        else:
            namespace = getattr(instance, "name", "unknown")  # For collection-based databases
        span_name = f"{db_operation} {namespace}"
        
        with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
            start_time = time.time()
            response = wrapped(*args, **kwargs)
            
            try:
                # Process response and generate telemetry
                response = process_vectordb_response(
                    response, db_operation, server_address, server_port,
                    environment, application_name, metrics, start_time, span,
                    capture_message_content, disable_metrics, version, 
                    instance, args, endpoint=gen_ai_endpoint, **kwargs
                )
                
            except Exception as e:
                handle_exception(span, e)
                
            return response
    
    return wrapper

3. Simplified Telemetry Pattern

Based on ChromaDB learnings, use a simplified, consistent pattern:

def common_vectordb_logic(scope, environment, application_name, 
    metrics, capture_message_content, disable_metrics, version, 
    instance=None, endpoint=None):
    """
    Process vector database request and generate telemetry.
    """
    scope._end_time = time.time()

    # Set common database span attributes using helper
    common_db_span_attributes(scope, SemanticConvention.DB_SYSTEM_{DATABASE_NAME}, 
        scope._server_address, scope._server_port, environment, application_name, version)

    # Set DB operation specific attributes
    scope._span.set_attribute(SemanticConvention.DB_OPERATION_NAME, scope._db_operation)
    scope._span.set_attribute(SemanticConvention.DB_CLIENT_OPERATION_DURATION, 
        scope._end_time - scope._start_time)

    # Handle operations with endpoint differentiation
    if scope._db_operation == SemanticConvention.DB_OPERATION_GET:
        collection_name = getattr(instance, "name", "unknown")
        
        # Differentiate between different GET operations
        if endpoint == "{database}.get":
            query = scope._kwargs.get("ids", [])
            
            # Standard database attributes
            scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
            scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
            scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
            
            # Database-specific attributes
            scope._span.set_attribute(SemanticConvention.DB_FILTER, 
                str(scope._kwargs.get("where", "")))
            
            # Response metrics
            if scope._response and scope._response.get('ids'):
                returned_rows = object_count(scope._response['ids'])  # Adjust for response structure
                scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, returned_rows)
            
            scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
                f"{scope._db_operation} {collection_name} "
                f"ids={query} "
                f"limit={scope._kwargs.get('limit', 'None')}")
        
        elif endpoint == "{database}.query":
            # Handle query-specific parameters
            query_texts = scope._kwargs.get("query_texts", [])
            query_embeddings = scope._kwargs.get("query_embeddings", [])
            
            # Create query content representation
            if query_texts:
                query_content = f"texts={query_texts}"
            elif query_embeddings:
                query_content = f"embeddings={len(query_embeddings)} vectors"
            else:
                query_content = "no query provided"
            
            # Standard database attributes
            scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, query_content)
            scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
            scope._span.set_attribute(SemanticConvention.DB_VECTOR_QUERY_TOP_K, 
                scope._kwargs.get("n_results", 10))
            
            # Response metrics (adjust for nested response structure)
            if scope._response and scope._response.get('ids'):
                returned_rows = object_count(scope._response['ids'][0]) if scope._response['ids'] else 0
                scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, returned_rows)
            
            scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
                f"{scope._db_operation} {collection_name} "
                f"n_results={scope._kwargs.get('n_results', 10)} "
                f"{query_content}")

    elif scope._db_operation == SemanticConvention.DB_OPERATION_INSERT:
        collection_name = getattr(instance, "name", "unknown")
        query = scope._kwargs.get("ids", [])
        
        # Standard database attributes
        scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
        
        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
            f"{scope._db_operation} {collection_name} "
            f"ids={query} "
            f"documents={scope._kwargs.get('documents', 'None')}")

    # ... similar patterns for other operations

    scope._span.set_status(Status(StatusCode.OK))

    # Record metrics using helper
    if not disable_metrics:
        record_db_metrics(metrics, SemanticConvention.DB_SYSTEM_{DATABASE_NAME}, 
            scope._server_address, scope._server_port, environment, application_name, 
            scope._start_time, scope._end_time)

4. Key Telemetry Patterns

Essential Attributes for All Operations:

• DB_QUERY_TEXT: String representation of the query/operation
• DB_COLLECTION_NAME: Collection name (for collection-based databases)
• DB_VECTOR_COUNT: Count of primary query parameter
• DB_QUERY_SUMMARY: Comprehensive operation summary

Query Summary Format:

f"{scope._db_operation} {collection_name} key_param={value} optional_param={value or 'None'}"

Response Structure Handling:

# For flat arrays (like ChromaDB GET)
returned_rows = object_count(scope._response['ids'])

# For nested arrays (like ChromaDB QUERY)
returned_rows = object_count(scope._response['ids'][0]) if scope._response['ids'] else 0

Naming Conventions

1. File Names

• Use lowercase with underscores: {database_name}.py
• Async files: async_{database_name}.py (only if needed)
• No hyphens, use underscores

2. Function Names

• Use snake_case consistently
• Standard function names:
  • general_wrap
  • async_general_wrap (if needed)
  • common_vectordb_logic
  • process_vectordb_response

3. Variable Names

• Use snake_case for all variables
• Simplified naming: Use query for primary operation parameter
• Standard variable names:
  • server_address, server_port
  • start_time, end_time
  • db_operation
  • collection_name (for collection-based databases)

4. Class Names

• Use PascalCase for classes
• Instrumentor class: {DatabaseName}Instrumentor

Code Style Guidelines

1. Quotes

• Always use double quotes for strings: "example"
• Never use single quotes for strings: ❌ 'example'
• Consistent across all files including inline comments, error messages, and string literals

2. Indentation

• 4 spaces for indentation
• No tabs
• Multi-line definitions: Use only one level of indentation for continuation lines

  # ✅ Correct: One level indentation for multi-line function definitions
  def general_wrap(gen_ai_endpoint, version, environment, application_name,
      tracer, pricing_info, capture_message_content, metrics, disable_metrics):
  
  # ✅ Correct: One level indentation for multi-line function calls
  response = process_vectordb_response(
      response, db_operation, server_address, server_port,
      environment, application_name, metrics, start_time, span,
      capture_message_content, disable_metrics, version, 
      instance, args, endpoint=gen_ai_endpoint, **kwargs
  )
  
  # ❌ Wrong: Don't align with opening parenthesis
  def general_wrap(gen_ai_endpoint, version, environment, application_name,
                   tracer, pricing_info, capture_message_content, metrics, disable_metrics):
  

3. Comments

• Triple double quotes for docstrings:

  def function_name():
      """
      Function description.
      """
  

• Standard inline comments for specific actions:

  # Get server address and port using the standard helper
  server_address, server_port = set_server_address_and_port(instance)
  
  # Process response and generate telemetry
  response = process_vectordb_response(...)
  

4. Import Order

Follow this specific order:

# Standard library imports
import time

# OpenTelemetry imports
from opentelemetry.trace import Status, StatusCode, SpanKind
from opentelemetry import context as context_api

# OpenLIT imports
from openlit.__helpers import (
    common_db_span_attributes,
    record_db_metrics,
    handle_exception,
    set_server_address_and_port,
)
from openlit.semcov import SemanticConvention

5. Function Parameters

Standard parameter order:

def function_name(gen_ai_endpoint, version, environment, application_name, tracer, 
                 pricing_info, capture_message_content, metrics, disable_metrics):

Implementation Steps

Step 1: Create Directory Structure

mkdir -p sdk/python/src/openlit/instrumentation/{database_name}

Step 2: Create init.py (Compact Pattern)

For databases with many operations, use the compact data structure approach:

"""
OpenLIT {DatabaseName} Instrumentation
"""

from typing import Collection
import importlib.metadata
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from wrapt import wrap_function_wrapper

from openlit.instrumentation.{database_name}.{database_name} import general_wrap
# Only import async if the database supports it
# from openlit.instrumentation.{database_name}.async_{database_name} import async_general_wrap

_instruments = ("{database-package} >= {version}",)

# Operations to wrap for both sync and async clients
{DATABASE_NAME}_OPERATIONS = [
    ("method_name", "{database}.endpoint"),
    ("create_collection", "{database}.create_collection"),
    ("upsert", "{database}.upsert"),
    ("query", "{database}.query"),
    ("delete", "{database}.delete"),
    # ... add all operations here
]

class {DatabaseName}Instrumentor(BaseInstrumentor):
    """
    An instrumentor for {DatabaseName}'s client library.
    """

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs):
        version = importlib.metadata.version("{database-package}")
        environment = kwargs.get("environment", "default")
        application_name = kwargs.get("application_name", "default")
        tracer = kwargs.get("tracer")
        pricing_info = kwargs.get("pricing_info", {})
        capture_message_content = kwargs.get("capture_message_content", False)
        metrics = kwargs.get("metrics_dict")
        disable_metrics = kwargs.get("disable_metrics")

        # Wrap sync operations
        for method_name, endpoint in {DATABASE_NAME}_OPERATIONS:
            wrap_function_wrapper(
                "{database_module}",
                f"{ClientClass}.{method_name}",
                general_wrap(
                    endpoint, version, environment, application_name, tracer,
                    pricing_info, capture_message_content, metrics, disable_metrics
                ),
            )

        # Wrap async operations (only if database supports them)
        # for method_name, endpoint in {DATABASE_NAME}_OPERATIONS:
        #     wrap_function_wrapper(
        #         "{database_module}",
        #         f"{AsyncClientClass}.{method_name}",
        #         async_general_wrap(
        #             endpoint, version, environment, application_name, tracer,
        #             pricing_info, capture_message_content, metrics, disable_metrics
        #         ),
        #     )

    def _uninstrument(self, **kwargs):
        pass

For databases with few operations, use the direct approach:

# Direct wrapping for databases with < 10 operations
wrap_function_wrapper(
    "{database_module}",
    "{Class}.{method}",
    general_wrap(
        "{database}.{operation}",
        version, environment, application_name, tracer,
        pricing_info, capture_message_content, metrics, disable_metrics
    ),
)

Step 3: Create Wrapper Files

Create sync (and async if needed) wrapper files following the general_wrap pattern.

Step 4: Create utils.py

Include functions in standard order:

1. DB_OPERATION_MAP with proper semantic operation mapping
2. object_count helper function
3. set_server_address_and_port function for server address extraction
4. common_vectordb_logic with endpoint differentiation
5. process_vectordb_response function

Server Address Extraction Function:

def set_server_address_and_port(instance):
    """
    Extracts server address and port from {database} client instance.
    
    Args:
        instance: {Database} client instance
        
    Returns:
        tuple: (server_address, server_port)
    """
    server_address = "default-host"
    server_port = default_port
    
    # Database-specific extraction logic
    # For Qdrant: instance.init_options
    # For others: instance._client.base_url, instance.config.host, etc.
    
    return server_address, server_port

Usage in wrapper files:

# Replace generic helper with database-specific function
from openlit.instrumentation.{database}.utils import set_server_address_and_port

# In wrapper function
server_address, server_port = set_server_address_and_port(instance)

Step 5: Database-Specific Adaptations

Server Address Patterns:

• Pinecone: "pinecone.io", 443
• ChromaDB: "localhost", 8000
• Adjust for your database's default

Collection vs Namespace:

• Collection-based (ChromaDB): Use DB_COLLECTION_NAME
• Namespace-based (Pinecone): Use DB_NAMESPACE

Response Structure:

• Flat arrays: object_count(response['ids'])
• Nested arrays: object_count(response['ids'][0])

Testing

Create Test File

Create sdk/python/tests/test_{database_name}.py:

"""
Tests for {DatabaseName} functionality.
"""

import pytest
import openlit
from {database_package} import {ClientClass}

# Initialize client
client = {ClientClass}()

# Initialize OpenLIT
openlit.init(environment="openlit-python-testing", application_name="openlit-python-{database}-test")

def test_sync_{database}_operation():
    """
    Tests synchronous operation.
    """
    try:
        response = client.operation(
            # parameters
        )
        assert response  # Add appropriate assertions
        
    except Exception as e:
        if "rate limit" in str(e).lower():
            print("Rate limit exceeded:", e)
        else:
            raise

# Add async tests only if database supports async
# @pytest.mark.asyncio
# async def test_async_{database}_operation():
#     """
#     Tests asynchronous operation.
#     """
#     pass

Integration

Step 1: Add to Base Instrumentation

Add to sdk/python/src/openlit/instrumentation/__init__.py:

from openlit.instrumentation.{database_name} import {DatabaseName}Instrumentor

def auto_instrument(**kwargs):
    """
    Auto-instrument supported libraries
    """
    # ... existing instrumentors ...
    
    # {DatabaseName} Instrumentation
    try:
        {DatabaseName}Instrumentor().instrument(**kwargs)
    except Exception as e:
        logger.warning(f"Failed to instrument {DatabaseName}: {e}")

Step 2: Add Semantic Conventions

Add any new semantic conventions to sdk/python/src/openlit/semcov/__init__.py.

Consistency Requirements

CRITICAL: All vector database instrumentations must follow identical patterns for maintainability and developer experience:

Code Style Consistency:

• Quotes: Always use double quotes (") - never single quotes (')
• Comments: Use standardized inline comments (see Code Style Guidelines)
• Function Order: Follow the exact order specified in utils.py structure
• Import Order: Standard library → OpenTelemetry → OpenLIT imports
• Variable Names: Use snake_case consistently

Structural Consistency:

• Exception Handling: wrapped() call outside try block, telemetry processing inside
• Server Address: Use set_server_address_and_port(instance) function signature
• Span Creation: Always use tracer.start_as_current_span() - never tracer.start_span()
• Suppression Check: Always include context_api._SUPPRESS_INSTRUMENTATION_KEY check

Key Reminders

1. Always use tracer.start_as_current_span() - Never use tracer.start_span()
2. Include suppression check - Always check context_api._SUPPRESS_INSTRUMENTATION_KEY
3. Use simplified patterns - Focus on essential telemetry attributes
4. Handle endpoint differentiation - Use endpoint parameter for operations mapping to same semantic type
5. Semantic operation mapping - Multiple endpoints can map to same DB operation
6. Response structure awareness - Handle flat vs nested response arrays correctly
7. Collection vs namespace - Use appropriate semantic convention for your database type
8. Default values - Use "None" for missing values in query summaries (double quotes)
9. Test thoroughly - Include both sync and async tests (if applicable)

Example Implementation

For complete reference implementations, see:

Qdrant (Optimized Pattern):

• Structure: 4-file structure with async support
• Lines: 80 (init.py), 55 (sync), 55 (async), 328 (utils.py)
• Features: Compact init.py with data structure, set_server_address_and_port function
• Pattern: Collection-based, endpoint differentiation, comprehensive operation coverage

Pinecone (Optimized Pattern):

• Structure: 4-file structure with async support
• Lines: 78 (init.py), 56 (sync), 56 (async), 227 (utils.py)
• Features: Compact init.py with data structure, set_server_address_and_port function
• Pattern: Namespace-based, comprehensive telemetry

ChromaDB (Simplified Pattern):

• Structure: 3-file structure (no async support)
• Lines: 90 (init.py), 57 (sync), 273 (utils.py)
• Features: Collection-based, endpoint differentiation, set_server_address_and_port function
• Pattern: Direct wrapping approach for fewer operations

Recommended Approach:

1. < 10 operations: Use ChromaDB pattern (direct wrapping)
2. 10-20 operations: Use Pinecone pattern (traditional init.py)
3. 20+ operations: Use Qdrant pattern (compact data structure)

This guide ensures consistency across all vector database instrumentations while accommodating database-specific requirements and maintaining high quality standards.