Skip to content

Commit bda3637

Browse files
committed
feat: support expanded reasoning response in agent chat
Add the `data` field to `AgentReasoningItem` to surface tool-call details from the agent's reasoning trace. Introduces `ToolCallDetail`, `ReasoningDataItem`, `ToolCallReasoningDataItem`, and `UnknownReasoningDataItem` following the same type-dispatch pattern used by `MessageContent`.
1 parent 3326f55 commit bda3637

3 files changed

Lines changed: 146 additions & 10 deletions

File tree

cognite/client/data_classes/agents/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,16 @@
3838
Message,
3939
MessageContent,
4040
MessageList,
41+
ReasoningDataItem,
4142
TextContent,
43+
ToolCallDetail,
44+
ToolCallReasoningDataItem,
4245
ToolConfirmationCall,
4346
ToolConfirmationResult,
4447
UnknownAction,
4548
UnknownActionCall,
4649
UnknownContent,
50+
UnknownReasoningDataItem,
4751
)
4852

4953
__all__ = [
@@ -81,14 +85,18 @@
8185
"QueryKnowledgeGraphAgentToolUpsert",
8286
"QueryTimeSeriesDatapointsAgentTool",
8387
"QueryTimeSeriesDatapointsAgentToolUpsert",
88+
"ReasoningDataItem",
8489
"SummarizeDocumentAgentTool",
8590
"SummarizeDocumentAgentToolUpsert",
8691
"TextContent",
92+
"ToolCallDetail",
93+
"ToolCallReasoningDataItem",
8794
"ToolConfirmationCall",
8895
"ToolConfirmationResult",
8996
"UnknownAction",
9097
"UnknownActionCall",
9198
"UnknownAgentTool",
9299
"UnknownAgentToolUpsert",
93100
"UnknownContent",
101+
"UnknownReasoningDataItem",
94102
]

cognite/client/data_classes/agents/chat.py

Lines changed: 109 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -464,25 +464,130 @@ def _load(cls, data: dict[str, Any]) -> AgentDataItem:
464464
return cls(type=item_type, data=item_data)
465465

466466

467+
@dataclass
468+
class ToolCallDetail(CogniteResource):
469+
"""Details of a tool call made during agent reasoning.
470+
471+
Args:
472+
id (str): The id of the tool call.
473+
name (str): The name of the tool that was called.
474+
tool_type (str): The type of the tool that was called.
475+
input (dict[str, Any]): The parameters that were passed to the tool.
476+
result (dict[str, Any]): The results that were returned by the tool.
477+
"""
478+
479+
id: str
480+
name: str
481+
tool_type: str
482+
input: dict[str, Any] = field(default_factory=dict)
483+
result: dict[str, Any] = field(default_factory=dict)
484+
485+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
486+
key = "toolType" if camel_case else "tool_type"
487+
return {"id": self.id, "name": self.name, key: self.tool_type, "input": self.input, "result": self.result}
488+
489+
@classmethod
490+
def _load(cls, data: dict[str, Any]) -> ToolCallDetail:
491+
return cls(
492+
id=data["id"],
493+
name=data["name"],
494+
tool_type=data["toolType"],
495+
input=data.get("input", {}),
496+
result=data.get("result", {}),
497+
)
498+
499+
500+
@dataclass
501+
class ReasoningDataItem(CogniteResource, ABC):
502+
"""Base class for reasoning data item types."""
503+
504+
_type: ClassVar[str]
505+
506+
@classmethod
507+
def _load(cls, data: dict[str, Any]) -> ReasoningDataItem:
508+
item_type = data.get("type", "")
509+
klass = _REASONING_DATA_CLS_BY_TYPE.get(item_type, UnknownReasoningDataItem)
510+
return klass._load_item(data)
511+
512+
@classmethod
513+
@abstractmethod
514+
def _load_item(cls, data: dict[str, Any]) -> ReasoningDataItem: ...
515+
516+
517+
@dataclass
518+
class ToolCallReasoningDataItem(ReasoningDataItem):
519+
"""Reasoning data item for a tool call.
520+
521+
Args:
522+
tool_call (ToolCallDetail | None): Details of the tool call.
523+
"""
524+
525+
_type: ClassVar[str] = "toolCall"
526+
tool_call: ToolCallDetail | None = None
527+
528+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
529+
key = "toolCall" if camel_case else "tool_call"
530+
result: dict[str, Any] = {"type": self._type}
531+
if self.tool_call is not None:
532+
result[key] = self.tool_call.dump(camel_case=camel_case)
533+
return result
534+
535+
@classmethod
536+
def _load_item(cls, data: dict[str, Any]) -> ToolCallReasoningDataItem:
537+
tool_call_data = data.get("toolCall")
538+
return cls(tool_call=ToolCallDetail._load(tool_call_data) if tool_call_data else None)
539+
540+
541+
@dataclass
542+
class UnknownReasoningDataItem(ReasoningDataItem):
543+
"""Unknown reasoning data item type for forward compatibility.
544+
545+
Args:
546+
type (str): The item type.
547+
data (dict[str, Any]): The raw item data.
548+
"""
549+
550+
type: str = ""
551+
data: dict[str, Any] = field(default_factory=dict)
552+
553+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
554+
result = self.data.copy()
555+
result["type"] = self.type
556+
return result
557+
558+
@classmethod
559+
def _load_item(cls, data: dict[str, Any]) -> UnknownReasoningDataItem:
560+
return cls(type=data.get("type", ""), data=data)
561+
562+
563+
_REASONING_DATA_CLS_BY_TYPE: dict[str, type[ReasoningDataItem]] = {
564+
ToolCallReasoningDataItem._type: ToolCallReasoningDataItem,
565+
}
566+
567+
467568
@dataclass
468569
class AgentReasoningItem(CogniteResource):
469570
"""Reasoning item in agent response.
470571
471572
Args:
472573
content (list[MessageContent]): The reasoning content.
574+
data (list[ReasoningDataItem] | None): The data of the reasoning.
473575
"""
474576

475577
content: list[MessageContent]
578+
data: list[ReasoningDataItem] | None = None
476579

477580
def dump(self, camel_case: bool = True) -> dict[str, Any]:
478-
return {
479-
"content": [item.dump(camel_case=camel_case) for item in self.content],
480-
}
581+
result: dict[str, Any] = {"content": [item.dump(camel_case=camel_case) for item in self.content]}
582+
if self.data is not None:
583+
result["data"] = [item.dump(camel_case=camel_case) for item in self.data]
584+
return result
481585

482586
@classmethod
483587
def _load(cls, data: dict[str, Any]) -> AgentReasoningItem:
484588
content = [MessageContent._load(item) for item in data.get("content", [])]
485-
return cls(content=content)
589+
data_items = [ReasoningDataItem._load(item) for item in data.get("data", [])] or None
590+
return cls(content=content, data=data_items)
486591

487592

488593
@dataclass

tests/tests_unit/test_api/test_agents_chat.py

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
AgentMessage,
1515
AgentReasoningItem,
1616
TextContent,
17+
ToolCallReasoningDataItem,
1718
)
1819
from tests.utils import get_url, jsgz_load
1920

@@ -51,7 +52,29 @@ def chat_response_body() -> dict:
5152
"text": "The user is asking about capabilities",
5253
"type": "text",
5354
}
54-
]
55+
],
56+
"data": [
57+
{
58+
"type": "toolCall",
59+
"toolCall": {
60+
"id": "tc_1",
61+
"name": "search_instances",
62+
"toolType": "query",
63+
"input": {
64+
"view_space": "cdf_cdm",
65+
"view_external_id": "CogniteAsset",
66+
"view_version": "v1",
67+
"query": "pump",
68+
"operator": "AND",
69+
"return_properties": ["name", "externalId"],
70+
},
71+
"result": {
72+
"result": {"items": [{"space": "my_space", "externalId": "pump_1"}], "count": 1},
73+
"error": None,
74+
},
75+
},
76+
}
77+
],
5578
}
5679
],
5780
"role": "agent",
@@ -113,11 +136,11 @@ def test_chat_simple_message(
113136
# Check reasoning
114137
assert agent_msg.reasoning is not None
115138
assert len(agent_msg.reasoning) == 1
116-
assert isinstance(agent_msg.reasoning[0], AgentReasoningItem)
117-
assert len(agent_msg.reasoning[0].content) == 1
118-
content = agent_msg.reasoning[0].content[0]
119-
assert isinstance(content, TextContent)
120-
assert content.text == "The user is asking about capabilities"
139+
reasoning_item = agent_msg.reasoning[0]
140+
assert isinstance(reasoning_item, AgentReasoningItem)
141+
assert isinstance(reasoning_item.content[0], TextContent)
142+
assert reasoning_item.data is not None
143+
assert isinstance(reasoning_item.data[0], ToolCallReasoningDataItem)
121144

122145
# Test convenience properties
123146
assert response.text == "I can help you with various tasks related to your industrial data."

0 commit comments

Comments
 (0)