diff --git a/cognite/client/data_classes/agents/agent_tools.py b/cognite/client/data_classes/agents/agent_tools.py
index d19d95eff3..94ee2f6a95 100644
--- a/cognite/client/data_classes/agents/agent_tools.py
+++ b/cognite/client/data_classes/agents/agent_tools.py
@@ -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(
@@ -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]:
@@ -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:
diff --git a/tests/tests_integration/test_api/test_agents.py b/tests/tests_integration/test_api/test_agents.py
index 405e245702..f7167b5a5b 100644
--- a/tests/tests_integration/test_api/test_agents.py
+++ b/tests/tests_integration/test_api/test_agents.py
@@ -41,7 +41,9 @@ def permanent_agent(cognite_client: CogniteClient) -> Agent:
                             version="v1",
                             view_external_ids=["CogniteAsset"],
                         )
-                    ]
+                    ],
+                    instance_spaces=None,
+                    version="v2",
                 ),
             )
         ],
@@ -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(
diff --git a/tests/tests_unit/test_api/test_agents.py b/tests/tests_unit/test_api/test_agents.py
index b05cfdad3f..bf853e69bd 100644
--- a/tests/tests_unit/test_api/test_agents.py
+++ b/tests/tests_unit/test_api/test_agents.py
@@ -138,6 +138,7 @@ def test_upsert_full(self, cognite_client: CogniteClient, mock_agent_upsert_resp
                             )
                         ],
                         instance_spaces=InstanceSpaces(type="all"),
+                        version="v2",
                     ),
                 )
             ],