diff --git a/cognite/client/_api/agents/agents.py b/cognite/client/_api/agents/agents.py index 03869378aa..5884cd8974 100644 --- a/cognite/client/_api/agents/agents.py +++ b/cognite/client/_api/agents/agents.py @@ -77,7 +77,39 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A ... labels=["published"], ... tools=[find_assets_tool], ... ) - >>> client.agents.upsert(agents=[agent]) + >>> client.agents.upsert(agent) + + Create an agent with the query tool: + + >>> from cognite.client.data_classes.agents import ( + ... AgentUpsert, + ... QueryAgentToolUpsert, + ... QueryAgentToolConfiguration, + ... DataModelInfo, + ... InstanceSpaces, + ... ) + >>> query_tool = QueryAgentToolUpsert( + ... name="explore data", + ... description="Run flexible queries against your data model", + ... configuration=QueryAgentToolConfiguration( + ... data_models=[ + ... DataModelInfo( + ... space="cdf_idm", + ... external_id="CogniteProcessIndustries", + ... version="v1", + ... ) + ... ], + ... instance_spaces=InstanceSpaces(type="all"), + ... ), + ... ) + >>> agent = AgentUpsert( + ... external_id="my_agent", + ... name="My Agent", + ... labels=["published"], + ... runtime_version="1.1.2-preview", + ... tools=[query_tool], + ... ) + >>> client.agents.upsert(agent) Create an agent with multiple different tools: @@ -159,7 +191,7 @@ async def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | A ... ts_tool, ... ], ... ) - >>> client.agents.upsert(agents=[agent]) + >>> client.agents.upsert(agent) """ diff --git a/cognite/client/_sync_api/agents/agents.py b/cognite/client/_sync_api/agents/agents.py index 01824b8959..b68bc08754 100644 --- a/cognite/client/_sync_api/agents/agents.py +++ b/cognite/client/_sync_api/agents/agents.py @@ -1,6 +1,6 @@ """ =============================================================================== -063c42ab744021733ccbdc455b150b2c +178ce7222985b04d03174af7b7e0f525 This file is auto-generated from the Async API modules, - do not edit manually! =============================================================================== """ @@ -75,7 +75,39 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi ... labels=["published"], ... tools=[find_assets_tool], ... ) - >>> client.agents.upsert(agents=[agent]) + >>> client.agents.upsert(agent) + + Create an agent with the query tool: + + >>> from cognite.client.data_classes.agents import ( + ... AgentUpsert, + ... QueryAgentToolUpsert, + ... QueryAgentToolConfiguration, + ... DataModelInfo, + ... InstanceSpaces, + ... ) + >>> query_tool = QueryAgentToolUpsert( + ... name="explore data", + ... description="Run flexible queries against your data model", + ... configuration=QueryAgentToolConfiguration( + ... data_models=[ + ... DataModelInfo( + ... space="cdf_idm", + ... external_id="CogniteProcessIndustries", + ... version="v1", + ... ) + ... ], + ... instance_spaces=InstanceSpaces(type="all"), + ... ), + ... ) + >>> agent = AgentUpsert( + ... external_id="my_agent", + ... name="My Agent", + ... labels=["published"], + ... runtime_version="1.1.2-preview", + ... tools=[query_tool], + ... ) + >>> client.agents.upsert(agent) Create an agent with multiple different tools: @@ -157,7 +189,7 @@ def upsert(self, agents: AgentUpsert | Sequence[AgentUpsert]) -> Agent | AgentLi ... ts_tool, ... ], ... ) - >>> client.agents.upsert(agents=[agent]) + >>> client.agents.upsert(agent) """ return run_sync(self.__async_client.agents.upsert(agents=agents)) diff --git a/cognite/client/data_classes/agents/__init__.py b/cognite/client/data_classes/agents/__init__.py index f0d39aebc9..9531002d5e 100644 --- a/cognite/client/data_classes/agents/__init__.py +++ b/cognite/client/data_classes/agents/__init__.py @@ -9,6 +9,9 @@ AskDocumentAgentToolUpsert, DataModelInfo, InstanceSpaces, + QueryAgentTool, + QueryAgentToolConfiguration, + QueryAgentToolUpsert, QueryKnowledgeGraphAgentTool, QueryKnowledgeGraphAgentToolConfiguration, QueryKnowledgeGraphAgentToolUpsert, @@ -70,6 +73,9 @@ "Message", "MessageContent", "MessageList", + "QueryAgentTool", + "QueryAgentToolConfiguration", + "QueryAgentToolUpsert", "QueryKnowledgeGraphAgentTool", "QueryKnowledgeGraphAgentToolConfiguration", "QueryKnowledgeGraphAgentToolUpsert", diff --git a/cognite/client/data_classes/agents/agent_tools.py b/cognite/client/data_classes/agents/agent_tools.py index 8957b208cf..8d01531034 100644 --- a/cognite/client/data_classes/agents/agent_tools.py +++ b/cognite/client/data_classes/agents/agent_tools.py @@ -194,6 +194,45 @@ def as_write(self) -> QueryKnowledgeGraphAgentToolConfiguration: return self +@dataclass +class QueryAgentToolConfiguration(WriteableCogniteResource): + """Configuration for query agent tools. + + Args: + data_models (Sequence[DataModelInfo]): The data models to query. + instance_spaces (InstanceSpaces | None): The instance spaces to query. + """ + + data_models: Sequence[DataModelInfo] + instance_spaces: InstanceSpaces | None = None + + @classmethod + def _load(cls, resource: dict[str, Any]) -> QueryAgentToolConfiguration: + dm_config = resource["dataModels"] + data_models = [DataModelInfo._load(dm) for dm in dm_config.get("dataModels", [])] + return cls( + data_models=data_models, + instance_spaces=InstanceSpaces._load_if(resource.get("instanceSpaces")), + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result: dict[str, Any] = {} + key = "dataModels" if camel_case else "data_models" + result[key] = { + # TODO(ks93): Clarify whether "type" can take other values (cf. InstanceSpaces.type); + # if so, make it a real field + "type": "manual", + key: [dm.dump(camel_case=camel_case) for dm in self.data_models], + } + if self.instance_spaces: + key = "instanceSpaces" if camel_case else "instance_spaces" + result[key] = self.instance_spaces.dump(camel_case=camel_case) + return result + + def as_write(self) -> QueryAgentToolConfiguration: + return self + + @dataclass class SummarizeDocumentAgentTool(AgentTool): """Agent tool for summarizing documents. @@ -389,6 +428,69 @@ def _load(cls, resource: dict[str, Any]) -> QueryTimeSeriesDatapointsAgentToolUp ) +@dataclass +class QueryAgentTool(AgentTool): + """Agent tool for running flexible queries against data models. + + Args: + name (str): The name of the agent tool. Used by the agent to decide when to use this tool. + description (str): The description of the agent tool. Used by the agent to decide when to use this tool. + configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool. + """ + + _type: ClassVar[str] = "query" + configuration: QueryAgentToolConfiguration | None = None + + @classmethod + def _load_tool(cls, resource: dict[str, Any]) -> QueryAgentTool: + return cls( + name=resource["name"], + description=resource["description"], + configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")), + ) + + def as_write(self) -> QueryAgentToolUpsert: + return QueryAgentToolUpsert( + name=self.name, + description=self.description, + configuration=self.configuration, + ) + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case=camel_case) + if self.configuration: + result["configuration"] = self.configuration.dump(camel_case=camel_case) + return result + + +@dataclass +class QueryAgentToolUpsert(AgentToolUpsert): + """Upsert version of query agent tool. + + Args: + name (str): The name of the agent tool. Used by the agent to decide when to use this tool. + description (str): The description of the agent tool. Used by the agent to decide when to use this tool. + configuration (QueryAgentToolConfiguration | None): The configuration of the query agent tool. + """ + + _type: ClassVar[str] = "query" + configuration: QueryAgentToolConfiguration | None = None + + def dump(self, camel_case: bool = True) -> dict[str, Any]: + result = super().dump(camel_case=camel_case) + if self.configuration: + result["configuration"] = self.configuration.dump(camel_case=camel_case) + return result + + @classmethod + def _load(cls, resource: dict[str, Any]) -> QueryAgentToolUpsert: + return cls( + name=resource["name"], + description=resource["description"], + configuration=QueryAgentToolConfiguration._load_if(resource.get("configuration")), + ) + + @dataclass class UnknownAgentTool(AgentTool): """Agent tool for unknown/unrecognized tool types. diff --git a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py index d0215314b6..358fd4cc5b 100644 --- a/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py +++ b/tests/tests_unit/test_data_classes/test_agents/test_agent_tools.py @@ -5,6 +5,8 @@ from cognite.client.data_classes.agents.agent_tools import ( AgentTool, AskDocumentAgentTool, + QueryAgentTool, + QueryAgentToolConfiguration, QueryKnowledgeGraphAgentTool, QueryKnowledgeGraphAgentToolConfiguration, QueryTimeSeriesDatapointsAgentTool, @@ -42,6 +44,32 @@ "description": "Query the time series datapoints", } +query_example = { + "name": "queryExample", + "type": "query", + "description": "Run flexible queries against your data model", + "configuration": { + "dataModels": { + "type": "manual", + "dataModels": [ + { + "space": "cdf_idm", + "externalId": "CogniteProcessIndustries", + "version": "v1", + "viewExternalIds": ["CogniteAsset"], + } + ], + }, + "instanceSpaces": {"type": "manual", "spaces": ["my_space"]}, + }, +} + +query_no_config_example = { + "name": "queryNoConfigExample", + "type": "query", + "description": "Run flexible queries against your data model", +} + unknown_example = { "name": "unknownExample", "type": "yolo", # This is not a known tool type @@ -58,9 +86,19 @@ class TestAgentToolLoad: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), + (query_no_config_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], - ids=["queryKnowledgeGraph", "askDocument", "summarizeDocument", "queryTimeSeriesDatapoints", "somethingElse"], + ids=[ + "queryKnowledgeGraph", + "askDocument", + "summarizeDocument", + "queryTimeSeriesDatapoints", + "query", + "queryNoConfig", + "somethingElse", + ], ) def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected_type: type[AgentTool]) -> None: """Test that AgentTool._load() returns the correct subtype based on the tool type.""" @@ -81,6 +119,10 @@ def test_agent_tool_load_returns_correct_subtype(self, tool_data: dict, expected assert isinstance(loaded_tool, QueryKnowledgeGraphAgentTool) assert isinstance(loaded_tool.configuration, QueryKnowledgeGraphAgentToolConfiguration) assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"] + elif expected_type is QueryAgentTool: + assert isinstance(loaded_tool, QueryAgentTool) + assert isinstance(loaded_tool.configuration, QueryAgentToolConfiguration) + assert loaded_tool.configuration.dump(camel_case=True) == tool_data["configuration"] elif expected_type is UnknownAgentTool: assert isinstance(loaded_tool, UnknownAgentTool) assert loaded_tool.configuration == tool_data["configuration"] @@ -104,6 +146,7 @@ class TestAgentToolDump: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) @@ -150,6 +193,7 @@ class TestAgentToolUpsert: (ask_document_example, AskDocumentAgentTool), (summarize_document_example, SummarizeDocumentAgentTool), (query_time_series_datapoints_example, QueryTimeSeriesDatapointsAgentTool), + (query_example, QueryAgentTool), (unknown_example, UnknownAgentTool), ], ) @@ -166,3 +210,10 @@ def test_agent_tool_upsert_returns_correct_type(self, tool_data: dict, expected_ assert dumped_tool["name"] == tool_data["name"] assert dumped_tool["description"] == tool_data["description"] + + +class TestQueryAgentTool: + def test_load_without_configuration(self) -> None: + loaded = AgentTool._load(query_no_config_example) + assert isinstance(loaded, QueryAgentTool) + assert loaded.configuration is None