agent-infra-sdk 0.2.0

Unified Rust SDK for Gateway-backed and local Agent Infra APIs
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

use crate::transport::{
    CallOptions, ClientOptions, HttpTransport, InfraClientError, ServiceEndpoint,
};
use agent_trace_contract::{
    AppendEventsRequest, AppendResponse, EVENTS_APPEND_PATH, EVENTS_SEARCH_PATH, EventsResponse,
    RuntimeEventRecord, SearchEventsRequest, TraceScope, redact_payload,
};
use reqwest::Client;

/// Client for the Trace service (`agent-trace-server`).
#[derive(Clone, Debug)]
pub struct TraceClient {
    transport: HttpTransport,
}

impl TraceClient {
    pub(crate) fn new_with_endpoint(
        http: Client,
        endpoint: ServiceEndpoint,
        options: ClientOptions,
    ) -> Self {
        Self {
            transport: HttpTransport::new_with_options(
                http,
                agent_trace_contract::SERVICE_NAME,
                endpoint,
                options,
            ),
        }
    }

    /// Append trace events to a run. Idempotent: duplicate IDs are skipped.
    pub async fn append_events(
        &self,
        events: Vec<RuntimeEventRecord>,
    ) -> Result<AppendResponse, InfraClientError> {
        self.append_events_with_options(events, CallOptions::default())
            .await
    }

    pub async fn append_events_with_options(
        &self,
        mut events: Vec<RuntimeEventRecord>,
        options: CallOptions,
    ) -> Result<AppendResponse, InfraClientError> {
        for event in &mut events {
            redact_payload(&mut event.payload);
            redact_payload(&mut event.metadata);
        }
        self.transport
            .post_json_with_options(
                EVENTS_APPEND_PATH,
                &AppendEventsRequest { events },
                options.idempotent(true),
            )
            .await
    }

    /// Search trace events in a run.
    pub async fn search_events(
        &self,
        conversation_id: &str,
        run_id: &str,
        event_type: Option<&str>,
        limit: Option<usize>,
    ) -> Result<Vec<RuntimeEventRecord>, InfraClientError> {
        let request = SearchEventsRequest {
            scope: TraceScope::run(conversation_id, run_id),
            event_type: event_type.map(ToOwned::to_owned),
            limit,
        };
        let response: EventsResponse = self
            .transport
            .post_json(EVENTS_SEARCH_PATH, &request)
            .await?;
        Ok(response.items)
    }
}