openlayer-python/src/openlayer/lib/__init__.py at 6742ac914b7316924d98771f437947be31e53c40 · openlayer-ai/openlayer-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
"""Openlayer lib."""

__all__ = [
    "trace",
    "trace_anthropic",
    "trace_openai",
    "trace_openai_assistant_thread_run",
    "trace_mistral",
    "trace_groq",
]

# ---------------------------------- Tracing --------------------------------- #
from .tracing import tracer

trace = tracer.trace


def trace_anthropic(client):
    """Trace Anthropic chat completions."""
    # pylint: disable=import-outside-toplevel
    import anthropic

    from .integrations import anthropic_tracer

    if not isinstance(client, anthropic.Anthropic):
        raise ValueError("Invalid client. Please provide an Anthropic client.")
    return anthropic_tracer.trace_anthropic(client)


def trace_openai(client):
    """Trace OpenAI chat completions."""
    # pylint: disable=import-outside-toplevel
    import openai

    from .integrations import openai_tracer

    if not isinstance(client, (openai.Client, openai.AzureOpenAI)):
        raise ValueError("Invalid client. Please provide an OpenAI client.")
    return openai_tracer.trace_openai(client)


def trace_async_openai(client):
    """Trace OpenAI chat completions."""
    # pylint: disable=import-outside-toplevel
    import openai

    from .integrations import async_openai_tracer

    if not isinstance(client, (openai.AsyncOpenAI, openai.AsyncAzureOpenAI)):
        raise ValueError("Invalid client. Please provide an OpenAI client.")
    return async_openai_tracer.trace_async_openai(client)


def trace_openai_assistant_thread_run(client, run):
    """Trace OpenAI Assistant thread run."""
    # pylint: disable=import-outside-toplevel
    from .integrations import openai_tracer

    return openai_tracer.trace_openai_assistant_thread_run(client, run)


def trace_mistral(client):
    """Trace Mistral chat completions."""
    # pylint: disable=import-outside-toplevel
    import mistralai

    from .integrations import mistral_tracer

    if not isinstance(client, mistralai.Mistral):
        raise ValueError("Invalid client. Please provide a Mistral client.")
    return mistral_tracer.trace_mistral(client)


def trace_groq(client):
    """Trace Groq queries."""
    # pylint: disable=import-outside-toplevel
    import groq

    from .integrations import groq_tracer

    if not isinstance(client, groq.Groq):
        raise ValueError("Invalid client. Please provide a Groq client.")
    return groq_tracer.trace_groq(client)