Skip to content

feat(agents): Add ClientToolAction.from_function() for automatic action generation #2351

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

feat(agents): Add ClientToolAction.from_function() for automatic action generation #2351

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d1b9e60
feat(agents): add support for client-side actions in agent chat
ks93 Oct 1, 2025
a6ebeed
refactor(agents): make ActionResult a base class with ClientToolResult
ks93 Oct 1, 2025
55efb96
test(agents): add comprehensive unit tests for actions feature
ks93 Oct 1, 2025
1f6bc6a
refactor(tests): remove redundant action tests, keep 6 essential
ks93 Oct 1, 2025
4ca658f
feat(agents): add ClientToolAction.from_function() with function intr…
ks93 Oct 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 52 additions & 7 deletions cognite/client/_api/agents/agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, overload
from typing import TYPE_CHECKING, Any, overload

from cognite.client._api_client import APIClient
from cognite.client.data_classes.agents import Agent, AgentList, AgentUpsert
from cognite.client.data_classes.agents.chat import AgentChatResponse, Message, MessageList
from cognite.client.data_classes.agents.chat import (
Action,
ActionResult,
AgentChatResponse,
Message,
)
from cognite.client.utils._experimental import FeaturePreviewWarning
from cognite.client.utils._identifier import IdentifierSequence
from cognite.client.utils.useful_types import SequenceNotStr
Expand Down Expand Up @@ -243,8 +248,9 @@ def list(self) -> AgentList: # The API does not yet support limit or pagination
def chat(
self,
agent_id: str,
messages: Message | Sequence[Message],
messages: Message | ActionResult | Sequence[Message | ActionResult],
cursor: str | None = None,
actions: Sequence[Action] | None = None,
) -> AgentChatResponse:
"""`Chat with an agent. <https://api-docs.cognite.com/20230101-beta/tag/Agents/operation/agent_session_ai_agents_chat_post/>`_

Expand All @@ -253,9 +259,10 @@ def chat(

Args:
agent_id (str): External ID that uniquely identifies the agent.
messages (Message | Sequence[Message]): A list of one or many input messages to the agent.
messages (Message | ActionResult | Sequence[Message | ActionResult]): A list of one or many input messages to the agent. Can include regular messages and action results.
cursor (str | None): The cursor to use for continuation of a conversation. Use this to
create multi-turn conversations, as the cursor will keep track of the conversation state.
actions (Sequence[Action] | None): A list of client-side actions that can be called by the agent.

Returns:
AgentChatResponse: The response from the agent.
Expand Down Expand Up @@ -290,22 +297,60 @@ def chat(
... Message("Once you have found it, find related time series.")
... ]
... )

Chat with client-side actions:

>>> from cognite.client.data_classes.agents import ClientToolAction, ClientToolResult
>>> add_numbers_action = ClientToolAction(
... name="add",
... description="Add two numbers together",
... parameters={
... "type": "object",
... "properties": {
... "a": {"type": "number", "description": "First number"},
... "b": {"type": "number", "description": "Second number"},
... },
... "required": ["a", "b"]
... }
... )
>>> response = client.agents.chat(
... agent_id="my_agent",
... messages=Message("What is 42 plus 58?"),
... actions=[add_numbers_action]
... )
>>> if response.action_calls:
... for call in response.action_calls:
... # Execute the action
... result = call.arguments["a"] + call.arguments["b"]
... # Send result back
... response = client.agents.chat(
... agent_id="my_agent",
... messages=ClientToolResult(
... action_id=call.action_id,
... content=f"The result is {result}"
... ),
... cursor=response.cursor,
... actions=[add_numbers_action]
... )
"""
self._warnings.warn()

# Convert single message to list
if isinstance(messages, Message):
if isinstance(messages, (Message, ActionResult)):
messages = [messages]

# Build request body
body = {
body: dict[str, Any] = {
"agentId": agent_id,
"messages": MessageList(messages).dump(camel_case=True),
"messages": [msg.dump(camel_case=True) for msg in messages],
}

if cursor is not None:
body["cursor"] = cursor

if actions is not None:
body["actions"] = [action.dump(camel_case=True) for action in actions]

# Make the API call
response = self._post(
url_path=self._RESOURCE_PATH + "/chat",
Expand Down
18 changes: 18 additions & 0 deletions cognite/client/data_classes/agents/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,19 +21,31 @@
)
from cognite.client.data_classes.agents.agents import Agent, AgentList, AgentUpsert, AgentUpsertList
from cognite.client.data_classes.agents.chat import (
Action,
ActionCall,
ActionResult,
AgentChatResponse,
AgentDataItem,
AgentMessage,
AgentMessageList,
AgentReasoningItem,
ClientToolAction,
ClientToolCall,
ClientToolResult,
Message,
MessageContent,
MessageList,
TextContent,
UnknownAction,
UnknownActionCall,
UnknownActionResult,
UnknownContent,
)

__all__ = [
"Action",
"ActionCall",
"ActionResult",
"Agent",
"AgentChatResponse",
"AgentDataItem",
Expand All @@ -49,6 +61,9 @@
"AgentUpsertList",
"AskDocumentAgentTool",
"AskDocumentAgentToolUpsert",
"ClientToolAction",
"ClientToolCall",
"ClientToolResult",
"DataModelInfo",
"InstanceSpaces",
"Message",
Expand All @@ -62,6 +77,9 @@
"SummarizeDocumentAgentTool",
"SummarizeDocumentAgentToolUpsert",
"TextContent",
"UnknownAction",
"UnknownActionCall",
"UnknownActionResult",
"UnknownAgentTool",
"UnknownAgentToolUpsert",
"UnknownContent",
Expand Down
Loading