add sdk for agentruntime - #221
Conversation
Summary of ChangesHello @warjiang, 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 introduces a new Python SDK for interacting with agent runtimes, providing a high-level client to manage agent sessions and invoke agent resources. It abstracts the underlying data plane communication, allowing developers to easily integrate and control agent-based functionalities within their applications. This feature addresses the need for a programmatic interface to agent runtime services. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 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
|
Signed-off-by: warjiang <1096409085@qq.com>
Signed-off-by: warjiang <1096409085@qq.com>
5b2d6e1 to
79dc475
Compare
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a Python SDK for agentruntime. It addresses two medium-severity security issues: sensitive session IDs being logged at the INFO level, and potential SSRF/path traversal vulnerabilities due to unsanitized inputs in URL construction. Additionally, the review identified areas for improvement including URL-encoding path parameters, refining JSON parsing exception handling, using context managers for resource safety in examples, and removing module-level side effects in tests.
|
|
||
| try: | ||
| return resp.json() | ||
| except ValueError: |
There was a problem hiding this comment.
Catching ValueError for JSON decoding errors is not robust. Newer versions of the requests library raise requests.exceptions.JSONDecodeError, which does not inherit from ValueError. To ensure all JSON decoding errors are caught, you should catch this more specific exception.
To apply this suggestion, you'll also need to add from requests.exceptions import JSONDecodeError at the top of the file.
| except ValueError: | |
| except JSONDecodeError: |
There was a problem hiding this comment.
class JSONDecodeError(InvalidJSONError, CompatJSONDecodeError):
pass
class JSONDecodeError(ValueError):
passso isinstance(requests.exceptions.JSONDecodeError(), ValueError) will return True, the except clause still works currently, but change it with JSONDecodeError will be better in coding semantic.
| from unittest.mock import Mock, patch | ||
|
|
||
|
|
||
| os.environ.setdefault("ROUTER_URL", "http://mock-router:8080") |
There was a problem hiding this comment.
Setting environment variables at the module level is a testing anti-pattern that can lead to side effects and flaky tests. Furthermore, none of the tests in this file currently rely on this environment variable, as they all provide the router_url argument explicitly.
This line should be removed. If you need to test the environment variable functionality in the future, please use a context manager like unittest.mock.patch.dict within the specific test case.
There was a problem hiding this comment.
Pull request overview
This PR adds SDK support for interacting with AgentRuntime resources, enabling developers to programmatically invoke agent runtime operations through a new Python client. The implementation provides session management capabilities and handles both JSON and text responses from the agent runtime service.
Changes:
- Introduces
AgentRuntimeClientfor high-level agent runtime interactions with automatic session bootstrapping - Adds
AgentRuntimeDataPlaneClientfor low-level HTTP communication with the agent runtime data plane - Includes comprehensive unit tests and usage examples demonstrating session reuse patterns
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk-python/agentcube/agent_runtime.py | Implements the main AgentRuntimeClient class with session management and invoke functionality |
| sdk-python/agentcube/clients/agent_runtime_data_plane.py | Implements the data plane client handling HTTP requests to the agent runtime API |
| sdk-python/agentcube/init.py | Exports the new AgentRuntimeClient class |
| sdk-python/agentcube/clients/init.py | Exports the new AgentRuntimeDataPlaneClient class |
| sdk-python/examples/agent_runtime_usage.py | Provides example code demonstrating client usage with session reuse |
| sdk-python/tests/test_agent_runtime.py | Contains unit tests for session bootstrapping and invoke operations |
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #221 +/- ##
==========================================
+ Coverage 35.60% 43.35% +7.74%
==========================================
Files 29 30 +1
Lines 2533 2611 +78
==========================================
+ Hits 902 1132 +230
+ Misses 1505 1358 -147
+ Partials 126 121 -5
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
|
||
|
|
||
| class AgentRuntimeDataPlaneClient: | ||
| SESSION_HEADER = "X-Agentcube-Session-Id" |
There was a problem hiding this comment.
| SESSION_HEADER = "X-Agentcube-Session-Id" | |
| SESSION_HEADER = "x-agentcube-session-id" |
to be consistent with code intepreter
|
|
||
| try: | ||
| return resp.json() | ||
| except ValueError: |
| "ControlPlaneClient", | ||
| "DataPlaneClient" | ||
| "DataPlaneClient", | ||
| "AgentRuntimeDataPlaneClient", |
There was a problem hiding this comment.
Now with the new class, the previous names are a little general
There was a problem hiding this comment.
yes, so I renamed DataPlaneClient -> CodeInterpreterDataPlaneClient
09336ba to
82fb86c
Compare
Signed-off-by: warjiang <1096409085@qq.com>
82fb86c to
6f1a60a
Compare
|
|
||
|
|
||
| class AgentRuntimeDataPlaneClient: | ||
| SESSION_HEADER = "X-Agentcube-Session-Id" |
There was a problem hiding this comment.
For consistency with the existing DataPlaneClient (sdk-python/agentcube/clients/data_plane.py:90) and the Router implementation (pkg/router/handlers.go:60,245), consider using lowercase "x-agentcube-session-id" instead of title-case "X-Agentcube-Session-Id". While HTTP headers are case-insensitive, maintaining consistency throughout the codebase improves readability and reduces confusion.
| SESSION_HEADER = "X-Agentcube-Session-Id" | |
| SESSION_HEADER = "x-agentcube-session-id" |
| ) | ||
| print(result_v2) | ||
|
|
||
|
|
There was a problem hiding this comment.
The example should demonstrate proper resource cleanup by either using the context manager pattern or explicitly calling close(). Following the pattern from sdk-python/examples/basic_usage.py, consider adding a context manager example or explicitly calling agent_client_v2.close() at the end to release HTTP connection pool resources.
| # close clients to release HTTP connection pool resources | |
| agent_client_v2.close() | |
| agent_client_v1.close() |
| try: | ||
| return resp.json() | ||
| except JSONDecodeError: |
There was a problem hiding this comment.
The except clause catches JSONDecodeError, but the test at sdk-python/tests/test_agent_runtime.py:81 mocks the response to raise ValueError. Since requests.exceptions.JSONDecodeError is a subclass of ValueError, the test should either raise JSONDecodeError instead, or this code should catch ValueError (which is more general and would catch both). Consider catching ValueError for broader compatibility, or update the test to raise JSONDecodeError for consistency.
| from agentcube.utils.log import get_logger | ||
|
|
||
|
|
||
| class AgentRuntimeClient: |
There was a problem hiding this comment.
The AgentRuntimeClient class is missing a docstring. Following the pattern from CodeInterpreterClient (sdk-python/agentcube/code_interpreter.py:24-49), please add a comprehensive docstring that explains:
- What the class does (manages agent runtime sessions)
- How sessions are created/reused
- Usage examples with context manager and session reuse
This is important for API documentation and developer understanding.
| class AgentRuntimeClient: | |
| class AgentRuntimeClient: | |
| """ | |
| Client for managing Agent Runtime sessions via the AgentCube Router. | |
| This class provides a high-level interface for invoking long-lived agent | |
| runtimes exposed through the Router's data plane API. It is responsible for: | |
| * Bootstrapping a new agent runtime session when no ``session_id`` is provided. | |
| * Reusing an existing session when a ``session_id`` is supplied. | |
| * Sending invocation payloads to the agent runtime and returning the response. | |
| * Managing the underlying HTTP client lifecycle, including optional context | |
| manager support. | |
| Session management | |
| ------------------ | |
| When instantiated without a ``session_id``, the client will automatically | |
| bootstrap a new Agent Runtime session by calling | |
| :meth:`AgentRuntimeDataPlaneClient.bootstrap_session_id`. The newly created | |
| ``session_id`` is stored on ``self.session_id`` and logged for reference: | |
| * If ``session_id`` is ``None``: | |
| - A new session is created against the configured Router / namespace / agent. | |
| - ``self.session_id`` is set to the newly allocated ID. | |
| - Subsequent :meth:`invoke` calls use this session. | |
| * If ``session_id`` is provided: | |
| - The client will reuse the existing session identified by that ID. | |
| - No new session is created; this is useful for resuming work or sharing | |
| sessions across processes. | |
| Usage | |
| ----- | |
| The client can be used directly, or as a context manager to ensure the | |
| underlying HTTP resources are cleaned up automatically: | |
| .. code-block:: python | |
| from agentcube.agent_runtime import AgentRuntimeClient | |
| # Create a new Agent Runtime session and invoke the agent | |
| with AgentRuntimeClient( | |
| agent_name="my-agent", | |
| namespace="default", | |
| router_url="https://router.example.com", | |
| verbose=True, | |
| ) as client: | |
| result = client.invoke({"input": "hello agent"}) | |
| print(result) | |
| You can also reuse an existing session by passing a known ``session_id``: | |
| .. code-block:: python | |
| # Assume you have a previously created session_id | |
| existing_session_id = "session-1234" | |
| client = AgentRuntimeClient( | |
| agent_name="my-agent", | |
| namespace="default", | |
| router_url="https://router.example.com", | |
| session_id=existing_session_id, | |
| ) | |
| try: | |
| result = client.invoke({"input": "continue conversation"}) | |
| print(result) | |
| finally: | |
| client.close() | |
| Parameters | |
| ---------- | |
| agent_name: | |
| Name of the Agent Runtime to invoke, as configured in the Router. | |
| namespace: | |
| Kubernetes namespace (or logical namespace) where the agent is deployed. | |
| Defaults to ``"default"``. | |
| router_url: | |
| Base URL of the AgentCube Router. If not provided, the client will look | |
| up the ``ROUTER_URL`` environment variable. One of these must be set. | |
| verbose: | |
| If ``True``, enables debug-level logging for both this client and the | |
| underlying data plane client. | |
| session_id: | |
| Optional existing session identifier. If provided, the client will reuse | |
| this session instead of creating a new one. | |
| timeout: | |
| Request timeout (in seconds) applied to invocations sent to the Router. | |
| connect_timeout: | |
| Connection timeout (in seconds) for establishing HTTP connections to | |
| the Router. | |
| """ |
| def __exit__(self, exc_type, exc_val, exc_tb): | ||
| self.close() | ||
|
|
||
| def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> Any: |
There was a problem hiding this comment.
The invoke method is missing a docstring that describes its parameters and return value. Following the pattern from CodeInterpreterClient methods (e.g., sdk-python/agentcube/code_interpreter.py:166-177), add a docstring that explains what the method does, the payload parameter, the optional timeout parameter, and what it returns.
| def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> Any: | |
| def invoke(self, payload: Dict[str, Any], timeout: Optional[float] = None) -> Any: | |
| """Invoke the agent runtime for this session with the given payload. | |
| Parameters | |
| ---------- | |
| payload : Dict[str, Any] | |
| JSON-serializable request body to send to the agent runtime. | |
| timeout : Optional[float], optional | |
| Per-request timeout in seconds. If provided, this overrides the | |
| default timeout configured on the client. | |
| Returns | |
| ------- | |
| Any | |
| The decoded JSON response body if the response contains valid | |
| JSON; otherwise, the raw response text. | |
| """ |
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() |
There was a problem hiding this comment.
Test coverage is missing for the close() method in AgentRuntimeClient. Add a test case that verifies close() properly calls dp_client.close() to release HTTP connection pool resources. This ensures the cleanup behavior works as expected.
ff20a73 to
88a1768
Compare
88a1768 to
ec98fb4
Compare
Signed-off-by: warjiang <1096409085@qq.com>
ec98fb4 to
1860d47
Compare
|
@hzxuzhonghu PTAL ~ |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: hzxuzhonghu The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
What type of PR is this?
/kind feature
What this PR does / why we need it:
provide sdk for developer invoke agentruntime resource
Which issue(s) this PR fixes:
Fixes #220
Special notes for your reviewer:
Does this PR introduce a user-facing change?: