comet-ml--opik/sdks/python/tests/testlib/fake_message_factory.py

import dataclasses
import random
import uuid
from datetime import datetime, timedelta
from typing import List, Optional

from opik.message_processing import messages
from opik.types import ErrorInfoDict


@dataclasses.dataclass
class LongStr:
    value: str

    def __str__(self) -> str:
        return self.value[1] + ".." + self.value[-1]

    def __repr__(self) -> str:
        return str(self)


ONE_KILOBYTE = 1024
ONE_MEGABYTE = ONE_KILOBYTE * ONE_KILOBYTE


def fake_create_trace_message_batch(
    count: int = 1000,
    approximate_trace_size: int = ONE_MEGABYTE,
    has_ended: Optional[bool] = None,
) -> List[messages.CreateTraceMessage]:
    """
    Factory method to create a batch with a specified number of
    CreateTraceMessage objects initialized with fake data.

    Args:
        approximate_trace_size: The approximate size of each trace in megabytes
        count: Number of CreateTraceMessage objects to include in the batch (default: 1000)
        has_ended: the flag to indicate if the trace has ended. If None, the trace will
            be randomly decided to be ended or not.

    Returns:
        CreateTraceBatchMessage containing the specified number of fake CreateTraceMessage objects
    """
    dummy_traces = []

    for i in range(count):
        # Generate a unique trace ID
        trace_id = str(uuid.uuid4())

        # Create a random start time within the last 24 hours
        start_time = datetime.now() - timedelta(
            hours=random.randint(0, 23),
            minutes=random.randint(0, 59),
            seconds=random.randint(0, 59),
        )

        # Randomly decide if the trace has ended
        if has_ended is None:
            has_ended = random.choice([True, False])

        if has_ended:
            end_time = start_time + timedelta(seconds=random.randint(1, 3600))
            last_updated_at = end_time
        else:
            end_time = None
            last_updated_at = start_time

        # Generate dummy input data
        input_data = {
            "prompt": f"This is a dummy prompt #{i}",
            "parameters": {
                "temperature": round(random.uniform(0.1, 1.0), 2),
                "max_tokens": random.randint(10, 1000),
                "long_string": LongStr("a" * approximate_trace_size),
            },
        }

        # Generate dummy output data if the trace has ended
        output_data = (
            {
                "response": f"This is a dummy response for prompt #{i}",
                "tokens_used": random.randint(10, 500),
            }
            if has_ended
            else None
        )

        # Generate dummy metadata
        metadata = {
            "model": random.choice(["gpt-3.5-turbo", "gpt-4", "claude-2", "llama-2"]),
            "environment": random.choice(["production", "staging", "development"]),
            "client_id": f"client-{random.randint(1000, 9999)}",
        }

        # Generate random tags
        available_tags = [
            "important",
            "experiment",
            "production",
            "test",
            "debug",
            "high-priority",
            "low-priority",
        ]
        tags = random.sample(
            available_tags, k=random.randint(0, min(3, len(available_tags)))
        )

        # Randomly decide if there's an error
        has_error = random.random() < 0.1  # 10% chance of error
        error_info = (
            ErrorInfoDict(
                exception_type=random.choice(
                    ["TimeoutError", "ValidationError", "AuthenticationError"]
                ),
                traceback=f"Dummy stacktrace for error in trace #{i}",
            )
            if has_error
            else None
        )

        # Generate a thread ID for some traces
        thread_id = (
            str(uuid.uuid4()) if random.random() < 0.7 else None
        )  # 70% chance of having a thread ID

        # Create the trace message
        trace_message = messages.CreateTraceMessage(
            trace_id=trace_id,
            project_name="dummy-project",
            name=f"Dummy Trace #{i}",
            start_time=start_time,
            end_time=end_time,
            input=input_data,
            output=output_data,
            metadata=metadata,
            tags=tags,
            error_info=error_info,
            thread_id=thread_id,
            last_updated_at=last_updated_at,
            source="sdk",
        )

        dummy_traces.append(trace_message)

    return dummy_traces


def fake_span_create_message_batch(
    count: int = 1000,
    approximate_span_size: int = ONE_MEGABYTE,
    has_ended: Optional[bool] = None,
) -> List[messages.CreateSpanMessage]:
    """
    Factory method to create a list with a specified number of
    CreateSpanMessage objects initialized with fake data.

    Args:
        approximate_span_size: The approximate size of each span in megabytes
        count: Number of CreateSpanMessage objects to include in the batch (default: 1000)
        has_ended: the flag to indicate if the span has ended. If None, the span will
            be randomly decided to be ended or not.

    Returns:
        CreateSpansBatchMessage containing the specified number of fake CreateSpanMessage objects
    """
    dummy_spans = []

    for i in range(count):
        # Generate a unique span ID
        span_id = str(uuid.uuid4())

        # Create a random start time within the last 24 hours
        start_time = datetime.now() - timedelta(
            hours=random.randint(0, 23),
            minutes=random.randint(0, 59),
            seconds=random.randint(0, 59),
        )

        # Randomly decide if the span has ended
        if has_ended is None:
            has_ended = random.choice([True, False])

        if has_ended:
            end_time = start_time + timedelta(seconds=random.randint(1, 3600))
            last_updated_at = end_time
        else:
            end_time = None
            last_updated_at = start_time

        # Generate dummy input data
        input_data = {
            "prompt": f"This is a dummy prompt #{i}",
            "parameters": {
                "temperature": round(random.uniform(0.1, 1.0), 2),
                "max_tokens": random.randint(10, 1000),
                "long_string": LongStr("a" * approximate_span_size),
            },
        }

        # Generate dummy output data if the span has ended
        output_data = (
            {
                "response": f"This is a dummy response for prompt #{i}",
                "tokens_used": random.randint(10, 500),
            }
            if has_ended
            else None
        )

        # Generate dummy metadata
        metadata = {
            "model": random.choice(["gpt-3.5-turbo", "gpt-4", "claude-2", "llama-2"]),
            "environment": random.choice(["production", "staging", "development"]),
            "client_id": f"client-{random.randint(1000, 9999)}",
        }

        # Generate random tags
        available_tags = [
            "important",
            "experiment",
            "production",
            "test",
            "debug",
            "high-priority",
            "low-priority",
        ]
        tags = random.sample(
            available_tags, k=random.randint(0, min(3, len(available_tags)))
        )

        # Randomly decide if there's an error
        has_error = random.random() < 0.1  # 10% chance of error
        error_info = (
            ErrorInfoDict(
                exception_type=random.choice(
                    ["TimeoutError", "ValidationError", "AuthenticationError"]
                ),
                traceback=f"Dummy stacktrace for error in trace #{i}",
            )
            if has_error
            else None
        )

        # Create the span message
        span_message = messages.CreateSpanMessage(
            span_id=span_id,
            trace_id=str(uuid.uuid4()),
            parent_span_id=span_id,  # This is wrong, but it's okay for dummy data
            project_name="dummy-project",
            name=f"Dummy Span #{i}",
            start_time=start_time,
            end_time=end_time,
            input=input_data,
            output=output_data,
            metadata=metadata,
            tags=tags,
            error_info=error_info,
            type="general",
            usage=None,
            model=metadata["model"],
            provider=None,
            total_cost=random.random() * 0.01,
            last_updated_at=last_updated_at,
            source="sdk",
        )

        dummy_spans.append(span_message)

    return dummy_spans