Introduce new models for graph execution and added Start_delay #348
Conversation
NiveditJain
commented
Sep 4, 2025
NiveditJain
commented
- Incremented the version in _version.py to 0.0.2b5.
- Added new models in models.py for GraphNode, Unites, RetryPolicy, and StoreConfig, enhancing the structure and validation of graph execution parameters.
- Refactored StateManager methods to utilize the new models, improving type safety and clarity in the trigger and upsert_graph functions.
- Updated trigger_graph to include a start_delay parameter for delayed execution, enhancing flexibility in graph triggering.
- Incremented the version in _version.py to 0.0.2b5. - Added new models in models.py for GraphNode, Unites, RetryPolicy, and StoreConfig, enhancing the structure and validation of graph execution parameters. - Refactored StateManager methods to utilize the new models, improving type safety and clarity in the trigger and upsert_graph functions. - Updated trigger_graph to include a start_delay parameter for delayed execution, enhancing flexibility in graph triggering.
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. Warning Rate limit exceeded@NiveditJain has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 14 minutes and 3 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds typed Pydantic models to the SDK, replaces TriggerState with model-driven APIs, adds start_delay/run_id to trigger flow and persists enqueue_after for delayed scheduling in the state manager, and bumps SDK version to 0.0.2b5. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor Client
participant SDK as SDK StateManager
participant API as State Manager API
participant Ctrl as Trigger Controller
participant DB as State Store
participant Q as Queue/Scheduler
Note over Client,SDK: trigger(graph_name, inputs?, store?, start_delay=0)
Client->>SDK: trigger(..., start_delay)
SDK->>API: POST /trigger {graph_name, inputs, store, start_delay}
API->>Ctrl: Validate & process
Ctrl->>DB: Create State {..., enqueue_after = now + start_delay}
DB-->>Ctrl: State saved (with run_id)
alt start_delay == 0
Ctrl->>Q: Enqueue immediately
else start_delay > 0
Ctrl->>Q: Schedule for enqueue_after
end
Ctrl-->>API: {run_id, ...}
API-->>SDK: response
SDK-->>Client: run_id
sequenceDiagram
autonumber
actor Dev
participant SDK as SDK StateManager
participant API as State Manager API
Note over Dev,SDK: upsert_graph with typed models
Dev->>SDK: upsert_graph(graph_name, [GraphNodeModel...], secrets, retry_policy?, store_config?)
SDK->>SDK: Serialize models via model_dump
SDK->>API: POST /graphs/upsert {graph_nodes, retry_policy, store_config, ...}
API-->>SDK: ack/validation response
SDK-->>Dev: result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @NiveditJain, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly refines the handling of graph execution parameters by introducing robust Pydantic models for various configurations, including graph nodes, retry policies, and store settings. This foundational change improves data validation and type safety across the system. Additionally, it introduces a new capability for scheduling graph executions with a specified delay, providing greater control over workflow initiation.
Highlights
- New Data Models: Introduced new Pydantic models (GraphNode, Unites, RetryPolicy, StoreConfig) to enhance the structure and validation of graph execution parameters, improving type safety and clarity.
- Refactored StateManager: Updated StateManager methods to leverage the newly introduced models, specifically for
triggerandupsert_graphfunctions, ensuring better data integrity and readability. - Delayed Graph Execution: Added a
start_delayparameter to thetrigger_graphfunction, allowing for delayed execution of graphs, which enhances scheduling flexibility. - Version Update: Incremented the project version to 0.0.2b5.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Code Review
This pull request introduces new Pydantic models for graph execution, enhancing type safety and validation, and adds a start_delay parameter for delayed execution. The refactoring to use these models in StateManager is a great improvement. My review includes several suggestions for the new models.py file to improve the validation logic, making it more concise and robust.
- Removed TriggerState from exports and added new models: UnitesStrategyEnum, UnitesModel, GraphNodeModel, RetryStrategyEnum, RetryPolicyModel, and StoreConfigModel. - Updated __all__ to reflect the changes in exported components, enhancing the module's structure and usability.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (9)
state-manager/app/models/trigger_model.py (1)
9-11: Use UUID type for run_id for stronger validation and clearer OpenAPI.This keeps wire format identical (string) but adds input/output validation.
Apply:
-from pydantic import BaseModel, Field +from pydantic import BaseModel, Field +from uuid import UUID @@ -class TriggerGraphResponseModel(BaseModel): +class TriggerGraphResponseModel(BaseModel): status: StateStatusEnum = Field(..., description="Status of the states") - run_id: str = Field(..., description="Unique run ID generated for this graph execution") + run_id: UUID = Field(..., description="Unique run ID for this graph execution")state-manager/app/controller/trigger_graph.py (3)
18-24: Guard against missing store_config on templates.If
graph_template.store_configcan beNone, this raises. Make the check tolerant.Apply:
-def check_required_store_keys(graph_template: GraphTemplate, store: dict[str, str]) -> None: - required_keys = set(graph_template.store_config.required_keys) +def check_required_store_keys(graph_template: GraphTemplate, store: dict[str, str]) -> None: + cfg = getattr(graph_template, "store_config", None) + required_keys = set(cfg.required_keys) if cfg else set() @@ - if missing_keys: - raise HTTPException(status_code=400, detail=f"Missing store keys: {missing_keys}") + if missing_keys: + missing = ", ".join(sorted(missing_keys)) + raise HTTPException(status_code=400, detail=f"Missing store keys: {missing}")
113-115: Log the exception with stack trace.Currently the error log loses traceback context.
Apply:
- logger.error(f"Error triggering graph {graph_name} for namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + logger.error( + f"Error triggering graph {graph_name} for namespace {namespace_name}", + x_exosphere_request_id=x_exosphere_request_id, + exc_info=True + )
94-106: Index enqueue_after for efficient delayed dispatch scans.Add a DB index on
State.enqueue_afterto avoid full scans when selecting ready states.If using MongoDB: ensure an index like
{ enqueue_after: 1, status: 1 }. If SQL: add a B-tree index on(status, enqueue_after).python-sdk/exospherehost/statemanager.py (5)
87-91: Add client timeouts to all HTTP calls.External calls lack timeouts; this can hang indefinitely.
Apply:
- async with aiohttp.ClientSession() as session: + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30)) as session: @@ - async with aiohttp.ClientSession() as session: + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=30)) as session: @@ - async with aiohttp.ClientSession() as session: + async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=60)) as session:Tune totals as needed.
Also applies to: 123-126, 170-173
73-77: Validate that inputs and store are str->str maps before sending.Prevents 422s from the API and makes errors local.
Apply:
if inputs is None: inputs = {} if store is None: store = {} + if not all(isinstance(k, str) and isinstance(v, str) for k, v in inputs.items()): + raise TypeError("inputs must be a dict[str, str]") + if not all(isinstance(k, str) and isinstance(v, str) for k, v in store.items()): + raise TypeError("store must be a dict[str, str]")
177-186: Validation polling loop is correct; consider jitter to reduce thundering herd.Optional: add small random jitter to
polling_interval.
1-5: Consider reusing an aiohttp session per StateManager instance.Creating a session per call adds overhead; reuse with a context manager or aclose in
__aenter__/__aexit__.
1-1: Restore theTriggerStateclass instatemanager.py.
TriggerStateis imported by__init__.pyand used throughout the tests but no longer defined instatemanager.py. Add or re-export the missing class to satisfy those imports and keep CI passing.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
python-sdk/exospherehost/_version.py(1 hunks)python-sdk/exospherehost/models.py(1 hunks)python-sdk/exospherehost/statemanager.py(7 hunks)state-manager/app/controller/trigger_graph.py(2 hunks)state-manager/app/models/trigger_model.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
python-sdk/exospherehost/statemanager.py (1)
python-sdk/exospherehost/models.py (3)
GraphNodeModel(16-75)RetryPolicyModel(92-97)StoreConfigModel(100-158)
🪛 GitHub Actions: Python SDK Tests
python-sdk/exospherehost/statemanager.py
[error] 1-1: TriggerState is not defined or exported in statemanager.py, causing ImportError during import.
🔇 Additional comments (7)
python-sdk/exospherehost/_version.py (1)
1-1: Semver check: this bump likely includes a breaking API removal.
TriggerStatewas removed fromstatemanager.py(tests fail). Either add a deprecation shim (preferred) or bump a breaking version once tests and docs are updated.python-sdk/exospherehost/statemanager.py (1)
128-168: Minor: model_dump usage looks correct and aligned with Pydantic v2.Serialization of nodes/retry/store_config is sound.
python-sdk/exospherehost/models.py (5)
6-9: LGTM! Clean enum implementation.The
UnitesStrategyEnumis well-defined with clear, descriptive values for graph execution strategies.
11-14: LGTM! Well-structured model with appropriate defaults.The
UnitesModelproperly uses Pydantic's Field with descriptions and sets a sensible default strategy.
78-90: LGTM! Comprehensive retry strategy options.The
RetryStrategyEnumprovides a complete set of retry strategies with various jitter options, following industry best practices for retry mechanisms.
92-98: LGTM! Well-designed retry policy with proper constraints.The
RetryPolicyModelhas sensible defaults and appropriate validation constraints (ge=0, gt=0) for each field.
154-154: Coercion ofdefault_valuesto strings is appropriate. The field is declared asdict[str, str]and all current tests and code paths expect string values, so converting viastr(value)matches intended behavior.
- Removed instances of TriggerState from tests, replacing them with a dictionary structure for state representation. - Updated test assertions to use GraphNodeModel for graph node definitions, enhancing clarity and consistency in the test suite. - Adjusted related test cases to ensure compatibility with the new state representation, improving overall test robustness.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
python-sdk/tests/test_package_init.py (2)
87-101: Test currently doesn’t import the package between measurements; it can’t detect side effects.Import (and reload) the package to assert idempotency of import-time logging setup.
Apply this diff:
def test_import_without_side_effects(): """Test that importing the package doesn't cause side effects.""" - import logging + import logging, importlib @@ - # Store initial state - initial_handlers = len(logging.getLogger().handlers) - - # Check that logging handlers weren't added unexpectedly - # (The package might add handlers during import, which is expected) - current_handlers = len(logging.getLogger().handlers) - - # The package should either not add handlers or add them consistently - assert current_handlers >= initial_handlers + initial_handlers = len(logging.getLogger().handlers) + import exospherehost # first import + first_import_handlers = len(logging.getLogger().handlers) + # Reload should not keep stacking handlers + importlib.reload(exospherehost) + second_import_handlers = len(logging.getLogger().handlers) + # Allow one-time handler addition on first import, but ensure idempotency + assert second_import_handlers == first_import_handlers
26-41: Also assert new model re-exports are importable from the top level.This guards against accidental regressions in init.py.
Apply this diff to add a small test:
def test_runtime_class_import(): @@ assert Runtime is RuntimeDirect + + +def test_model_exports_top_level(): + from exospherehost import GraphNodeModel, RetryPolicyModel, StoreConfigModel + assert GraphNodeModel is not None + assert RetryPolicyModel is not None + assert StoreConfigModel is not Nonepython-sdk/tests/test_statemanager_comprehensive.py (1)
95-113: Expand coverage: start_delay propagation and optional model fields on upsert.
- Trigger tests validate success/error paths but don’t assert that a non-zero
start_delayis forwarded.- Upsert tests don’t cover serialization of
retry_policyandstore_config.Add the following tests to close these gaps.
Apply this diff:
@@ class TestStateManagerTrigger: @@ async def test_trigger_single_state_success(self, state_manager_config): @@ assert result == {"status": "success"} + + @pytest.mark.asyncio + async def test_trigger_includes_nonzero_start_delay(self, state_manager_config): + with patch('exospherehost.statemanager.aiohttp.ClientSession') as mock_session_class: + mock_session, mock_post_response, *_ = create_mock_aiohttp_session() + mock_post_response.status = 200 + mock_post_response.json = AsyncMock(return_value={"status": "success"}) + mock_session_class.return_value = mock_session + sm = StateManager(**state_manager_config) + await sm.trigger("test_graph", inputs={"k": "v"}, start_delay=2500) + mock_session.post.assert_called_once() + _, kwargs = mock_session.post.call_args + assert kwargs["json"]["start_delay"] == 2500 @@ class TestStateManagerUpsertGraph: @@ async def test_upsert_graph_success_201(self, state_manager_config): @@ assert result["name"] == "test_graph" + + @pytest.mark.asyncio + async def test_upsert_graph_serializes_retry_policy_and_store_config(self, state_manager_config): + with patch('exospherehost.statemanager.aiohttp.ClientSession') as mock_session_class, \ + patch('exospherehost.statemanager.StateManager.get_graph') as mock_get_graph: + mock_session, _, _, mock_put_response = create_mock_aiohttp_session() + mock_put_response.status = 201 + mock_put_response.json = AsyncMock(return_value={"name": "g", "validation_status": "PENDING"}) + mock_session_class.return_value = mock_session + mock_get_graph.side_effect = [{"validation_status": "VALID", "name": "g"}] + + sm = StateManager(**state_manager_config) + graph_nodes = [GraphNodeModel(node_name="n1", namespace="ns", identifier="n1", inputs={}, next_nodes=None, unites=None)] + from exospherehost.models import RetryPolicyModel, RetryStrategyEnum, StoreConfigModel + retry = RetryPolicyModel(max_retries=3, strategy=RetryStrategyEnum.EXPONENTIAL, backoff_factor=0.5) + store = StoreConfigModel(required_keys=["cursor"], default_values={"cursor": "0"}) + + result = await sm.upsert_graph("g", graph_nodes, secrets={}, retry_policy=retry, store_config=store) + assert result["validation_status"] == "VALID" + # Ensure payload included both fields + mock_session.put.assert_called_once() + _, kwargs = mock_session.put.call_args + body = kwargs["json"] + assert "retry_policy" in body and "store_config" in bodyAlso applies to: 188-260, 300-356, 378-397
python-sdk/tests/test_integration.py (2)
163-166: Tests rely on private_node_mapping; prefer a public helper to register state->node mapping.Depending on a private attribute is brittle and may break with internal refactors. If feasible, expose a tiny helper like
Runtime._bind_state_node_for_tests(state_id, node_cls)or public configure hook gated behind a test flag.I can open a follow-up PR to add a small testing hook if you’d like.
Also applies to: 310-312, 374-376, 499-505
392-398: Unifystate_manager_versionnaming across Runtime and tests
Rename thestate_manage_versionparameter inRuntime.__init__(and its docstring/assignment) tostate_manager_version, and update all callers (tests and code) to usestate_manager_versionrather thanstate_manage_version.
- In
python-sdk/exospherehost/runtime.py
• Change the__init__signature at line 83 tostate_manager_version: str = "v0"
• Update the docstring parameter name at line 71 and the assignment at line 96 to match- In test files, replace all
state_manage_version=withstate_manager_version=(e.g. lines 396 & 414 intests/test_integration.py, occurrences intest_signals_and_runtime_functions.pyandtest_runtime_comprehensive.py)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
python-sdk/exospherehost/__init__.py(1 hunks)python-sdk/tests/test_coverage_additions.py(1 hunks)python-sdk/tests/test_integration.py(3 hunks)python-sdk/tests/test_package_init.py(1 hunks)python-sdk/tests/test_statemanager_comprehensive.py(10 hunks)
🧰 Additional context used
🧬 Code graph analysis (4)
python-sdk/exospherehost/__init__.py (2)
python-sdk/exospherehost/statemanager.py (1)
StateManager(9-188)python-sdk/exospherehost/models.py (4)
UnitesModel(11-13)GraphNodeModel(16-75)RetryStrategyEnum(78-89)StoreConfigModel(100-158)
python-sdk/tests/test_statemanager_comprehensive.py (1)
python-sdk/exospherehost/statemanager.py (2)
StateManager(9-188)trigger(37-91)
python-sdk/tests/test_integration.py (2)
python-sdk/exospherehost/statemanager.py (3)
StateManager(9-188)upsert_graph(128-188)trigger(37-91)python-sdk/exospherehost/models.py (1)
GraphNodeModel(16-75)
python-sdk/tests/test_package_init.py (1)
python-sdk/exospherehost/statemanager.py (1)
StateManager(9-188)
🔇 Additional comments (3)
python-sdk/exospherehost/__init__.py (1)
40-46: API surface change looks consistent; verify docs and migration note.
- Re-exports and all align with the new models; looks good.
- The top-level docstring still shows
runtime.connect([...]); runtime.start(). If the current Runtime API expectsnodes=[...]in the constructor (as used across tests), refresh the example to avoid confusion.Do you want me to open a follow-up PR to (a) update the docstring example to the constructor-based API and (b) add a migration note for removing
TriggerStatefrom the public API?python-sdk/tests/test_package_init.py (1)
2-2: Exports updated to match new public API — looks good.The test now tracks the added model exports and removal of
TriggerState; aligns with init.py.Also applies to: 17-17
python-sdk/tests/test_integration.py (1)
208-217: LGTM: switching to GraphNodeModel for upserts.The model-based approach improves validation and readability.
- Updated the `validate_required_keys` and `validate_default_values` methods in `StoreConfigModel` to be class methods, enhancing their functionality and consistency with the model's design. - Introduced a new test file for comprehensive testing of `GraphNodeModel`, `StoreConfigModel`, and `StateManager`, including validation checks and default behavior, improving test coverage and reliability.
…y to streamline the test file.
1 participant
