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