unitycatalog_client/codegen/agents/
builders.rs1#![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::*;
12pub struct ListAgentsBuilder {
14 client: AgentServiceClient,
15 request: ListAgentsRequest,
16}
17impl ListAgentsBuilder {
18 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 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 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 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 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}
81pub struct CreateAgentBuilder {
83 client: AgentServiceClient,
84 request: CreateAgentRequest,
85}
86impl CreateAgentBuilder {
87 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 pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
109 self.request.description = description.into();
110 self
111 }
112 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 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 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}
140pub struct GetAgentBuilder {
142 client: AgentServiceClient,
143 request: GetAgentRequest,
144}
145impl GetAgentBuilder {
146 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 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}
171pub struct UpdateAgentBuilder {
173 client: AgentServiceClient,
174 request: UpdateAgentRequest,
175}
176impl UpdateAgentBuilder {
177 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 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 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 pub fn with_endpoint(mut self, endpoint: impl Into<Option<String>>) -> Self {
201 self.request.endpoint = endpoint.into();
202 self
203 }
204 pub fn with_description(mut self, description: impl Into<Option<String>>) -> Self {
206 self.request.description = description.into();
207 self
208 }
209 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 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 pub fn with_comment(mut self, comment: impl Into<Option<String>>) -> Self {
224 self.request.comment = comment.into();
225 self
226 }
227 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}
242pub struct DeleteAgentBuilder {
244 client: AgentServiceClient,
245 request: DeleteAgentRequest,
246}
247impl DeleteAgentBuilder {
248 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}