1use crate::llm::actions::{
4 client_trait::{Client, ClientActionResult},
5 protocol_trait::Protocol,
6 ActionDefinition, Parameter, ParameterDefinition,
7};
8use crate::protocol::EventType;
9use crate::state::app_state::AppState;
10use anyhow::{Context, Result};
11use serde_json::json;
12use std::sync::LazyLock;
13
14pub static DYNAMODB_CLIENT_CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
16 EventType::new(
17 "dynamodb_connected",
18 "DynamoDB client initialized and ready to send requests",
19 json!({"type": "put_item", "table_name": "Users", "item": {"id": {"S": "user456"}, "name": {"S": "Bob"}}}),
20 )
21 .with_parameters(vec![
22 Parameter {
23 name: "region".to_string(),
24 type_hint: "string".to_string(),
25 description: "AWS region for DynamoDB".to_string(),
26 required: true,
27 },
28 Parameter {
29 name: "endpoint".to_string(),
30 type_hint: "string".to_string(),
31 description: "DynamoDB endpoint URL".to_string(),
32 required: false,
33 },
34 ])
35});
36
37pub static DYNAMODB_CLIENT_RESPONSE_RECEIVED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
39 EventType::new(
40 "dynamodb_response_received",
41 "DynamoDB response received from server",
42 json!({"type": "query", "table_name": "Users", "key_condition_expression": "id = :id"}),
43 )
44 .with_parameters(vec![
45 Parameter {
46 name: "operation".to_string(),
47 type_hint: "string".to_string(),
48 description: "DynamoDB operation performed".to_string(),
49 required: true,
50 },
51 Parameter {
52 name: "success".to_string(),
53 type_hint: "boolean".to_string(),
54 description: "Whether the operation succeeded".to_string(),
55 required: true,
56 },
57 Parameter {
58 name: "data".to_string(),
59 type_hint: "object".to_string(),
60 description: "Response data from DynamoDB".to_string(),
61 required: false,
62 },
63 Parameter {
64 name: "error".to_string(),
65 type_hint: "string".to_string(),
66 description: "Error message if operation failed".to_string(),
67 required: false,
68 },
69 ])
70});
71
72pub struct DynamoDbClientProtocol;
74
75impl DynamoDbClientProtocol {
76 pub fn new() -> Self {
77 Self
78 }
79}
80
81impl Protocol for DynamoDbClientProtocol {
83 fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
84 vec![
85 ParameterDefinition {
86 name: "region".to_string(),
87 description: "AWS region (e.g., us-east-1). Defaults to us-east-1".to_string(),
88 type_hint: "string".to_string(),
89 required: false,
90 example: json!("us-east-1"),
91 default: None,
92 },
93 ParameterDefinition {
94 name: "endpoint_url".to_string(),
95 description: "Custom DynamoDB endpoint URL (for local testing with DynamoDB Local or LocalStack)".to_string(),
96 type_hint: "string".to_string(),
97 required: false,
98 example: json!("http://localhost:8000"),
99 default: None,
100 },
101 ParameterDefinition {
102 name: "access_key_id".to_string(),
103 description: "AWS access key ID (defaults to environment variable)".to_string(),
104 type_hint: "string".to_string(),
105 required: false,
106 example: json!("AKIAIOSFODNN7EXAMPLE"),
107 default: None,
108 },
109 ParameterDefinition {
110 name: "secret_access_key".to_string(),
111 description: "AWS secret access key (defaults to environment variable)".to_string(),
112 type_hint: "string".to_string(),
113 required: false,
114 example: json!("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"),
115 default: None,
116 },
117 ]
118 }
119 fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
120 vec![
121 ActionDefinition {
122 name: "put_item".to_string(),
123 description: "Put an item into a DynamoDB table".to_string(),
124 parameters: vec![
125 Parameter {
126 name: "table_name".to_string(),
127 type_hint: "string".to_string(),
128 description: "Name of the DynamoDB table".to_string(),
129 required: true,
130 },
131 Parameter {
132 name: "item".to_string(),
133 type_hint: "object".to_string(),
134 description: "Item to put (map of attribute names to values with types)"
135 .to_string(),
136 required: true,
137 },
138 ],
139 example: json!({
140 "type": "put_item",
141 "table_name": "Users",
142 "item": {
143 "id": {"S": "user123"},
144 "name": {"S": "Alice"},
145 "age": {"N": "30"}
146 }
147 }),
148 log_template: None,
149 },
150 ActionDefinition {
151 name: "get_item".to_string(),
152 description: "Get an item from a DynamoDB table by primary key".to_string(),
153 parameters: vec![
154 Parameter {
155 name: "table_name".to_string(),
156 type_hint: "string".to_string(),
157 description: "Name of the DynamoDB table".to_string(),
158 required: true,
159 },
160 Parameter {
161 name: "key".to_string(),
162 type_hint: "object".to_string(),
163 description:
164 "Primary key of the item (map of attribute names to values with types)"
165 .to_string(),
166 required: true,
167 },
168 ],
169 example: json!({
170 "type": "get_item",
171 "table_name": "Users",
172 "key": {
173 "id": {"S": "user123"}
174 }
175 }),
176 log_template: None,
177 },
178 ActionDefinition {
179 name: "query".to_string(),
180 description: "Query items from a DynamoDB table".to_string(),
181 parameters: vec![
182 Parameter {
183 name: "table_name".to_string(),
184 type_hint: "string".to_string(),
185 description: "Name of the DynamoDB table".to_string(),
186 required: true,
187 },
188 Parameter {
189 name: "key_condition_expression".to_string(),
190 type_hint: "string".to_string(),
191 description: "Query condition expression".to_string(),
192 required: true,
193 },
194 Parameter {
195 name: "expression_attribute_values".to_string(),
196 type_hint: "object".to_string(),
197 description: "Values for the expression".to_string(),
198 required: false,
199 },
200 ],
201 example: json!({
202 "type": "query",
203 "table_name": "Users",
204 "key_condition_expression": "id = :id",
205 "expression_attribute_values": {
206 ":id": {"S": "user123"}
207 }
208 }),
209 log_template: None,
210 },
211 ActionDefinition {
212 name: "scan".to_string(),
213 description: "Scan all items in a DynamoDB table".to_string(),
214 parameters: vec![
215 Parameter {
216 name: "table_name".to_string(),
217 type_hint: "string".to_string(),
218 description: "Name of the DynamoDB table".to_string(),
219 required: true,
220 },
221 Parameter {
222 name: "filter_expression".to_string(),
223 type_hint: "string".to_string(),
224 description: "Optional filter expression".to_string(),
225 required: false,
226 },
227 Parameter {
228 name: "expression_attribute_values".to_string(),
229 type_hint: "object".to_string(),
230 description: "Values for the filter expression".to_string(),
231 required: false,
232 },
233 ],
234 example: json!({
235 "type": "scan",
236 "table_name": "Users",
237 "filter_expression": "age > :min_age",
238 "expression_attribute_values": {
239 ":min_age": {"N": "21"}
240 }
241 }),
242 log_template: None,
243 },
244 ActionDefinition {
245 name: "update_item".to_string(),
246 description: "Update an item in a DynamoDB table".to_string(),
247 parameters: vec![
248 Parameter {
249 name: "table_name".to_string(),
250 type_hint: "string".to_string(),
251 description: "Name of the DynamoDB table".to_string(),
252 required: true,
253 },
254 Parameter {
255 name: "key".to_string(),
256 type_hint: "object".to_string(),
257 description: "Primary key of the item to update".to_string(),
258 required: true,
259 },
260 Parameter {
261 name: "update_expression".to_string(),
262 type_hint: "string".to_string(),
263 description: "Update expression".to_string(),
264 required: true,
265 },
266 Parameter {
267 name: "expression_attribute_values".to_string(),
268 type_hint: "object".to_string(),
269 description: "Values for the update expression".to_string(),
270 required: false,
271 },
272 ],
273 example: json!({
274 "type": "update_item",
275 "table_name": "Users",
276 "key": {
277 "id": {"S": "user123"}
278 },
279 "update_expression": "SET age = :age",
280 "expression_attribute_values": {
281 ":age": {"N": "31"}
282 }
283 }),
284 log_template: None,
285 },
286 ActionDefinition {
287 name: "delete_item".to_string(),
288 description: "Delete an item from a DynamoDB table".to_string(),
289 parameters: vec![
290 Parameter {
291 name: "table_name".to_string(),
292 type_hint: "string".to_string(),
293 description: "Name of the DynamoDB table".to_string(),
294 required: true,
295 },
296 Parameter {
297 name: "key".to_string(),
298 type_hint: "object".to_string(),
299 description: "Primary key of the item to delete".to_string(),
300 required: true,
301 },
302 ],
303 example: json!({
304 "type": "delete_item",
305 "table_name": "Users",
306 "key": {
307 "id": {"S": "user123"}
308 }
309 }),
310 log_template: None,
311 },
312 ActionDefinition {
313 name: "disconnect".to_string(),
314 description: "Disconnect from DynamoDB".to_string(),
315 parameters: vec![],
316 example: json!({
317 "type": "disconnect"
318 }),
319 log_template: None,
320 },
321 ]
322 }
323 fn get_sync_actions(&self) -> Vec<ActionDefinition> {
324 vec![
325 ActionDefinition {
326 name: "put_item".to_string(),
327 description: "Put another item in response to received data".to_string(),
328 parameters: vec![
329 Parameter {
330 name: "table_name".to_string(),
331 type_hint: "string".to_string(),
332 description: "Name of the DynamoDB table".to_string(),
333 required: true,
334 },
335 Parameter {
336 name: "item".to_string(),
337 type_hint: "object".to_string(),
338 description: "Item to put".to_string(),
339 required: true,
340 },
341 ],
342 example: json!({
343 "type": "put_item",
344 "table_name": "Users",
345 "item": {
346 "id": {"S": "user456"},
347 "name": {"S": "Bob"}
348 }
349 }),
350 log_template: None,
351 },
352 ActionDefinition {
353 name: "query".to_string(),
354 description: "Query items in response to received data".to_string(),
355 parameters: vec![
356 Parameter {
357 name: "table_name".to_string(),
358 type_hint: "string".to_string(),
359 description: "Name of the DynamoDB table".to_string(),
360 required: true,
361 },
362 Parameter {
363 name: "key_condition_expression".to_string(),
364 type_hint: "string".to_string(),
365 description: "Query condition expression".to_string(),
366 required: true,
367 },
368 ],
369 example: json!({
370 "type": "query",
371 "table_name": "Users",
372 "key_condition_expression": "id = :id"
373 }),
374 log_template: None,
375 },
376 ]
377 }
378 fn protocol_name(&self) -> &'static str {
379 "DynamoDB"
380 }
381 fn get_event_types(&self) -> Vec<EventType> {
382 vec![
383 EventType::new(
384 "dynamodb_connected",
385 "Triggered when DynamoDB client is initialized",
386 json!({"type": "put_item", "table_name": "Users", "item": {"id": {"S": "user456"}, "name": {"S": "Bob"}}}),
387 ),
388 EventType::new(
389 "dynamodb_response_received",
390 "Triggered when DynamoDB client receives a response",
391 json!({"type": "query", "table_name": "Users", "key_condition_expression": "id = :id"}),
392 ),
393 ]
394 }
395 fn stack_name(&self) -> &'static str {
396 "ETH>IP>TCP>TLS>HTTP>DynamoDB"
397 }
398 fn keywords(&self) -> Vec<&'static str> {
399 vec!["dynamodb", "dynamo", "aws dynamodb", "connect to dynamodb"]
400 }
401 fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
402 use crate::protocol::metadata::{DevelopmentState, ProtocolMetadataV2};
403
404 ProtocolMetadataV2::builder()
405 .state(DevelopmentState::Experimental)
406 .implementation("AWS SDK for DynamoDB client library")
407 .llm_control("Full control over DynamoDB operations (PutItem, GetItem, Query, Scan, UpdateItem, DeleteItem)")
408 .e2e_testing("DynamoDB Local or LocalStack")
409 .build()
410 }
411 fn description(&self) -> &'static str {
412 "DynamoDB client for interacting with AWS DynamoDB or local instances"
413 }
414 fn example_prompt(&self) -> &'static str {
415 "Connect to DynamoDB and put an item in the Users table"
416 }
417 fn group_name(&self) -> &'static str {
418 "Database"
419 }
420
421 fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
422 use crate::llm::actions::StartupExamples;
423 use serde_json::json;
424
425 StartupExamples::new(
428 json!({
430 "type": "open_client",
431 "remote_addr": "localhost:8000",
432 "base_stack": "dynamodb",
433 "startup_params": {
434 "access_key_id": "local",
435 "secret_access_key": "local"
436 },
437 "instruction": "Scan the Users table and report all items with age greater than 21"
438 }),
439 json!({
441 "type": "open_client",
442 "remote_addr": "localhost:8000",
443 "base_stack": "dynamodb",
444 "startup_params": {
445 "access_key_id": "local",
446 "secret_access_key": "local"
447 },
448 "event_handlers": [{
449 "event_pattern": "dynamodb_response_received",
450 "handler": {
451 "type": "script",
452 "language": "python",
453 "code": "<dynamodb_client_handler>"
454 }
455 }]
456 }),
457 json!({
459 "type": "open_client",
460 "remote_addr": "localhost:8000",
461 "base_stack": "dynamodb",
462 "startup_params": {
463 "access_key_id": "local",
464 "secret_access_key": "local"
465 },
466 "event_handlers": [
467 {
468 "event_pattern": "dynamodb_connected",
469 "handler": {
470 "type": "static",
471 "actions": [{
472 "type": "scan",
473 "table_name": "Users"
474 }]
475 }
476 },
477 {
478 "event_pattern": "dynamodb_response_received",
479 "handler": {
480 "type": "static",
481 "actions": [{
482 "type": "disconnect"
483 }]
484 }
485 }
486 ]
487 }),
488 )
489 }
490}
491
492impl Client for DynamoDbClientProtocol {
494 fn connect(
495 &self,
496 ctx: crate::protocol::ConnectContext,
497 ) -> std::pin::Pin<
498 Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
499 > {
500 Box::pin(async move {
501 use crate::client::dynamodb::DynamoDbClient;
502 DynamoDbClient::connect_with_llm_actions(
503 ctx.remote_addr,
504 ctx.llm_client,
505 ctx.state,
506 ctx.status_tx,
507 ctx.client_id,
508 ctx.startup_params,
509 )
510 .await
511 })
512 }
513 fn execute_action(&self, action: serde_json::Value) -> Result<ClientActionResult> {
514 let action_type = action
515 .get("type")
516 .and_then(|v| v.as_str())
517 .context("Missing 'type' field in action")?;
518
519 match action_type {
520 "put_item" => {
521 let table_name = action
522 .get("table_name")
523 .and_then(|v| v.as_str())
524 .context("Missing 'table_name' field")?
525 .to_string();
526
527 let item = action
528 .get("item")
529 .and_then(|v| v.as_object())
530 .context("Missing 'item' field")?
531 .clone();
532
533 Ok(ClientActionResult::Custom {
534 name: "put_item".to_string(),
535 data: json!({
536 "table_name": table_name,
537 "item": item,
538 }),
539 })
540 }
541 "get_item" => {
542 let table_name = action
543 .get("table_name")
544 .and_then(|v| v.as_str())
545 .context("Missing 'table_name' field")?
546 .to_string();
547
548 let key = action
549 .get("key")
550 .and_then(|v| v.as_object())
551 .context("Missing 'key' field")?
552 .clone();
553
554 Ok(ClientActionResult::Custom {
555 name: "get_item".to_string(),
556 data: json!({
557 "table_name": table_name,
558 "key": key,
559 }),
560 })
561 }
562 "query" => {
563 let table_name = action
564 .get("table_name")
565 .and_then(|v| v.as_str())
566 .context("Missing 'table_name' field")?
567 .to_string();
568
569 let key_condition_expression = action
570 .get("key_condition_expression")
571 .and_then(|v| v.as_str())
572 .context("Missing 'key_condition_expression' field")?
573 .to_string();
574
575 let expression_attribute_values = action
576 .get("expression_attribute_values")
577 .and_then(|v| v.as_object())
578 .cloned();
579
580 Ok(ClientActionResult::Custom {
581 name: "query".to_string(),
582 data: json!({
583 "table_name": table_name,
584 "key_condition_expression": key_condition_expression,
585 "expression_attribute_values": expression_attribute_values,
586 }),
587 })
588 }
589 "scan" => {
590 let table_name = action
591 .get("table_name")
592 .and_then(|v| v.as_str())
593 .context("Missing 'table_name' field")?
594 .to_string();
595
596 let filter_expression = action
597 .get("filter_expression")
598 .and_then(|v| v.as_str())
599 .map(|s| s.to_string());
600
601 let expression_attribute_values = action
602 .get("expression_attribute_values")
603 .and_then(|v| v.as_object())
604 .cloned();
605
606 Ok(ClientActionResult::Custom {
607 name: "scan".to_string(),
608 data: json!({
609 "table_name": table_name,
610 "filter_expression": filter_expression,
611 "expression_attribute_values": expression_attribute_values,
612 }),
613 })
614 }
615 "update_item" => {
616 let table_name = action
617 .get("table_name")
618 .and_then(|v| v.as_str())
619 .context("Missing 'table_name' field")?
620 .to_string();
621
622 let key = action
623 .get("key")
624 .and_then(|v| v.as_object())
625 .context("Missing 'key' field")?
626 .clone();
627
628 let update_expression = action
629 .get("update_expression")
630 .and_then(|v| v.as_str())
631 .context("Missing 'update_expression' field")?
632 .to_string();
633
634 let expression_attribute_values = action
635 .get("expression_attribute_values")
636 .and_then(|v| v.as_object())
637 .cloned();
638
639 Ok(ClientActionResult::Custom {
640 name: "update_item".to_string(),
641 data: json!({
642 "table_name": table_name,
643 "key": key,
644 "update_expression": update_expression,
645 "expression_attribute_values": expression_attribute_values,
646 }),
647 })
648 }
649 "delete_item" => {
650 let table_name = action
651 .get("table_name")
652 .and_then(|v| v.as_str())
653 .context("Missing 'table_name' field")?
654 .to_string();
655
656 let key = action
657 .get("key")
658 .and_then(|v| v.as_object())
659 .context("Missing 'key' field")?
660 .clone();
661
662 Ok(ClientActionResult::Custom {
663 name: "delete_item".to_string(),
664 data: json!({
665 "table_name": table_name,
666 "key": key,
667 }),
668 })
669 }
670 "disconnect" => Ok(ClientActionResult::Disconnect),
671 _ => Err(anyhow::anyhow!(
672 "Unknown DynamoDB client action: {}",
673 action_type
674 )),
675 }
676 }
677}