├── .github └── workflows │ └── test.yaml ├── .gitignore ├── LICENSE ├── README.md ├── asimov ├── __init__.py ├── asimov_base.py ├── caches │ ├── cache.py │ ├── mock_redis_cache.py │ └── redis_cache.py ├── constants.py ├── data │ └── postgres │ │ └── manager.py ├── executors │ ├── __init__.py │ └── executor.py ├── graph │ ├── __init__.py │ ├── agent_directed_flow.py │ └── tasks.py ├── py.typed ├── services │ └── inference_clients.py └── utils │ ├── models.py │ ├── token_buffer.py │ ├── token_counter.py │ └── visualize.py ├── bismuth.toml ├── docs ├── basic_agent.md └── llm_agent.md ├── examples ├── __init__.py ├── agent_directed_flow.py ├── basic_agent.py └── llm_agent.py ├── pyproject.toml ├── pytest.ini ├── stack.dev.yaml └── tests ├── __init__.py ├── test_agent.py ├── test_anthropic_inference_client.py ├── test_basic_agent.py ├── test_cache.py ├── test_llm_agent.py └── test_readme_examples.py /.github/workflows/test.yaml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: 4 | pull_request: 5 | push: 6 | 7 | jobs: 8 | Test: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - name: Checkout code 12 | uses: actions/checkout@v4 13 | 14 | - name: Set up Python 15 | uses: actions/setup-python@v5 16 | with: 17 | python-version: "3.12" 18 | 19 | - name: Install hatch 20 | run: | 21 | pip install hatch 22 | 23 | - name: Run tests 24 | run: | 25 | docker run --rm -d -p 6379:6379 redis 26 | hatch test -- -vv -s --log-cli-level DEBUG 27 | 28 | mypy: 29 | runs-on: ubuntu-latest 30 | steps: 31 | - name: Checkout code 32 | uses: actions/checkout@v4 33 | 34 | - name: Set up Python 35 | uses: actions/setup-python@v5 36 | with: 37 | python-version: "3.12" 38 | 39 | - name: Install hatch 40 | run: | 41 | pip install hatch 42 | 43 | - name: Run mypy 44 | run: | 45 | hatch run types:check 46 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/* 2 | **/*/__pycache__/* 3 | **/__pycache__/* 4 | .vscode/* 5 | .torb_buildstate/* 6 | dist/* 7 | venv/* 8 | .venv/* 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![PyPI - Version](https://img.shields.io/pypi/v/asimov_agents.svg)](https://pypi.org/project/asimov_agents) 2 | [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/asimov_agents.svg)](https://pypi.org/project/asimov_agents) 3 | 4 | [![](https://dcbadge.limes.pink/api/server/https://discord.gg/bismuthai)](https://discord.gg/bismuthai) 5 | 6 | # Asimov Agents 7 | 8 | A Python framework for building AI agent systems with robust task management in the form of a graph execution engine, inference capabilities, and caching. 9 | 10 | We support advanced features like State Snapshotting, Middleware, Agent Directed Graph Execution, Open Telemetry Integrations and more. 11 | 12 | 🔮 Asimov is the foundation of [bismuth.sh](https://waitlist.bismuth.sh) an end to end AI software developer that can handle many tasks autonomously. Check us out! 🔮 13 | 14 | ## Quickstart 15 | ```bash 16 | pip install asimov_agents 17 | ``` 18 | 19 | Checkout [these docs](https://github.com/BismuthCloud/asimov/tree/main/docs) which show off two basic examples that should be enough to get you experimenting! 20 | 21 | Further documentation greatly appreciated in PRs! 22 | 23 | ## System Overview 24 | 25 | Asimov Agents is composed of three main components: 26 | 27 | 1. **Graph Primitives** 28 | - Manages task execution flow and dependencies 29 | 30 | 2. **Inference Clients** 31 | - Supports multiple LLM providers: 32 | - Anthropic Claude (via API) 33 | - AWS Bedrock 34 | - OpenAI (Including local models) 35 | - Vertex 36 | - OpenRouter 37 | - Features: 38 | - Streaming responses 39 | - Tool/function calling capabilities 40 | - Token usage tracking 41 | - OpenTelemetry instrumentation 42 | - Prompt caching support 43 | 44 | 3. **Caching System** 45 | - Abstract Cache interface with Redis implementation 46 | - Features: 47 | - Key-value storage with JSON serialization 48 | - Prefix/suffix namespacing 49 | - Pub/sub messaging via mailboxes 50 | - Bulk operations (get_all, clear) 51 | - Async interface 52 | 53 | ## Component Interactions 54 | 55 | ### Task Management 56 | - Tasks are created and tracked using the `Task` class 57 | - Each task has: 58 | - Unique ID 59 | - Type and objective 60 | - Parameters dictionary 61 | - Status tracking 62 | - Result/error storage 63 | 64 | ### Graph System Architecture 65 | - **Module Types** 66 | - `SUBGRAPH`: Nodes composes of other nodes. 67 | - `EXECUTOR`: Task execution modules 68 | - `FLOW_CONTROL`: Execution flow control modules 69 | 70 | - **Node Configuration** 71 | ```python 72 | node_config = NodeConfig( 73 | parallel=True, # Enable parallel execution 74 | condition="task.ready", # Conditional execution 75 | retry_on_failure=True, # Enable retry mechanism 76 | max_retries=3, # Maximum retry attempts 77 | max_visits=5, # Maximum node visits 78 | inputs=["data"], # Required inputs 79 | outputs=["result"] # Expected outputs 80 | ) 81 | ``` 82 | 83 | - **Flow Control Features** 84 | - Conditional branching based on task state 85 | - Dynamic node addition during execution 86 | - Dependency chain management 87 | - Automatic cleanup of completed nodes 88 | - Execution state tracking and recovery 89 | - LLM directed flow for complex decisisons 90 | 91 | - **Snapshot System** 92 | - State preservation modes: 93 | - `NEVER`: No snapshots 94 | - `ONCE`: Single snapshot 95 | - `ALWAYS`: Continuous snapshots 96 | - Captures: 97 | - Agent state 98 | - Cache contents 99 | - Task status 100 | - Execution history 101 | - Configurable storage location via `ASIMOV_SNAPSHOT` 102 | 103 | - **Error Handling** 104 | - Automatic retry mechanisms 105 | - Partial completion states 106 | - Failed chain tracking 107 | - Detailed error reporting 108 | - Timeout management 109 | 110 | ### Inference Pipeline 111 | 1. Messages are formatted with appropriate roles (SYSTEM, USER, ASSISTANT, TOOL_RESULT) 112 | 2. Inference clients handle: 113 | - Message formatting 114 | - API communication 115 | - Response streaming 116 | - Token accounting 117 | - Error handling 118 | 119 | ### Caching Layer 120 | - Redis cache provides: 121 | - Fast key-value storage 122 | - Message queuing 123 | - Namespace management 124 | - Atomic operations 125 | 126 | ## Agent Primitives 127 | 128 | The Asimov Agents framework is built around several core primitives that enable flexible and powerful agent architectures: 129 | 130 | ### Module Types 131 | The framework supports different types of modules through the `ModuleType` enum: 132 | - `SUBGRAPH`: Nodes composes of other nodes. 133 | - `EXECUTOR`: Task execution and action implementation 134 | - `FLOW_CONTROL`: Execution flow and routing control 135 | 136 | ### Agent Module 137 | The `AgentModule` is the base class for all agent components: 138 | ```python 139 | class AgentModule: 140 | name: str # Unique module identifier 141 | type: ModuleType # Module type classification 142 | config: ModuleConfig # Module configuration 143 | dependencies: List[str] # Module dependencies 144 | input_mailboxes: List[str] # Input communication channels 145 | output_mailbox: str # Output communication channel 146 | trace: bool # OpenTelemetry tracing flag 147 | ``` 148 | 149 | ### Node Configuration 150 | Nodes can be configured with various parameters through `NodeConfig`: 151 | ```python 152 | class NodeConfig: 153 | parallel: bool = False # Enable parallel execution 154 | condition: Optional[str] = None # Execution condition 155 | retry_on_failure: bool = True # Auto-retry on failures 156 | max_retries: int = 3 # Maximum retry attempts 157 | max_visits: int = 5 # Maximum node visits 158 | inputs: List[str] = [] # Required inputs 159 | outputs: List[str] = [] # Expected outputs 160 | ``` 161 | 162 | ### Flow Control 163 | Flow control enables dynamic execution paths: 164 | ```python 165 | class FlowDecision: 166 | next_node: str # Target node 167 | condition: Optional[str] = None # Jump condition 168 | cleanup_on_jump: bool = False # Cleanup on transition 169 | 170 | class FlowControlConfig: 171 | decisions: List[FlowDecision] # Decision rules 172 | default: Optional[str] = None # Default node 173 | cleanup_on_default: bool = True # Cleanup on default 174 | ``` 175 | 176 | ### Agent Directed Flow Control 177 | 178 | Agent Directed Flow Control is a powerful feature that enables intelligent routing of tasks based on LLM decision making. It allows the system to: 179 | 180 | - Dynamically route tasks to specialized modules based on content analysis 181 | - Use example-based learning for routing decisions 182 | - Support multiple voters for consensus-based routing 183 | - Handle fallback cases with error handlers 184 | 185 | Example configuration: 186 | ```python 187 | flow_control = Node( 188 | name="flow_control", 189 | type=ModuleType.FLOW_CONTROL, 190 | modules=[ 191 | AgentDirectedFlowControl( 192 | name="ContentFlowControl", 193 | type=ModuleType.FLOW_CONTROL, 194 | voters=3, # Number of voters for consensus 195 | inference_client=inference_client, 196 | system_description="A system that handles various content creation tasks", 197 | flow_config=AgentDrivenFlowControlConfig( 198 | decisions=[ 199 | AgentDrivenFlowDecision( 200 | next_node="blog_writer", 201 | metadata={"description": "Writes blog posts on technical topics"}, 202 | examples=[ 203 | Example( 204 | message="Write a blog post about AI agents", 205 | choices=[ 206 | {"choice": "blog_writer", "description": "Writes blog posts"}, 207 | {"choice": "code_writer", "description": "Writes code"} 208 | ], 209 | choice="blog_writer", 210 | reasoning="The request is specifically for blog content" 211 | ) 212 | ] 213 | ), 214 | AgentDrivenFlowDecision( 215 | next_node="code_writer", 216 | metadata={"description": "Writes code examples and tutorials"}, 217 | examples=[ 218 | Example( 219 | message="Create a Python script for data processing", 220 | choices=[ 221 | {"choice": "blog_writer", "description": "Writes blog posts"}, 222 | {"choice": "code_writer", "description": "Writes code"} 223 | ], 224 | choice="code_writer", 225 | reasoning="The request is for code creation" 226 | ) 227 | ] 228 | ) 229 | ], 230 | default="error_handler" # Fallback node for unmatched requests 231 | ) 232 | ) 233 | ] 234 | ) 235 | ``` 236 | 237 | Key features: 238 | - Example-based routing decisions with clear reasoning 239 | - Multiple voter support (configurable number of voters) for robust decision making 240 | - Specialized executor modules for different content types (e.g., blog posts, code) 241 | - Metadata-enriched routing configuration for better decision context 242 | - Fallback error handling for unmatched requests 243 | - Cached message passing between nodes using Redis 244 | - Asynchronous execution with semaphore control 245 | - Comprehensive error handling and reporting 246 | 247 | For a complete working example of Agent Directed Flow Control, check out the `examples/agent_directed_flow.py` file which demonstrates a content creation system that intelligently routes tasks between blog writing and code generation modules. 248 | 249 | ### Middleware System 250 | Middleware allows for processing interception: 251 | ```python 252 | class Middleware: 253 | async def process(self, data: Dict[str, Any], cache: Cache) -> Dict[str, Any]: 254 | return data # Process or transform data 255 | ``` 256 | 257 | ### Execution State 258 | The framework maintains execution state through: 259 | ```python 260 | class ExecutionState: 261 | execution_index: int # Current execution position 262 | current_plan: ExecutionPlan # Active execution plan 263 | execution_history: List[ExecutionPlan] # Historical plans 264 | total_iterations: int # Total execution iterations 265 | ``` 266 | 267 | ### Snapshot Control 268 | State persistence is managed through `SnapshotControl`: 269 | - `NEVER`: No snapshots taken 270 | - `ONCE`: Single snapshot capture 271 | - `ALWAYS`: Continuous state capture 272 | 273 | ## Setup and Configuration 274 | 275 | ### Redis Cache Setup 276 | ```python 277 | cache = RedisCache( 278 | host="localhost", # Redis host 279 | port=6379, # Redis port 280 | db=0, # Database number 281 | password=None, # Optional password 282 | default_prefix="" # Optional key prefix 283 | ) 284 | ``` 285 | 286 | ### Inference Client Setup 287 | ```python 288 | # Anthropic Client 289 | client = AnthropicInferenceClient( 290 | model="claude-3-5-sonnet-20241022", 291 | api_key="your-api-key", 292 | api_url="https://api.anthropic.com/v1/messages" 293 | ) 294 | 295 | # AWS Bedrock Client 296 | client = BedrockInferenceClient( 297 | model="anthropic.claude-3-5-sonnet-20241022-v2:0", 298 | region_name="us-east-1" 299 | ) 300 | ``` 301 | 302 | There is similar set up for VertexAI and OpenAI 303 | 304 | ### Task and Graph Setup 305 | ```python 306 | # Create a task 307 | task = Task( 308 | type="processing", 309 | objective="Process data", 310 | params={"input": "data"} 311 | ) 312 | 313 | # Create nodes with different module types 314 | executor_node = Node( 315 | name="executor", 316 | type=ModuleType.EXECUTOR, 317 | modules=[ExecutorModule()], 318 | dependencies=["planner"] 319 | ) 320 | 321 | flow_control = Node( 322 | name="flow_control", 323 | type=ModuleType.FLOW_CONTROL, 324 | modules=[FlowControlModule( 325 | flow_config=FlowControlConfig( 326 | decisions=[ 327 | FlowDecision( 328 | next_node="executor", 329 | condition="task.ready == true" # Conditions are small lua scripts that get run based on current state. 330 | ) 331 | ], 332 | default="planner" 333 | ) 334 | )] 335 | ) 336 | 337 | # Set up the agent 338 | agent = Agent( 339 | cache=RedisCache(), 340 | max_concurrent_tasks=5, 341 | max_total_iterations=100 342 | ) 343 | 344 | # Add nodes to the agent 345 | agent.add_multiple_nodes([executor_node, flow_control]) 346 | 347 | # Run the task 348 | await agent.run_task(task) 349 | ``` 350 | 351 | ## Advanced Features 352 | 353 | ### Middleware System 354 | ```python 355 | class LoggingMiddleware(Middleware): 356 | async def process(self, data: Dict[str, Any], cache: Cache) -> Dict[str, Any]: 357 | print(f"Processing data: {data}") 358 | return data 359 | 360 | node = Node( 361 | name="executor", 362 | type=ModuleType.EXECUTOR, 363 | modules=[ExecutorModule()], 364 | config=ModuleConfig( 365 | middlewares=[LoggingMiddleware()], 366 | timeout=30.0 367 | ) 368 | ) 369 | ``` 370 | 371 | ### Execution State Management 372 | - Tracks execution history 373 | - Supports execution plan compilation 374 | - Enables dynamic plan modification 375 | - Provides state restoration capabilities 376 | ```python 377 | # Access execution state 378 | current_plan = agent.execution_state.current_plan 379 | execution_history = agent.execution_state.execution_history 380 | total_iterations = agent.execution_state.total_iterations 381 | 382 | # Compile execution plans 383 | full_plan = agent.compile_execution_plan() 384 | partial_plan = agent.compile_execution_plan_from("specific_node") 385 | 386 | # Restore from snapshot 387 | await agent.run_from_snapshot(snapshot_dir) 388 | ``` 389 | 390 | ### OpenTelemetry Integration 391 | - Automatic span creation for nodes 392 | - Execution tracking 393 | - Performance monitoring 394 | - Error tracing 395 | ```python 396 | node = Node( 397 | name="traced_node", 398 | type=ModuleType.EXECUTOR, 399 | modules=[ExecutorModule()], 400 | trace=True # Enable OpenTelemetry tracing 401 | ) 402 | ``` 403 | 404 | ## Performance Considerations 405 | 406 | ### Caching 407 | - Use appropriate key prefixes/suffixes for namespace isolation 408 | - Consider timeout settings for blocking operations 409 | - Monitor Redis memory usage 410 | - Use raw mode when bypassing JSON serialization 411 | 412 | ### Inference 413 | - Token usage is tracked automatically 414 | - Streaming reduces time-to-first-token 415 | - Tool calls support iteration limits 416 | - Prompt caching can improve response times 417 | 418 | ### Task Management 419 | - Tasks support partial failure states 420 | - Use UUIDs for guaranteed uniqueness 421 | - Status transitions are atomic 422 | 423 | ## Development 424 | 425 | ### Running Tests 426 | ```bash 427 | pytest tests/ 428 | ``` 429 | 430 | ### Required Dependencies 431 | - Redis server (If using caching) 432 | - Python 3.12+ 433 | - See requirements.txt for Python packages 434 | 435 | ## License 436 | 437 | ApacheV2 438 | -------------------------------------------------------------------------------- /asimov/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.6.0" 2 | -------------------------------------------------------------------------------- /asimov/asimov_base.py: -------------------------------------------------------------------------------- 1 | from pydantic import ConfigDict, BaseModel 2 | from abc import ABC 3 | 4 | 5 | class AsimovBase(BaseModel, ABC): 6 | model_config = ConfigDict(arbitrary_types_allowed=True) 7 | 8 | pass 9 | -------------------------------------------------------------------------------- /asimov/caches/cache.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from contextvars import ContextVar 3 | from typing import Dict, Any, Set, Optional 4 | from pydantic import ConfigDict, Field 5 | from asimov.asimov_base import AsimovBase 6 | from contextlib import asynccontextmanager 7 | 8 | 9 | class Cache(AsimovBase, ABC): 10 | model_config = ConfigDict(arbitrary_types_allowed=True) 11 | # TODO: this is treated more like a "namespace" - perhaps rename? 12 | default_prefix: str = Field(default="") 13 | default_suffix: str = Field(default="") 14 | affix_sep: str = ":" 15 | _prefix: ContextVar[str] 16 | _suffix: ContextVar[str] 17 | 18 | def __init__(self, **kwargs): 19 | super().__init__(**kwargs) 20 | self._prefix = ContextVar("prefix", default=self.default_prefix) 21 | self._suffix = ContextVar("suffix", default=self.default_suffix) 22 | 23 | async def get_prefix(self) -> str: 24 | return self._prefix.get() 25 | 26 | async def get_suffix(self) -> str: 27 | return self._suffix.get() 28 | 29 | async def apply_key_modifications(self, key: str) -> str: 30 | prefix = await self.get_prefix() 31 | suffix = await self.get_suffix() 32 | 33 | if prefix: 34 | key = f"{prefix}{self.affix_sep}{key}" 35 | if suffix: 36 | key = f"{key}{self.affix_sep}{suffix}" 37 | return key 38 | 39 | @asynccontextmanager 40 | async def with_prefix(self, prefix: str): 41 | token = self._prefix.set(prefix) 42 | try: 43 | yield self 44 | finally: 45 | self._prefix.reset(token) 46 | 47 | @asynccontextmanager 48 | async def with_suffix(self, suffix: str): 49 | token = self._suffix.set(suffix) 50 | try: 51 | yield self 52 | finally: 53 | self._suffix.reset(token) 54 | 55 | 56 | def __getitem__(self, key: str): 57 | return self.get(key) 58 | 59 | @abstractmethod 60 | async def get(self, key: str, default: Optional[Any] = None, raw: bool = False): 61 | pass 62 | 63 | @abstractmethod 64 | async def set(self, key: str, value, raw: bool = False): 65 | pass 66 | 67 | @abstractmethod 68 | async def delete(self, key: str): 69 | pass 70 | 71 | @abstractmethod 72 | async def clear(self): 73 | pass 74 | 75 | @abstractmethod 76 | async def get_all(self) -> Dict[str, Any]: 77 | pass 78 | 79 | @abstractmethod 80 | async def publish_to_mailbox(self, mailbox_id: str, value): 81 | pass 82 | 83 | @abstractmethod 84 | async def get_message(self, mailbox: str, timeout: Optional[float] = None): 85 | pass 86 | 87 | @abstractmethod 88 | async def close(self): 89 | pass 90 | 91 | @abstractmethod 92 | async def keys(self) -> Set[str]: 93 | pass 94 | -------------------------------------------------------------------------------- /asimov/caches/mock_redis_cache.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from collections import defaultdict 3 | from asimov.caches.cache import Cache 4 | from pydantic import Field 5 | from typing import Any, Optional, Dict, Set 6 | from asimov.caches.cache import Cache 7 | from pydantic import Field, field_validator 8 | import copy 9 | 10 | 11 | RAISE_ON_NONE = object() 12 | 13 | 14 | class MockRedisCache(Cache): 15 | data: dict = Field(default_factory=dict) 16 | mailboxes: Dict[str, list[str]] = Field(default_factory=dict) 17 | 18 | @field_validator("mailboxes", mode="before") 19 | def set_mailboxes(cls, v): 20 | return defaultdict(list, v) 21 | 22 | def __init__(self): 23 | super().__init__() 24 | self.mailboxes = defaultdict(list) 25 | 26 | async def clear(self): 27 | self.data = {} 28 | self.mailboxes = defaultdict(list) 29 | 30 | async def get( 31 | self, key: str, default: Any = RAISE_ON_NONE, raw: bool = False 32 | ) -> Optional[Any]: 33 | modified_key = key 34 | 35 | if not raw: 36 | modified_key = await self.apply_key_modifications(key) 37 | 38 | out = self.data.get(modified_key, default) 39 | if out is RAISE_ON_NONE: 40 | raise KeyError(key) 41 | return copy.deepcopy(out) 42 | 43 | async def get_all(self) -> Dict[str, Any]: 44 | prefix = await self.get_prefix() 45 | return {k: v for k, v in self.data.items() if k.startswith(prefix)} 46 | 47 | async def set(self, key: str, value: Any, raw: bool = False) -> None: 48 | modified_key = key 49 | 50 | if not raw: 51 | modified_key = await self.apply_key_modifications(key) 52 | 53 | self.data[modified_key] = value 54 | 55 | async def delete(self, key: str) -> None: 56 | modified_key = key 57 | self.data.pop(modified_key, None) 58 | 59 | async def peek_mailbox(self, mailbox_id: str) -> list: 60 | return self.mailboxes[mailbox_id][:] 61 | 62 | async def peek_message(self, mailbox_id: str) -> str: 63 | return self.mailboxes[mailbox_id][0] 64 | 65 | async def get_message(self, mailbox: str, timeout: Optional[float] = None): 66 | async def _get(): 67 | while True: 68 | if len(self.mailboxes[mailbox]) > 0: 69 | return self.mailboxes[mailbox].pop(0) 70 | await asyncio.sleep(0.1) 71 | 72 | try: 73 | return await asyncio.wait_for(_get(), timeout=timeout) 74 | except asyncio.TimeoutError: 75 | return None 76 | 77 | async def publish_to_mailbox(self, mailbox: str, message: Any): 78 | self.mailboxes[mailbox].append(message) 79 | 80 | async def keys(self) -> Set[str]: 81 | prefix = await self.get_prefix() 82 | suffix = await self.get_suffix() 83 | return set( 84 | k for k in self.data.keys() if k.startswith(prefix) and k.endswith(suffix) 85 | ) 86 | 87 | async def close(self) -> None: 88 | # No-op for this mock implementation 89 | pass 90 | -------------------------------------------------------------------------------- /asimov/caches/redis_cache.py: -------------------------------------------------------------------------------- 1 | import redis.asyncio 2 | import redis.exceptions 3 | import jsonpickle 4 | from typing import Dict, Any, Optional, Set 5 | 6 | from asimov.caches.cache import Cache 7 | 8 | 9 | RAISE_ON_NONE = object() 10 | 11 | 12 | class RedisCache(Cache): 13 | host: str = "localhost" 14 | port: int = 6379 15 | db: int = 0 16 | password: str | None = None 17 | _client: redis.asyncio.Redis 18 | 19 | def __init__(self, **kwargs): 20 | super().__init__(**kwargs) 21 | self._client = redis.asyncio.Redis( 22 | host=self.host, 23 | port=self.port, 24 | db=self.db, 25 | password=self.password, 26 | ) 27 | 28 | async def get(self, key: str, default=RAISE_ON_NONE, raw=False): 29 | modified_key = key 30 | 31 | if not raw: 32 | modified_key = await self.apply_key_modifications(key) 33 | 34 | value = await self._client.get(modified_key) 35 | if value is None and default is RAISE_ON_NONE: 36 | raise KeyError(key) 37 | return jsonpickle.decode(value) if value is not None else default 38 | 39 | async def set(self, key: str, value, raw: bool = False): 40 | modified_key = key 41 | 42 | if not raw: 43 | modified_key = await self.apply_key_modifications(key) 44 | 45 | await self._client.set(modified_key, jsonpickle.encode(value)) 46 | 47 | async def delete(self, key: str): 48 | modified_key = await self.apply_key_modifications(key) 49 | await self._client.delete(modified_key) 50 | 51 | async def clear(self): 52 | prefix = await self.get_prefix() 53 | all_keys = await self._client.keys(f"{prefix}{self.affix_sep}*") 54 | if all_keys: 55 | await self._client.delete(*all_keys) 56 | 57 | async def get_all(self) -> Dict[str, Any]: 58 | prefix = await self.get_prefix() 59 | all_keys = await self._client.keys(f"{prefix}{self.affix_sep}*") 60 | result = {} 61 | for key in all_keys: 62 | try: 63 | value = await self.get(key.decode("utf-8"), raw=True) 64 | except redis.exceptions.ResponseError: 65 | # Attempt to GET a non-normal key, e.g. a mailbox list 66 | continue 67 | result[key.decode("utf-8")] = value 68 | return result 69 | 70 | async def publish_to_mailbox(self, mailbox_id: str, value): 71 | modified_mailbox_id = await self.apply_key_modifications(mailbox_id) 72 | await self._client.rpush(modified_mailbox_id, jsonpickle.encode(value)) # type: ignore 73 | 74 | async def get_message(self, mailbox_id: str, timeout: Optional[float] = None): 75 | modified_mailbox_id = await self.apply_key_modifications(mailbox_id) 76 | res = await self._client.blpop([modified_mailbox_id], timeout=timeout) # type: ignore 77 | if res is None: 78 | return None 79 | _key, message = res 80 | return jsonpickle.decode(message) 81 | 82 | async def keys(self) -> Set[str]: 83 | keys: Set[str] = set() 84 | 85 | cursor = 0 86 | prefix = await self.get_prefix() 87 | suffix = await self.get_suffix() 88 | key_string = f"*" 89 | 90 | if prefix: 91 | key_string = f"{prefix}{self.affix_sep}{key_string}" 92 | if suffix: 93 | key_string = f"{key_string}{self.affix_sep}{suffix}" 94 | 95 | while True: 96 | cursor, partial_keys = await self._client.scan( 97 | cursor=cursor, match=key_string, count=1000 98 | ) 99 | 100 | keys.update([k.decode("utf-8") for k in partial_keys]) 101 | if cursor == 0: 102 | break 103 | 104 | return keys 105 | 106 | async def close(self): 107 | if self._client: 108 | await self._client.aclose() 109 | -------------------------------------------------------------------------------- /asimov/constants.py: -------------------------------------------------------------------------------- 1 | import enum 2 | 3 | ANTHROPIC_MODELS = [ 4 | "anthropic.claude-3-5-sonnet-20240620-v1:0", 5 | "anthropic.claude-3-5-sonnet-20241022-v2:0", 6 | "claude-3-5-sonnet-20241022", 7 | "claude-3-5-sonnet-20240620", 8 | "claude-3-5-haiku-20241022", 9 | ] 10 | 11 | 12 | OAI_REASONING_MODELS = ["o1-mini", "o1-preview"] 13 | 14 | OAI_GPT_MODELS = ["gpt-4o", "gpt-4o-turbo"] 15 | 16 | 17 | LLAMA_MODELS = ["hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4"] 18 | 19 | 20 | class ModelFamily(enum.Enum): 21 | OAI_REASONING = "OAI_REASONING" 22 | OAI_GPT = "OAI_GPT" 23 | ANTHROPIC = "ANTHROPIC" 24 | LLAMA = "LLAMA" 25 | -------------------------------------------------------------------------------- /asimov/data/postgres/manager.py: -------------------------------------------------------------------------------- 1 | import psycopg2 2 | from threading import Lock 3 | import contextlib 4 | from psycopg2.extras import RealDictCursor 5 | import psycopg2.pool 6 | import opentelemetry.instrumentation.psycopg2 7 | 8 | opentelemetry.instrumentation.psycopg2.Psycopg2Instrumentor().instrument() 9 | 10 | 11 | class DatabaseManager: 12 | _instances: dict[str, dict] = {} 13 | _lock = Lock() 14 | _initialized = False 15 | 16 | def __new__(cls, dsn: str): 17 | with cls._lock: 18 | if cls._instances.get(dsn) is None: 19 | cls._instances[dsn] = { 20 | "instance": super(DatabaseManager, cls).__new__(cls), 21 | "initialized": False, 22 | } 23 | return cls._instances[dsn]["instance"] 24 | 25 | def __init__(self, dsn: str): 26 | if not DatabaseManager._instances[dsn]["initialized"]: 27 | DatabaseManager._instances[dsn]["instance"].initialize(dsn) 28 | 29 | def initialize(self, dsn: str = ""): 30 | self.connection_pool = psycopg2.pool.ThreadedConnectionPool( 31 | minconn=1, 32 | maxconn=20, # Adjust based on your needs 33 | dsn=dsn, 34 | ) 35 | 36 | DatabaseManager._instances[dsn]["initialized"] = True 37 | 38 | @contextlib.contextmanager 39 | def get_connection(self): 40 | conn = self.connection_pool.getconn() 41 | try: 42 | yield conn 43 | finally: 44 | self.connection_pool.putconn(conn) 45 | 46 | @contextlib.contextmanager 47 | def get_cursor(self, commit=True): 48 | with self.get_connection() as conn: 49 | cursor = conn.cursor(cursor_factory=RealDictCursor) 50 | try: 51 | yield cursor 52 | finally: 53 | if commit: 54 | conn.commit() 55 | else: 56 | conn.rollback() 57 | cursor.close() 58 | 59 | def execute_query(self, query, params=None, cursor=None): 60 | with self.get_cursor() as cur: 61 | if cursor is not None: 62 | cur = cursor 63 | cur.execute(query, params) 64 | if cur.description: 65 | return cur.fetchall() 66 | return None 67 | 68 | def execute_and_fetch_one(self, query, params=None, cursor=None): 69 | with self.get_cursor() as cur: 70 | if cursor is not None: 71 | cur = cursor 72 | cur.execute(query, params) 73 | return cur.fetchone() 74 | 75 | def execute_and_return_id(self, query, params=None, cursor=None): 76 | with self.get_cursor() as cur: 77 | if cursor is not None: 78 | cur = cursor 79 | cur.execute(query, params) 80 | return cur.fetchone()["id"] 81 | -------------------------------------------------------------------------------- /asimov/executors/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BismuthCloud/asimov/ba2b9e28878406a9f81738718e7f094a5fbbdbe3/asimov/executors/__init__.py -------------------------------------------------------------------------------- /asimov/executors/executor.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import Any 3 | from pydantic import Field 4 | import asyncio 5 | 6 | from asimov.caches.cache import Cache 7 | from asimov.graph import AgentModule, ModuleType 8 | 9 | 10 | class Executor(AgentModule, ABC): 11 | type: ModuleType = Field(default=ModuleType.EXECUTOR) 12 | 13 | @abstractmethod 14 | async def process( 15 | self, cache: Cache, semaphore: asyncio.Semaphore 16 | ) -> dict[str, Any]: 17 | raise NotImplementedError 18 | -------------------------------------------------------------------------------- /asimov/graph/__init__.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | from enum import Enum 3 | import jsonpickle 4 | import os 5 | import pathlib 6 | import pickle 7 | import re 8 | import logging 9 | import opentelemetry.trace 10 | from typing import Awaitable, List, Dict, Any, AsyncGenerator, Optional, Sequence 11 | from typing_extensions import TypedDict 12 | from pydantic import ( 13 | Field, 14 | PrivateAttr, 15 | model_validator, 16 | ) 17 | 18 | from lupa import LuaRuntime 19 | from collections import defaultdict 20 | 21 | from asimov.graph.tasks import Task, TaskStatus 22 | from asimov.caches.cache import Cache 23 | from asimov.asimov_base import AsimovBase 24 | from asimov.caches.redis_cache import RedisCache 25 | 26 | lua = LuaRuntime(unpack_returned_tuples=True) # type: ignore[call-arg] 27 | 28 | tracer = opentelemetry.trace.get_tracer("asimov") 29 | 30 | 31 | class NonRetryableException(Exception): 32 | """ 33 | An exception which, when raised within a module, will prevent the module from being retried 34 | regardless of the retry_on_failure configuration. 35 | """ 36 | 37 | pass 38 | 39 | 40 | class ModuleType(Enum): 41 | EXECUTOR = "executor" 42 | SUBGRAPH = "subgraph" 43 | FLOW_CONTROL = "flow_control" 44 | 45 | 46 | class Middleware(AsimovBase): 47 | async def process(self, data: Dict[str, Any], cache: Cache) -> Dict[str, Any]: 48 | return data 49 | 50 | 51 | class ModuleConfig(AsimovBase): 52 | stop_on_success: bool = False 53 | middlewares: List[Middleware] = Field(default_factory=list) 54 | timeout: float = 60.0 55 | context: Dict[str, Any] = Field(default_factory=dict) 56 | 57 | 58 | class NodeConfig(AsimovBase): 59 | parallel: bool = False 60 | condition: Optional[str] = None 61 | retry_on_failure: bool = True 62 | retry_from: Optional[str] = None 63 | max_retries: int = 3 64 | max_visits: int = Field(default=5) 65 | inputs: List[str] = Field(default_factory=list) 66 | outputs: List[str] = Field(default_factory=list) 67 | 68 | 69 | class AgentModule(AsimovBase): 70 | name: str 71 | type: ModuleType 72 | config: ModuleConfig = Field(default_factory=ModuleConfig) 73 | dependencies: List[str] = Field(default_factory=list) 74 | input_mailboxes: List[str] = Field(default_factory=list) 75 | output_mailbox: str = Field(default="") 76 | container: Optional["CompositeModule"] = None 77 | executions: int = Field(default=0) 78 | trace: bool = Field(default=False) 79 | _generator: Optional[AsyncGenerator] = None 80 | 81 | async def run(self, cache: Cache, semaphore: asyncio.Semaphore) -> dict[str, Any]: 82 | try: 83 | if self._generator: 84 | output = await self._generator.__anext__() 85 | else: 86 | task = self.process(cache, semaphore) 87 | 88 | if isinstance(task, AsyncGenerator): 89 | self._generator = task 90 | 91 | output = await task.__anext__() 92 | else: 93 | output = await task 94 | except StopAsyncIteration as e: 95 | output = {"status": "success", "result": "Generator Exhausted"} 96 | except StopIteration as e: 97 | output = {"status": "success", "result": "Generator Exhausted"} 98 | 99 | if self.output_mailbox: 100 | await cache.publish_to_mailbox(self.output_mailbox, output) 101 | 102 | return output 103 | 104 | def process( 105 | self, cache: Cache, semaphore: asyncio.Semaphore 106 | ) -> Awaitable[dict[str, Any]] | AsyncGenerator: 107 | raise NotImplementedError 108 | 109 | def is_success(self, result: Dict[str, Any]) -> bool: 110 | return result.get("status") == "success" 111 | 112 | 113 | class FlowDecision(AsimovBase): 114 | next_node: str 115 | condition: Optional[str] = None 116 | metadata: Dict[str, Any] = Field(default_factory=dict) 117 | condition_variables: List[str] = Field(default_factory=list) 118 | cleanup_on_jump: bool = Field(default=False) 119 | 120 | 121 | class FlowControlConfig(AsimovBase): 122 | decisions: Sequence[FlowDecision] 123 | default: Optional[str] = None 124 | cleanup_on_default: bool = True 125 | 126 | 127 | class FlowControlModule(AgentModule): 128 | flow_config: FlowControlConfig 129 | 130 | async def run(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 131 | # Get keys once 132 | cache_keys = await cache.keys() 133 | 134 | for decision in self.flow_config.decisions: 135 | if decision.condition: 136 | if await self.evaluate_condition(decision.condition, cache, cache_keys): 137 | # Unset variables used in the condition 138 | for var in decision.condition_variables: 139 | await cache.delete(var) 140 | lua_globals = lua.globals() 141 | lua_globals[var] = None 142 | 143 | return { 144 | "status": "success", 145 | "cleanup": decision.cleanup_on_jump, 146 | "decision": decision.next_node, 147 | "metadata": decision.metadata, 148 | } 149 | else: 150 | # Unconditional jump 151 | return { 152 | "status": "success", 153 | "cleanup": decision.cleanup_on_jump, 154 | "decision": decision.next_node, 155 | "metadata": decision.metadata, 156 | } 157 | 158 | if self.flow_config.default: 159 | return { 160 | "status": "success", 161 | "decision": self.flow_config.default, 162 | "cleanup": self.flow_config.cleanup_on_default, 163 | "metadata": {}, 164 | } 165 | else: 166 | # If no decisions were met, fall through 167 | return { 168 | "status": "success", 169 | "decision": None, # Indicates fall-through 170 | "cleanup": True, 171 | "metadata": {}, 172 | } 173 | 174 | async def _apply_cache_affixes_condition( 175 | self, condition: str, cache: Cache, cache_keys: set[str] 176 | ) -> str: 177 | parts = condition.split(" ") 178 | 179 | try: 180 | keys = {k.split(cache.affix_sep)[1] for k in cache_keys} 181 | except IndexError: 182 | print("No affixes, returning raw keys.") 183 | keys = cache_keys 184 | 185 | new_parts = [] 186 | 187 | for part in parts: 188 | new_part = part 189 | 190 | if new_part in keys: 191 | new_part = await cache.apply_key_modifications(part) 192 | new_part = "v_" + new_part.replace(cache.affix_sep, "_") 193 | 194 | new_parts.append(new_part) 195 | 196 | return " ".join(new_parts) 197 | 198 | async def evaluate_condition( 199 | self, condition: str, cache: Cache, cache_keys: set[str] 200 | ) -> bool: 201 | # TODO: these are never cleaned up 202 | lua_globals = lua.globals() 203 | 204 | lua_vars = re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]*\b", condition) 205 | for orig_var in lua_vars: 206 | renamed_var = await cache.apply_key_modifications(orig_var) 207 | if renamed_var in cache_keys: 208 | lua_safe_key = "v_" + renamed_var.replace(cache.affix_sep, "_") 209 | lua_globals[lua_safe_key] = await cache.get(orig_var) 210 | 211 | modified_condition = await self._apply_cache_affixes_condition( 212 | condition, cache, cache_keys 213 | ) 214 | return lua.eval(modified_condition) 215 | 216 | 217 | class CompositeModule(AgentModule): 218 | modules: List[AgentModule] 219 | node_config: NodeConfig 220 | 221 | async def run(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 222 | if self.node_config.parallel: 223 | result = await self.run_parallel_modules(cache, semaphore) 224 | else: 225 | result = await self.run_sequential_modules(cache, semaphore) 226 | 227 | if self.output_mailbox: 228 | await cache.publish_to_mailbox(self.output_mailbox, result) 229 | 230 | return result 231 | 232 | async def apply_middlewares( 233 | self, middlewares: List[Middleware], data: Dict[str, Any], cache: Cache 234 | ) -> Dict[str, Any]: 235 | for middleware in middlewares: 236 | data = await middleware.process(data, cache) 237 | return data 238 | 239 | async def run_sequential_modules( 240 | self, cache: Cache, semaphore: asyncio.Semaphore 241 | ) -> Dict[str, Any]: 242 | results = [] 243 | for module in self.modules: 244 | module.container = self 245 | result = await self.run_module_with_concurrency_control( 246 | module, cache, semaphore 247 | ) 248 | if not self.is_success(result): 249 | return result 250 | 251 | results.append(result) 252 | 253 | if self.config.stop_on_success and self.is_success(result): 254 | break 255 | return {"status": "success", "results": results} 256 | 257 | async def run_parallel_modules( 258 | self, cache: Cache, semaphore: asyncio.Semaphore 259 | ) -> Dict[str, Any]: 260 | tasks = [] 261 | for module in self.modules: 262 | module.container = self 263 | task = asyncio.create_task( 264 | self.run_module_with_concurrency_control(module, cache, semaphore) 265 | ) 266 | tasks.append(task) 267 | 268 | results = await asyncio.gather(*tasks) 269 | for result in results: 270 | if not self.is_success(result): 271 | result["all_results"] = results 272 | return result 273 | 274 | return {"status": "success", "results": results} 275 | 276 | async def run_module_with_concurrency_control( 277 | self, 278 | module: AgentModule, 279 | cache: Cache, 280 | semaphore: asyncio.Semaphore, 281 | ) -> Dict[str, Any]: 282 | async with semaphore: 283 | try: 284 | async with asyncio.timeout(module.config.timeout): 285 | if module.trace: 286 | with tracer.start_as_current_span("run_module") as span: 287 | span.set_attribute("module_name", module.name) 288 | result = await module.run(cache, semaphore) 289 | else: 290 | result = await module.run(cache, semaphore) 291 | return await self.apply_middlewares( 292 | module.config.middlewares, result, cache 293 | ) 294 | except asyncio.TimeoutError: 295 | print( 296 | f"Module {module.name} execution timed out after {module.config.timeout} seconds" 297 | ) 298 | return { 299 | "status": "error", 300 | "error": f"Module {module.name} execution timed out after {module.config.timeout} seconds", 301 | "module": module.name, 302 | } 303 | 304 | 305 | class SnapshotControl(Enum): 306 | NEVER = "never" 307 | ONCE = "once" 308 | ALWAYS = "always" 309 | 310 | 311 | class Node(CompositeModule): 312 | _nodes: Dict[str, "Node"] = PrivateAttr() 313 | modules: List[AgentModule] = Field(default_factory=list) 314 | node_config: NodeConfig = Field(default_factory=NodeConfig) 315 | type: ModuleType = Field(default=ModuleType.SUBGRAPH) 316 | dependencies: List[str] = Field(default_factory=list) 317 | snapshot: SnapshotControl = Field(default=SnapshotControl.NEVER) 318 | 319 | def set_nodes(self, nodes: Dict[str, "Node"]): 320 | self._nodes = nodes 321 | 322 | def get_nodes(self): 323 | return self._nodes 324 | 325 | async def cancel_gen(self, agen): 326 | await agen.aclose() 327 | 328 | async def cleanup(self): 329 | for module in self.modules: 330 | if module._generator: 331 | await self.cancel_gen(module._generator) 332 | 333 | module._generator = None 334 | 335 | 336 | class ExecutionStep(TypedDict): 337 | executed: List[bool] 338 | nodes: List[str] 339 | skipped: List[bool] 340 | 341 | 342 | ExecutionPlan = List[ExecutionStep] 343 | 344 | 345 | class ExecutionState(AsimovBase): 346 | execution_index: int = Field(default=0) 347 | current_plan: ExecutionPlan = Field(default_factory=list) 348 | execution_history: List[ExecutionPlan] = Field(default_factory=list) 349 | total_iterations: int = Field(default=0) 350 | 351 | def mark_executed(self, step_index: int, node_index: int): 352 | if 0 <= step_index < len(self.current_plan): 353 | step = self.current_plan[step_index] 354 | if 0 <= node_index < len(step["nodes"]): 355 | step["executed"][node_index] = True 356 | 357 | def was_executed(self, node_name: str) -> bool: 358 | for step in self.current_plan: 359 | if node_name in step["nodes"]: 360 | index = step["nodes"].index(node_name) 361 | return step["executed"][index] 362 | return False 363 | 364 | def was_skipped(self, node_name: str) -> bool: 365 | for step in self.current_plan: 366 | if node_name in step["nodes"]: 367 | index = step["nodes"].index(node_name) 368 | return step["skipped"][index] 369 | return False 370 | 371 | 372 | def create_redis_cache(): 373 | return RedisCache() 374 | 375 | 376 | class Agent(AsimovBase): 377 | cache: Cache = Field(default_factory=create_redis_cache) 378 | nodes: Dict[str, Node] = Field(default_factory=dict) 379 | graph: Dict[str, List[str]] = Field(default_factory=dict) 380 | max_total_iterations: int = Field(default=100) 381 | max_concurrent_tasks: int = Field(default=5) 382 | _semaphore: asyncio.Semaphore = PrivateAttr() 383 | node_results: Dict[str, Any] = Field(default_factory=dict) 384 | output_mailbox: str = "agent_output" 385 | error_mailbox: str = "agent_error" 386 | execution_state: ExecutionState = Field(default_factory=ExecutionState) 387 | auto_snapshot: bool = False 388 | _logger: logging.Logger = PrivateAttr() 389 | _snapshot_generation: int = 0 390 | 391 | @model_validator(mode="after") 392 | def set_semaphore(self): 393 | self._semaphore = asyncio.Semaphore(self.max_concurrent_tasks) 394 | self._logger = logging.getLogger(__name__) 395 | 396 | return self 397 | 398 | def __getstate__(self) -> Dict[Any, Any]: 399 | # N.B. This does _not_ use pydantic's __getstate__ so unpickling results in a half-initialized object 400 | # This is fine for snapshot purposes though since we only need to copy some state out. 401 | state = self.__dict__.copy() 402 | del state["cache"] 403 | # Don't serialize nodes since those can contain unserializable things like callbacks 404 | del state["nodes"] 405 | return state 406 | 407 | def __setstate__(self, state: Dict[Any, Any]) -> None: 408 | self.__dict__.update(state) 409 | 410 | def is_success(self, result): 411 | return result["status"] == "success" 412 | 413 | def add_node(self, node: Node): 414 | if node.name in self.nodes: 415 | raise ValueError(f"Node with name {node.name} already exists") 416 | self.nodes[node.name] = node 417 | self.graph[node.name] = list(node.dependencies) 418 | 419 | node.set_nodes(self.nodes) 420 | 421 | def add_multiple_nodes(self, nodes: List[Node]): 422 | for node in nodes: 423 | self.add_node(node) 424 | 425 | async def run_task(self, task: Task) -> None: 426 | task.status = TaskStatus.EXECUTING 427 | 428 | self.execution_state.current_plan = self.compile_execution_plan() 429 | self.execution_state.execution_history.append(self.execution_state.current_plan) 430 | 431 | with tracer.start_as_current_span("run_task") as span: 432 | span.set_attribute("task_id", str(task.id)) 433 | await self._run_task(task) 434 | 435 | async def _run_task(self, task: Task) -> None: 436 | await self.cache.set("task", task) 437 | 438 | failed_chains: set[tuple[str, ...]] = set() 439 | node_visit_count: defaultdict[str, int] = defaultdict(int) 440 | 441 | while self.execution_state.current_plan: 442 | while self.execution_state.execution_index < len( 443 | self.execution_state.current_plan 444 | ): 445 | self.execution_state.total_iterations += 1 446 | if self.execution_state.total_iterations > self.max_total_iterations: 447 | self._logger.warning( 448 | f"Graph execution exceeded maximum total iterations ({self.max_total_iterations})" 449 | ) 450 | await self.cache.publish_to_mailbox( 451 | self.error_mailbox, 452 | { 453 | "status": "error", 454 | "error": f"Graph execution exceeded maximum total iterations ({self.max_total_iterations})", 455 | }, 456 | ) 457 | 458 | task.status = TaskStatus.FAILED 459 | await self.cache.publish_to_mailbox( 460 | self.output_mailbox, 461 | { 462 | "status": TaskStatus.FAILED, 463 | "failed_chains": list(failed_chains), 464 | "result": self.node_results, 465 | }, 466 | ) 467 | 468 | return 469 | 470 | parallel_group = self.execution_state.current_plan[ 471 | self.execution_state.execution_index 472 | ] 473 | 474 | if self.auto_snapshot and any( 475 | self.nodes[n].snapshot == SnapshotControl.ALWAYS 476 | or ( 477 | self.nodes[n].snapshot == SnapshotControl.ONCE 478 | and node_visit_count[n] == 0 479 | ) 480 | for n in parallel_group["nodes"] 481 | ): 482 | dir = await self.snapshot( 483 | task, 484 | "-".join( 485 | sorted( 486 | n 487 | for n in parallel_group["nodes"] 488 | if self.nodes[n].snapshot != SnapshotControl.NEVER 489 | ) 490 | ), 491 | ) 492 | self._logger.info(f"Snapshot: {dir}") 493 | 494 | tasks = [] 495 | 496 | # print(parallel_group["nodes"]) 497 | from pprint import pprint 498 | 499 | # pprint(self.execution_state.current_plan) 500 | 501 | for i, node_name in enumerate(parallel_group["nodes"]): 502 | if ( 503 | not self.execution_state.was_executed(node_name) 504 | and not self.execution_state.was_skipped(node_name) 505 | and not any( 506 | dep in chain 507 | for chain in failed_chains 508 | for dep in self.nodes[node_name].dependencies 509 | ) 510 | ): 511 | # print(node_name) 512 | node_visit_count[node_name] += 1 513 | if ( 514 | self.nodes[node_name].node_config.max_visits > 0 515 | and node_visit_count[node_name] 516 | > self.nodes[node_name].node_config.max_visits 517 | ): 518 | self._logger.warning( 519 | f"Node {node_name} exceeded maximum visits ({self.nodes[node_name].node_config.max_visits})" 520 | ) 521 | await self.cache.publish_to_mailbox( 522 | self.error_mailbox, 523 | { 524 | "status": "error", 525 | "error": f"Node {node_name} exceeded maximum visits ({self.nodes[node_name].node_config.max_visits})", 526 | }, 527 | ) 528 | failed_chains.add(self.get_dependent_chains(node_name)) 529 | else: 530 | tasks.append((i, self.run_node(node_name))) 531 | 532 | results = await asyncio.gather( 533 | *[task[1] for task in tasks], return_exceptions=True 534 | ) 535 | 536 | flow_control_executed = False 537 | new_nodes_added = False 538 | for (i, _), result in zip(tasks, results): 539 | node_name = parallel_group["nodes"][i] 540 | 541 | self.node_results[node_name] = result 542 | 543 | if not self.is_success(result): 544 | dependent_chain = self.get_dependent_chains(node_name) 545 | failed_chains.add(dependent_chain) 546 | self._logger.warning(f"Node '{node_name}' error '{result}'") 547 | await self.cache.publish_to_mailbox( 548 | self.error_mailbox, 549 | { 550 | "status": "error", 551 | "node": node_name, 552 | "error": str(result), 553 | }, 554 | ) 555 | continue 556 | 557 | self.execution_state.mark_executed( 558 | self.execution_state.execution_index, i 559 | ) 560 | 561 | if isinstance(result, dict) and "results" in result: 562 | for module_result in result["results"]: 563 | if isinstance(module_result, dict): 564 | if "new_nodes" in module_result: 565 | for new_node in module_result["new_nodes"]: 566 | self.add_node(new_node) 567 | self.update_dependencies( 568 | node_name, new_node.name 569 | ) 570 | new_nodes_added = True 571 | elif "decision" in module_result: 572 | next_node = module_result["decision"] 573 | if module_result["cleanup"]: 574 | # TODO This probably isn't entirely correct so look into it. 575 | for step in self.execution_state.current_plan: 576 | if next_node in step["nodes"]: 577 | await self.nodes[next_node].cleanup() 578 | break 579 | 580 | for node in step["nodes"]: 581 | await self.nodes[node].cleanup() 582 | 583 | if ( 584 | next_node is not None 585 | ): # A specific next node was chosen 586 | if next_node in self.nodes: 587 | if next_node not in failed_chains: 588 | new_plan = ( 589 | self.compile_execution_plan_from( 590 | next_node 591 | ) 592 | ) 593 | 594 | self.execution_state.current_plan = ( 595 | new_plan 596 | ) 597 | self.execution_state.execution_index = 0 598 | self.execution_state.execution_history.append( 599 | new_plan 600 | ) 601 | flow_control_executed = True 602 | else: 603 | self._logger.warning( 604 | f"Flow control attempted to jump to failed node: {next_node}" 605 | ) 606 | await self.cache.publish_to_mailbox( 607 | self.error_mailbox, 608 | { 609 | "status": "error", 610 | "error": f"Flow control attempted to jump to failed node: {next_node}", 611 | }, 612 | ) 613 | # Mark the chain containing this flow control node as failed 614 | failed_chains.add( 615 | self.get_dependent_chains(node_name) 616 | ) 617 | break 618 | else: 619 | self._logger.warning( 620 | f"Invalid next node: {next_node}" 621 | ) 622 | await self.cache.publish_to_mailbox( 623 | self.error_mailbox, 624 | { 625 | "status": "error", 626 | "error": f"Invalid next node: {next_node}", 627 | }, 628 | ) 629 | else: # Fall-through case 630 | flow_control_executed = False # Continue with the next node in the current plan 631 | 632 | if new_nodes_added: 633 | new_plan = self.compile_execution_plan_from( 634 | module_result["new_nodes"][0].name 635 | ) 636 | self.execution_state.current_plan = new_plan 637 | self.execution_state.execution_index = 0 638 | self.execution_state.execution_history.append(new_plan) 639 | elif not flow_control_executed: 640 | self.execution_state.execution_index += 1 641 | 642 | if self.execution_state.execution_index >= len( 643 | self.execution_state.current_plan 644 | ): 645 | # Current plan completed, prepare for the next plan if any 646 | self.execution_state.current_plan = [] 647 | self.execution_state.execution_index = 0 648 | 649 | if failed_chains: 650 | failed_count = 0 651 | for chain in list(failed_chains): 652 | failed_count += len(chain) 653 | 654 | if failed_count >= len(self.nodes): 655 | task.status = TaskStatus.FAILED 656 | await self.cache.publish_to_mailbox( 657 | self.output_mailbox, 658 | { 659 | "status": TaskStatus.FAILED, 660 | "failed_chains": list(failed_chains), 661 | "result": self.node_results, 662 | }, 663 | ) 664 | else: 665 | task.status = TaskStatus.PARTIAL 666 | await self.cache.publish_to_mailbox( 667 | self.output_mailbox, 668 | { 669 | "status": TaskStatus.PARTIAL, 670 | "failed_chains": list(failed_chains), 671 | "result": self.node_results, 672 | }, 673 | ) 674 | else: 675 | task.status = TaskStatus.COMPLETE 676 | await self.cache.publish_to_mailbox( 677 | self.output_mailbox, 678 | {"status": TaskStatus.COMPLETE, "result": self.node_results}, 679 | ) 680 | 681 | await self.cache.set("task", task) 682 | 683 | async def run_node(self, node_name: str) -> Dict[str, Any]: 684 | node = self.nodes[node_name] 685 | node.executions += 1 686 | retries = 0 687 | 688 | result = None 689 | last_exception = None 690 | 691 | while ( 692 | retries < node.node_config.max_retries and node.node_config.retry_on_failure 693 | ): 694 | try: 695 | if node.trace: 696 | with tracer.start_as_current_span("run_node") as span: 697 | span.set_attribute("node_name", node_name) 698 | span.set_attribute("retry", retries) 699 | result = await node.run( 700 | cache=self.cache, semaphore=self._semaphore 701 | ) 702 | span.set_attribute("success", self.is_success(result)) 703 | else: 704 | result = await node.run(cache=self.cache, semaphore=self._semaphore) 705 | 706 | if not self.is_success(result): 707 | retries += 1 708 | continue 709 | 710 | return result 711 | except Exception as e: 712 | last_exception = e 713 | 714 | self._logger.exception( 715 | f"Error in node {node_name} (attempt {retries + 1}/{node.node_config.max_retries})" 716 | ) 717 | if isinstance(e, NonRetryableException): 718 | break 719 | retries += 1 720 | 721 | print(f"Node {node_name} failed after {retries} attempts") 722 | 723 | if not result: 724 | result = {"status": "error", "result": str(last_exception)} 725 | 726 | return result 727 | 728 | def get_dependent_chains(self, node_name: str) -> tuple[str, ...]: 729 | dependent_chains: set[str] = set() 730 | dependent_chains.add(node_name) 731 | queue = [node_name] 732 | while queue: 733 | current = queue.pop(0) 734 | for name, node in self.nodes.items(): 735 | if current in node.dependencies: 736 | dependent_chains.add(name) 737 | queue.append(name) 738 | return tuple(list(dependent_chains)) 739 | 740 | def update_dependencies(self, parent_node: str, new_node: str): 741 | # Find all nodes that depend on the parent node 742 | dependent_nodes = [ 743 | node 744 | for node, deps in self.graph.items() 745 | if parent_node in deps and node != new_node 746 | ] 747 | 748 | # Update their dependencies to include the new node instead of the parent 749 | for node in dependent_nodes: 750 | # Update graph 751 | self.graph[node] = [ 752 | new_node if dep == parent_node else dep for dep in self.graph[node] 753 | ] 754 | 755 | # Update node dependencies 756 | self.nodes[node].dependencies = [ 757 | new_node if dep == parent_node else dep 758 | for dep in self.nodes[node].dependencies 759 | ] 760 | 761 | # Set the new node's dependency to the parent node 762 | self.graph[new_node] = [parent_node] 763 | self.nodes[new_node].dependencies = [parent_node] 764 | 765 | # Ensure the parent node is not dependent on the new node 766 | if new_node in self.graph[parent_node]: 767 | self.graph[parent_node].remove(new_node) 768 | if new_node in self.nodes[parent_node].dependencies: 769 | self.nodes[parent_node].dependencies.remove(new_node) 770 | 771 | def slice_before_node(self, execution_plan: List[dict], node: str) -> List[dict]: 772 | for index, step in enumerate(execution_plan): 773 | if node in step["nodes"]: 774 | return execution_plan[index:] 775 | return execution_plan # or return an empty list if the node isn't found 776 | 777 | def mark_node_and_deps_as_skipped(self, node_name: str): 778 | for step in self.execution_state.current_plan: 779 | if node_name in step["nodes"]: 780 | index = step["nodes"].index(node_name) 781 | step["skipped"][index] = True 782 | 783 | # Mark all dependencies as skipped 784 | deps_to_skip = set(self.get_dependent_chains(node_name)) 785 | for dep_step in self.execution_state.current_plan: 786 | for i, dep_node in enumerate(dep_step["nodes"]): 787 | if dep_node in deps_to_skip: 788 | dep_step["skipped"][i] = True 789 | 790 | break 791 | 792 | def _topological_sort(self, graph): 793 | visited = set() 794 | stack = [] 795 | 796 | def dfs(node): 797 | if node in visited: 798 | return 799 | visited.add(node) 800 | for neighbor in graph.get(node, []): 801 | dfs(neighbor) 802 | stack.append(node) 803 | 804 | for node in graph: 805 | if node not in visited: 806 | dfs(node) 807 | 808 | # Reverse the stack to get the correct topological order 809 | stack.reverse() 810 | 811 | # Group nodes that can be executed in parallel 812 | result = [] 813 | while stack: 814 | parallel_group = [] 815 | next_stack = [] 816 | for node in stack: 817 | if all(dep not in stack for dep in graph.get(node, [])): 818 | parallel_group.append(node) 819 | else: 820 | next_stack.append(node) 821 | result.append(parallel_group) 822 | stack = next_stack 823 | 824 | return result 825 | 826 | def compile_execution_plan(self) -> ExecutionPlan: 827 | plan = self._topological_sort(self.graph) 828 | return [ 829 | ExecutionStep( 830 | executed=[False] * len(group), nodes=group, skipped=[False] * len(group) 831 | ) 832 | for group in plan 833 | ] 834 | 835 | def compile_execution_plan_from(self, start_node: str) -> ExecutionPlan: 836 | full_plan = self.compile_execution_plan() 837 | start_index = next( 838 | (i for i, step in enumerate(full_plan) if start_node in step["nodes"]), None 839 | ) 840 | 841 | if start_index is None: 842 | return full_plan # Return the full plan if start_node is not found 843 | 844 | new_plan = full_plan[start_index:] 845 | 846 | # Mark nodes as skipped in the new plan 847 | for step_index, step in enumerate(new_plan): 848 | if start_node in step["nodes"]: 849 | # We've found the start node 850 | # Mark all previous nodes as skipped 851 | for prev_step in new_plan[:step_index]: 852 | prev_step["skipped"] = [True] * len(prev_step["nodes"]) 853 | 854 | # Mark siblings in this parallel group as skipped 855 | for i, node in enumerate(step["nodes"]): 856 | if node != start_node: 857 | step["skipped"][i] = True 858 | # Mark all dependents of this sibling as skipped 859 | dependents = self.get_dependent_chains(node) 860 | for future_step in new_plan[step_index + 1 :]: 861 | for j, future_node in enumerate(future_step["nodes"]): 862 | if future_node in dependents: 863 | future_step["skipped"][j] = True 864 | 865 | # We unmark as skipped any shared dependents of the node we're jumping to. 866 | # Users have to be aware of this fact in docs as it can footgun them. 867 | dependents = self.get_dependent_chains(start_node) 868 | for future_step in new_plan[step_index + 1 :]: 869 | for j, future_node in enumerate(future_step["nodes"]): 870 | if future_node in dependents: 871 | future_step["skipped"][j] = False 872 | # No need to process further steps 873 | return new_plan 874 | 875 | return new_plan 876 | 877 | def _is_reachable(self, start: str, end: str) -> bool: 878 | visited = set() 879 | stack = [start] 880 | while stack: 881 | node = stack.pop() 882 | if node not in visited: 883 | if node == end: 884 | return True 885 | visited.add(node) 886 | stack.extend(self.graph[node]) 887 | return False 888 | 889 | async def snapshot(self, task: Task, name_hint: str) -> pathlib.Path: 890 | out_dir = pathlib.Path( 891 | os.environ.get("ASIMOV_SNAPSHOT", "/tmp/asimov_snapshot") 892 | ) 893 | out_dir = out_dir / str(task.id) 894 | out_dir = out_dir / f"{self._snapshot_generation}_{name_hint}" 895 | self._snapshot_generation += 1 896 | out_dir.mkdir(parents=True, exist_ok=True) 897 | with open(out_dir / "agent.pkl", "wb") as f: 898 | pickle.dump(self, f) 899 | with open(out_dir / "cache.json", "w") as f: 900 | f.write(jsonpickle.encode(await self.cache.get_all())) 901 | with open(out_dir / "task.pkl", "wb") as f: 902 | pickle.dump(task, f) 903 | return out_dir 904 | 905 | async def run_from_snapshot(self, snapshot_dir: pathlib.Path): 906 | with open(snapshot_dir / "agent.pkl", "rb") as f: 907 | agent = pickle.load(f) 908 | self.graph = agent.graph 909 | self.max_total_iterations = agent.max_total_iterations 910 | self.max_concurrent_tasks = agent.max_concurrent_tasks 911 | self.node_results = agent.node_results 912 | self.output_mailbox = agent.output_mailbox 913 | self.error_mailbox = agent.error_mailbox 914 | self.execution_state = agent.execution_state 915 | with open(snapshot_dir / "cache.json", "r") as f: 916 | await self.cache.clear() 917 | cache = jsonpickle.decode(f.read()) 918 | for k, v in cache.items(): 919 | # get_all returns keys with prefixes, so set with raw=True 920 | await self.cache.set(k, v, raw=True) 921 | with open(snapshot_dir / "task.pkl", "rb") as f: 922 | task = pickle.load(f) 923 | await self._run_task(task) 924 | 925 | # Need to build after CompositeModule is defined 926 | AgentModule.model_rebuild() 927 | -------------------------------------------------------------------------------- /asimov/graph/agent_directed_flow.py: -------------------------------------------------------------------------------- 1 | import textwrap 2 | import logging 3 | from typing import List, Dict, Any, Sequence, Optional 4 | from pydantic import Field, PrivateAttr, model_validator 5 | import json 6 | import asyncio 7 | 8 | from asimov.caches.cache import Cache 9 | from asimov.graph import ( 10 | FlowControlModule, 11 | FlowDecision, 12 | FlowControlConfig, 13 | CompositeModule, # noqa 14 | Node, # noqa 15 | ) 16 | from asimov.services.inference_clients import InferenceClient, ChatMessage, ChatRole 17 | 18 | MAKE_DECISION_SCHEMA = { 19 | "name": "make_decision", 20 | "description": "Make decision about what choice to make from the provided options in the prompt.", 21 | "input_schema": { 22 | "type": "object", 23 | "properties": { 24 | "thoughts": { 25 | "type": "string", 26 | "description": "Your thoughts about why you are making the decsiion you are mkaing regarding the possible choices available.", 27 | }, 28 | "decision": { 29 | "type": "string", 30 | "description": "A phrase to search for exact case sensitive matches over the codebase.", 31 | }, 32 | }, 33 | "required": ["decision"], 34 | }, 35 | } 36 | 37 | 38 | class Example: 39 | def __init__( 40 | self, 41 | message: str, 42 | choices: List[Dict[str, str]], 43 | choice: str, 44 | reasoning: str = "", 45 | ): 46 | self.message = message 47 | self.choices = choices 48 | self.choice = choice 49 | self.reasoning = reasoning 50 | 51 | 52 | class AgentDrivenFlowDecision(FlowDecision): 53 | examples: List[Example] = Field(default_factory=list) 54 | 55 | 56 | class AgentDrivenFlowControlConfig(FlowControlConfig): # type: ignore[override] 57 | decisions: Sequence[AgentDrivenFlowDecision] 58 | default: Optional[str] = None 59 | 60 | 61 | class AgentDirectedFlowControl(FlowControlModule): # type: ignore[override] 62 | inference_client: InferenceClient 63 | system_description: str 64 | flow_config: AgentDrivenFlowControlConfig 65 | input_var: str = Field(default="input_message") 66 | use_historicals: bool = Field(default=True) 67 | voters: int = 1 68 | _prompt: str = PrivateAttr() 69 | _logger: logging.Logger = PrivateAttr() 70 | 71 | @model_validator(mode="after") 72 | def set_attrs(self): 73 | examples = [] 74 | 75 | for decision in self.flow_config.decisions: 76 | prompt = "" 77 | for _, example in enumerate(decision.examples, 1): 78 | prompt += textwrap.dedent( 79 | f""" 80 | Message: {example.message} 81 | Choices: 82 | {json.dumps(example.choices, indent=2)} 83 | Choice: {example.choice} 84 | """ 85 | ).strip() 86 | 87 | if example.reasoning != "": 88 | prompt += f"\nReasoning: {example.reasoning}" 89 | 90 | examples.append(prompt) 91 | 92 | self._prompt = f"""You are going to be provided with some input for a user, you need to decide based on the available options which will also be provided what option best fits the message the user sent. 93 | Users may ask for tasks to be completed and not every task may be possible given the available options. You'll be provided examples between and and the choices between and 94 | You may also be provided with context, such as past messages, to help you make your decision. If this context is available it will be between and Your response must be valid json. 95 | 96 | Description of the system: {self.system_description} 97 | 98 | 99 | {"\n".join(examples)} 100 | """ 101 | 102 | return self 103 | 104 | def most_common_string(self, string_list): 105 | return max(set(string_list), key=string_list.count) 106 | 107 | async def gen(self, generation_input, cache): 108 | async def make_decision(resp): 109 | print(dict(resp)) 110 | async with cache.with_suffix(self.name): 111 | votes = await cache.get("votes", []) 112 | votes.append(resp["decision"].lower()) 113 | 114 | await cache.set("votes", votes) 115 | 116 | return dict(resp) 117 | 118 | await self.inference_client.tool_chain( 119 | messages=generation_input, 120 | top_p=0.9, 121 | tool_choice="any", 122 | temperature=0, 123 | max_iterations=1, 124 | tools=[ 125 | (make_decision, MAKE_DECISION_SCHEMA), 126 | ], 127 | ) 128 | 129 | async def run(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 130 | self._logger = logging.getLogger(__name__).getChild( 131 | await cache.get("request_id") 132 | ) 133 | 134 | message = await cache.get(self.input_var) 135 | 136 | choices = [] 137 | for decision in self.flow_config.decisions: 138 | choice = { 139 | "choice": decision.next_node, 140 | "description": decision.metadata.get("description"), 141 | } 142 | choices.append(choice) 143 | 144 | prompt = self._prompt 145 | 146 | if self.use_historicals: 147 | context = await cache.get("formatted_chat_messages", []) 148 | context_string = "\n".join([ctx.model_dump_json() for ctx in context]) 149 | prompt = ( 150 | self._prompt 151 | + textwrap.dedent( 152 | f""" 153 | 154 | {context_string} 155 | 156 | """ 157 | ).strip() 158 | ) 159 | 160 | prompt = ( 161 | prompt 162 | + "\n" 163 | + textwrap.dedent( 164 | f""" 165 | 166 | {json.dumps(choices, indent=2)} 167 | 168 | """ 169 | ).strip() 170 | ) 171 | 172 | generation_input = [ 173 | ChatMessage( 174 | role=ChatRole.SYSTEM, 175 | content=prompt, 176 | ), 177 | ChatMessage( 178 | role=ChatRole.USER, 179 | content=message, 180 | ), 181 | ] 182 | 183 | tasks = [] 184 | for _ in range(0, self.voters): 185 | tasks.append(self.gen(generation_input, cache)) 186 | 187 | await asyncio.gather(*tasks) 188 | 189 | async with cache.with_suffix(self.name): 190 | votes = await cache.get("votes", []) 191 | 192 | print(votes) 193 | if votes: 194 | voted_choice = self.most_common_string(votes) 195 | 196 | for decision in self.flow_config.decisions: 197 | if decision.next_node.lower() == voted_choice.lower(): 198 | self._logger.info(f"agent directed flow decision: {decision.next_node}") 199 | return { 200 | "status": "success", 201 | "cleanup": decision.cleanup_on_jump, 202 | "decision": decision.next_node, 203 | "metadata": decision.metadata, 204 | } 205 | 206 | if self.flow_config.default: 207 | self._logger.info(f"Using default decision: {self.flow_config.default}") 208 | return { 209 | "status": "success", 210 | "cleanup": True, 211 | "decision": self.flow_config.default, 212 | "metadata": {}, 213 | } 214 | 215 | # If no decisions were met, fall through 216 | self._logger.warning("No agent directed flow decision was met, falling through") 217 | return { 218 | "status": "success", 219 | "decision": None, # Indicates fall-through 220 | "cleanup": True, 221 | "metadata": {}, 222 | } 223 | -------------------------------------------------------------------------------- /asimov/graph/tasks.py: -------------------------------------------------------------------------------- 1 | from enum import Enum 2 | from uuid import uuid4, UUID 3 | from typing import Dict, Any, Optional 4 | from asimov.asimov_base import AsimovBase 5 | from pydantic import ConfigDict, Field 6 | 7 | 8 | class TaskStatus(Enum): 9 | WAITING = "waiting" 10 | EXECUTING = "executing" 11 | COMPLETE = "complete" 12 | FAILED = "failed" 13 | PARTIAL = "partial_failure" 14 | 15 | 16 | class Task(AsimovBase): 17 | id: UUID = Field(default_factory=uuid4) 18 | type: str 19 | objective: str 20 | params: Dict[str, Any] = Field(default_factory=dict) 21 | status: TaskStatus = TaskStatus.WAITING 22 | result: Optional[Any] = None 23 | error: Optional[str] = None 24 | model_config = ConfigDict(arbitrary_types_allowed=True) 25 | 26 | def update_status(self, status: TaskStatus) -> None: 27 | self.status = status 28 | 29 | def set_result(self, result: Any) -> None: 30 | self.result = result 31 | self.status = TaskStatus.COMPLETE 32 | 33 | def set_error(self, error: str) -> None: 34 | self.error = error 35 | self.status = TaskStatus.FAILED 36 | -------------------------------------------------------------------------------- /asimov/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BismuthCloud/asimov/ba2b9e28878406a9f81738718e7f094a5fbbdbe3/asimov/py.typed -------------------------------------------------------------------------------- /asimov/utils/models.py: -------------------------------------------------------------------------------- 1 | from asimov.constants import ( 2 | ModelFamily, 3 | OAI_REASONING_MODELS, 4 | OAI_GPT_MODELS, 5 | ANTHROPIC_MODELS, 6 | LLAMA_MODELS, 7 | ) 8 | from asimov.services.inference_clients import InferenceClient, ChatMessage 9 | 10 | 11 | def get_model_family(client: InferenceClient): 12 | if client.model in OAI_GPT_MODELS: 13 | return ModelFamily.OAI_GPT 14 | elif client.model in OAI_REASONING_MODELS: 15 | return ModelFamily.OAI_REASONING 16 | elif client.model in ANTHROPIC_MODELS: 17 | return ModelFamily.ANTHROPIC 18 | elif client.model in LLAMA_MODELS: 19 | return ModelFamily.LLAMA 20 | else: 21 | return None 22 | 23 | 24 | def is_model_family(client: InferenceClient, families: list[ModelFamily]) -> bool: 25 | model_family = get_model_family(client) 26 | 27 | return model_family in families 28 | 29 | 30 | def prepare_model_generation_input( 31 | messages: list[ChatMessage], client: InferenceClient 32 | ): 33 | final_output = [] 34 | client_model_family = get_model_family(client) 35 | 36 | for message in messages: 37 | if not message.model_families or client_model_family in message.model_families: 38 | final_output.append(message) 39 | 40 | return final_output 41 | -------------------------------------------------------------------------------- /asimov/utils/token_buffer.py: -------------------------------------------------------------------------------- 1 | import threading 2 | import logging 3 | 4 | 5 | class TokenBuffer: 6 | def __init__(self, size): 7 | self.size = size 8 | self.buffer = [None] * size 9 | self.current_index = 0 10 | self.mutex = threading.Lock() 11 | self.logger = logging.getLogger(self.__class__.__name__) 12 | 13 | def add_locking(self, token): 14 | with self.mutex: 15 | self.buffer[self.current_index] = token 16 | self.current_index = (self.current_index + 1) % self.size 17 | 18 | def clear_locking(self): 19 | with self.mutex: 20 | self.buffer = [None] * self.size 21 | 22 | def __str__(self): 23 | return "".join( 24 | filter( 25 | None, 26 | self.buffer[self.current_index :] + self.buffer[: self.current_index], 27 | ) 28 | ) 29 | 30 | def get_current_buffer(self): 31 | return self.buffer 32 | -------------------------------------------------------------------------------- /asimov/utils/token_counter.py: -------------------------------------------------------------------------------- 1 | # utils/approx_token_count.py 2 | from typing import Any, Dict, List 3 | import math 4 | 5 | AVG_CHARS_PER_TOKEN = 4 # heuristic—you can tweak if your data skews long/short 6 | TOKENS_PER_MSG = 4 # ChatML fixed overhead (role, separators, etc.) 7 | TOKENS_PER_NAME = -1 # spec quirk: “name” field shaves one token 8 | END_OF_REQ_TOKENS = 2 # every request implicitly ends with: 9 | 10 | def approx_tokens_from_serialized_messages( 11 | serialized_messages: List[Dict[str, Any]], 12 | avg_chars_per_token: int = AVG_CHARS_PER_TOKEN, 13 | ) -> int: 14 | """ 15 | Fast, model-agnostic token estimate for a ChatML message array. 16 | 17 | Parameters 18 | ---------- 19 | serialized_messages : list[dict] 20 | Your [{role, content:[{type,text}]}] structure. 21 | avg_chars_per_token : int, optional 22 | How many characters you assume map to one token (default 4). 23 | 24 | Returns 25 | ------- 26 | int 27 | Estimated prompt token count. 28 | """ 29 | total_tokens = 0 30 | 31 | for msg in serialized_messages: 32 | total_tokens += TOKENS_PER_MSG 33 | 34 | # role string itself 35 | total_tokens += math.ceil(len(msg["role"]) / avg_chars_per_token) 36 | 37 | if "name" in msg: 38 | total_tokens += TOKENS_PER_NAME 39 | 40 | for part in msg["content"]: 41 | if part["type"] == "text": 42 | total_tokens += math.ceil(len(part["text"]) / avg_chars_per_token) 43 | else: 44 | # non-text parts: fall back to raw length heuristic 45 | total_tokens += math.ceil(len(str(part)) / avg_chars_per_token) 46 | 47 | total_tokens += END_OF_REQ_TOKENS 48 | return max(total_tokens, 0) 49 | -------------------------------------------------------------------------------- /asimov/utils/visualize.py: -------------------------------------------------------------------------------- 1 | import graphviz 2 | 3 | from asimov.graph import Agent, FlowControlModule 4 | 5 | 6 | def create_agent_graph(agent: Agent, output_file="agent_graph"): 7 | """ 8 | Creates a DOT graph visualization of an Agent's nodes and their relationships. 9 | 10 | Args: 11 | agent: The Agent object containing nodes and their relationships 12 | output_file: The name of the output file (without extension) 13 | 14 | Returns: 15 | The path to the generated DOT file 16 | """ 17 | # Create a new directed graph 18 | dot = graphviz.Digraph(comment="Agent Node Graph") 19 | dot.attr(rankdir="LR") # Left to right layout 20 | 21 | # Add all nodes first 22 | for node in agent.nodes.values(): 23 | if node.modules: 24 | with dot.subgraph(name=f"cluster_{node.name}") as sub: 25 | sub.attr(label=node.name) 26 | for mod in node.modules: 27 | sub.node(f"{node.name}__{mod.name}", label=mod.name) 28 | for a, b in zip(node.modules, node.modules[1:]): 29 | sub.edge(f"{node.name}__{a.name}", f"{node.name}__{b.name}") 30 | else: 31 | dot.node(node.name) 32 | 33 | # Add edges for dependencies 34 | for node in agent.nodes.values(): 35 | # Add dependency edges 36 | for dep in node.dependencies: 37 | dot.edge( 38 | ( 39 | f"{dep}__{agent.nodes[dep].modules[-1].name}" 40 | if agent.nodes[dep].modules 41 | else dep 42 | ), 43 | f"{node.name}__{node.modules[0].name}" if node.modules else node.name, 44 | ) 45 | 46 | # Add flow control edges 47 | for mod in node.modules: 48 | if isinstance(mod, FlowControlModule): 49 | for decision in mod.flow_config.decisions: 50 | dot.edge( 51 | f"{node.name}__{mod.name}", 52 | ( 53 | f"{decision.next_node}__{agent.nodes[decision.next_node].modules[0].name}" 54 | if agent.nodes[decision.next_node].modules 55 | else decision.next_node 56 | ), 57 | label=decision.condition, 58 | ) 59 | if mod.flow_config.default: 60 | dot.edge( 61 | f"{node.name}__{mod.name}", 62 | ( 63 | f"{mod.flow_config.default}__{agent.nodes[mod.flow_config.default].modules[0].name}" 64 | if agent.nodes[mod.flow_config.default].modules 65 | else mod.flow_config.default 66 | ), 67 | color="blue", 68 | label="default", 69 | ) 70 | 71 | # Save the graph 72 | try: 73 | # Save as both .dot and rendered format 74 | dot.save(output_file) 75 | # Also render as PNG for visualization 76 | dot.render(output_file, format="png") 77 | return f"{output_file}.dot" 78 | except Exception as e: 79 | print(f"Error saving graph: {e}") 80 | return None 81 | -------------------------------------------------------------------------------- /bismuth.toml: -------------------------------------------------------------------------------- 1 | [chat] 2 | # Timeout in seconds for commands run by the agent. 3 | # If omitted, defaults to 60. 4 | command_timeout = 320 5 | -------------------------------------------------------------------------------- /docs/basic_agent.md: -------------------------------------------------------------------------------- 1 | # Building a Basic Agent 2 | 3 | This walkthrough will guide you through creating a simple agent using the Asimov Agents framework. We'll build a text processing agent that demonstrates the core concepts of the framework. 4 | 5 | ## Prerequisites 6 | 7 | - Python 3.12+ 8 | - Redis server running (see Redis Setup below) 9 | - Asimov Agents package installed 10 | 11 | ## Redis Setup with Docker 12 | 13 | To run Redis using Docker: 14 | 15 | ```bash 16 | # Pull the official Redis image 17 | docker pull redis:latest 18 | 19 | # Run Redis container 20 | docker run --name asimov-redis -d -p 6379:6379 redis:latest 21 | 22 | # Verify Redis is running 23 | docker ps | grep asimov-redis 24 | ``` 25 | 26 | This will start Redis on the default port 6379. The container will run in the background and restart automatically unless explicitly stopped. 27 | 28 | ## Concepts Covered 29 | 30 | 1. Agent Module Types 31 | 2. Node Configuration 32 | 3. Flow Control 33 | 4. Task Management 34 | 5. Cache Usage 35 | 36 | ## Step-by-Step Guide 37 | 38 | ### 1. Setting Up the Project Structure 39 | 40 | Create a new Python file `basic_agent.py` with the following imports: 41 | 42 | ```python 43 | import asyncio 44 | from typing import Dict, Any 45 | from asimov.graph import ( 46 | Agent, 47 | AgentModule, 48 | ModuleType, 49 | Node, 50 | NodeConfig, 51 | FlowControlModule, 52 | FlowControlConfig, 53 | FlowDecision, 54 | ) 55 | from asimov.graph.tasks import Task 56 | from asimov.caches.redis_cache import RedisCache 57 | ``` 58 | 59 | ### 2. Creating the Planning Executor Module 60 | 61 | The planning executor module is responsible for analyzing the task and creating a plan: 62 | 63 | ```python 64 | class TextPlannerModule(AgentModule): 65 | """Plans text processing operations.""" 66 | 67 | name = "text_planner" 68 | type = ModuleType.EXECUTOR 69 | 70 | async def process(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 71 | # Get the task from cache 72 | task = await cache.get("task") 73 | text = task.params.get("text", "") 74 | 75 | # Create a simple plan 76 | plan = { 77 | "operations": [ 78 | {"type": "count_words", "text": text}, 79 | {"type": "calculate_stats", "text": text} 80 | ] 81 | } 82 | 83 | # Store the plan in cache 84 | await cache.set("plan", plan) 85 | 86 | return { 87 | "status": "success", 88 | "result": "Plan created successfully" 89 | } 90 | ``` 91 | 92 | Key points: 93 | - The module inherits from `AgentModule` 94 | - It has a unique name and is an EXECUTOR type 95 | - The `process` method contains the core logic 96 | - Uses cache for state management 97 | - Returns a standardized response format 98 | 99 | ### 3. Creating the Executor Module 100 | 101 | The executor module implements the actual processing logic: 102 | 103 | ```python 104 | class TextExecutorModule(AgentModule): 105 | """Executes text processing operations.""" 106 | 107 | name = "text_executor" 108 | type = ModuleType.EXECUTOR 109 | 110 | async def process(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 111 | # Get the plan and task 112 | plan = await cache.get("plan") 113 | task = await cache.get("task") 114 | 115 | results = [] 116 | for operation in plan["operations"]: 117 | if operation["type"] == "count_words": 118 | word_count = len(operation["text"].split()) 119 | results.append({ 120 | "operation": "count_words", 121 | "result": word_count 122 | }) 123 | elif operation["type"] == "calculate_stats": 124 | char_count = len(operation["text"]) 125 | line_count = len(operation["text"].splitlines()) 126 | results.append({ 127 | "operation": "calculate_stats", 128 | "result": { 129 | "characters": char_count, 130 | "lines": line_count 131 | } 132 | }) 133 | 134 | return { 135 | "status": "success", 136 | "result": results 137 | } 138 | ``` 139 | 140 | Key points: 141 | - Implements specific processing operations 142 | - Retrieves state from cache 143 | - Processes each operation in the plan 144 | - Returns structured results 145 | 146 | ### 4. Setting Up Flow Control 147 | 148 | Configure how the agent moves between nodes: 149 | 150 | ```python 151 | flow_control = Node( 152 | name="flow_control", 153 | type=ModuleType.FLOW_CONTROL, 154 | modules=[FlowControlModule( 155 | flow_config=FlowControlConfig( 156 | decisions=[ 157 | FlowDecision( 158 | next_node="executor", 159 | condition="plan ~= null" # Conditions are lua. 160 | ) 161 | ], 162 | default="planner" 163 | ) 164 | )] 165 | ) 166 | ``` 167 | 168 | Key points: 169 | - Uses conditions to determine execution flow 170 | - Provides a default node 171 | - Can be extended with multiple decision paths 172 | 173 | ### 5. Creating and Configuring the Agent 174 | 175 | Put everything together: 176 | 177 | ```python 178 | async def main(): 179 | # Create the agent 180 | agent = Agent( 181 | cache=RedisCache(), 182 | max_concurrent_tasks=1, 183 | max_total_iterations=10 184 | ) 185 | 186 | # Create nodes 187 | planner_node = Node( 188 | name="planner", 189 | type=ModuleType.EXECUTOR, 190 | modules=[TextPlannerModule()], 191 | node_config=NodeConfig( 192 | parallel=False, 193 | max_retries=3 194 | ) 195 | ) 196 | 197 | executor_node = Node( 198 | name="executor", 199 | type=ModuleType.EXECUTOR, 200 | modules=[TextExecutorModule()], 201 | dependencies=["planner"] 202 | ) 203 | 204 | flow_control = Node( 205 | name="flow_control", 206 | type=ModuleType.FLOW_CONTROL, 207 | modules=[FlowControlModule( 208 | flow_config=FlowControlConfig( 209 | decisions=[ 210 | FlowDecision( 211 | next_node="executor", 212 | condition="plan ~= null" # Conditions are lua. 213 | ) 214 | ], 215 | default="planner" 216 | ) 217 | )] 218 | ) 219 | 220 | # Add nodes to agent 221 | agent.add_multiple_nodes([planner_node, executor_node, flow_control]) 222 | 223 | # Create and run a task 224 | task = Task( 225 | type="text_processing", 226 | objective="Process sample text", 227 | params={ 228 | "text": "Hello world!\nThis is a sample text." 229 | } 230 | ) 231 | 232 | # Run the task 233 | await agent.run_task(task) 234 | 235 | # Get the results 236 | results = agent.node_results.get("executor", {}).get("result", []) 237 | print("\nProcessing Results:") 238 | for result in results: 239 | print(f"\nOperation: {result['operation']}") 240 | print(f"Result: {result['result']}") 241 | 242 | if __name__ == "__main__": 243 | asyncio.run(main()) 244 | ``` 245 | 246 | Key points: 247 | - Configure agent with cache and execution limits 248 | - Create and configure nodes 249 | - Set up dependencies between nodes 250 | - Create and run a task 251 | - Handle results 252 | 253 | ## Understanding the Flow 254 | 255 | 1. The agent starts with the task in the planner node 256 | 2. The planner creates a processing plan 257 | 3. Flow control checks if a plan exists 258 | 4. If a plan exists, execution moves to the executor node 259 | 5. The executor processes the text according to the plan 260 | 6. Results are stored and can be retrieved from node_results 261 | 262 | ## Common Patterns 263 | 264 | 1. **State Management** 265 | - Use cache for sharing state between modules 266 | - Store intermediate results and plans 267 | - Track progress and status 268 | 269 | 2. **Module Communication** 270 | - Modules communicate through the cache 271 | - Use standardized response formats 272 | - Handle errors and status consistently 273 | 274 | 3. **Flow Control** 275 | - Use conditions to control execution flow 276 | - Provide default paths 277 | - Handle edge cases and errors 278 | 279 | 4. **Error Handling** 280 | - Use retry mechanisms 281 | - Return appropriate status codes 282 | - Include error details in responses 283 | 284 | ## Next Steps 285 | 286 | 1. Add more complex processing operations 287 | 2. Implement error handling and recovery 288 | 3. Add monitoring and logging 289 | 4. Explore parallel processing capabilities 290 | 5. Integrate with external services 291 | 292 | ## Best Practices 293 | 294 | 1. **Module Design** 295 | - Keep modules focused and single-purpose 296 | - Use clear naming conventions 297 | - Document module behavior and requirements 298 | 299 | 2. **State Management** 300 | - Use meaningful cache keys 301 | - Clean up temporary state 302 | - Handle cache failures gracefully 303 | 304 | 3. **Flow Control** 305 | - Keep conditions simple and clear 306 | - Plan for all possible paths 307 | - Use meaningful node names 308 | 309 | 4. **Testing** 310 | - Test modules in isolation 311 | - Test different execution paths 312 | - Verify error handling 313 | - Test with different input types and sizes 314 | -------------------------------------------------------------------------------- /docs/llm_agent.md: -------------------------------------------------------------------------------- 1 | # Building an LLM-Based Agent 2 | 3 | This walkthrough will guide you through creating an agent that uses Large Language Models (LLMs) for task planning and execution. This type of agent is particularly useful for complex tasks that require natural language understanding and generation. 4 | 5 | ## Prerequisites 6 | 7 | - Python 3.12+ 8 | - Redis server running (see Redis Setup below) 9 | - Asimov Agents package installed 10 | - API key for an LLM provider (e.g., Anthropic, AWS Bedrock) set in the ANTHROPIC_API_KEY environment variable (see Environment Setup below) 11 | 12 | ## Environment Setup 13 | 14 | ### API Key Configuration 15 | 16 | Before running the agent, you need to set up your Anthropic API key as an environment variable: 17 | 18 | ```bash 19 | # Linux/macOS 20 | export ANTHROPIC_API_KEY="your-api-key" 21 | 22 | # Windows (Command Prompt) 23 | set ANTHROPIC_API_KEY=your-api-key 24 | 25 | # Windows (PowerShell) 26 | $env:ANTHROPIC_API_KEY="your-api-key" 27 | ``` 28 | 29 | The agent will automatically use this environment variable for authentication with the Anthropic API. 30 | 31 | > **Security Note**: Never hardcode API keys in your source code. Always use environment variables or secure configuration management systems to handle sensitive credentials. This helps prevent accidental exposure of API keys in version control systems or logs. 32 | 33 | ### Redis Setup with Docker 34 | 35 | To run Redis using Docker: 36 | 37 | ```bash 38 | # Pull the official Redis image 39 | docker pull redis:latest 40 | 41 | # Run Redis container 42 | docker run --name asimov-redis -d -p 6379:6379 redis:latest 43 | 44 | # Verify Redis is running 45 | docker ps | grep asimov-redis 46 | ``` 47 | 48 | This will start Redis on the default port 6379. The container will run in the background and restart automatically unless explicitly stopped. 49 | 50 | ## Concepts Covered 51 | 52 | 1. LLM Integration 53 | 2. Complex Task Planning 54 | 3. Dynamic Execution 55 | 4. Result Validation 56 | 5. Decision Making with LLMs 57 | 58 | ## Step-by-Step Guide 59 | 60 | ### 1. Setting Up the Project 61 | 62 | Create a new Python file `llm_agent.py` and import the necessary modules: 63 | 64 | ```python 65 | import os 66 | import asyncio 67 | from typing import Dict, Any, List 68 | from asimov.graph import ( 69 | Agent, 70 | AgentModule, 71 | ModuleType, 72 | Node, 73 | NodeConfig, 74 | FlowControlModule, 75 | FlowControlConfig, 76 | FlowDecision, 77 | ) 78 | from asimov.graph.tasks import Task 79 | from asimov.caches.redis_cache import RedisCache 80 | from asimov.services.inference_clients import AnthropicInferenceClient 81 | ``` 82 | 83 | ### 2. Creating the Planning Executor Module 84 | 85 | The planning executor module uses an LLM to analyze tasks and create execution plans: 86 | 87 | ```python 88 | class LLMPlannerModule(AgentModule): 89 | """Uses LLM to plan task execution.""" 90 | 91 | name: str = "llm_planner" 92 | type: ModuleType = ModuleType.EXECUTOR 93 | 94 | client: AnthropicInferenceClient = None 95 | 96 | def __init__(self): 97 | super().__init__() 98 | api_key = os.getenv("ANTHROPIC_API_KEY") 99 | if not api_key: 100 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 101 | self.client = AnthropicInferenceClient( 102 | model="claude-3-5-sonnet-20241022", api_key=api_key 103 | ) 104 | 105 | async def process( 106 | self, cache: Cache, semaphore: asyncio.Semaphore 107 | ) -> Dict[str, Any]: 108 | print(f"{self.name}: Starting planning process") 109 | print(await cache.keys()) 110 | task = await cache.get("task") 111 | print(f"{self.name}: Retrieved task: {task.objective}") 112 | 113 | # Create a planning prompt 114 | prompt = f""" 115 | Task Objective: {task.objective} 116 | Parameters: {task.params} 117 | 118 | Create a step-by-step plan to accomplish this task. 119 | Format the response as a JSON array of steps, where each step has: 120 | - description: what needs to be done 121 | - requirements: any input needed 122 | - validation: how to verify the step was successful 123 | """ 124 | 125 | # Get plan from LLM 126 | try: 127 | print(f"{self.name}: Sending planning request to LLM") 128 | response = await asyncio.wait_for( 129 | self.client.get_generation([ChatMessage(role=ChatRole.USER, content=prompt)]), timeout=30.0 130 | ) 131 | print(f"{self.name}: Received plan from LLM") 132 | except asyncio.TimeoutError: 133 | print(f"{self.name}: Timeout waiting for LLM response") 134 | raise 135 | 136 | # Store the plan 137 | await cache.set("plan", json.loads(response)) 138 | await cache.set("current_step", 0) 139 | 140 | return {"status": "success", "result": "Plan created successfully"} 141 | ``` 142 | 143 | Key points: 144 | - Initializes LLM client in constructor with API key from environment variable 145 | - Creates structured prompts for the LLM 146 | - Stores plan and progress in cache 147 | - Uses standardized response format 148 | - Implements secure credential handling 149 | 150 | ### 3. Creating the LLM Executor Module 151 | 152 | The executor module uses LLM to perform task steps: 153 | 154 | ```python 155 | class LLMExecutorModule(AgentModule): 156 | """Executes steps using LLM guidance.""" 157 | 158 | name = "llm_executor" 159 | type = ModuleType.EXECUTOR 160 | 161 | def __init__(self): 162 | super().__init__() 163 | api_key = os.getenv("ANTHROPIC_API_KEY") 164 | if not api_key: 165 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 166 | self.client = AnthropicInferenceClient( 167 | model="claude-3-5-sonnet-20241022", 168 | api_key=api_key 169 | ) 170 | 171 | async def process(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 172 | plan = await cache.get("plan") 173 | current_step = await cache.get("current_step") 174 | task = await cache.get("task") 175 | 176 | if current_step >= len(plan): 177 | return { 178 | "status": "success", 179 | "result": "All steps completed" 180 | } 181 | 182 | step = plan[current_step] 183 | 184 | # Create execution prompt 185 | prompt = f""" 186 | Task: {task.objective} 187 | Current Step: {step['description']} 188 | Requirements: {step['requirements']} 189 | 190 | Execute this step and provide the results. 191 | Include: 192 | 1. The actions taken 193 | 2. The outcome 194 | 3. Any relevant output or artifacts 195 | """ 196 | 197 | # Execute step with LLM 198 | response = await self.client.complete(prompt) 199 | result = response.choices[0].message.content 200 | 201 | # Validate step 202 | validation_prompt = f""" 203 | Step: {step['description']} 204 | Validation Criteria: {step['validation']} 205 | Result: {result} 206 | 207 | Evaluate if the step was completed successfully. 208 | Return either "success" or "failure" with a brief explanation. 209 | """ 210 | 211 | validation = await self.client.complete(validation_prompt) 212 | validation_result = validation.choices[0].message.content 213 | 214 | if "success" in validation_result.lower(): 215 | current_step += 1 216 | await cache.set("current_step", current_step) 217 | status = "success" 218 | else: 219 | status = "error" 220 | 221 | return { 222 | "status": status, 223 | "result": { 224 | "step": step['description'], 225 | "execution_result": result, 226 | "validation": validation_result 227 | } 228 | } 229 | ``` 230 | 231 | Key points: 232 | - Executes one step at a time 233 | - Uses LLM for both execution and validation 234 | - Maintains progress through steps 235 | - Includes detailed results and validation 236 | 237 | ### 4. Flow Control Options 238 | 239 | The framework provides two types of flow control modules: 240 | 241 | 1. Basic Flow Control Module 242 | 2. Agent-Directed Flow Control Module 243 | 244 | #### Basic Flow Control Module 245 | 246 | The basic flow control module uses Lua conditions to make routing decisions: 247 | 248 | ```python 249 | flow_control = Node( 250 | name="flow_control", 251 | type=ModuleType.FLOW_CONTROL, 252 | modules=[FlowControlModule( 253 | flow_config=FlowControlConfig( 254 | decisions=[ 255 | FlowDecision( 256 | next_node="executor", 257 | condition="plan ~= nil and current_step < #plan" # Lua conditions 258 | ), 259 | FlowDecision( 260 | next_node="flow_control", 261 | condition="execution_history ~= nil" 262 | ) 263 | ], 264 | default="planner" 265 | ) 266 | )] 267 | ) 268 | ``` 269 | 270 | Key points: 271 | - Uses Lua conditions for decision making 272 | - Supports multiple decision paths 273 | - Can access cache variables in conditions 274 | - Has a default fallback path 275 | 276 | #### Agent-Directed Flow Control Module 277 | 278 | The agent-directed flow control module uses LLMs to make intelligent routing decisions: 279 | 280 | ```python 281 | from asimov.graph.agent_directed_flow import ( 282 | AgentDirectedFlowControl, 283 | AgentDrivenFlowControlConfig, 284 | AgentDrivenFlowDecision, 285 | Example 286 | ) 287 | 288 | flow_control = Node( 289 | name="flow_control", 290 | type=ModuleType.FLOW_CONTROL, 291 | modules=[AgentDirectedFlowControl( 292 | inference_client=AnthropicInferenceClient(...), 293 | system_description="A system that handles various content creation tasks", 294 | flow_config=AgentDrivenFlowControlConfig( 295 | decisions=[ 296 | AgentDrivenFlowDecision( 297 | next_node="blog_writer", 298 | metadata={"description": "Writes blog posts on technical topics"}, 299 | examples=[ 300 | Example( 301 | message="Write a blog post about AI agents", 302 | choices=[ 303 | {"choice": "blog_writer", "description": "Writes blog posts"}, 304 | {"choice": "code_writer", "description": "Writes code"} 305 | ], 306 | choice="blog_writer", 307 | reasoning="The request is specifically for blog content" 308 | ) 309 | ] 310 | ), 311 | AgentDrivenFlowDecision( 312 | next_node="code_writer", 313 | metadata={"description": "Writes code examples and tutorials"}, 314 | examples=[ 315 | Example( 316 | message="Create a Python script for data processing", 317 | choices=[ 318 | {"choice": "blog_writer", "description": "Writes blog posts"}, 319 | {"choice": "code_writer", "description": "Writes code"} 320 | ], 321 | choice="code_writer", 322 | reasoning="The request is for code creation" 323 | ) 324 | ] 325 | ) 326 | ], 327 | default="error_handler" 328 | ) 329 | )] 330 | ) 331 | ``` 332 | 333 | Key points: 334 | - Uses LLM for intelligent decision making 335 | - Supports example-based learning 336 | - Can include reasoning for decisions 337 | - Maintains context awareness 338 | - Handles complex routing scenarios 339 | - Supports metadata for rich decision context 340 | 341 | ### 5. Configuring Flow Control 342 | 343 | Choose the appropriate flow control based on your needs: 344 | 345 | ```python 346 | flow_control = Node( 347 | name="flow_control", 348 | type=ModuleType.FLOW_CONTROL, 349 | modules=[FlowControlModule( 350 | flow_config=FlowControlConfig( 351 | decisions=[ 352 | FlowDecision( 353 | next_node="executor", 354 | condition="plan ~= nil and current_step < #plan" # Conditions are lua 355 | ), 356 | FlowDecision( 357 | next_node="flow_control", 358 | condition="execution_history ~= nil" # Conditions are lua 359 | ) 360 | ], 361 | default="planner" 362 | ) 363 | )] 364 | ) 365 | ``` 366 | 367 | ### 6. Example Implementations 368 | 369 | #### Basic Task Planning Agent 370 | 371 | Create a basic agent that uses LLM for task planning and execution: 372 | 373 | ```python 374 | async def main(): 375 | # Create the agent 376 | agent = Agent( 377 | cache=RedisCache(), 378 | max_concurrent_tasks=1, 379 | max_total_iterations=20 380 | ) 381 | 382 | # Create and add nodes 383 | planner_node = Node( 384 | name="planner", 385 | type=ModuleType.EXECUTOR, 386 | modules=[LLMPlannerModule()], 387 | node_config=NodeConfig( 388 | parallel=False, 389 | max_retries=3 390 | ) 391 | ) 392 | 393 | executor_node = Node( 394 | name="executor", 395 | type=ModuleType.EXECUTOR, 396 | modules=[LLMExecutorModule()], 397 | dependencies=["planner", "flow_control"] 398 | ) 399 | 400 | flow_control_node = Node( 401 | name="flow_control", 402 | type=ModuleType.FLOW_CONTROL, 403 | modules=[LLMFlowControlModule()], 404 | dependencies=["executor"] 405 | ) 406 | 407 | # Add nodes to agent 408 | agent.add_multiple_nodes([ 409 | planner_node, 410 | executor_node, 411 | flow_control_node 412 | ]) 413 | 414 | # Create and run a task 415 | task = Task( 416 | type="content_creation", 417 | objective="Write a blog post about AI agents", 418 | params={ 419 | "topic": "AI Agents in Production", 420 | "length": "1000 words", 421 | "style": "technical but accessible", 422 | "key_points": [ 423 | "Definition of AI agents", 424 | "Common architectures", 425 | "Real-world applications", 426 | "Future trends" 427 | ] 428 | } 429 | ) 430 | 431 | await agent.run_task(task) 432 | ``` 433 | 434 | #### Intelligent Content Router 435 | 436 | Create an agent that uses AgentDirectedFlow to route content creation tasks: 437 | 438 | ```python 439 | async def main(): 440 | # Create the agent 441 | agent = Agent( 442 | cache=RedisCache(default_prefix=str(uuid4())), 443 | max_concurrent_tasks=1, 444 | max_total_iterations=20 445 | ) 446 | 447 | # Set up specialized content nodes 448 | blog_writer = Node( 449 | name="blog_writer", 450 | type=ModuleType.EXECUTOR, 451 | modules=[BlogWriterModule()] 452 | ) 453 | 454 | code_writer = Node( 455 | name="code_writer", 456 | type=ModuleType.EXECUTOR, 457 | modules=[CodeWriterModule()] 458 | ) 459 | 460 | error_handler = Node( 461 | name="error_handler", 462 | type=ModuleType.EXECUTOR, 463 | modules=[ErrorHandlerModule()] 464 | ) 465 | 466 | # Create flow control with agent-directed decisions 467 | flow_control = Node( 468 | name="flow_control", 469 | type=ModuleType.FLOW_CONTROL, 470 | modules=[AgentDirectedFlowControl( 471 | inference_client=AnthropicInferenceClient(...), 472 | system_description="A system that handles various content creation tasks", 473 | flow_config=AgentDrivenFlowControlConfig( 474 | decisions=[ 475 | AgentDrivenFlowDecision( 476 | next_node="blog_writer", 477 | metadata={"description": "Writes blog posts"}, 478 | examples=[ 479 | Example( 480 | message="Write a blog post about AI", 481 | choices=[ 482 | {"choice": "blog_writer", "description": "Writes blogs"}, 483 | {"choice": "code_writer", "description": "Writes code"} 484 | ], 485 | choice="blog_writer", 486 | reasoning="Request is for blog content" 487 | ) 488 | ] 489 | ), 490 | AgentDrivenFlowDecision( 491 | next_node="code_writer", 492 | metadata={"description": "Writes code"}, 493 | examples=[ 494 | Example( 495 | message="Create a sorting algorithm", 496 | choices=[ 497 | {"choice": "blog_writer", "description": "Writes blogs"}, 498 | {"choice": "code_writer", "description": "Writes code"} 499 | ], 500 | choice="code_writer", 501 | reasoning="Request is for code implementation" 502 | ) 503 | ] 504 | ) 505 | ], 506 | default="error_handler" 507 | ) 508 | )] 509 | ) 510 | 511 | # Add all nodes to agent 512 | agent.add_multiple_nodes([ 513 | blog_writer, 514 | code_writer, 515 | error_handler, 516 | flow_control 517 | ]) 518 | 519 | # Example tasks 520 | tasks = [ 521 | Task( 522 | type="content_request", 523 | objective="Write about AI", 524 | params={"input_message": "Write a blog post about AI"} 525 | ), 526 | Task( 527 | type="content_request", 528 | objective="Implement quicksort", 529 | params={"input_message": "Create a quicksort implementation"} 530 | ) 531 | ] 532 | 533 | # Run tasks 534 | for task in tasks: 535 | await agent.run_task(task) 536 | 537 | if __name__ == "__main__": 538 | asyncio.run(main()) 539 | ``` 540 | 541 | ## Best Practices for LLM Agents 542 | 543 | 1. **Prompt Engineering** 544 | - Be specific and structured in prompts 545 | - Include relevant context 546 | - Request structured outputs 547 | - Validate LLM responses 548 | 549 | 2. **Error Handling** 550 | - Handle LLM API errors 551 | - Implement retries for failed requests 552 | - Validate LLM outputs 553 | - Have fallback strategies 554 | 555 | 3. **State Management** 556 | - Track progress through steps 557 | - Store intermediate results 558 | - Maintain execution history 559 | - Clean up temporary state 560 | 561 | 4. **Cost Management** 562 | - Cache LLM responses when appropriate 563 | - Use smaller models for simpler tasks 564 | - Batch similar requests 565 | - Monitor token usage 566 | 567 | 5. **Performance** 568 | - Use streaming for long responses 569 | - Implement timeouts 570 | - Consider parallel processing 571 | - Cache frequently used prompts 572 | 573 | ## Common Patterns 574 | 575 | 1. **Plan-Execute-Validate** 576 | - Create detailed plans 577 | - Execute steps systematically 578 | - Validate results 579 | - Adjust plans based on feedback 580 | 581 | 2. **Progressive Refinement** 582 | - Start with high-level plans 583 | - Refine details progressively 584 | - Validate intermediate results 585 | - Adjust based on outcomes 586 | 587 | 3. **Decision Making** 588 | - Use LLMs for complex decisions 589 | - Provide context and history 590 | - Include decision criteria 591 | - Document reasoning 592 | 593 | 4. **Error Recovery** 594 | - Detect errors early 595 | - Implement retry strategies 596 | - Consider alternative approaches 597 | - Learn from failures 598 | 599 | ## Next Steps 600 | 601 | 1. Implement more sophisticated planning strategies 602 | 2. Add support for multiple LLM providers 603 | 3. Implement caching for LLM responses 604 | 4. Add monitoring and logging 605 | 5. Implement cost tracking and optimization 606 | 6. Add support for streaming responses 607 | 7. Implement parallel processing for suitable tasks 608 | -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- 1 | # Make examples directory a proper Python module -------------------------------------------------------------------------------- /examples/agent_directed_flow.py: -------------------------------------------------------------------------------- 1 | """ 2 | Agent-Directed Flow Example 3 | 4 | This example demonstrates how to use the AgentDirectedFlow module to create an intelligent 5 | routing system for content creation tasks. The system will: 6 | 1. Analyze incoming requests using LLM 7 | 2. Route tasks to appropriate specialized modules 8 | 3. Handle multiple types of content creation tasks 9 | """ 10 | 11 | import os 12 | import asyncio 13 | from typing import Dict, Any 14 | from uuid import uuid4 15 | 16 | from asimov.graph import ( 17 | Agent, 18 | AgentModule, 19 | ModuleType, 20 | Node, 21 | NodeConfig, 22 | Cache, 23 | ) 24 | from asimov.graph.agent_directed_flow import ( 25 | AgentDirectedFlowControl, 26 | AgentDrivenFlowControlConfig, 27 | AgentDrivenFlowDecision, 28 | Example, 29 | ) 30 | from asimov.graph.tasks import Task 31 | from asimov.caches.redis_cache import RedisCache 32 | from asimov.services.inference_clients import ( 33 | AnthropicInferenceClient, 34 | ChatMessage, 35 | ChatRole, 36 | ) 37 | 38 | 39 | class BlogWriterModule(AgentModule): 40 | """Specialized module for writing blog posts.""" 41 | 42 | name: str = "blog_writer" 43 | type: ModuleType = ModuleType.EXECUTOR 44 | 45 | client: AnthropicInferenceClient = None 46 | 47 | def __init__(self): 48 | super().__init__() 49 | api_key = os.getenv("ANTHROPIC_API_KEY") 50 | if not api_key: 51 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 52 | self.client = AnthropicInferenceClient( 53 | model="claude-3-5-sonnet-20241022", api_key=api_key 54 | ) 55 | 56 | async def process( 57 | self, cache: Cache, semaphore: asyncio.Semaphore 58 | ) -> Dict[str, Any]: 59 | message = await cache.get("input_message") 60 | 61 | prompt = f""" 62 | Write a blog post based on the following request: 63 | {message} 64 | 65 | Format the post with: 66 | 1. A compelling title 67 | 2. An introduction 68 | 3. Main content sections 69 | 4. A conclusion 70 | """ 71 | 72 | response = await self.client.get_generation( 73 | [ChatMessage(role=ChatRole.USER, content=prompt)] 74 | ) 75 | 76 | return {"status": "success", "result": response} 77 | 78 | 79 | class CodeWriterModule(AgentModule): 80 | """Specialized module for writing code examples.""" 81 | 82 | name: str = "code_writer" 83 | type: ModuleType = ModuleType.EXECUTOR 84 | 85 | client: AnthropicInferenceClient = None 86 | 87 | def __init__(self): 88 | super().__init__() 89 | api_key = os.getenv("ANTHROPIC_API_KEY") 90 | if not api_key: 91 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 92 | self.client = AnthropicInferenceClient( 93 | model="claude-3-5-sonnet-20241022", api_key=api_key 94 | ) 95 | 96 | async def process( 97 | self, cache: Cache, semaphore: asyncio.Semaphore 98 | ) -> Dict[str, Any]: 99 | message = await cache.get("input_message") 100 | 101 | prompt = f""" 102 | Create code based on the following request: 103 | {message} 104 | 105 | Provide: 106 | 1. The complete code implementation 107 | 2. Comments explaining key parts 108 | 3. Usage examples 109 | 4. Any necessary setup instructions 110 | """ 111 | 112 | response = await self.client.get_generation( 113 | [ChatMessage(role=ChatRole.USER, content=prompt)] 114 | ) 115 | 116 | return {"status": "success", "result": response} 117 | 118 | 119 | class ErrorHandlerModule(AgentModule): 120 | """Handles cases where no other module is appropriate.""" 121 | 122 | name: str = "error_handler" 123 | type: ModuleType = ModuleType.EXECUTOR 124 | 125 | async def process( 126 | self, cache: Cache, semaphore: asyncio.Semaphore 127 | ) -> Dict[str, Any]: 128 | message = await cache.get("input_message") 129 | 130 | return { 131 | "status": "error", 132 | "result": f"Unable to process request: {message}. Please try a different type of request.", 133 | } 134 | 135 | 136 | async def main(): 137 | # Create the agent 138 | cache = RedisCache(default_prefix=str(uuid4())) 139 | agent = Agent( 140 | cache=cache, 141 | max_concurrent_tasks=1, 142 | max_total_iterations=20, 143 | ) 144 | 145 | # Set up the inference client for flow control 146 | api_key = os.getenv("ANTHROPIC_API_KEY") 147 | if not api_key: 148 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 149 | inference_client = AnthropicInferenceClient( 150 | model="claude-3-5-sonnet-20241022", api_key=api_key 151 | ) 152 | 153 | # Create specialized content nodes 154 | blog_writer = Node( 155 | name="blog_writer", 156 | type=ModuleType.EXECUTOR, 157 | modules=[BlogWriterModule()], 158 | dependencies=["flow_control"], 159 | ) 160 | 161 | code_writer = Node( 162 | name="code_writer", 163 | type=ModuleType.EXECUTOR, 164 | modules=[CodeWriterModule()], 165 | dependencies=["flow_control"], 166 | ) 167 | 168 | error_handler = Node( 169 | name="error_handler", 170 | type=ModuleType.EXECUTOR, 171 | modules=[ErrorHandlerModule()], 172 | dependencies=["flow_control"], 173 | ) 174 | 175 | # Create flow control node with agent-directed decisions 176 | flow_control = Node( 177 | name="flow_control", 178 | type=ModuleType.FLOW_CONTROL, 179 | modules=[ 180 | AgentDirectedFlowControl( 181 | name="ContentFlowControl", 182 | type=ModuleType.FLOW_CONTROL, 183 | voters=3, 184 | inference_client=inference_client, 185 | system_description="A system that handles various content creation tasks", 186 | flow_config=AgentDrivenFlowControlConfig( 187 | decisions=[ 188 | AgentDrivenFlowDecision( 189 | next_node="blog_writer", 190 | metadata={ 191 | "description": "Writes blog posts on technical topics" 192 | }, 193 | examples=[ 194 | Example( 195 | message="Write a blog post about AI agents", 196 | choices=[ 197 | { 198 | "choice": "blog_writer", 199 | "description": "Writes blog posts", 200 | }, 201 | { 202 | "choice": "code_writer", 203 | "description": "Writes code", 204 | }, 205 | ], 206 | choice="blog_writer", 207 | reasoning="The request is specifically for blog content", 208 | ), 209 | Example( 210 | message="Create an article explaining machine learning basics", 211 | choices=[ 212 | { 213 | "choice": "blog_writer", 214 | "description": "Writes blog posts", 215 | }, 216 | { 217 | "choice": "code_writer", 218 | "description": "Writes code", 219 | }, 220 | ], 221 | choice="blog_writer", 222 | reasoning="The request is for an explanatory article", 223 | ), 224 | ], 225 | ), 226 | AgentDrivenFlowDecision( 227 | next_node="code_writer", 228 | metadata={ 229 | "description": "Writes code examples and tutorials" 230 | }, 231 | examples=[ 232 | Example( 233 | message="Create a Python script for data processing", 234 | choices=[ 235 | { 236 | "choice": "blog_writer", 237 | "description": "Writes blog posts", 238 | }, 239 | { 240 | "choice": "code_writer", 241 | "description": "Writes code", 242 | }, 243 | ], 244 | choice="code_writer", 245 | reasoning="The request is for code creation", 246 | ), 247 | Example( 248 | message="Show me how to implement a binary search tree", 249 | choices=[ 250 | { 251 | "choice": "blog_writer", 252 | "description": "Writes blog posts", 253 | }, 254 | { 255 | "choice": "code_writer", 256 | "description": "Writes code", 257 | }, 258 | ], 259 | choice="code_writer", 260 | reasoning="The request is for a code implementation", 261 | ), 262 | ], 263 | ), 264 | ], 265 | default="error_handler", 266 | ), 267 | ) 268 | ], 269 | ) 270 | 271 | # Add all nodes to the agent 272 | agent.add_multiple_nodes([blog_writer, code_writer, error_handler, flow_control]) 273 | 274 | # Example tasks to demonstrate routing 275 | tasks = [ 276 | Task( 277 | type="content_request", 278 | objective="Write a blog post about the future of AI", 279 | params={"input_message": "Write a blog post about the future of AI"}, 280 | ), 281 | Task( 282 | type="content_request", 283 | objective="Create a Python implementation of quicksort", 284 | params={"input_message": "Create a Python implementation of quicksort"}, 285 | ), 286 | Task( 287 | type="content_request", 288 | objective="Generate a haiku about programming", 289 | params={"input_message": "Generate a haiku about programming"}, 290 | ), 291 | ] 292 | 293 | # Run each task and show results 294 | for task in tasks: 295 | await cache.set("input_message", task.params["input_message"]) 296 | print(f"\nProcessing task: {task.objective}") 297 | result = await agent.run_task(task) 298 | print(f"Result: {result}") 299 | print("Node results:") 300 | for node, node_result in agent.node_results.items(): 301 | print(f"{node}: {node_result}") 302 | 303 | 304 | if __name__ == "__main__": 305 | asyncio.run(main()) 306 | -------------------------------------------------------------------------------- /examples/basic_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Basic Agent Example 3 | 4 | This example demonstrates how to create a simple agent that processes text using 5 | a planner and executor pattern. The agent will: 6 | 1. Plan how to process the text (e.g., identify operations needed) 7 | 2. Execute the planned operations 8 | 3. Use flow control to manage the execution flow 9 | """ 10 | 11 | import asyncio 12 | from typing import Dict, Any 13 | from asimov.graph import ( 14 | Agent, 15 | AgentModule, 16 | ModuleType, 17 | Node, 18 | NodeConfig, 19 | FlowControlModule, 20 | FlowControlConfig, 21 | FlowDecision, 22 | Cache, 23 | ) 24 | from asimov.graph.tasks import Task 25 | from asimov.caches.redis_cache import RedisCache 26 | 27 | 28 | class TextPlannerModule(AgentModule): 29 | """Plans text processing operations.""" 30 | 31 | name: str = "text_planner" 32 | type: ModuleType = ModuleType.EXECUTOR 33 | 34 | async def process(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 35 | print(f"{self.name}: Starting planning process") 36 | # Get the task from cache 37 | task = await cache.get("task") 38 | text = task.params.get("text", "") 39 | print(f"{self.name}: Retrieved task with text length {len(text)}") 40 | 41 | # Create a simple plan 42 | plan = { 43 | "operations": [ 44 | {"type": "count_words", "text": text}, 45 | {"type": "calculate_stats", "text": text} 46 | ] 47 | } 48 | 49 | # Store the plan in cache 50 | await cache.set("plan", plan) 51 | 52 | return { 53 | "status": "success", 54 | "result": "Plan created successfully" 55 | } 56 | 57 | 58 | class TextExecutorModule(AgentModule): 59 | """Executes text processing operations.""" 60 | 61 | name: str = "text_executor" 62 | type: ModuleType = ModuleType.EXECUTOR 63 | 64 | async def process(self, cache: Cache, semaphore: asyncio.Semaphore) -> Dict[str, Any]: 65 | print(f"{self.name}: Starting execution process") 66 | # Get the plan and task 67 | plan = await cache.get("plan") 68 | task = await cache.get("task") 69 | print(f"{self.name}: Retrieved plan with {len(plan['operations'])} operations") 70 | 71 | results = [] 72 | for operation in plan["operations"]: 73 | if operation["type"] == "count_words": 74 | word_count = len(operation["text"].split()) 75 | results.append({ 76 | "operation": "count_words", 77 | "result": word_count 78 | }) 79 | elif operation["type"] == "calculate_stats": 80 | char_count = len(operation["text"]) 81 | line_count = len(operation["text"].splitlines()) 82 | results.append({ 83 | "operation": "calculate_stats", 84 | "result": { 85 | "characters": char_count, 86 | "lines": line_count 87 | } 88 | }) 89 | 90 | return { 91 | "status": "success", 92 | "result": results 93 | } 94 | 95 | 96 | async def main(): 97 | print("Starting basic agent example") 98 | # Create the agent 99 | agent = Agent( 100 | cache=RedisCache(), 101 | max_concurrent_tasks=1, 102 | max_total_iterations=10 103 | ) 104 | print("Agent created with Redis cache") 105 | 106 | # Create nodes 107 | planner_node = Node( 108 | name="planner", 109 | type=ModuleType.EXECUTOR, 110 | modules=[TextPlannerModule()], 111 | node_config=NodeConfig( 112 | parallel=False, 113 | max_retries=3 114 | ) 115 | ) 116 | 117 | executor_node = Node( 118 | name="executor", 119 | type=ModuleType.EXECUTOR, 120 | modules=[TextExecutorModule()], 121 | dependencies=["planner"] 122 | ) 123 | 124 | flow_control = Node( 125 | name="flow_control", 126 | type=ModuleType.FLOW_CONTROL, 127 | modules=[FlowControlModule( 128 | name="flow_control_module", 129 | type=ModuleType.FLOW_CONTROL, 130 | flow_config=FlowControlConfig( 131 | decisions=[ 132 | FlowDecision( 133 | next_node="executor", 134 | condition="plan ~= null" 135 | ) 136 | ], 137 | default="planner" 138 | ) 139 | )] 140 | ) 141 | 142 | # Add nodes to agent 143 | agent.add_multiple_nodes([planner_node, executor_node, flow_control]) 144 | 145 | # Create and run a task 146 | task = Task( 147 | type="text_processing", 148 | objective="Process sample text", 149 | params={ 150 | "text": "Hello world!\nThis is a sample text.\nIt demonstrates the basic agent functionality." 151 | } 152 | ) 153 | 154 | # Run the task 155 | await agent.run_task(task) 156 | 157 | # Get the final results 158 | results = agent.node_results.get("executor", {}).get("results", [])[0]['result'] 159 | print("\nProcessing Results:") 160 | for result in results: 161 | print(f"\nOperation: {result['operation']}") 162 | print(f"Result: {result['result']}") 163 | 164 | 165 | if __name__ == "__main__": 166 | asyncio.run(main()) 167 | -------------------------------------------------------------------------------- /examples/llm_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Cake Baking LLM Agent Example 3 | 4 | This example demonstrates how to create an agent that uses Large Language Models (LLMs) 5 | for baking a cake. The agent will: 6 | 1. Use an LLM to plan the cake baking process 7 | 2. Execute the plan by following recipe steps and monitoring progress 8 | 3. Use flow control to manage the baking process and validate the results 9 | """ 10 | 11 | import json 12 | import os 13 | import asyncio 14 | from typing import Dict, Any 15 | from asimov.graph import ( 16 | Agent, 17 | AgentModule, 18 | ModuleType, 19 | Node, 20 | NodeConfig, 21 | FlowControlModule, 22 | FlowControlConfig, 23 | FlowDecision, 24 | Cache, 25 | ) 26 | from asimov.graph.tasks import Task 27 | from asimov.caches.redis_cache import RedisCache 28 | from asimov.services.inference_clients import ( 29 | AnthropicInferenceClient, 30 | ChatMessage, 31 | ChatRole, 32 | ) 33 | 34 | from uuid import uuid4 35 | 36 | 37 | class LLMPlannerModule(AgentModule): 38 | """Uses LLM to plan cake baking process.""" 39 | 40 | name: str = "llm_planner" 41 | type: ModuleType = ModuleType.EXECUTOR 42 | 43 | client: AnthropicInferenceClient = None 44 | 45 | def __init__(self): 46 | super().__init__() 47 | api_key = os.getenv("ANTHROPIC_API_KEY") 48 | if not api_key: 49 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 50 | self.client = AnthropicInferenceClient( 51 | model="claude-3-5-sonnet-20241022", api_key=api_key 52 | ) 53 | 54 | async def process( 55 | self, cache: Cache, semaphore: asyncio.Semaphore 56 | ) -> Dict[str, Any]: 57 | print(f"{self.name}: Starting planning process") 58 | task = await cache.get("task") 59 | print(f"{self.name}: Retrieved task: {task.objective}") 60 | 61 | # Create a planning prompt 62 | prompt = f""" 63 | Task Objective: {task.objective} 64 | Parameters: {task.params} 65 | 66 | Create a step-by-step plan to accomplish this task. 67 | Format the response as a JSON array of steps, where each step has: 68 | - description: what needs to be done 69 | - requirements: any input needed 70 | - validation: how to verify the step was successful 71 | """ 72 | 73 | # Get plan from LLM 74 | try: 75 | print(f"{self.name}: Sending planning request to LLM") 76 | response = await asyncio.wait_for( 77 | self.client.get_generation( 78 | [ChatMessage(role=ChatRole.USER, content=prompt)] 79 | ), 80 | timeout=30.0, 81 | ) 82 | 83 | try: 84 | loaded_response = json.loads(response) 85 | response_content = loaded_response["steps"] 86 | except json.JSONDecodeError as e: 87 | return {"status": "error", "result": str(e)} 88 | 89 | print(f"{self.name}: Received plan from LLM") 90 | except asyncio.TimeoutError: 91 | print(f"{self.name}: Timeout waiting for LLM response") 92 | raise 93 | 94 | # Store the plan 95 | await cache.set("plan", response_content) # Store raw JSON string 96 | await cache.set("current_step", 0) 97 | 98 | return {"status": "success", "result": "Plan created successfully"} 99 | 100 | 101 | class LLMExecutorModule(AgentModule): 102 | """Executes cake baking steps using LLM guidance.""" 103 | 104 | name: str = "llm_executor" 105 | type: ModuleType = ModuleType.EXECUTOR 106 | 107 | client: AnthropicInferenceClient = None 108 | 109 | def __init__(self): 110 | super().__init__() 111 | api_key = os.getenv("ANTHROPIC_API_KEY") 112 | if not api_key: 113 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 114 | self.client = AnthropicInferenceClient( 115 | model="claude-3-5-sonnet-20241022", api_key=api_key 116 | ) 117 | 118 | async def process( 119 | self, cache: Cache, semaphore: asyncio.Semaphore 120 | ) -> Dict[str, Any]: 121 | print(f"{self.name}: Starting execution process") 122 | try: 123 | # Note the real cache and not the mock testing cache uses jsonpickle and handles serialization and deserialization so you don't need to do this. 124 | plan_json = await cache.get("plan") 125 | plan = ( 126 | json.loads(plan_json) 127 | if isinstance(plan_json, (str, bytes)) 128 | else plan_json 129 | ) 130 | current_step = await cache.get("current_step") 131 | task = await cache.get("task") 132 | print(f"{self.name}: Retrieved plan and current step {current_step}") 133 | except asyncio.TimeoutError: 134 | print(f"{self.name}: Timeout retrieving task data") 135 | raise 136 | 137 | if current_step >= len(plan): 138 | return {"status": "success", "result": "All steps completed"} 139 | 140 | step = plan[current_step] 141 | 142 | print(step) 143 | 144 | # Create execution prompt 145 | prompt = f""" 146 | Task: {task.objective} 147 | Current Step: {step['description']} 148 | Requirements: {step['requirements']} 149 | 150 | Execute this step and provide the results. 151 | Include: 152 | 1. The actions taken 153 | 2. The outcome 154 | 3. Any relevant output or artifacts 155 | """ 156 | 157 | # Execute step with LLM 158 | try: 159 | print(f"{self.name}: Sending execution request to LLM") 160 | response = await asyncio.wait_for( 161 | self.client.get_generation( 162 | [ChatMessage(role=ChatRole.USER, content=prompt)] 163 | ), 164 | timeout=30.0, 165 | ) 166 | print(f"{self.name}: Received execution result from LLM") 167 | except asyncio.TimeoutError: 168 | print(f"{self.name}: Timeout waiting for LLM execution response") 169 | raise 170 | 171 | # Validate step 172 | validation_prompt = f""" 173 | Step: {step['description']} 174 | Validation Criteria: {step['validation']} 175 | Result: {response} 176 | 177 | Evaluate if the step was completed successfully. 178 | Return either "success" or "failure" with a brief explanation. 179 | """ 180 | 181 | try: 182 | print(f"{self.name}: Sending validation request to LLM") 183 | validation_result = await asyncio.wait_for( 184 | self.client.get_generation( 185 | [ChatMessage(role=ChatRole.USER, content=validation_prompt)] 186 | ), 187 | timeout=30.0, 188 | ) 189 | print(f"{self.name}: Received validation result from LLM") 190 | except asyncio.TimeoutError: 191 | print(f"{self.name}: Timeout waiting for LLM validation response") 192 | raise 193 | 194 | if "success" in validation_result.lower(): 195 | current_step += 1 196 | await cache.set("current_step", current_step) 197 | status = "success" 198 | else: 199 | status = "error" 200 | 201 | # Create execution record 202 | execution_record = { 203 | "step": step["description"], 204 | "execution_result": response, 205 | "validation": validation_result, 206 | "status": status, 207 | "timestamp": str(asyncio.get_event_loop().time()), 208 | } 209 | 210 | # Update execution history 211 | execution_history = await cache.get("execution_history", []) 212 | execution_history.append(execution_record) 213 | await cache.set("execution_history", execution_history) 214 | 215 | return { 216 | "status": status, 217 | "result": execution_record, 218 | } 219 | 220 | 221 | class LLMFlowControlModule(AgentModule): 222 | """Makes decisions about cake baking flow based on LLM analysis.""" 223 | 224 | name: str = "llm_flow_control" 225 | type: ModuleType = ModuleType.FLOW_CONTROL 226 | 227 | client: AnthropicInferenceClient = None 228 | 229 | def __init__(self): 230 | super().__init__() 231 | api_key = os.getenv("ANTHROPIC_API_KEY") 232 | if not api_key: 233 | raise ValueError("ANTHROPIC_API_KEY environment variable must be set") 234 | self.client = AnthropicInferenceClient( 235 | model="claude-3-5-sonnet-20241022", api_key=api_key 236 | ) 237 | 238 | async def process( 239 | self, cache: Cache, semaphore: asyncio.Semaphore 240 | ) -> Dict[str, Any]: 241 | print(f"{self.name}: Starting flow control process") 242 | plan_json = await cache.get("plan") 243 | plan = ( 244 | json.loads(plan_json) if isinstance(plan_json, (str, bytes)) else plan_json 245 | ) 246 | current_step = await cache.get("current_step") 247 | execution_history = await cache.get("execution_history", []) 248 | print( 249 | f"{self.name}: Retrieved plan and history with {len(execution_history)} entries" 250 | ) 251 | 252 | if not execution_history: 253 | return { 254 | "status": "success", 255 | "result": {"decision": "continue", "reason": "No history to analyze"}, 256 | } 257 | 258 | # Create analysis prompt 259 | prompt = f""" 260 | Execution History: {execution_history} 261 | Current Step: {current_step} of {len(plan)} steps 262 | 263 | Analyze the execution history and determine if we should: 264 | 1. continue: proceed with the next step 265 | 2. retry: retry the current step 266 | 3. replan: create a new plan 267 | 4. abort: stop execution 268 | 269 | Provide your decision and reasoning. 270 | """ 271 | 272 | try: 273 | print(f"{self.name}: Sending analysis request to LLM") 274 | analysis = await asyncio.wait_for( 275 | self.client.get_generation( 276 | [ChatMessage(role=ChatRole.USER, content=prompt)] 277 | ), 278 | timeout=30.0, 279 | ) 280 | print(f"{self.name}: Received analysis result from LLM") 281 | except asyncio.TimeoutError: 282 | print(f"{self.name}: Timeout waiting for LLM analysis response") 283 | raise 284 | 285 | return { 286 | "status": "success", 287 | "result": { 288 | "analysis": analysis, 289 | "decision": "continue", # Extract actual decision from analysis 290 | }, 291 | } 292 | 293 | 294 | async def main(): 295 | print("Starting LLM agent example") 296 | # Create the agent 297 | 298 | agent = Agent( 299 | cache=RedisCache(default_prefix=str(uuid4())), 300 | max_concurrent_tasks=1, 301 | max_total_iterations=20, 302 | ) 303 | print("Agent created with Redis cache") 304 | 305 | # Create nodes 306 | planner_node = Node( 307 | name="planner", 308 | type=ModuleType.EXECUTOR, 309 | modules=[LLMPlannerModule()], 310 | node_config=NodeConfig(parallel=False, max_retries=3), 311 | ) 312 | 313 | executor_node = Node( 314 | name="executor", 315 | type=ModuleType.EXECUTOR, 316 | modules=[LLMExecutorModule()], 317 | dependencies=["planner"], 318 | ) 319 | 320 | flow_control = Node( 321 | name="flow_control", 322 | type=ModuleType.FLOW_CONTROL, 323 | modules=[ 324 | FlowControlModule( 325 | name="flow_control", 326 | type=ModuleType.FLOW_CONTROL, 327 | flow_config=FlowControlConfig( 328 | decisions=[ 329 | FlowDecision( 330 | next_node="executor", 331 | condition="plan ~= nil and current_step < #plan", 332 | ), 333 | FlowDecision( 334 | next_node="flow_control", 335 | condition="execution_history ~= nil", 336 | ), 337 | ], 338 | default="planner", 339 | ), 340 | ) 341 | ], 342 | ) 343 | 344 | # Add nodes to agent 345 | agent.add_multiple_nodes([planner_node, executor_node, flow_control]) 346 | 347 | # Create and run a task 348 | task = Task( 349 | type="cake_baking", 350 | objective="Bake a chocolate cake", 351 | params={ 352 | "cake_type": "Chocolate", 353 | "servings": "8-10", 354 | "difficulty": "intermediate", 355 | "requirements": [ 356 | "Moist and fluffy texture", 357 | "Rich chocolate flavor", 358 | "Professional presentation", 359 | "Even baking throughout", 360 | ], 361 | }, 362 | ) 363 | 364 | # Run the task 365 | await agent.run_task(task) 366 | 367 | # Print results 368 | print("\nTask Execution Results:") 369 | for node, result in agent.node_results.items(): 370 | print(f"\nNode: {node}") 371 | print(f"Result: {result}") 372 | 373 | 374 | if __name__ == "__main__": 375 | asyncio.run(main()) 376 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "asimov-agents" 7 | dynamic = ["version"] 8 | description = "A library of primitives for building agentic flows." 9 | readme = "README.md" 10 | requires-python = ">=3.12" 11 | license = "Apache-2.0" 12 | keywords = ["git", "agent", "bismuth", "ai"] 13 | authors = [ 14 | { name = "Ian Butler", email = "ian@bismuth.cloud" }, 15 | ] 16 | classifiers = [ 17 | "Development Status :: 4 - Beta", 18 | "Programming Language :: Python", 19 | "Programming Language :: Python :: 3.12", 20 | "Programming Language :: Python :: Implementation :: CPython", 21 | "Programming Language :: Python :: Implementation :: PyPy", 22 | ] 23 | dependencies = [ 24 | "gitpython>=3.1.0,<4.0.0", 25 | "httpx>=0.23.0,<1.0.0", 26 | "opentelemetry-instrumentation-httpx", 27 | "aioboto3>=13.0.0", 28 | "pydantic>=2.0.0", 29 | "psycopg2-binary>=2.9.0", 30 | "opentelemetry-instrumentation-psycopg2", 31 | "lupa==2.2.0", 32 | "jsonpickle>=3.0.0", 33 | "redis>=5.0.0", 34 | "opentelemetry-api>=1.27.0", 35 | "google-cloud-aiplatform==1.69.0", 36 | "backoff>=2.2.0", 37 | "google-generativeai==0.8.3", 38 | "google-genai", 39 | ] 40 | 41 | [project.optional-dependencies] 42 | dev = [ 43 | "black>=23.1.0", 44 | "mypy>=1.0.0", 45 | "pytest>=7.0.0", 46 | "pytest-cov>=4.0.0", 47 | "pytest-asyncio>=0.24.0", 48 | "graphviz", 49 | ] 50 | 51 | [project.urls] 52 | Homepage = "https://github.com/BismuthCloud/asimov" 53 | URL = "https://github.com/BismuthCloud/asimov" 54 | Documentation = "https://github.com/BismuthCloud/asimov" 55 | Issues = "https://github.com/BismuthCloud/asimov/issues" 56 | Source = "https://github.com/BismuthCloud/asimov" 57 | 58 | [tool.hatch.version] 59 | path = "asimov/__init__.py" 60 | 61 | [[tool.hatch.envs.test.matrix]] 62 | python = ["312"] 63 | 64 | [tool.hatch.envs.hatch-test] 65 | extra-dependencies = [ 66 | "pytest-asyncio>=0.24.0", 67 | "pytest-timeout>=2.0.0", 68 | ] 69 | 70 | [tool.hatch.envs.types] 71 | extra-dependencies = [ 72 | "mypy>=1.0.0", 73 | # types-lupa is wrong in a bunch of places 74 | "types-aioboto3>=13.0.0", 75 | ] 76 | 77 | [tool.hatch.envs.types.scripts] 78 | check = "mypy --install-types --non-interactive {args:asimov}" 79 | 80 | [tool.mypy] 81 | ignore_missing_imports = true 82 | 83 | [tool.coverage.run] 84 | branch = true 85 | parallel = true 86 | omit = [ 87 | "asimov/__init__.py" 88 | ] 89 | 90 | [tool.coverage.report] 91 | exclude_lines = [ 92 | "no cov", 93 | "if __name__ == .__main__.:", 94 | "if TYPE_CHECKING:", 95 | ] 96 | 97 | [tool.hatch.build.targets.wheel] 98 | packages = ["asimov/"] 99 | 100 | [tool.hatch.build.targets.sdist] 101 | exclude = [ 102 | "/.github", 103 | "/docs", 104 | "/tests", 105 | ] 106 | 107 | [tool.pytest.ini_options] 108 | timeout = 10 109 | timeout_method = "thread" 110 | asyncio_default_fixture_loop_scope = "function" 111 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | pythonpath = . 3 | addopts = -v -------------------------------------------------------------------------------- /stack.dev.yaml: -------------------------------------------------------------------------------- 1 | version: v1.0.0 2 | kind: stack 3 | name: "Asimov Dev" 4 | description: "" 5 | namespace: "asimov-dev" 6 | release: "dev" 7 | services: 8 | redis_1: 9 | service: "redis" 10 | values: 11 | cluster: 12 | enable: false 13 | master: 14 | service: 15 | type: NodePort 16 | auth: 17 | password: "abcd1234" 18 | repository: "oci://registry-1.docker.io/bitnamicharts" 19 | chart: "redis" 20 | expedient: true 21 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BismuthCloud/asimov/ba2b9e28878406a9f81738718e7f094a5fbbdbe3/tests/__init__.py -------------------------------------------------------------------------------- /tests/test_anthropic_inference_client.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from unittest.mock import AsyncMock, MagicMock, patch 3 | from asimov.services.inference_clients import ( 4 | AnthropicInferenceClient, 5 | ChatMessage, 6 | ChatRole, 7 | ) 8 | 9 | 10 | @pytest.mark.asyncio 11 | async def test_tool_chain_streaming(): 12 | # Mock response chunks that simulate the streaming response 13 | response_chunks = [ 14 | 'data: {"type":"message_start","message":{"id":"msg_123","type":"message","role":"assistant","model":"claude-3-5-sonnet-20241022","usage":{"input_tokens":10,"output_tokens":0},"content":[],"stop_reason":null}}\n', 15 | 'data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n', 16 | 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Let me check"}}\n', 17 | 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" the weather"}}\n', 18 | 'data: {"type":"content_block_stop","index":0}\n', 19 | 'data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_123","name":"get_weather","input":{}}}\n', 20 | 'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\\"location\\": \\""}}\n', 21 | 'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"San Francisco"}}\n', 22 | 'data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\\"}"}}\n', 23 | 'data: {"type":"content_block_stop","index":1}\n', 24 | 'data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}\n', 25 | 'data: {"type":"message_stop"}\n', 26 | ] 27 | 28 | # Mock the weather tool function 29 | get_weather = AsyncMock(return_value={"temperature": 72, "conditions": "sunny"}) 30 | 31 | tools = [ 32 | ( 33 | get_weather, 34 | {"name": "get_weather", "description": "Get weather for location"}, 35 | ) 36 | ] 37 | 38 | # Create test messages 39 | messages = [ 40 | ChatMessage(role=ChatRole.SYSTEM, content="You are a helpful assistant"), 41 | ChatMessage( 42 | role=ChatRole.USER, content="What's the weather like in San Francisco?" 43 | ), 44 | ] 45 | 46 | # Mock the HTTP client's stream method 47 | mock_aiter = MagicMock(name="aiter") 48 | mock_aiter.__aiter__.return_value = response_chunks 49 | 50 | mock_response = AsyncMock() 51 | mock_response.status_code = 200 52 | mock_response.aiter_lines = MagicMock(return_value=mock_aiter) 53 | mock_response.__aenter__.return_value = mock_response 54 | mock_response.__aexit__.return_value = MagicMock() 55 | 56 | mock_client = AsyncMock() 57 | mock_client.stream = MagicMock(return_value=mock_response) 58 | mock_client.__aenter__.return_value = mock_client 59 | mock_client.__aexit__.return_value = MagicMock() 60 | 61 | # Create the client and patch httpx.AsyncClient 62 | client = AnthropicInferenceClient("claude-3-5-sonnet-20241022", "test-key") 63 | 64 | with patch("httpx.AsyncClient") as mock_http_client: 65 | mock_http_client.return_value = mock_client 66 | 67 | # Call tool_chain 68 | result = await client.tool_chain( 69 | messages=messages, tools=tools, max_iterations=1 70 | ) 71 | 72 | # Verify the request was made correctly 73 | mock_client.stream.assert_called_once() 74 | call_args = mock_client.stream.call_args 75 | assert call_args[0][0] == "POST" # First positional arg is the method 76 | assert call_args[1]["headers"]["anthropic-version"] == "2023-06-01" 77 | assert call_args[1]["json"]["stream"] == True 78 | assert call_args[1]["json"]["tools"] == [tools[0][1]] 79 | 80 | # Verify the result contains the expected messages 81 | print(result) 82 | assert len(result) == 3 # user, assistant, tool result 83 | assert result[0]["role"] == "user" 84 | assert result[1]["role"] == "assistant" 85 | assert result[1]["content"][0]["type"] == "text" 86 | assert result[1]["content"][0]["text"] == "Let me check the weather" 87 | assert result[1]["content"][1]["type"] == "tool_use" 88 | assert result[1]["content"][1]["name"] == "get_weather" 89 | assert result[1]["content"][1]["input"] == {"location": "San Francisco"} 90 | assert result[2]["role"] == "user" 91 | assert result[2]["content"][0]["type"] == "tool_result" 92 | assert result[2]["content"][0]["content"] == "{'temperature': 72, 'conditions': 'sunny'}" 93 | 94 | get_weather.assert_awaited_once_with({"location": "San Francisco"}) 95 | -------------------------------------------------------------------------------- /tests/test_basic_agent.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import asyncio 3 | from typing import Dict, Any 4 | 5 | from asimov.graph import ( 6 | Agent, 7 | ModuleType, 8 | Node, 9 | NodeConfig, 10 | FlowControlModule, 11 | FlowControlConfig, 12 | FlowDecision, 13 | ) 14 | from asimov.graph.tasks import Task 15 | from asimov.caches.mock_redis_cache import MockRedisCache 16 | from examples.basic_agent import TextPlannerModule, TextExecutorModule 17 | 18 | 19 | @pytest.fixture 20 | def mock_cache(): 21 | return MockRedisCache() 22 | 23 | 24 | @pytest.fixture 25 | def basic_agent(mock_cache): 26 | """Setup and cleanup for basic agent tests.""" 27 | print("Setting up basic agent for test") 28 | # Setup 29 | agent = Agent(cache=mock_cache, max_concurrent_tasks=1, max_total_iterations=10) 30 | 31 | # Create nodes 32 | planner_node = Node( 33 | name="planner", 34 | type=ModuleType.EXECUTOR, 35 | modules=[TextPlannerModule()], 36 | node_config=NodeConfig(parallel=False, max_retries=3), 37 | ) 38 | 39 | executor_node = Node( 40 | name="executor", 41 | type=ModuleType.EXECUTOR, 42 | modules=[TextExecutorModule()], 43 | dependencies=["planner"], 44 | ) 45 | 46 | flow_control = Node( 47 | name="flow_control", 48 | type=ModuleType.FLOW_CONTROL, 49 | dependencies=["executor"], 50 | modules=[ 51 | FlowControlModule( 52 | name="flow_control", 53 | type=ModuleType.FLOW_CONTROL, 54 | flow_config=FlowControlConfig( 55 | decisions=[ 56 | FlowDecision(next_node="executor", condition="plan ~= null") 57 | ], 58 | default="planner", 59 | ), 60 | ) 61 | ], 62 | ) 63 | 64 | # Add nodes to agent 65 | agent.add_multiple_nodes([planner_node, executor_node, flow_control]) 66 | 67 | return agent 68 | 69 | 70 | @pytest.mark.asyncio 71 | async def test_basic_agent_initialization(basic_agent): 72 | """Test that the basic agent is initialized with correct components.""" 73 | assert len(basic_agent.nodes) == 3 74 | assert "planner" in basic_agent.nodes 75 | assert "executor" in basic_agent.nodes 76 | assert "flow_control" in basic_agent.nodes 77 | 78 | # Verify node types 79 | assert basic_agent.nodes["planner"].type == ModuleType.EXECUTOR 80 | assert basic_agent.nodes["executor"].type == ModuleType.EXECUTOR 81 | assert basic_agent.nodes["flow_control"].type == ModuleType.FLOW_CONTROL 82 | 83 | 84 | @pytest.mark.asyncio 85 | async def test_text_planning(basic_agent, mock_cache): 86 | """Test the text planning functionality.""" 87 | task = Task( 88 | type="text_processing", 89 | objective="Test planning", 90 | params={"text": "Hello world"}, 91 | ) 92 | 93 | await mock_cache.set("task", task) 94 | 95 | # Run just the planner node 96 | planner_node = basic_agent.nodes["planner"] 97 | await planner_node.run(mock_cache, asyncio.Semaphore()) 98 | 99 | # Verify plan was created and stored 100 | plan = await mock_cache.get("plan") 101 | assert plan is not None 102 | assert "operations" in plan 103 | assert len(plan["operations"]) == 2 104 | assert plan["operations"][0]["type"] == "count_words" 105 | assert plan["operations"][1]["type"] == "calculate_stats" 106 | 107 | 108 | @pytest.mark.asyncio 109 | async def test_text_execution(basic_agent, mock_cache): 110 | """Test the text execution functionality.""" 111 | # Setup test data 112 | test_plan = { 113 | "operations": [ 114 | {"type": "count_words", "text": "Hello world"}, 115 | {"type": "calculate_stats", "text": "Hello world"}, 116 | ] 117 | } 118 | test_task = Task( 119 | type="text_processing", 120 | objective="Test execution", 121 | params={"text": "Hello world"}, 122 | ) 123 | 124 | # Store test data in cache 125 | await mock_cache.set("plan", test_plan) 126 | await mock_cache.set("task", test_task) 127 | 128 | # Run executor node 129 | executor_node = basic_agent.nodes["executor"] 130 | result = await executor_node.run(mock_cache, asyncio.Semaphore()) 131 | 132 | # Verify results 133 | result = result["results"][0] 134 | assert result["status"] == "success" 135 | assert len(result) == 2 136 | word_count_result = result["result"][0] 137 | assert word_count_result["operation"] == "count_words" 138 | assert word_count_result["result"] == 2 # "Hello world" has 2 words 139 | 140 | stats_result = result["result"][1] 141 | assert stats_result["operation"] == "calculate_stats" 142 | assert stats_result["result"]["characters"] == 11 # Length of "Hello world" 143 | assert stats_result["result"]["lines"] == 1 144 | 145 | 146 | import asyncio 147 | from functools import partial 148 | 149 | 150 | @pytest.mark.asyncio 151 | @pytest.mark.timeout(5) # 5 second timeout 152 | async def test_end_to_end_processing(basic_agent, mock_cache): 153 | # Add logging 154 | print("Starting end-to-end test") 155 | """Test complete end-to-end text processing.""" 156 | task = Task( 157 | type="text_processing", 158 | objective="Process sample text", 159 | params={ 160 | "text": "Hello world!\nThis is a sample text.\nIt demonstrates the basic agent functionality." 161 | }, 162 | ) 163 | 164 | # Run the task 165 | print(f"Running task: {task.objective}") 166 | try: 167 | await asyncio.wait_for( 168 | basic_agent.run_task(task), timeout=4.0 169 | ) # 4 second timeout 170 | print("Task completed successfully") 171 | except asyncio.TimeoutError: 172 | print("Task execution timed out") 173 | raise 174 | except Exception as e: 175 | print(f"Task execution failed: {str(e)}") 176 | raise 177 | 178 | print("Verifying executor results") 179 | # Get final results from the executor node 180 | result = basic_agent.node_results.get("executor", {}) 181 | assert result.get("status") == "success" 182 | 183 | results = result.get("results", [])[0]["result"] 184 | assert len(results) == 2 185 | 186 | # Verify word count operation 187 | word_count_result = results[0] 188 | assert word_count_result["operation"] == "count_words" 189 | assert word_count_result["result"] == 13 # Count of words in the test text 190 | 191 | # Verify stats operation 192 | stats_result = results[1] 193 | assert stats_result["operation"] == "calculate_stats" 194 | assert stats_result["result"]["lines"] == 3 # Number of lines in test text 195 | assert stats_result["result"]["characters"] > 0 # Should have characters 196 | 197 | 198 | @pytest.mark.asyncio 199 | async def test_error_handling(basic_agent, mock_cache): 200 | """Test error handling with invalid input.""" 201 | task = Task( 202 | type="text_processing", 203 | objective="Process invalid input", 204 | params={}, # Missing required 'text' parameter 205 | ) 206 | 207 | # Run the task and expect it to complete (even with empty text) 208 | print(f"Running error handling test with task: {task.objective}") 209 | try: 210 | await asyncio.wait_for( 211 | basic_agent.run_task(task), timeout=4.0 212 | ) # 4 second timeout 213 | print("Error handling test completed") 214 | except asyncio.TimeoutError: 215 | print("Error handling test timed out") 216 | raise 217 | except Exception as e: 218 | print(f"Error handling test failed: {str(e)}") 219 | raise 220 | 221 | print("Verifying error handling results") 222 | # Verify results still contain expected structure 223 | result = basic_agent.node_results.get("executor", {}) 224 | assert result.get("status") == "success" 225 | 226 | results = result.get("results", [])[0]["result"] 227 | assert len(results) == 2 228 | 229 | # Both operations should handle empty text gracefully 230 | word_count_result = results[0] 231 | assert word_count_result["result"] == 0 # Empty text has 0 words 232 | 233 | stats_result = results[1] 234 | assert stats_result["result"]["characters"] == 0 235 | assert stats_result["result"]["lines"] == 0 236 | -------------------------------------------------------------------------------- /tests/test_cache.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import pytest_asyncio 3 | from asimov.caches.redis_cache import RedisCache 4 | 5 | 6 | @pytest_asyncio.fixture 7 | async def redis_cache(): 8 | cache = RedisCache(default_prefix="test") 9 | await cache.clear() 10 | yield cache 11 | await cache.close() 12 | 13 | 14 | @pytest.mark.asyncio 15 | async def test_mailbox(redis_cache): 16 | MBOX = "test_mailbox" 17 | 18 | assert (await redis_cache.get_message(MBOX, timeout=0.5)) is None 19 | await redis_cache.publish_to_mailbox(MBOX, {"k": "v"}) 20 | assert (await redis_cache.get_message(MBOX, timeout=0.5)) == {"k": "v"} 21 | 22 | 23 | @pytest.mark.asyncio 24 | async def test_get_all(redis_cache): 25 | """ 26 | Mainly just regression test that get_all works with (i.e. excludes) list types used for mailboxes 27 | """ 28 | prefix = redis_cache.default_prefix 29 | 30 | await redis_cache.publish_to_mailbox("mailbox", {"k": "v"}) 31 | await redis_cache.set("key1", "value1") 32 | await redis_cache.set("key2", 2) 33 | assert (await redis_cache.get_all()) == { 34 | f"{prefix}:key1": "value1", 35 | f"{prefix}:key2": 2, 36 | } 37 | -------------------------------------------------------------------------------- /tests/test_llm_agent.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import pytest_asyncio 3 | import asyncio 4 | import os 5 | from typing import Dict, Any, List 6 | from unittest.mock import MagicMock, AsyncMock, patch 7 | import json 8 | 9 | from asimov.graph import ( 10 | Agent, 11 | ModuleType, 12 | Node, 13 | NodeConfig, 14 | FlowControlModule, 15 | FlowControlConfig, 16 | FlowDecision, 17 | ) 18 | from asimov.graph.tasks import Task 19 | from asimov.caches.mock_redis_cache import MockRedisCache 20 | from examples.llm_agent import LLMPlannerModule, LLMExecutorModule, LLMFlowControlModule 21 | 22 | 23 | class MockAnthropicClient: 24 | """Mock client for testing LLM interactions.""" 25 | 26 | def __init__(self, responses=None): 27 | self.responses = responses or {} 28 | self.default_response = {} 29 | 30 | async def get_generation(self, messages: List[Any]) -> str: 31 | # Extract prompt from messages 32 | prompt = messages[-1].content if messages else "" 33 | # Return predefined response if available, otherwise default 34 | for key, response in self.responses.items(): 35 | if key.lower() in prompt.lower(): 36 | return response 37 | return json.dumps(self.default_response) 38 | 39 | 40 | @pytest.fixture(autouse=True) 41 | def setup_env(): 42 | """Set up test environment variables.""" 43 | os.environ["ANTHROPIC_API_KEY"] = "test-api-key" 44 | yield 45 | if "ANTHROPIC_API_KEY" in os.environ: 46 | del os.environ["ANTHROPIC_API_KEY"] 47 | 48 | 49 | @pytest.fixture 50 | def mock_cache(): 51 | return MockRedisCache() 52 | 53 | 54 | @pytest.fixture 55 | def mock_anthropic_client(): 56 | return MockAnthropicClient( 57 | { 58 | "Create a step-by-step plan": '{"steps": [{"description": "Gather ingredients and equipment", "requirements": "Recipe ingredients list and tools", "validation": "All ingredients and tools are present and measured"}, {"description": "Prepare cake batter", "requirements": "Ingredients and mixing equipment", "validation": "Batter has proper consistency and ingredients are well combined"}]}', 59 | "Execute this step": "Step executed successfully with the following results:\n1. Actions: Gathered and measured all ingredients\n2. Outcome: All ingredients and tools ready for baking\n3. Output: Ingredients measured and organized according to recipe", 60 | "Evaluate if the step": "success - all ingredients are properly measured and equipment is ready", 61 | "Analyze the execution history": "Analysis complete. Decision: continue - ingredient preparation is complete and accurate", 62 | } 63 | ) 64 | 65 | 66 | @pytest.fixture 67 | def llm_agent(mock_cache, mock_anthropic_client): 68 | # Create modules with mock client 69 | planner = LLMPlannerModule() 70 | planner.client = mock_anthropic_client 71 | 72 | executor = LLMExecutorModule() 73 | executor.client = mock_anthropic_client 74 | 75 | flow_control_module = LLMFlowControlModule() 76 | flow_control_module.client = mock_anthropic_client 77 | 78 | # Create agent 79 | agent = Agent(cache=mock_cache, max_concurrent_tasks=1, max_total_iterations=20) 80 | 81 | # Create nodes 82 | planner_node = Node( 83 | name="planner", 84 | type=ModuleType.EXECUTOR, 85 | modules=[planner], 86 | node_config=NodeConfig(parallel=False, max_retries=3), 87 | ) 88 | 89 | executor_node = Node( 90 | name="executor", 91 | type=ModuleType.EXECUTOR, 92 | modules=[executor], 93 | dependencies=["planner"], 94 | ) 95 | 96 | flow_control_node = Node( 97 | name="flow_control_llm", 98 | type=ModuleType.FLOW_CONTROL, 99 | modules=[flow_control_module], 100 | dependencies=["executor"], 101 | ) 102 | 103 | # Add nodes to agent 104 | agent.add_multiple_nodes([planner_node, executor_node, flow_control_node]) 105 | return agent 106 | 107 | 108 | @pytest.mark.asyncio 109 | async def test_llm_agent_initialization(llm_agent): 110 | """Test that the LLM agent is initialized with correct components.""" 111 | assert len(llm_agent.nodes) == 3 112 | assert "planner" in llm_agent.nodes 113 | assert "executor" in llm_agent.nodes 114 | assert "flow_control_llm" in llm_agent.nodes 115 | 116 | # Verify node types 117 | assert llm_agent.nodes["planner"].type == ModuleType.EXECUTOR 118 | assert llm_agent.nodes["executor"].type == ModuleType.EXECUTOR 119 | assert llm_agent.nodes["flow_control_llm"].type == ModuleType.FLOW_CONTROL 120 | 121 | 122 | @pytest.mark.asyncio 123 | async def test_llm_planning(llm_agent, mock_cache): 124 | """Test the LLM-based planning functionality.""" 125 | task = Task( 126 | type="cake_baking", 127 | objective="Bake a chocolate cake", 128 | params={"cake_type": "Chocolate", "servings": "8-10", "difficulty": "intermediate"}, 129 | ) 130 | 131 | await mock_cache.set("task", task) 132 | 133 | # Run planner node 134 | planner_node = llm_agent.nodes["planner"] 135 | result = await planner_node.run(mock_cache, asyncio.Semaphore()) 136 | 137 | # Verify planning result 138 | assert result["status"] == "success" 139 | 140 | print(result) 141 | assert "Plan created successfully" in result["results"][0]["result"] 142 | 143 | # Verify plan was stored in cache 144 | plan = await mock_cache.get("plan") 145 | 146 | print(plan) 147 | assert isinstance(plan, list) 148 | assert len(plan) > 0 149 | assert all(key in plan[0] for key in ["description", "requirements", "validation"]) 150 | 151 | 152 | @pytest.mark.asyncio 153 | async def test_llm_execution(llm_agent, mock_cache): 154 | """Test the LLM-based execution functionality.""" 155 | # Setup test data 156 | task = Task( 157 | type="cake_baking", 158 | objective="Bake a chocolate cake", 159 | params={"cake_type": "Chocolate", "servings": "8-10", "difficulty": "intermediate"}, 160 | ) 161 | plan = json.dumps( 162 | [ 163 | { 164 | "description": "Gather ingredients and equipment", 165 | "requirements": "Recipe ingredients list and tools", 166 | "validation": "All ingredients and tools are present and measured", 167 | } 168 | ] 169 | ) 170 | 171 | # Store test data in cache 172 | await mock_cache.set("plan", plan) 173 | await mock_cache.set("current_step", 0) 174 | await mock_cache.set("task", task) 175 | 176 | # Run executor node 177 | executor_node = llm_agent.nodes["executor"] 178 | result = await executor_node.run(mock_cache, asyncio.Semaphore()) 179 | 180 | result = result["results"][0] 181 | 182 | # Verify execution results 183 | assert result["status"] == "success" 184 | assert "step" in result["result"] 185 | assert "execution_result" in result["result"] 186 | assert "validation" in result["result"] 187 | assert "success" in result["result"]["validation"].lower() 188 | 189 | # Verify execution history was updated 190 | execution_history = await mock_cache.get("execution_history") 191 | assert execution_history is not None 192 | assert len(execution_history) == 1 # Should have one entry after first execution 193 | 194 | history_entry = execution_history[0] 195 | assert "step" in history_entry 196 | assert "execution_result" in history_entry 197 | assert "validation" in history_entry 198 | assert "status" in history_entry 199 | assert "timestamp" in history_entry 200 | assert history_entry["status"] == "success" 201 | assert history_entry["step"] == "Gather ingredients and equipment" 202 | 203 | 204 | @pytest.mark.asyncio 205 | @pytest.mark.timeout(30) # 30 second timeout 206 | async def test_end_to_end_processing(llm_agent, mock_cache): 207 | print("Starting end-to-end LLM agent test") 208 | """Test complete end-to-end cake baking process.""" 209 | task = Task( 210 | type="cake_baking", 211 | objective="Bake a chocolate cake", 212 | params={ 213 | "cake_type": "Chocolate", 214 | "servings": "8-10", 215 | "difficulty": "intermediate", 216 | "style": "classic recipe", 217 | }, 218 | ) 219 | 220 | # Run the task 221 | print(f"Running task: {task.objective}") 222 | try: 223 | await asyncio.wait_for( 224 | llm_agent.run_task(task), timeout=25.0 225 | ) # 25 second timeout 226 | print("Task completed successfully") 227 | except asyncio.TimeoutError: 228 | print("Task execution timed out") 229 | raise 230 | except Exception as e: 231 | print(f"Task execution failed: {str(e)}") 232 | raise 233 | 234 | print("Verifying node results") 235 | # Verify results from each node 236 | planner_result = llm_agent.node_results.get("planner", {}) 237 | assert planner_result.get("status") == "success" 238 | 239 | executor_result = llm_agent.node_results.get("executor", {}) 240 | assert executor_result.get("status") == "success" 241 | result = executor_result.get("results", [])[0]["result"] 242 | assert "step" in result 243 | 244 | 245 | @pytest.mark.asyncio 246 | async def test_missing_api_key(): 247 | """Test that modules properly handle missing API key.""" 248 | if "ANTHROPIC_API_KEY" in os.environ: 249 | del os.environ["ANTHROPIC_API_KEY"] 250 | 251 | with pytest.raises( 252 | ValueError, match="ANTHROPIC_API_KEY environment variable must be set" 253 | ): 254 | LLMPlannerModule() 255 | 256 | with pytest.raises( 257 | ValueError, match="ANTHROPIC_API_KEY environment variable must be set" 258 | ): 259 | LLMExecutorModule() 260 | 261 | with pytest.raises( 262 | ValueError, match="ANTHROPIC_API_KEY environment variable must be set" 263 | ): 264 | LLMFlowControlModule() 265 | 266 | 267 | @pytest.mark.asyncio 268 | async def test_error_handling(llm_agent, mock_cache): 269 | """Test error handling with problematic LLM responses.""" 270 | # Create a client that simulates errors 271 | error_client = MockAnthropicClient( 272 | { 273 | "Create a step-by-step plan": "Invalid JSON response", 274 | "Execute this step": "Error: Unable to process step", 275 | "Evaluate if the step": "failure - validation criteria not met", 276 | } 277 | ) 278 | 279 | # Update modules with error client 280 | llm_agent.nodes["planner"].modules[0].client = error_client 281 | llm_agent.nodes["executor"].modules[0].client = error_client 282 | 283 | task = Task( 284 | type="cake_baking", 285 | objective="Bake a chocolate cake", 286 | params={"cake_type": "Chocolate", "servings": "8-10", "difficulty": "intermediate"}, 287 | ) 288 | 289 | # Run the task and expect it to handle errors gracefully 290 | await llm_agent.run_task(task) 291 | 292 | # Verify error handling in results 293 | planner_result = llm_agent.node_results.get("planner", {}) 294 | 295 | print(llm_agent.node_results) 296 | assert planner_result.get("status") in ["success", "error"] 297 | 298 | # If planning succeeded despite invalid JSON, executor should handle the error 299 | if "executor" in llm_agent.node_results: 300 | executor_result = llm_agent.node_results["executor"] 301 | assert executor_result.get("status") in ["success", "error"] 302 | -------------------------------------------------------------------------------- /tests/test_readme_examples.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | from unittest.mock import patch, MagicMock 3 | import asyncio 4 | from asimov.graph import AgentModule, ModuleType 5 | from asimov.caches.mock_redis_cache import MockRedisCache 6 | from asimov.services.inference_clients import InferenceClient, ChatMessage, ChatRole 7 | 8 | class TestReadmeExamples(unittest.TestCase): 9 | def setUp(self): 10 | self.cache = MockRedisCache() 11 | 12 | async def test_basic_graph_setup(self): 13 | """Test the basic setup example that will be shown in README""" 14 | graph = Graph() 15 | cache = self.cache 16 | client = InferenceClient(cache=cache) 17 | 18 | # Verify components are initialized correctly 19 | self.assertIsInstance(graph, Graph) 20 | self.assertIsInstance(cache, MockRedisCache) 21 | self.assertIsInstance(client, InferenceClient) 22 | 23 | # Verify cache connection 24 | self.assertTrue(await cache.ping()) 25 | 26 | async def test_cache_operations(self): 27 | """Test cache operations example for README""" 28 | cache = self.cache 29 | 30 | # Basic cache operations 31 | await cache.set("test_key", {"value": 42}) 32 | result = await cache.get("test_key") 33 | self.assertEqual(result, {"value": 42}) 34 | 35 | # Mailbox functionality 36 | await cache.publish_to_mailbox("test_mailbox", {"message": "hello"}) 37 | message = await cache.get_message("test_mailbox", timeout=1) 38 | self.assertEqual(message, {"message": "hello"}) 39 | 40 | async def test_graph_module_execution(self): 41 | """Test graph module execution example for README""" 42 | class TestModule(AgentModule): 43 | name = "test_module" 44 | type = ModuleType.EXECUTOR 45 | 46 | async def process(self, cache, semaphore): 47 | return {"status": "success", "result": "test completed"} 48 | 49 | graph = Graph() 50 | cache = self.cache 51 | module = TestModule() 52 | 53 | # Execute module 54 | semaphore = asyncio.Semaphore(1) 55 | result = await module.run(cache, semaphore) 56 | self.assertEqual(result["status"], "success") 57 | self.assertEqual(result["result"], "test completed") 58 | 59 | async def test_inference_client_usage(self): 60 | """Test inference client usage example for README""" 61 | cache = self.cache 62 | client = InferenceClient(cache=cache) 63 | 64 | messages = [ 65 | ChatMessage(role=ChatRole.SYSTEM, content="You are a helpful assistant."), 66 | ChatMessage(role=ChatRole.USER, content="Hello!") 67 | ] 68 | 69 | # Mock the generation method since we can't make real API calls in tests 70 | with patch.object(client, 'get_generation') as mock_generate: 71 | mock_generate.return_value = "Hello! How can I help you today?" 72 | response = await client.get_generation(messages) 73 | self.assertEqual(response, "Hello! How can I help you today?") 74 | self.assertTrue(await cache.ping()) 75 | 76 | async def test_cache_operations(self): 77 | """Test cache operations example for README""" 78 | cache = self.cache 79 | 80 | # Basic cache operations 81 | await cache.set("test_key", {"value": 42}) 82 | result = await cache.get("test_key") 83 | self.assertEqual(result, {"value": 42}) 84 | 85 | # Mailbox functionality 86 | await cache.publish_to_mailbox("test_mailbox", {"message": "hello"}) 87 | message = await cache.get_message("test_mailbox", timeout=1) 88 | self.assertEqual(message, {"message": "hello"}) 89 | 90 | async def test_graph_module_execution(self): 91 | """Test graph module execution example for README""" 92 | class TestModule(AgentModule): 93 | name = "test_module" 94 | type = ModuleType.EXECUTOR 95 | 96 | async def process(self, cache, semaphore): 97 | return {"status": "success", "result": "test completed"} 98 | 99 | graph = Graph() 100 | cache = self.cache 101 | module = TestModule() 102 | 103 | # Execute module 104 | semaphore = asyncio.Semaphore(1) 105 | result = await module.run(cache, semaphore) 106 | self.assertEqual(result["status"], "success") 107 | self.assertEqual(result["result"], "test completed") 108 | 109 | async def test_inference_client_usage(self): 110 | """Test inference client usage example for README""" 111 | cache = self.cache 112 | client = InferenceClient(cache=cache) 113 | 114 | messages = [ 115 | ChatMessage(role=ChatRole.SYSTEM, content="You are a helpful assistant."), 116 | ChatMessage(role=ChatRole.USER, content="Hello!") 117 | ] 118 | 119 | # Mock the generation method since we can't make real API calls in tests 120 | with patch.object(client, 'get_generation') as mock_generate: 121 | mock_generate.return_value = "Hello! How can I help you today?" 122 | response = await client.get_generation(messages) 123 | self.assertEqual(response, "Hello! How can I help you today?") --------------------------------------------------------------------------------