erdo 0.1.4__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of erdo might be problematic. Click here for more details.

erdo/types.py ADDED
@@ -0,0 +1,1142 @@
1
+ """
2
+ Erdo Agent SDK - Core Types
3
+
4
+ This file contains the main SDK classes for building AI agents.
5
+ Auto-generated types are imported from the generated module.
6
+ """
7
+
8
+ import os
9
+ from enum import Enum
10
+ from pathlib import Path
11
+ from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple, Union, cast
12
+
13
+ from pydantic import BaseModel, Field
14
+
15
+ from ._generated.types import (
16
+ ExecutionModeType,
17
+ HandlerType,
18
+ OutputContentType,
19
+ OutputVisibility,
20
+ ParameterDefinition,
21
+ ParameterHydrationBehaviour,
22
+ Tool,
23
+ )
24
+ from .template import TemplateString
25
+
26
+ # Type aliases for complex types that are json.RawMessage in Go
27
+ ExecutionCondition = Dict[str, Any] # Complex execution condition configuration
28
+
29
+
30
+ # Protocols for better type safety
31
+ class ActionProtocol(Protocol):
32
+ """Protocol for action objects that can be converted to parameters."""
33
+
34
+ name: str
35
+
36
+ def model_dump(self) -> Dict[str, Any]: ...
37
+
38
+
39
+ class StepLike(Protocol):
40
+ """Protocol for step-like objects."""
41
+
42
+ key: Optional[str]
43
+ id: Optional[str]
44
+
45
+
46
+ class ConditionProtocol(Protocol):
47
+ """Protocol for condition objects."""
48
+
49
+ def to_dict(self) -> Dict[str, Any]: ...
50
+
51
+
52
+ class ExecutionModeProtocol(Protocol):
53
+ """Protocol for execution mode objects."""
54
+
55
+ def to_dict(self) -> Dict[str, Any]: ...
56
+
57
+
58
+ class PythonFile(BaseModel):
59
+ """Reference to a Python file for code execution.
60
+
61
+ Can be used in code_files parameter to reference local Python files.
62
+ When used, the file content will be automatically loaded during export.
63
+ """
64
+
65
+ filename: str = Field(
66
+ ..., description="Path to the file (e.g., 'analyze_file_files/analyze.py')"
67
+ )
68
+
69
+ def to_dict(self) -> Dict[str, Any]:
70
+ """Convert to dict format for serialization."""
71
+ return {
72
+ "filename": self.filename,
73
+ "_type": "PythonFile", # Marker to identify this as a file reference
74
+ }
75
+
76
+ def resolve_content(self, base_path: Optional[str] = None) -> Dict[str, str]:
77
+ """Resolve the file content for inclusion in code_files.
78
+
79
+ Args:
80
+ base_path: Base directory path to resolve relative paths
81
+
82
+ Returns:
83
+ Dict with filename and content
84
+ """
85
+ if base_path is None:
86
+ base_path = os.getcwd()
87
+
88
+ file_path = Path(base_path) / self.filename
89
+
90
+ try:
91
+ with open(file_path, "r", encoding="utf-8") as f:
92
+ content = f.read()
93
+ # Return just the base filename, not the full path
94
+ base_filename = Path(self.filename).name
95
+ return {"filename": base_filename, "content": content}
96
+ except FileNotFoundError:
97
+ raise FileNotFoundError(f"Python file not found: {file_path}")
98
+ except Exception as e:
99
+ raise RuntimeError(f"Error reading Python file {file_path}: {e}")
100
+
101
+
102
+ class StepMetadata(BaseModel):
103
+ """Metadata for workflow steps, containing configuration and execution parameters."""
104
+
105
+ model_config = {"arbitrary_types_allowed": True}
106
+
107
+ key: Optional[str] = None
108
+ depends_on: Union[List[Union[Any, str]], None] = (
109
+ None # Can be Step objects or strings
110
+ )
111
+ execution_mode: Optional[Union[Any, Dict[str, Any]]] = (
112
+ None # Can be ExecutionMode object or dict
113
+ )
114
+ output_behavior: Optional[Dict[str, Any]] = None
115
+ output_channels: List[str] = Field(default_factory=list)
116
+ output_content_type: OutputContentType = OutputContentType.TEXT
117
+ history_content_type: Optional[str] = None
118
+ ui_content_type: Optional[str] = None
119
+ user_output_visibility: OutputVisibility = OutputVisibility.VISIBLE
120
+ bot_output_visibility: OutputVisibility = OutputVisibility.HIDDEN
121
+ running_message: Optional[str] = None
122
+ finished_message: Optional[str] = None
123
+ parameter_hydration_behaviour: ParameterHydrationBehaviour = (
124
+ ParameterHydrationBehaviour.NONE
125
+ )
126
+
127
+
128
+ class Step(StepMetadata):
129
+ """A single step in an agent workflow."""
130
+
131
+ agent: Optional[Any] = None # Reference to the agent this step belongs to
132
+ result_handler: Optional[Any] = (
133
+ None # Reference to the result handler this step belongs to
134
+ )
135
+ action: Union[Any, Dict[str, Any]] # Function call like codeexec.execute(...)
136
+ parameters: Dict[str, Any] = Field(default_factory=dict)
137
+ result_handlers: List["ResultHandler"] = Field(default_factory=list)
138
+
139
+ def __init__(self, action: Any = None, step: Optional["Step"] = None, **data: Any):
140
+ # Validate that only one of action or step is provided
141
+ if action is not None and step is not None:
142
+ raise ValueError("Cannot specify both 'action' and 'step' parameters")
143
+
144
+ if action is None and step is None:
145
+ raise ValueError("Must specify either 'action' or 'step' parameter")
146
+
147
+ if step is not None:
148
+ # Copy all fields from the provided step
149
+ # Type check is redundant due to type annotation, but kept for runtime safety
150
+ if not hasattr(step, "action"):
151
+ raise ValueError("'step' parameter must be a Step object")
152
+
153
+ # Copy fields directly to preserve object types (especially action)
154
+ # Don't use model_dump() as it serializes objects to dicts
155
+ data.update(
156
+ {
157
+ "action": step.action,
158
+ "key": step.key,
159
+ "parameters": step.parameters,
160
+ "depends_on": step.depends_on,
161
+ "execution_mode": step.execution_mode,
162
+ "output_behavior": step.output_behavior,
163
+ "result_handlers": step.result_handlers,
164
+ "output_channels": step.output_channels,
165
+ "output_content_type": step.output_content_type,
166
+ "history_content_type": step.history_content_type,
167
+ "ui_content_type": step.ui_content_type,
168
+ "user_output_visibility": step.user_output_visibility,
169
+ "bot_output_visibility": step.bot_output_visibility,
170
+ "running_message": step.running_message,
171
+ "finished_message": step.finished_message,
172
+ "parameter_hydration_behaviour": step.parameter_hydration_behaviour,
173
+ }
174
+ )
175
+ else:
176
+ # Standard action-based step
177
+ data["action"] = action
178
+
179
+ super().__init__(**data)
180
+ # For enhanced syntax support - use private field not tracked by pydantic
181
+ self._decorator_handlers: List[Tuple[Any, Callable[..., Any]]] = []
182
+
183
+ # Automatically register with agent if provided and this is not a nested step
184
+ if self.agent is not None and self.result_handler is None:
185
+ self.agent.add_step(self)
186
+
187
+ def extract_action_parameters(self) -> Dict[str, Any]:
188
+ """Extract parameters from the action object safely."""
189
+ if not self.action:
190
+ return {}
191
+
192
+ # Handle nested Step objects (for result handlers)
193
+ if isinstance(self.action, Step):
194
+ return self.action.extract_action_parameters()
195
+
196
+ # If action is a Pydantic model, get its dict representation
197
+ if hasattr(self.action, "model_dump") and callable(
198
+ getattr(self.action, "model_dump")
199
+ ):
200
+ try:
201
+ action_obj = cast(ActionProtocol, self.action)
202
+ params = action_obj.model_dump()
203
+ # Remove the redundant 'name' field since it's already known from action type
204
+ params.pop("name", None)
205
+ # Map Python SDK field names to backend field names
206
+ if "json_data" in params:
207
+ params["json"] = params.pop("json_data")
208
+ # Remove None values for optional parameters (matches action function behavior)
209
+ params = {k: v for k, v in params.items() if v is not None}
210
+ # Sort parameters for deterministic output
211
+ return dict(sorted(params.items()))
212
+ except Exception:
213
+ return {}
214
+ elif hasattr(self.action, "dict") and callable(getattr(self.action, "dict")):
215
+ try:
216
+ # Legacy pydantic v1 support
217
+ action_dict_method = getattr(self.action, "dict")
218
+ params = action_dict_method()
219
+ # Remove the redundant 'name' field since it's already known from action type
220
+ params.pop("name", None)
221
+ # Map Python SDK field names to backend field names
222
+ if "json_data" in params:
223
+ params["json"] = params.pop("json_data")
224
+ # Remove None values for optional parameters (matches action function behavior)
225
+ params = {k: v for k, v in params.items() if v is not None}
226
+ # Sort parameters for deterministic output
227
+ return dict(sorted(params.items()))
228
+ except Exception:
229
+ return {}
230
+ elif isinstance(self.action, dict):
231
+ # Map Python SDK field names to backend field names
232
+ dict_params: Dict[str, Any] = dict(self.action)
233
+ # Remove the redundant 'name' field since it's already known from action type
234
+ dict_params.pop("name", None)
235
+ if "json_data" in dict_params:
236
+ dict_params["json"] = dict_params.pop("json_data")
237
+ # Sort parameters for deterministic output
238
+ return dict(sorted(dict_params.items()))
239
+ else:
240
+ return {}
241
+
242
+ def get_action_type(self) -> str:
243
+ """Get the action type string safely."""
244
+ if not self.action:
245
+ raise ValueError("Action is not set")
246
+
247
+ # Handle nested Step objects (for result handlers)
248
+ if isinstance(self.action, Step):
249
+ return self.action.get_action_type()
250
+
251
+ # Action objects should have a name attribute with the action type
252
+ if hasattr(self.action, "name"):
253
+ name_attr = getattr(self.action, "name")
254
+ if isinstance(name_attr, str):
255
+ return name_attr
256
+ elif hasattr(name_attr, "__str__"):
257
+ return str(name_attr)
258
+
259
+ # Fallback for unexpected cases
260
+ return str(type(self.action).__name__).lower()
261
+
262
+ def get_depends_on_keys(self) -> Optional[List[str]]:
263
+ """Get dependency keys safely without circular references, preserving None vs [] distinction."""
264
+ # CRITICAL: Preserve None vs [] distinction for 100% roundtrip parity
265
+ if self.depends_on is None:
266
+ return None # Explicitly return None for null dependencies
267
+
268
+ if not self.depends_on: # Empty list case
269
+ return []
270
+
271
+ # Depends_on is guaranteed to be a non-empty list at this point
272
+ depends_list = self.depends_on
273
+
274
+ result: List[str] = []
275
+ for dep in depends_list:
276
+ if isinstance(dep, str):
277
+ # Already a string identifier
278
+ result.append(dep)
279
+ elif hasattr(dep, "key") and getattr(dep, "key", None):
280
+ # Prefer the step key if available
281
+ key_val = getattr(dep, "key")
282
+ result.append(str(key_val) if key_val is not None else "")
283
+ elif hasattr(dep, "id") and getattr(dep, "id", None):
284
+ # Fall back to step ID if no key
285
+ id_val = getattr(dep, "id")
286
+ result.append(str(id_val) if id_val is not None else "")
287
+ else:
288
+ # Fail if we can't get a valid identifier
289
+ raise ValueError(f"Step dependency has no key or id: {dep}")
290
+ return result
291
+
292
+ @property
293
+ def output(self) -> "StepOutput":
294
+ """Get typed output reference for this step."""
295
+ return StepOutput(step_key=self.key or f"step_{id(self)}")
296
+
297
+ def when(
298
+ self, condition: Any
299
+ ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
300
+ """Decorator for adding result handlers with conditions.
301
+
302
+ Usage:
303
+ @step.when(IsSuccess() & GreaterThan("confidence", 0.8))
304
+ def handle_high_confidence(result):
305
+ return store_analysis(result)
306
+ """
307
+
308
+ def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
309
+ # Store the handler function and condition for later processing
310
+ self._decorator_handlers.append((condition, func))
311
+ return func
312
+
313
+ return decorator
314
+
315
+ def on(
316
+ self,
317
+ condition: Any,
318
+ *actions: Any,
319
+ handler_type: HandlerType = HandlerType.INTERMEDIATE,
320
+ ) -> "Step":
321
+ """Add a result handler with condition and action(s).
322
+
323
+ Usage:
324
+ # Single action
325
+ step.on(IsSuccess(), utils.store_analysis(data=state.analyze_step.output))
326
+
327
+ # Multiple actions (variadic arguments)
328
+ step.on(And(IsError(), LessThan(number='r"{{coalesce "code_retry_loops?" 0}}"', value="2")),
329
+ send_status(status="retrying", message="Code execution failed, attempting to fix..."),
330
+ utils.echo(data={"code_retry_loops": 'r"{{incrementCounter "code_retry_loops"}}'}),
331
+ raise_action(status="go to step", message="code", parameters={...})
332
+ )
333
+ """
334
+ if not actions:
335
+ raise ValueError("Must specify at least one action")
336
+
337
+ # Create steps for each action, handling Step objects correctly
338
+ step_actions: List[Step] = []
339
+ for action in actions:
340
+ if isinstance(action, Step):
341
+ # If it's already a Step, use step= parameter
342
+ step_actions.append(Step(step=action))
343
+ else:
344
+ # If it's an action function, use action= parameter
345
+ step_actions.append(Step(action=action))
346
+
347
+ # Create a result handler and add it to this step
348
+ handler = ResultHandler(
349
+ type=handler_type,
350
+ if_conditions=condition,
351
+ steps=step_actions,
352
+ )
353
+ self.result_handlers.append(handler)
354
+ return self
355
+
356
+ def to_dict(self) -> Dict[str, Any]:
357
+ """Convert to dict format expected by backend."""
358
+ # Process decorator handlers before serialization
359
+ self._process_decorator_handlers()
360
+
361
+ # Use model_dump but exclude problematic circular reference fields
362
+ result = self.model_dump(
363
+ exclude={
364
+ "agent",
365
+ "result_handler",
366
+ "_decorator_handlers",
367
+ "action",
368
+ "depends_on",
369
+ },
370
+ )
371
+
372
+ # Filter out empty/None optional fields to match export behavior
373
+ if result.get("ui_content_type") in [None, ""]:
374
+ result.pop("ui_content_type", None)
375
+ if result.get("history_content_type") in [None, ""]:
376
+ result.pop("history_content_type", None)
377
+
378
+ # Add action information without circular references
379
+ result["action_type"] = self.get_action_type()
380
+ result["parameters"] = self.extract_action_parameters()
381
+
382
+ # Add depends_on only if there are actual dependencies
383
+ if self.depends_on:
384
+ # Convert Step objects to their keys
385
+ deps = []
386
+ for dep in self.depends_on:
387
+ if isinstance(dep, Step):
388
+ # If it's a Step object, use its key
389
+ if dep.key:
390
+ deps.append(dep.key)
391
+ elif isinstance(dep, str):
392
+ # If it's already a string (step key), use it directly
393
+ deps.append(dep)
394
+ result["depends_on"] = deps
395
+
396
+ # Always include result_handlers, even if they were excluded by exclude_unset
397
+ # This ensures that result handlers added via .on() method are included
398
+ result["result_handlers"] = []
399
+ if self.result_handlers:
400
+ converted_handlers: List[Any] = []
401
+ for handler in self.result_handlers:
402
+ if hasattr(handler, "to_dict"):
403
+ converted_handlers.append(handler.to_dict())
404
+ elif isinstance(handler, dict):
405
+ # Already a dict, convert conditions if needed
406
+ handler_dict = cast(Dict[str, Any], handler)
407
+ if handler_dict.get("if_conditions") and hasattr(
408
+ handler_dict["if_conditions"], "to_dict"
409
+ ):
410
+ handler_dict["if_conditions"] = handler_dict[
411
+ "if_conditions"
412
+ ].to_dict()
413
+ converted_handlers.append(handler_dict)
414
+ else:
415
+ converted_handlers.append(handler)
416
+ result["result_handlers"] = converted_handlers
417
+
418
+ # Convert ExecutionMode objects to dictionaries
419
+ if self.execution_mode is not None:
420
+ if hasattr(self.execution_mode, "to_dict") and callable(
421
+ getattr(self.execution_mode, "to_dict")
422
+ ):
423
+ exec_mode_obj = cast(ExecutionModeProtocol, self.execution_mode)
424
+ result["execution_mode"] = exec_mode_obj.to_dict()
425
+ elif isinstance(self.execution_mode, dict):
426
+ result["execution_mode"] = self.execution_mode
427
+ else:
428
+ # Handle string mode types or other formats
429
+ result["execution_mode"] = self.execution_mode
430
+
431
+ # Recursively convert any other condition objects to dictionaries
432
+ def convert_conditions(obj: Any) -> Any:
433
+ """Recursively convert condition objects to dictionaries."""
434
+
435
+ if obj is None or isinstance(obj, (str, int, float, bool)):
436
+ return obj
437
+ elif isinstance(obj, Enum):
438
+ # Convert enum values to their string values for JSON serialization
439
+ return obj.value
440
+ elif hasattr(obj, "to_dict"):
441
+ return obj.to_dict()
442
+ elif isinstance(obj, list):
443
+ return [convert_conditions(item) for item in obj]
444
+ elif isinstance(obj, dict):
445
+ # Handle enum keys in dictionaries
446
+ converted_dict: Dict[str, Any] = {}
447
+ obj_dict = cast(Dict[Any, Any], obj)
448
+ for key, value in obj_dict.items():
449
+ # Convert enum keys to strings
450
+ str_key: str
451
+ if isinstance(key, Enum):
452
+ str_key = str(key.value)
453
+ else:
454
+ str_key = str(key)
455
+ converted_dict[str_key] = convert_conditions(value)
456
+ return converted_dict
457
+ else:
458
+ return obj
459
+
460
+ # Apply condition conversion to all fields
461
+ for key, value in result.items():
462
+ result[key] = convert_conditions(value)
463
+
464
+ # Recursively sort all dictionaries for deterministic serialization
465
+ def sort_dict_recursively(obj: Any) -> Any:
466
+ """Recursively sort all dictionaries to ensure deterministic JSON output."""
467
+ if isinstance(obj, dict):
468
+ # Convert keys to strings for sorting to handle mixed types (enums + strings)
469
+ def sort_key(item: Tuple[Any, Any]) -> str:
470
+ k, v = item
471
+ if hasattr(k, "value"): # Enum
472
+ return str(k.value)
473
+ return str(k)
474
+
475
+ obj_dict = cast(Dict[Any, Any], obj)
476
+ return dict(
477
+ sorted(
478
+ ((k, sort_dict_recursively(v)) for k, v in obj_dict.items()),
479
+ key=sort_key,
480
+ )
481
+ )
482
+ elif isinstance(obj, list):
483
+ return [sort_dict_recursively(item) for item in obj]
484
+ else:
485
+ return obj
486
+
487
+ result = sort_dict_recursively(result)
488
+ return cast(Dict[str, Any], result)
489
+
490
+ def _process_decorator_handlers(self) -> None:
491
+ """Process decorator handlers and convert them to ResultHandler objects."""
492
+ if not hasattr(self, "_decorator_handlers") or not self._decorator_handlers:
493
+ return
494
+
495
+ # Get current result_handlers list or create new one
496
+ current_handlers = list(self.result_handlers) if self.result_handlers else []
497
+
498
+ for condition, _ in self._decorator_handlers:
499
+ # Create a simple handler that calls the function
500
+ # For now, we'll create a basic handler structure
501
+ # In a full implementation, you'd want to analyze the function and create appropriate steps
502
+ handler = ResultHandler(
503
+ type=HandlerType.FINAL,
504
+ if_conditions=condition,
505
+ output_content_type=OutputContentType.TEXT,
506
+ steps=[], # Would need to convert function to steps
507
+ )
508
+ current_handlers.append(handler)
509
+
510
+ # Update the result_handlers field
511
+ self.result_handlers = current_handlers
512
+
513
+ # Clear processed handlers
514
+ self._decorator_handlers.clear()
515
+
516
+
517
+ class StepOutput(BaseModel):
518
+ """Represents the output of a step for type-safe access."""
519
+
520
+ step_key: str
521
+
522
+ def __init__(self, **data: Any):
523
+ super().__init__(**data)
524
+ # Use object attribute instead of Pydantic field
525
+ object.__setattr__(self, "_expected_fields", {})
526
+
527
+ def __getitem__(self, key: str) -> str:
528
+ """Allow bracket notation for accessing output fields."""
529
+ # Track that this field is being accessed for validation
530
+ if hasattr(self, "_expected_fields"):
531
+ expected_fields = getattr(self, "_expected_fields", {})
532
+ expected_fields[key] = (
533
+ None # Will be populated with actual type during execution
534
+ )
535
+ return f"{{{{{self.step_key}.{key}}}}}"
536
+
537
+ def __getattr__(self, name: str) -> str:
538
+ """Allow dot notation for accessing output fields."""
539
+ # Avoid infinite recursion for Pydantic internal attributes
540
+ if name.startswith("_") or name in {"step_key", "model_fields", "model_config"}:
541
+ raise AttributeError(
542
+ f"'{type(self).__name__}' object has no attribute '{name}'"
543
+ )
544
+
545
+ # Track that this field is being accessed for validation
546
+ if hasattr(self, "_expected_fields"):
547
+ expected_fields = getattr(self, "_expected_fields", {})
548
+ expected_fields[name] = (
549
+ None # Will be populated with actual type during execution
550
+ )
551
+ return f"{{{{{self.step_key}.{name}}}}}"
552
+
553
+ def get_expected_fields(self) -> Dict[str, Any]:
554
+ """Get the fields that have been accessed on this step output."""
555
+ return getattr(self, "_expected_fields", {})
556
+
557
+ def validate_field_access(self, available_fields: Dict[str, Any]) -> List[str]:
558
+ """Validate that all accessed fields are available in the step result."""
559
+ errors = []
560
+ expected = self.get_expected_fields()
561
+
562
+ for field_name in expected.keys():
563
+ if field_name not in available_fields:
564
+ errors.append(
565
+ f"Step '{self.step_key}' output field '{field_name}' is not available"
566
+ )
567
+
568
+ return errors
569
+
570
+
571
+ class ResultHandler(BaseModel):
572
+ """Result handler definition that matches the Go backend structure."""
573
+
574
+ type: HandlerType = HandlerType.FINAL
575
+ if_conditions: Optional[Any] = None # Can be condition objects or None
576
+ output_content_type: OutputContentType = OutputContentType.TEXT
577
+ history_content_type: Optional[str] = None
578
+ ui_content_type: Optional[str] = None
579
+ steps: List[Step] = Field(default_factory=list)
580
+
581
+ def to_dict(self) -> Dict[str, Any]:
582
+ """Convert to dict format expected by backend."""
583
+ # Convert Step objects in steps list to dictionaries BEFORE model_dump()
584
+ # because model_dump() will convert Step objects to dicts but won't call to_dict()
585
+ converted_steps = []
586
+ if self.steps:
587
+ for step in self.steps:
588
+ if hasattr(step, "to_dict"):
589
+ step_dict = step.to_dict()
590
+ # DO NOT auto-generate keys for result handler steps
591
+ # Only include keys if explicitly set to preserve roundtrip parity
592
+
593
+ # Ensure bot_output_visibility is preserved (defaults to hidden for result handler steps)
594
+ if "bot_output_visibility" not in step_dict:
595
+ step_dict["bot_output_visibility"] = "hidden"
596
+
597
+ converted_steps.append(step_dict)
598
+ else:
599
+ # Convert step to dict if it doesn't have to_dict method
600
+ if isinstance(step, dict):
601
+ converted_steps.append(step)
602
+ else:
603
+ # Fallback: convert to dict using model_dump if available
604
+ if hasattr(step, "model_dump"):
605
+ converted_steps.append(step.model_dump())
606
+ else:
607
+ converted_steps.append({})
608
+
609
+ # Use model_dump but exclude steps, then add our properly converted steps
610
+ result = self.model_dump(exclude={"steps"})
611
+ result["steps"] = converted_steps
612
+
613
+ # Filter out empty/None optional fields to match export behavior
614
+ if result.get("ui_content_type") in [None, ""]:
615
+ result.pop("ui_content_type", None)
616
+ if result.get("history_content_type") in [None, ""]:
617
+ result.pop("history_content_type", None)
618
+
619
+ # Convert condition objects to dictionaries if needed
620
+ if self.if_conditions is not None and hasattr(self.if_conditions, "to_dict"):
621
+ result["if_conditions"] = self.if_conditions.to_dict()
622
+
623
+ # Convert enum values to their string values (both keys and values)
624
+ converted_result = {}
625
+ for key, value in result.items():
626
+ # Convert enum keys to strings
627
+ if isinstance(key, Enum):
628
+ key = key.value
629
+ # Convert enum values to strings
630
+ if isinstance(value, Enum):
631
+ value = value.value
632
+ converted_result[key] = value
633
+ result = converted_result
634
+
635
+ return result
636
+
637
+
638
+ class SecretsDict(Dict[str, Any]):
639
+ """A dictionary wrapper that provides .get() method for secrets access."""
640
+
641
+ def __init__(self, secrets_data: Optional[Dict[str, Any]] = None):
642
+ super().__init__(secrets_data or {})
643
+
644
+ def get(self, key: str, default: Any = None) -> Any:
645
+ """Get decrypted secrets for a specific resource/service key."""
646
+ return super().get(key, default)
647
+
648
+
649
+ class ParametersDict(Dict[str, Any]):
650
+ """A dictionary wrapper that provides .get() method for parameters access."""
651
+
652
+ def __init__(self, parameters_data: Optional[Dict[str, Any]] = None):
653
+ super().__init__(parameters_data or {})
654
+
655
+ def get(self, key: str, default: Any = None) -> Any:
656
+ """Get a step parameter with optional default."""
657
+ return super().get(key, default)
658
+
659
+
660
+ class StepContext(BaseModel):
661
+ """Context available to step functions with type-safe access
662
+
663
+ Provides access to:
664
+ - User query and parameters
665
+ - Previous step results (state)
666
+ - Available resources (datasets, APIs)
667
+ - Encrypted secrets
668
+ - System information
669
+ """
670
+
671
+ # Core user input
672
+ query: Optional[str] = None
673
+
674
+ # Previous step results and state
675
+ state: Dict[str, Any] = Field(default_factory=dict)
676
+ steps: Dict[str, Any] = Field(default_factory=dict) # Alias for state.steps
677
+
678
+ # Resources and data access
679
+ resources: List[Dict[str, Any]] = Field(default_factory=list)
680
+ resource_definitions: Optional[Dict[str, Any]] = None
681
+
682
+ # Security and credentials - raw data
683
+ secrets_data: Dict[str, Any] = Field(default_factory=dict)
684
+
685
+ # Step-specific parameters - raw data
686
+ parameters_data: Dict[str, Any] = Field(default_factory=dict)
687
+
688
+ # System context
689
+ system: Dict[str, Any] = Field(default_factory=dict)
690
+
691
+ # Raw context for advanced users
692
+ raw: Dict[str, Any] = Field(default_factory=dict)
693
+
694
+ def __init__(self, **data: Any):
695
+ # Extract secrets and parameters from input data
696
+ secrets_data = data.pop("secrets", {})
697
+ parameters_data = data.pop("parameters", {})
698
+
699
+ # Set the internal data fields
700
+ data["secrets_data"] = secrets_data
701
+ data["parameters_data"] = parameters_data
702
+
703
+ super().__init__(**data)
704
+
705
+ # Create wrapper objects for secrets and parameters
706
+ self._secrets_wrapper = SecretsDict(self.secrets_data)
707
+ self._parameters_wrapper = ParametersDict(self.parameters_data)
708
+
709
+ @property
710
+ def secrets(self) -> SecretsDict:
711
+ """Get secrets wrapper that supports .get() method."""
712
+ return self._secrets_wrapper
713
+
714
+ @property
715
+ def parameters(self) -> ParametersDict:
716
+ """Get parameters wrapper that supports .get() method."""
717
+ return self._parameters_wrapper
718
+
719
+ def __getitem__(self, key: str) -> Any:
720
+ """Allow bracket notation access to state"""
721
+ try:
722
+ state_value = super().__getattribute__("state")
723
+ return state_value.get(key) if isinstance(state_value, dict) else None
724
+ except AttributeError:
725
+ return None
726
+
727
+ def __getattr__(self, name: str) -> Any:
728
+ """Allow dot notation access to state"""
729
+ # Avoid infinite recursion for Pydantic internal attributes
730
+ if name.startswith("_") or name in {
731
+ "state",
732
+ "steps",
733
+ "resources",
734
+ "secrets",
735
+ "parameters",
736
+ "system",
737
+ "raw",
738
+ "query",
739
+ "resource_definitions",
740
+ }:
741
+ raise AttributeError(
742
+ f"'{self.__class__.__name__}' object has no attribute '{name}'"
743
+ )
744
+
745
+ try:
746
+ state_value = super().__getattribute__("state")
747
+ if isinstance(state_value, dict) and name in state_value:
748
+ return state_value[name]
749
+ except AttributeError:
750
+ pass
751
+ raise AttributeError(
752
+ f"'{self.__class__.__name__}' object has no attribute '{name}'"
753
+ )
754
+
755
+ def get_resource(self, key: str) -> Optional[Dict[str, Any]]:
756
+ """Get a specific resource by key"""
757
+ for resource in self.resources:
758
+ if resource.get("key") == key:
759
+ return resource
760
+ return None
761
+
762
+ def get_secret(self, key: str) -> Optional[Any]:
763
+ """Get decrypted secrets for a specific resource/service"""
764
+ return self.secrets.get(key)
765
+
766
+ def get_parameter(self, key: str, default: Any = None) -> Any:
767
+ """Get a step parameter with optional default"""
768
+ return self.parameters.get(key, default)
769
+
770
+
771
+ class StepResult(BaseModel):
772
+ """Result from executing a workflow step."""
773
+
774
+ success: bool
775
+ data: Any = None
776
+ error: Optional[str] = None
777
+ step_name: str
778
+
779
+
780
+ class Agent(BaseModel):
781
+ """Main agent class for defining AI workflows."""
782
+
783
+ name: str
784
+ description: Optional[str] = None
785
+ persona: Optional[str] = None
786
+ visibility: str = "private"
787
+ running_message: Optional[str] = None
788
+ finished_message: Optional[str] = None
789
+ version: str = "1.0"
790
+ timeout: Optional[int] = None
791
+ retry_attempts: int = 0
792
+ tags: List[str] = Field(default_factory=list)
793
+ steps: List[Step] = Field(default_factory=list)
794
+ parameter_definitions: List["ParameterDefinition"] = Field(default_factory=list)
795
+
796
+ def step(
797
+ self,
798
+ action: Any,
799
+ key: Optional[str] = None,
800
+ depends_on: Optional[Union[Step, List[Step], str, List[str]]] = None,
801
+ **kwargs: Any,
802
+ ) -> Step:
803
+ """Create a step with cleaner syntax and better defaults.
804
+
805
+ Args:
806
+ action: The action to execute (e.g., codeexec.execute(...))
807
+ key: Optional step key, auto-generated if not provided
808
+ depends_on: Step(s) or key(s) this step depends on
809
+ **kwargs: Additional step configuration
810
+
811
+ Returns:
812
+ Step: The created step with sensible defaults
813
+ """
814
+ # Auto-generate key if not provided
815
+ if key is None:
816
+ key = f"step_{len(self.steps) + 1}"
817
+
818
+ # Set sensible defaults
819
+ step_config = {
820
+ "agent": self,
821
+ "key": key,
822
+ "action": action,
823
+ "depends_on": (
824
+ [depends_on]
825
+ if depends_on and not isinstance(depends_on, list)
826
+ else depends_on
827
+ ),
828
+ "user_output_visibility": OutputVisibility.VISIBLE,
829
+ "bot_output_visibility": OutputVisibility.HIDDEN,
830
+ "output_content_type": OutputContentType.TEXT,
831
+ "parameter_hydration_behaviour": ParameterHydrationBehaviour.NONE,
832
+ }
833
+
834
+ # Override with any provided kwargs
835
+ step_config.update(kwargs)
836
+
837
+ step = Step(**step_config)
838
+ return step
839
+
840
+ def exec(
841
+ self,
842
+ step_metadata: Optional[StepMetadata] = None,
843
+ **codeexec_params: Any,
844
+ ) -> Callable[..., Step]:
845
+ """Create a codeexec step using decorator syntax.
846
+
847
+ This method is designed to be used as a decorator for functions that implement
848
+ code execution logic. It creates a codeexec.execute action with the provided
849
+ parameters and step metadata.
850
+
851
+ Args:
852
+ step_metadata: Step configuration (key, depends_on, etc.)
853
+ **codeexec_params: Parameters for codeexec.execute (entrypoint, parameters, resources, etc.)
854
+
855
+ Returns:
856
+ Callable: Decorator function that creates and returns a Step
857
+ """
858
+
859
+ def decorator(func: Callable[..., Any]) -> Step:
860
+ # Import here to avoid circular imports
861
+ from .actions import codeexec
862
+
863
+ # Create the codeexec.execute action with the provided parameters
864
+ action = codeexec.execute(**codeexec_params)
865
+
866
+ # Extract step metadata or use defaults
867
+ if step_metadata:
868
+ step_config = {
869
+ "agent": self,
870
+ "action": action,
871
+ "key": step_metadata.key or func.__name__,
872
+ "depends_on": step_metadata.depends_on,
873
+ "execution_mode": step_metadata.execution_mode,
874
+ "output_behavior": step_metadata.output_behavior,
875
+ "output_channels": step_metadata.output_channels,
876
+ "output_content_type": step_metadata.output_content_type,
877
+ "history_content_type": step_metadata.history_content_type,
878
+ "ui_content_type": step_metadata.ui_content_type,
879
+ "user_output_visibility": step_metadata.user_output_visibility,
880
+ "bot_output_visibility": step_metadata.bot_output_visibility,
881
+ "running_message": step_metadata.running_message,
882
+ "finished_message": step_metadata.finished_message,
883
+ "parameter_hydration_behaviour": step_metadata.parameter_hydration_behaviour,
884
+ }
885
+ else:
886
+ step_config = {
887
+ "agent": self,
888
+ "action": action,
889
+ "key": func.__name__,
890
+ "user_output_visibility": OutputVisibility.VISIBLE,
891
+ "bot_output_visibility": OutputVisibility.HIDDEN,
892
+ "output_content_type": OutputContentType.TEXT,
893
+ "parameter_hydration_behaviour": ParameterHydrationBehaviour.NONE,
894
+ }
895
+
896
+ # Remove None values and extract action separately
897
+ filtered_config = {
898
+ k: v for k, v in step_config.items() if v is not None and k != "action"
899
+ }
900
+
901
+ # Create step with the codeexec action, not the function
902
+ step = Step(action=action, step=None, **filtered_config)
903
+
904
+ # Store the function name on the step for runtime extraction
905
+ if hasattr(step, "__dict__"):
906
+ step.__dict__["__name__"] = func.__name__
907
+
908
+ return step
909
+
910
+ return decorator
911
+
912
+ def add_step(self, step: Step) -> None:
913
+ """Add a step to this agent if it's not already present."""
914
+ # Only add if the step is not already in the list
915
+ if step in self.steps:
916
+ return
917
+
918
+ # Create a new list with the existing steps plus the new one
919
+ current_steps = list(self.steps) if self.steps else []
920
+ current_steps.append(step)
921
+ self.steps = current_steps
922
+
923
+ def get_step(self, key: str) -> Optional[Step]:
924
+ """Get a step by its key."""
925
+ for step in self.steps:
926
+ if step.key == key:
927
+ return step
928
+ return None
929
+
930
+ def to_json(self) -> str:
931
+ """Export to JSON format expected by Go backend."""
932
+ export_data = {
933
+ "bot": {
934
+ "Name": self.name,
935
+ "Description": self.description or "",
936
+ "Visibility": self.visibility,
937
+ "Persona": self.persona, # Export as string or null, not object
938
+ "RunningMessage": self.running_message,
939
+ "FinishedMessage": self.finished_message,
940
+ "Source": "python",
941
+ },
942
+ "parameter_definitions": self.parameter_definitions,
943
+ "steps": [step.to_dict() for step in self.steps],
944
+ }
945
+
946
+ import json
947
+
948
+ return json.dumps(export_data, indent=2)
949
+
950
+
951
+ class ExecutionMode(BaseModel):
952
+ """Execution mode matching Go backend structure."""
953
+
954
+ mode: ExecutionModeType = ExecutionModeType.ALL
955
+ data: Optional[Any] = None
956
+ if_condition: Optional[Any] = None
957
+
958
+ def to_dict(self) -> Dict[str, Any]:
959
+ """Convert to dict format expected by backend."""
960
+ result: Dict[str, Any] = {"mode": self.mode}
961
+ if self.data is not None:
962
+ result["data"] = self.data
963
+ # Always include if_condition to maintain parity with backend structure
964
+ result["if_condition"] = self.if_condition
965
+ return result
966
+
967
+
968
+ class Prompt(BaseModel):
969
+ """A prompt that can be loaded from a .prompt file or defined inline.
970
+
971
+ This class provides a clean way to manage prompts in agent code,
972
+ supporting both inline strings and external .prompt files.
973
+ """
974
+
975
+ content: str = Field(..., description="The prompt content")
976
+ filename: Optional[str] = Field(
977
+ None, description="Source filename if loaded from file"
978
+ )
979
+
980
+ def __init__(
981
+ self, content: Optional[str] = None, filename: Optional[str] = None, **data: Any
982
+ ):
983
+ """Initialize a Prompt.
984
+
985
+ Args:
986
+ content: Direct prompt content
987
+ filename: Path to .prompt file to load
988
+ **data: Additional model data
989
+ """
990
+ if content is not None and filename is not None:
991
+ raise ValueError("Cannot specify both 'content' and 'filename' parameters")
992
+
993
+ if content is None and filename is None:
994
+ raise ValueError("Must specify either 'content' or 'filename' parameter")
995
+
996
+ if filename is not None:
997
+ # Load content from file
998
+ content = self._load_from_file(filename)
999
+ data.update({"content": content, "filename": filename})
1000
+ else:
1001
+ data.update({"content": content})
1002
+
1003
+ super().__init__(**data)
1004
+
1005
+ @classmethod
1006
+ def from_file(cls, filename: str, base_path: Optional[str] = None) -> "Prompt":
1007
+ """Load a prompt from a .prompt file.
1008
+
1009
+ Args:
1010
+ filename: Path to the .prompt file
1011
+ base_path: Base directory to resolve relative paths (defaults to cwd)
1012
+
1013
+ Returns:
1014
+ Prompt: Loaded prompt instance
1015
+ """
1016
+ if base_path is None:
1017
+ base_path = os.getcwd()
1018
+
1019
+ file_path = Path(base_path) / filename
1020
+
1021
+ try:
1022
+ with open(file_path, "r", encoding="utf-8") as f:
1023
+ content = f.read().strip()
1024
+ return cls(content=content)
1025
+ except FileNotFoundError:
1026
+ raise FileNotFoundError(f"Prompt file not found: {file_path}")
1027
+ except Exception as e:
1028
+ raise RuntimeError(f"Error reading prompt file {file_path}: {e}")
1029
+
1030
+ def _load_from_file(self, filename: str) -> str:
1031
+ """Load content from a .prompt file."""
1032
+ return self.from_file(filename).content
1033
+
1034
+ def __str__(self) -> str:
1035
+ """Return the prompt content when used as a string."""
1036
+ return self.content
1037
+
1038
+ def __call__(self) -> str:
1039
+ """Allow the prompt to be called like a function, returning content."""
1040
+ return self.content
1041
+
1042
+ def __repr__(self) -> str:
1043
+ """Return a helpful representation."""
1044
+ if self.filename:
1045
+ return f"Prompt(filename='{self.filename}')"
1046
+ else:
1047
+ preview = (
1048
+ self.content[:50] + "..." if len(self.content) > 50 else self.content
1049
+ )
1050
+ return f"Prompt(content='{preview}')"
1051
+
1052
+ @classmethod
1053
+ def load_from_directory(
1054
+ cls, directory: str, base_path: Optional[str] = None
1055
+ ) -> Dict[str, "Prompt"]:
1056
+ """Load all .prompt files from a directory as Prompt objects.
1057
+
1058
+ Args:
1059
+ directory: Directory containing .prompt files
1060
+ base_path: Base directory to resolve relative paths (defaults to smart detection)
1061
+
1062
+ Returns:
1063
+ Dict[str, Prompt]: Mapping of filename (without extension) to Prompt objects
1064
+ """
1065
+ if base_path is None:
1066
+ # Use stack inspection to determine the calling file's directory
1067
+ # This allows agents to load prompts from their own directory
1068
+ # even when executed from a different working directory
1069
+ import inspect
1070
+
1071
+ frame = inspect.currentframe()
1072
+ try:
1073
+ # Get the caller's frame (skip current frame)
1074
+ if frame and frame.f_back:
1075
+ caller_frame = frame.f_back
1076
+ caller_file = caller_frame.f_code.co_filename
1077
+ caller_dir = Path(caller_file).parent
1078
+ # Check if the caller is in an agent directory structure
1079
+ # (i.e., the calling file is agent.py in a directory under erdo-agents)
1080
+ if caller_file.endswith("agent.py") and "erdo-agents" in str(
1081
+ caller_dir
1082
+ ):
1083
+ # The caller is an agent.py file, use its directory
1084
+ base_path = str(caller_dir)
1085
+ else:
1086
+ # Fall back to current working directory
1087
+ base_path = os.getcwd()
1088
+ else:
1089
+ # Fall back to current working directory
1090
+ base_path = os.getcwd()
1091
+ finally:
1092
+ del frame # Avoid reference cycles
1093
+
1094
+ dir_path = Path(base_path) / directory
1095
+
1096
+ if not dir_path.exists() or not dir_path.is_dir():
1097
+ # Provide a more helpful error message
1098
+ cwd = os.getcwd()
1099
+ caller_path = None
1100
+ import inspect
1101
+
1102
+ frame = inspect.currentframe()
1103
+ try:
1104
+ if frame and frame.f_back:
1105
+ caller_file = frame.f_back.f_code.co_filename
1106
+ caller_path = str(Path(caller_file).parent)
1107
+ finally:
1108
+ del frame
1109
+
1110
+ error_msg = f"Prompts directory not found: {dir_path}"
1111
+ if caller_path and caller_path != str(Path(base_path)):
1112
+ error_msg += f"\nCalling file: {caller_path}"
1113
+ error_msg += f"\nCurrent working directory: {cwd}"
1114
+ error_msg += f"\nBase path used: {base_path}"
1115
+ raise FileNotFoundError(error_msg)
1116
+
1117
+ prompts = {}
1118
+ for prompt_file in dir_path.glob("*.prompt"):
1119
+ prompt_name = prompt_file.stem # filename without extension
1120
+ relative_path = str(prompt_file.relative_to(base_path))
1121
+ prompts[prompt_name] = cls.from_file(relative_path, base_path)
1122
+
1123
+ return prompts
1124
+
1125
+
1126
+ # Re-export commonly used types for convenience
1127
+ __all__ = [
1128
+ "Agent",
1129
+ "Step",
1130
+ "StepMetadata",
1131
+ "StepContext",
1132
+ "StepResult",
1133
+ "StepOutput",
1134
+ "ResultHandler",
1135
+ "ExecutionMode",
1136
+ "Tool",
1137
+ "PythonFile",
1138
+ "Prompt",
1139
+ "TemplateString",
1140
+ # Complex types
1141
+ "ExecutionCondition",
1142
+ ]