unitycatalog_client/codegen/agents/
builders.rs1#![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::*;
18pub struct ListAgentsBuilder {
20 client: AgentServiceClient,
21 request: ListAgentsRequest,
22}
23impl ListAgentsBuilder {
24 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 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 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 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 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}
91pub struct CreateAgentBuilder {
93 client: AgentServiceClient,
94 request: CreateAgentRequest,
95}
96impl CreateAgentBuilder {
97 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 pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
119 self.request.description = description.into();
120 self
121 }
122 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 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 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}
150pub struct GetAgentBuilder {
152 client: AgentServiceClient,
153 request: GetAgentRequest,
154}
155impl GetAgentBuilder {
156 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 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}
181pub struct UpdateAgentBuilder {
183 client: AgentServiceClient,
184 request: UpdateAgentRequest,
185}
186impl UpdateAgentBuilder {
187 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 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 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 pub fn with_endpoint(mut self, endpoint: impl Into<Option<String>>) -> Self {
211 self.request.endpoint = endpoint.into();
212 self
213 }
214 pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
216 self.request.description = description.into();
217 self
218 }
219 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 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 pub fn with_comment(mut self, comment: impl Into<Option<String>>) -> Self {
234 self.request.comment = comment.into();
235 self
236 }
237 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}
252pub struct DeleteAgentBuilder {
254 client: AgentServiceClient,
255 request: DeleteAgentRequest,
256}
257impl DeleteAgentBuilder {
258 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}