galileo-agents/shared/otel.py at main · rungalileo/galileo-agents · 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
84
85
86
87
88
89
90
91
import argparse
import glob

import requests
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
    ExportTraceServiceRequest,
)


def parse_trace(body_bytes):
    reqtrace = ExportTraceServiceRequest()
    reqtrace.ParseFromString(body_bytes)
    return reqtrace


def main():
    parser = argparse.ArgumentParser(description="Send OTLP traces to Galileo API")
    parser.add_argument(
        "--api-key",
        required=True,
        help="Galileo API key",
    )
    parser.add_argument(
        "--project",
        help="Project name (at least one of --project or --projectid required)",
    )
    parser.add_argument(
        "--projectid",
        help="Project ID (at least one of --project or --projectid required)",
    )
    parser.add_argument(
        "--logstream",
        help="Log stream name (at least one of --logstream or --logstreamid required)",
    )
    parser.add_argument(
        "--logstreamid",
        help="Log stream ID (at least one of --logstream or --logstreamid required)",
    )
    parser.add_argument(
        "--url",
        default="https://api.galileo.ai/otel/v1/traces",
        help="API endpoint URL (default: https://api.galileo.ai/otel/v1/traces)",
    )
    parser.add_argument(
        "--directory",
        default="agents-langgraph/weather/otlp_trace",
        help="Directory containing .bin trace files (default: agents-langgraph/weather/otlp_trace)",
    )

    args = parser.parse_args()

    # Validate that at least one of project or projectid is provided
    if not args.project and not args.projectid:
        parser.error("At least one of --project or --projectid must be provided")

    # Validate that at least one of logstream or logstreamid is provided
    if not args.logstream and not args.logstreamid:
        parser.error("At least one of --logstream or --logstreamid must be provided")

    # Build headers
    headers = {
        "Galileo-API-Key": args.api_key,
        "Content-Type": "application/x-protobuf",
    }

    # Add project or projectid to headers (prefer project over projectid if both provided)
    if args.project:
        headers["project"] = args.project
    elif args.projectid:
        headers["projectid"] = args.projectid

    # Add logstream or logstreamid to headers (prefer logstream over logstreamid if both provided)
    if args.logstream:
        headers["logstream"] = args.logstream
    elif args.logstreamid:
        headers["logstreamid"] = args.logstreamid

    glob_files = glob.glob(f"{args.directory}/*.bin")
    glob_files.sort()
    for file in glob_files:
        print(f"Processing file: {file}")
        with open(file, "rb") as f:
            body_bytes = f.read()
            reqtrace = parse_trace(body_bytes)
        response = requests.post(args.url, headers=headers, data=reqtrace.SerializeToString())
        print(f"Status: {response.status_code}")
        print(f"Response: {response.text}")


if __name__ == "__main__":
    main()