Skip to main content

unitycatalog_client/codegen/agents/
builders.rs

1// @generated — do not edit by hand.
2#![allow(unused_mut)]
3#![allow(unused_imports)]
4#[cfg(not(target_arch = "wasm32"))]
5type BoxFut<'a, T> = ::futures::future::BoxFuture<'a, T>;
6#[cfg(target_arch = "wasm32")]
7type BoxFut<'a, T> = ::futures::future::LocalBoxFuture<'a, T>;
8#[cfg(not(target_arch = "wasm32"))]
9type BoxStr<'a, T> = ::futures::stream::BoxStream<'a, T>;
10#[cfg(target_arch = "wasm32")]
11type BoxStr<'a, T> = ::futures::stream::LocalBoxStream<'a, T>;
12use super::super::stream_paginated;
13use super::client::*;
14use crate::Result;
15use futures::{StreamExt, TryStreamExt};
16use std::future::IntoFuture;
17use unitycatalog_common::models::agents::v0alpha1::*;
18/// Builder for listing agents
19pub struct ListAgentsBuilder {
20    client: AgentServiceClient,
21    request: ListAgentsRequest,
22}
23impl ListAgentsBuilder {
24    /// Create a new builder instance.
25    /// Obtain via the corresponding method on `AgentServiceClient`.
26    pub(crate) fn new(
27        client: AgentServiceClient,
28        catalog_name: impl Into<String>,
29        schema_name: impl Into<String>,
30    ) -> Self {
31        let request = ListAgentsRequest {
32            catalog_name: catalog_name.into(),
33            schema_name: schema_name.into(),
34            ..Default::default()
35        };
36        Self { client, request }
37    }
38    /// The maximum number of results per page that should be returned.
39    pub fn with_max_results(mut self, max_results: impl Into<Option<i32>>) -> Self {
40        self.request.max_results = max_results.into();
41        self
42    }
43    /// Opaque pagination token to go to next page based on previous query.
44    pub fn with_page_token(mut self, page_token: impl Into<Option<String>>) -> Self {
45        self.request.page_token = page_token.into();
46        self
47    }
48    /** Whether to include agents in the response for which the principal can only
49    access selective metadata for.*/
50    pub fn with_include_browse(mut self, include_browse: impl Into<Option<bool>>) -> Self {
51        self.request.include_browse = include_browse.into();
52        self
53    }
54    /// Convert paginated request into stream of results
55    pub fn into_stream(self) -> BoxStr<'static, Result<Agent>> {
56        let remaining = self.request.max_results;
57        let stream = stream_paginated(
58            (self, remaining),
59            move |(mut builder, mut remaining), page_token| async move {
60                builder.request.page_token = page_token;
61                let res = builder.client.list_agents(&builder.request).await?;
62                if let Some(ref mut rem) = remaining {
63                    *rem -= res.agents.len() as i32;
64                }
65                let next_page_token = if remaining.is_some_and(|r| r <= 0) {
66                    None
67                } else {
68                    res.next_page_token.clone()
69                };
70                Ok((res, (builder, remaining), next_page_token))
71            },
72        )
73        .map_ok(|resp| futures::stream::iter(resp.agents.into_iter().map(Ok)))
74        .try_flatten();
75        #[cfg(not(target_arch = "wasm32"))]
76        let stream = stream.boxed();
77        #[cfg(target_arch = "wasm32")]
78        let stream = stream.boxed_local();
79        stream
80    }
81}
82impl IntoFuture for ListAgentsBuilder {
83    type Output = Result<ListAgentsResponse>;
84    type IntoFuture = BoxFut<'static, Self::Output>;
85    fn into_future(self) -> Self::IntoFuture {
86        let client = self.client;
87        let request = self.request;
88        Box::pin(async move { client.list_agents(&request).await })
89    }
90}
91/// Builder for creating a agent
92pub struct CreateAgentBuilder {
93    client: AgentServiceClient,
94    request: CreateAgentRequest,
95}
96impl CreateAgentBuilder {
97    /// Create a new builder instance.
98    /// Obtain via the corresponding method on `AgentServiceClient`.
99    pub(crate) fn new(
100        client: AgentServiceClient,
101        catalog_name: impl Into<String>,
102        schema_name: impl Into<String>,
103        name: impl Into<String>,
104        invocation_protocol: InvocationProtocol,
105        endpoint: impl Into<String>,
106    ) -> Self {
107        let request = CreateAgentRequest {
108            catalog_name: catalog_name.into(),
109            schema_name: schema_name.into(),
110            name: name.into(),
111            invocation_protocol: buffa::EnumValue::Known(invocation_protocol),
112            endpoint: endpoint.into(),
113            ..Default::default()
114        };
115        Self { client, request }
116    }
117    /// An LLM-readable description of what the agent does and the inputs it expects.
118    pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
119        self.request.description = description.into();
120        self
121    }
122    /// Capability identifiers advertised by the agent.
123    pub fn with_capabilities<I>(mut self, capabilities: I) -> Self
124    where
125        I: IntoIterator<Item = String>,
126    {
127        self.request.capabilities = capabilities.into_iter().collect();
128        self
129    }
130    /// A JSON Schema (encoded as a JSON string) describing the expected input.
131    pub fn with_input_schema(mut self, input_schema: impl Into<Option<String>>) -> Self {
132        self.request.input_schema = input_schema.into();
133        self
134    }
135    /// User-provided free-form text description.
136    pub fn with_comment(mut self, comment: impl Into<Option<String>>) -> Self {
137        self.request.comment = comment.into();
138        self
139    }
140}
141impl IntoFuture for CreateAgentBuilder {
142    type Output = Result<Agent>;
143    type IntoFuture = BoxFut<'static, Self::Output>;
144    fn into_future(self) -> Self::IntoFuture {
145        let client = self.client;
146        let request = self.request;
147        Box::pin(async move { client.create_agent(&request).await })
148    }
149}
150/// Builder for getting a agent
151pub struct GetAgentBuilder {
152    client: AgentServiceClient,
153    request: GetAgentRequest,
154}
155impl GetAgentBuilder {
156    /// Create a new builder instance.
157    /// Obtain via the corresponding method on `AgentServiceClient`.
158    pub(crate) fn new(client: AgentServiceClient, name: impl Into<String>) -> Self {
159        let request = GetAgentRequest {
160            name: name.into(),
161            ..Default::default()
162        };
163        Self { client, request }
164    }
165    /** Whether to include agents in the response for which the principal can only
166    access selective metadata for.*/
167    pub fn with_include_browse(mut self, include_browse: impl Into<Option<bool>>) -> Self {
168        self.request.include_browse = include_browse.into();
169        self
170    }
171}
172impl IntoFuture for GetAgentBuilder {
173    type Output = Result<Agent>;
174    type IntoFuture = BoxFut<'static, Self::Output>;
175    fn into_future(self) -> Self::IntoFuture {
176        let client = self.client;
177        let request = self.request;
178        Box::pin(async move { client.get_agent(&request).await })
179    }
180}
181/// Builder for updating a agent
182pub struct UpdateAgentBuilder {
183    client: AgentServiceClient,
184    request: UpdateAgentRequest,
185}
186impl UpdateAgentBuilder {
187    /// Create a new builder instance.
188    /// Obtain via the corresponding method on `AgentServiceClient`.
189    pub(crate) fn new(client: AgentServiceClient, name: impl Into<String>) -> Self {
190        let request = UpdateAgentRequest {
191            name: name.into(),
192            ..Default::default()
193        };
194        Self { client, request }
195    }
196    /// New name for the agent.
197    pub fn with_new_name(mut self, new_name: impl Into<Option<String>>) -> Self {
198        self.request.new_name = new_name.into();
199        self
200    }
201    /// The protocol a recipient uses to invoke the agent.
202    pub fn with_invocation_protocol(
203        mut self,
204        invocation_protocol: impl Into<Option<InvocationProtocol>>,
205    ) -> Self {
206        self.request.invocation_protocol = invocation_protocol.into().map(buffa::EnumValue::Known);
207        self
208    }
209    /// The agent's invocation endpoint URL.
210    pub fn with_endpoint(mut self, endpoint: impl Into<Option<String>>) -> Self {
211        self.request.endpoint = endpoint.into();
212        self
213    }
214    /// Updated LLM-readable description.
215    pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
216        self.request.description = description.into();
217        self
218    }
219    /// Updated capability identifiers advertised by the agent.
220    pub fn with_capabilities<I>(mut self, capabilities: I) -> Self
221    where
222        I: IntoIterator<Item = String>,
223    {
224        self.request.capabilities = capabilities.into_iter().collect();
225        self
226    }
227    /// Updated JSON Schema (encoded as a JSON string) describing the expected input.
228    pub fn with_input_schema(mut self, input_schema: impl Into<Option<String>>) -> Self {
229        self.request.input_schema = input_schema.into();
230        self
231    }
232    /// The comment attached to the agent.
233    pub fn with_comment(mut self, comment: impl Into<Option<String>>) -> Self {
234        self.request.comment = comment.into();
235        self
236    }
237    /// The identifier of the user who owns the agent.
238    pub fn with_owner(mut self, owner: impl Into<Option<String>>) -> Self {
239        self.request.owner = owner.into();
240        self
241    }
242}
243impl IntoFuture for UpdateAgentBuilder {
244    type Output = Result<Agent>;
245    type IntoFuture = BoxFut<'static, Self::Output>;
246    fn into_future(self) -> Self::IntoFuture {
247        let client = self.client;
248        let request = self.request;
249        Box::pin(async move { client.update_agent(&request).await })
250    }
251}
252/// Builder for deleting a agent
253pub struct DeleteAgentBuilder {
254    client: AgentServiceClient,
255    request: DeleteAgentRequest,
256}
257impl DeleteAgentBuilder {
258    /// Create a new builder instance.
259    /// Obtain via the corresponding method on `AgentServiceClient`.
260    pub(crate) fn new(client: AgentServiceClient, name: impl Into<String>) -> Self {
261        let request = DeleteAgentRequest {
262            name: name.into(),
263            ..Default::default()
264        };
265        Self { client, request }
266    }
267}
268impl IntoFuture for DeleteAgentBuilder {
269    type Output = Result<()>;
270    type IntoFuture = BoxFut<'static, Self::Output>;
271    fn into_future(self) -> Self::IntoFuture {
272        let client = self.client;
273        let request = self.request;
274        Box::pin(async move { client.delete_agent(&request).await })
275    }
276}