Skip to content

Commit dda47b6

Browse files
ks93ks93
committed
feat: add QueryAgentTool for the "query" tool type
The cog-ai backend and Fusion frontend already support the "query" tool type, but the Python SDK deserialized it as UnknownAgentTool. This adds first-class support with QueryAgentTool, QueryAgentToolUpsert, and QueryAgentToolConfiguration classes.
1 parent 8e291fc commit dda47b6

5 files changed

Lines changed: 243 additions & 2 deletions

File tree

‎cognite/client/_api/agents/agents.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,37 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A
7979
... )
8080
>>> client.agents.upsert(agents=[agent])
8181
82+
Create an agent with the query tool (Preview):
83+
84+
>>> from cognite.client.data_classes.agents import (
85+
... AgentUpsert,
86+
... QueryAgentToolUpsert,
87+
... QueryAgentToolConfiguration,
88+
... DataModelInfo,
89+
... InstanceSpaces,
90+
... )
91+
>>> query_tool = QueryAgentToolUpsert(
92+
... name="explore data",
93+
... description="Run flexible queries against your data model",
94+
... configuration=QueryAgentToolConfiguration(
95+
... data_models=[
96+
... DataModelInfo(
97+
... space="cdf_idm",
98+
... external_id="CogniteProcessIndustries",
99+
... version="v1",
100+
... )
101+
... ],
102+
... instance_spaces=InstanceSpaces(type="all"),
103+
... ),
104+
... )
105+
>>> agent = AgentUpsert(
106+
... external_id="my_agent",
107+
... name="My Agent",
108+
... labels=["published"],
109+
... tools=[query_tool],
110+
... )
111+
>>> client.agents.upsert(agents=[agent])
112+
82113
Create an agent with multiple different tools:
83114
84115
>>> from cognite.client.data_classes.agents import (

‎cognite/client/_sync_api/agents/agents.py‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
"""
22
===============================================================================
3-
bfc6535f11bf8b175836aebca409eb8b
3+
d88277929dfc151715f521eb7176622b
44
This file is auto-generated from the Async API modules, - do not edit manually!
55
===============================================================================
66
"""
@@ -77,6 +77,37 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi
7777
... )
7878
>>> client.agents.upsert(agents=[agent])
7979
80+
Create an agent with the query tool (Preview):
81+
82+
>>> from cognite.client.data_classes.agents import (
83+
... AgentUpsert,
84+
... QueryAgentToolUpsert,
85+
... QueryAgentToolConfiguration,
86+
... DataModelInfo,
87+
... InstanceSpaces,
88+
... )
89+
>>> query_tool = QueryAgentToolUpsert(
90+
... name="explore data",
91+
... description="Run flexible queries against your data model",
92+
... configuration=QueryAgentToolConfiguration(
93+
... data_models=[
94+
... DataModelInfo(
95+
... space="cdf_idm",
96+
... external_id="CogniteProcessIndustries",
97+
... version="v1",
98+
... )
99+
... ],
100+
... instance_spaces=InstanceSpaces(type="all"),
101+
... ),
102+
... )
103+
>>> agent = AgentUpsert(
104+
... external_id="my_agent",
105+
... name="My Agent",
106+
... labels=["published"],
107+
... tools=[query_tool],
108+
... )
109+
>>> client.agents.upsert(agents=[agent])
110+
80111
Create an agent with multiple different tools:
81112
82113
>>> from cognite.client.data_classes.agents import (

‎cognite/client/data_classes/agents/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99
AskDocumentAgentToolUpsert,
1010
DataModelInfo,
1111
InstanceSpaces,
12+
QueryAgentTool,
13+
QueryAgentToolConfiguration,
14+
QueryAgentToolUpsert,
1215
QueryKnowledgeGraphAgentTool,
1316
QueryKnowledgeGraphAgentToolConfiguration,
1417
QueryKnowledgeGraphAgentToolUpsert,
@@ -70,6 +73,9 @@
7073
"Message",
7174
"MessageContent",
7275
"MessageList",
76+
"QueryAgentTool",
77+
"QueryAgentToolConfiguration",
78+
"QueryAgentToolUpsert",
7379
"QueryKnowledgeGraphAgentTool",
7480
"QueryKnowledgeGraphAgentToolConfiguration",
7581
"QueryKnowledgeGraphAgentToolUpsert",

‎cognite/client/data_classes/agents/agent_tools.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,43 @@ def as_write(self) -> QueryKnowledgeGraphAgentToolConfiguration:
194194
return self
195195

196196

197+
@dataclass
198+
class QueryAgentToolConfiguration(WriteableCogniteResource):
199+
"""Configuration for query agent tools.
200+
201+
Args:
202+
data_models (Sequence[DataModelInfo]): The data models to query.
203+
instance_spaces (InstanceSpaces | None): The instance spaces to query.
204+
"""
205+
206+
data_models: Sequence[DataModelInfo]
207+
instance_spaces: InstanceSpaces | None = None
208+
209+
@classmethod
210+
def _load(cls, resource: dict[str, Any]) -> QueryAgentToolConfiguration:
211+
dm_config = resource.get("dataModels", {})
212+
data_models = [DataModelInfo._load(dm) for dm in dm_config.get("dataModels", [])]
213+
return cls(
214+
data_models=data_models,
215+
instance_spaces=InstanceSpaces._load_if(resource.get("instanceSpaces")),
216+
)
217+
218+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
219+
result: dict[str, Any] = {}
220+
key = "dataModels" if camel_case else "data_models"
221+
result[key] = {
222+
"type": "manual",
223+
key: [dm.dump(camel_case=camel_case) for dm in self.data_models],
224+
}
225+
if self.instance_spaces:
226+
key = "instanceSpaces" if camel_case else "instance_spaces"
227+
result[key] = self.instance_spaces.dump(camel_case=camel_case)
228+
return result
229+
230+
def as_write(self) -> QueryAgentToolConfiguration:
231+
return self
232+
233+
197234
@dataclass
198235
class SummarizeDocumentAgentTool(AgentTool):
199236
"""Agent tool for summarizing documents.
@@ -389,6 +426,69 @@ def _load(cls, resource: dict[str, Any]) -> QueryTimeSeriesDatapointsAgentToolUp
389426
)
390427

391428

429+
@dataclass
430+
class QueryAgentTool(AgentTool):
431+
"""Agent tool for running flexible queries against data models.
432+
433+
Args:
434+
name (str): The name of the agent tool. Used by the agent to decide when to use this tool.
435+
description (str): The description of the agent tool. Used by the agent to decide when to use this tool.
436+
configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool.
437+
"""
438+
439+
_type: ClassVar[str] = "query"
440+
configuration: QueryAgentToolConfiguration | None = None
441+
442+
@classmethod
443+
def _load_tool(cls, resource: dict[str, Any]) -> QueryAgentTool:
444+
return cls(
445+
name=resource["name"],
446+
description=resource["description"],
447+
configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")),
448+
)
449+
450+
def as_write(self) -> QueryAgentToolUpsert:
451+
return QueryAgentToolUpsert(
452+
name=self.name,
453+
description=self.description,
454+
configuration=self.configuration,
455+
)
456+
457+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
458+
result = super().dump(camel_case=camel_case)
459+
if self.configuration:
460+
result["configuration"] = self.configuration.dump(camel_case=camel_case)
461+
return result
462+
463+
464+
@dataclass
465+
class QueryAgentToolUpsert(AgentToolUpsert):
466+
"""Upsert version of query agent tool.
467+
468+
Args:
469+
name (str): The name of the agent tool. Used by the agent to decide when to use this tool.
470+
description (str): The description of the agent tool. Used by the agent to decide when to use this tool.
471+
configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool.
472+
"""
473+
474+
_type: ClassVar[str] = "query"
475+
configuration: QueryAgentToolConfiguration | None = None
476+
477+
def dump(self, camel_case: bool = True) -> dict[str, Any]:
478+
result = super().dump(camel_case=camel_case)
479+
if self.configuration:
480+
result["configuration"] = self.configuration.dump(camel_case=camel_case)
481+
return result
482+
483+
@classmethod
484+
def _load(cls, resource: dict[str, Any]) -> QueryAgentToolUpsert:
485+
return cls(
486+
name=resource["name"],
487+
description=resource["description"],
488+
configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")),
489+
)
490+
491+
392492
@dataclass
393493
class UnknownAgentTool(AgentTool):
394494
"""Agent tool for unknown/unrecognized tool types.

‎tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
from cognite.client.data_classes.agents.agent_tools import (
66
AgentTool,
77
AskDocumentAgentTool,
8+
QueryAgentTool,
9+
QueryAgentToolConfiguration,
810
QueryKnowledgeGraphAgentTool,
911
QueryKnowledgeGraphAgentToolConfiguration,
1012
QueryTimeSeriesDatapointsAgentTool,
@@ -42,6 +44,32 @@
4244
"description": "Query the time series datapoints",
4345
}
4446

47+
query_example = {
48+
"name": "queryExample",
49+
"type": "query",
50+
"description": "Run flexible queries against your data model",
51+
"configuration": {
52+
"dataModels": {
53+
"type": "manual",
54+
"dataModels": [
55+
{
56+
"space": "cdf_idm",
57+
"externalId": "CogniteProcessIndustries",
58+
"version": "v1",
59+
"viewExternalIds": ["CogniteAsset"],
60+
}
61+
],
62+
},
63+
"instanceSpaces": {"type": "manual", "spaces": ["my_space"]},
64+
},
65+
}
66+
67+
query_no_config_example = {
68+
"name": "queryNoConfigExample",
69+
"type": "query",
70+
"description": "Run flexible queries against your data model",
71+
}
72+
4573
unknown_example = {
4674
"name": "unknownExample",
4775
"type": "yolo", # This is not a known tool type
@@ -58,9 +86,19 @@ class TestAgentToolLoad:
5886
(ask_document_example, AskDocumentAgentTool),
5987
(summarize_document_example, SummarizeDocumentAgentTool),
6088
(query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool),
89+
(query_example, QueryAgentTool),
90+
(query_no_config_example, QueryAgentTool),
6191
(unknown_example, UnknownAgentTool),
6292
],
63-
ids=["queryKnowledgeGraph", "askDocument", "summarizeDocument", "queryTimeSeriesDatapoints", "somethingElse"],
93+
ids=[
94+
"queryKnowledgeGraph",
95+
"askDocument",
96+
"summarizeDocument",
97+
"queryTimeSeriesDatapoints",
98+
"query",
99+
"queryNoConfig",
100+
"somethingElse",
101+
],
64102
)
65103
def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected_type: type[AgentTool]) -> None:
66104
"""Test that AgentTool._load() returns the correct subtype based on the tool type."""
@@ -82,6 +120,9 @@ def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected
82120
assert isinstance(loaded_tool.configuration, QueryKnowledgeGraphAgentToolConfiguration)
83121
# Compare by serializing the structured object back to dict
84122
assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"]
123+
elif isinstance(loaded_tool, QueryAgentTool):
124+
assert isinstance(loaded_tool.configuration, QueryAgentToolConfiguration)
125+
assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"]
85126
elif isinstance(loaded_tool, UnknownAgentTool):
86127
# For other tools (like UnknownAgentTool), configuration should be a dict
87128
assert loaded_tool.configuration == tool_data["configuration"]
@@ -105,6 +146,7 @@ class TestAgentToolDump:
105146
(ask_document_example, AskDocumentAgentTool),
106147
(summarize_document_example, SummarizeDocumentAgentTool),
107148
(query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool),
149+
(query_example, QueryAgentTool),
108150
(unknown_example, UnknownAgentTool),
109151
],
110152
)
@@ -150,6 +192,7 @@ class TestAgentToolUpsert:
150192
(ask_document_example, AskDocumentAgentTool),
151193
(summarize_document_example, SummarizeDocumentAgentTool),
152194
(query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool),
195+
(query_example, QueryAgentTool),
153196
(unknown_example, UnknownAgentTool),
154197
],
155198
)
@@ -165,3 +208,33 @@ def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_
165208

166209
assert dumped_tool["name"] == tool_data["name"]
167210
assert dumped_tool["description"] == tool_data["description"]
211+
212+
213+
class TestQueryAgentTool:
214+
def test_load_with_configuration(self) -> None:
215+
loaded = AgentTool._load(query_example)
216+
assert isinstance(loaded, QueryAgentTool)
217+
assert loaded.configuration is not None
218+
assert isinstance(loaded.configuration, QueryAgentToolConfiguration)
219+
assert len(loaded.configuration.data_models) == 1
220+
assert loaded.configuration.data_models[0].space == "cdf_idm"
221+
assert loaded.configuration.data_models[0].external_id == "CogniteProcessIndustries"
222+
assert loaded.configuration.instance_spaces is not None
223+
assert loaded.configuration.instance_spaces.type == "manual"
224+
assert loaded.configuration.instance_spaces.spaces == ["my_space"]
225+
226+
def test_load_without_configuration(self) -> None:
227+
loaded = AgentTool._load(query_no_config_example)
228+
assert isinstance(loaded, QueryAgentTool)
229+
assert loaded.configuration is None
230+
231+
def test_round_trip(self) -> None:
232+
loaded = AgentTool._load(query_example)
233+
dumped = loaded.dump(camel_case=True)
234+
assert dumped == query_example
235+
236+
def test_as_write_round_trip(self) -> None:
237+
loaded = AgentTool._load(query_example)
238+
write = loaded.as_write()
239+
dumped = write.dump(camel_case=True)
240+
assert dumped == query_example

0 commit comments

Comments
 (0)