feat(agents): add actions support for custom tool calls in chat endpoint #2337
Conversation
- Add ClientTool, ClientToolParameters, and ClientToolList data classes - Update agents.chat() method to accept actions parameter - Support single ClientTool or sequence of ClientTool objects - Add comprehensive tests for new functionality - Update documentation with examples - Maintain backward compatibility with existing chat functionality The actions parameter allows clients to inject custom tool calls that agents can use during conversations, following the JSON Schema specification for parameter definitions.
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. |
|
/gemini review |
![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 adds support for custom tool calls in the agents chat endpoint by introducing ClientTool, ClientToolParameters, and ClientToolList data classes. The changes are well-structured and include comprehensive tests. My main feedback is to align ClientToolList with the existing CogniteResourceList pattern in the SDK for consistency and to leverage its functionality. This also requires updating the corresponding tests.
|
|
||
| from typing import Any | ||
|
|
||
| from cognite.client.data_classes._base import CogniteResource |
There was a problem hiding this comment.
To align ClientToolList with the SDK's pattern of using CogniteResourceList for resource collections, CogniteResourceList should be imported here. This is a prerequisite for the suggested change to the ClientToolList class.
| from cognite.client.data_classes._base import CogniteResource | |
| from cognite.client.data_classes._base import CogniteResource, CogniteResourceList |
| class ClientToolList: | ||
| """List of client tools.""" | ||
|
|
||
| def __init__(self, tools: list[ClientTool]) -> None: | ||
| self.tools = tools | ||
|
|
||
| def dump(self, camel_case: bool = True) -> list[dict[str, Any]]: | ||
| """Dump the client tools to a list of dictionaries.""" | ||
| return [tool.dump(camel_case=camel_case) for tool in self.tools] | ||
|
|
||
| @classmethod | ||
| def _load(cls, data: list[dict[str, Any]], cognite_client: Any = None) -> ClientToolList: | ||
| """Load client tools from a list of dictionaries.""" | ||
| tools = [ClientTool._load(tool_data) for tool_data in data] | ||
| return cls(tools=tools) |
There was a problem hiding this comment.
For consistency with other resource list classes in the SDK, ClientToolList should inherit from CogniteResourceList[ClientTool]. This provides useful methods like to_pandas(), get(), and proper handling of cognite_client propagation for free. The current implementation re-implements dump and _load and misses passing cognite_client during loading, which is a potential bug.1
class ClientToolList(CogniteResourceList[ClientTool]):
"""List of client tools."""
_RESOURCE = ClientToolStyle Guide References
Footnotes
-
Follow established patterns across the codebase. ↩
| class TestClientToolList: | ||
| def test_init(self) -> None: | ||
| """Test ClientToolList initialization.""" | ||
| tool1 = ClientTool(name="tool1", parameters=ClientToolParameters()) | ||
| tool2 = ClientTool(name="tool2", parameters=ClientToolParameters()) | ||
| tools = [tool1, tool2] | ||
|
|
||
| tool_list = ClientToolList(tools) | ||
| assert tool_list.tools == tools | ||
|
|
||
| def test_dump_camel_case(self) -> None: | ||
| """Test dumping ClientToolList with camel case.""" | ||
| tool1 = ClientTool( | ||
| name="tool1", | ||
| parameters=ClientToolParameters(properties={"param1": {"type": "string"}}), | ||
| ) | ||
| tool2 = ClientTool( | ||
| name="tool2", | ||
| parameters=ClientToolParameters(properties={"param2": {"type": "integer"}}), | ||
| ) | ||
| tool_list = ClientToolList([tool1, tool2]) | ||
|
|
||
| result = tool_list.dump(camel_case=True) | ||
| assert len(result) == 2 | ||
| assert result[0]["name"] == "tool1" | ||
| assert result[1]["name"] == "tool2" | ||
| assert result[0]["parameters"]["properties"]["param1"]["type"] == "string" | ||
| assert result[1]["parameters"]["properties"]["param2"]["type"] == "integer" | ||
|
|
||
| def test_load_from_list(self) -> None: | ||
| """Test loading ClientToolList from list of dictionaries.""" | ||
| data = [ | ||
| { | ||
| "name": "tool1", | ||
| "type": "clientTool", | ||
| "parameters": { | ||
| "type": "object", | ||
| "properties": {"param1": {"type": "string"}}, | ||
| "required": [], | ||
| "propertyOrdering": [], | ||
| }, | ||
| }, | ||
| { | ||
| "name": "tool2", | ||
| "type": "clientTool", | ||
| "parameters": { | ||
| "type": "object", | ||
| "properties": {"param2": {"type": "integer"}}, | ||
| "required": [], | ||
| "propertyOrdering": [], | ||
| }, | ||
| }, | ||
| ] | ||
|
|
||
| tool_list = ClientToolList._load(data) | ||
| assert len(tool_list.tools) == 2 | ||
| assert tool_list.tools[0].name == "tool1" | ||
| assert tool_list.tools[1].name == "tool2" | ||
| assert isinstance(tool_list.tools[0], ClientTool) | ||
| assert isinstance(tool_list.tools[1], ClientTool) |
There was a problem hiding this comment.
Following the change of ClientToolList to inherit from CogniteResourceList, these tests need to be updated. The list's items are now in the data attribute (from UserList) and can be accessed directly via indexing or iteration, not through a .tools attribute.
class TestClientToolList:
def test_init(self) -> None:
"""Test ClientToolList initialization."""
tool1 = ClientTool(name="tool1", parameters=ClientToolParameters())
tool2 = ClientTool(name="tool2", parameters=ClientToolParameters())
tools = [tool1, tool2]
tool_list = ClientToolList(tools)
assert tool_list.data == tools
def test_dump_camel_case(self) -> None:
"""Test dumping ClientToolList with camel case."""
tool1 = ClientTool(
name="tool1",
parameters=ClientToolParameters(properties={"param1": {"type": "string"}}),
)
tool2 = ClientTool(
name="tool2",
parameters=ClientToolParameters(properties={"param2": {"type": "integer"}}),
)
tool_list = ClientToolList([tool1, tool2])
result = tool_list.dump(camel_case=True)
assert len(result) == 2
assert result[0]["name"] == "tool1"
assert result[1]["name"] == "tool2"
assert result[0]["parameters"]["properties"]["param1"]["type"] == "string"
assert result[1]["parameters"]["properties"]["param2"]["type"] == "integer"
def test_load_from_list(self) -> None:
"""Test loading ClientToolList from list of dictionaries."""
data = [
{
"name": "tool1",
"type": "clientTool",
"parameters": {
"type": "object",
"properties": {"param1": {"type": "string"}},
"required": [],
"propertyOrdering": [],
},
},
{
"name": "tool2",
"type": "clientTool",
"parameters": {
"type": "object",
"properties": {"param2": {"type": "integer"}},
"required": [],
"propertyOrdering": [],
},
},
]
tool_list = ClientToolList._load(data)
assert len(tool_list) == 2
assert tool_list[0].name == "tool1"
assert tool_list[1].name == "tool2"
assert isinstance(tool_list[0], ClientTool)
assert isinstance(tool_list[1], ClientTool)- Make ClientToolList inherit from CogniteResourceList[ClientTool] - Update tests to work with CogniteResourceList interface - Maintains all existing functionality while following SDK patterns - Addresses code review feedback for consistency
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #2337 +/- ##
==========================================
- Coverage 90.66% 90.65% -0.01%
==========================================
Files 169 170 +1
Lines 25184 25224 +40
==========================================
+ Hits 22832 22866 +34
- Misses 2352 2358 +6
🚀 New features to boost your workflow:
|
The actions parameter allows clients to inject custom tool calls that agents can use during conversations, following the JSON Schema specification for parameter definitions.
Description
Checklist:
If a new method has been added it should be referenced in cognite.rst in order to generate docs based on its docstring.