Skip to content
9 changes: 7 additions & 2 deletions cognite/client/data_classes/agents/agent_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,10 +166,12 @@ class QueryKnowledgeGraphAgentToolConfiguration(WriteableCogniteResource):
Args:
data_models (Sequence[DataModelInfo] | None): The data models and views to query.
instance_spaces (InstanceSpaces | None): The instance spaces to query.
version (str | None): The version of the query generation strategy to use. A higher number does not necessarily mean a better query. Supported values are "v1" and "v2".
"""

data_models: Sequence[DataModelInfo] | None = None
instance_spaces: InstanceSpaces | None = None
data_models: Sequence[DataModelInfo] | None
instance_spaces: InstanceSpaces | None
version: Literal["v1", "v2"] | str | None

@classmethod
def _load(
Expand All @@ -186,6 +188,7 @@ def _load(
return cls(
data_models=data_models,
instance_spaces=instance_spaces,
version=resource.get("version"),
)

def dump(self, camel_case: bool = True) -> dict[str, Any]:
Expand All @@ -196,6 +199,8 @@ def dump(self, camel_case: bool = True) -> dict[str, Any]:
if self.instance_spaces:
key = "instanceSpaces" if camel_case else "instance_spaces"
result[key] = self.instance_spaces.dump(camel_case=camel_case)
if self.version:
result["version"] = self.version
return result

def as_write(self) -> QueryKnowledgeGraphAgentToolConfiguration:
Expand Down
8 changes: 6 additions & 2 deletions tests/tests_integration/test_api/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,9 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent:
version="v1",
view_external_ids=["CogniteAsset"],
)
]
],
instance_spaces=None,
version="v2",
),
)
],
Expand Down Expand Up @@ -80,7 +82,9 @@ def test_create_retrieve_update_delete_agent(self, cognite_client: CogniteClient
version="v1",
view_external_ids=["CogniteAsset"],
)
]
],
instance_spaces=None,
version="v2",
),
),
SummarizeDocumentAgentToolUpsert(
Expand Down
1 change: 1 addition & 0 deletions tests/tests_unit/test_api/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@ def test_upsert_full(self, cognite_client: CogniteClient, mock_agent_upsert_resp
)
],
instance_spaces=InstanceSpaces(type="all"),
version="v2",
),
)
],
Expand Down
Loading