Skip to content

Conversation

ks93
Copy link
Contributor

@ks93 ks93 commented Aug 19, 2025

Description

This PR introduces a higher-level abstraction for interacting with Cognite Agents, simplifying multi-turn conversations by automatically managing the chat cursor.

A new AgentSession class is added, which encapsulates the state of a conversation, including the agent_id and the current cursor. This allows users to initiate a session and then call session.chat() repeatedly without manually extracting and passing the cursor from previous responses.

The AgentSession can be started in two ways:

  1. Via client.agents.start_session(agent_id: str, cursor: str | None = None).
  2. Directly from an Agent instance: agent.start_session(cursor: str | None = None).

This change significantly improves the developer experience for building conversational applications with Cognite Agents.

Checklist:

  • Tests added/updated.
  • Documentation updated. Documentation is generated from docstrings - these must be updated according to your change.
    If a new method has been added it should be referenced in cognite.rst in order to generate docs based on its docstring.
  • Changelog updated in CHANGELOG.md.
  • Version bumped. If triggering a new release is desired, bump the version number in _version.py and pyproject.toml per semantic versioning.

Open in Cursor Open in Web

Copy link

cursor bot commented Aug 19, 2025

Cursor Agent can help with this pull request. Just @cursor in comments and I'll start working on changes in this branch.
Learn more about Cursor Agents

@ks93 ks93 marked this pull request as ready for review August 20, 2025 04:04
@ks93 ks93 requested review from a team as code owners August 20, 2025 04:04
@ks93 ks93 marked this pull request as draft August 20, 2025 04:17
Copy link

codecov bot commented Aug 20, 2025

Codecov Report

❌ Patch coverage is 95.65217% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 90.85%. Comparing base (a63de6d) to head (a969924).
⚠️ Report is 6 commits behind head on master.

Files with missing lines Patch % Lines
cognite/client/data_classes/agents/agents.py 83.33% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #2259      +/-   ##
==========================================
+ Coverage   90.84%   90.85%   +0.01%     
==========================================
  Files         167      167              
  Lines       24785    24865      +80     
==========================================
+ Hits        22516    22592      +76     
- Misses       2269     2273       +4     
Files with missing lines Coverage Δ
cognite/client/_api/agents/agents.py 100.00% <100.00%> (ø)
cognite/client/data_classes/agents/__init__.py 100.00% <ø> (ø)
cognite/client/data_classes/agents/chat.py 98.59% <100.00%> (+0.14%) ⬆️
cognite/client/data_classes/agents/agents.py 98.66% <83.33%> (-1.34%) ⬇️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@ks93 ks93 marked this pull request as ready for review August 21, 2025 18:47
@ks93 ks93 marked this pull request as draft August 21, 2025 18:48
@ks93
Copy link
Contributor Author

ks93 commented Aug 21, 2025

@cursor Hey! I still see some issues in the linting and static code checks:

Run pre-commit run --all-files --show-diff-on-failure
[INFO] Initializing environment for https://github.com/astral-sh/ruff-pre-commit.
[INFO] Initializing environment for https://github.com/pre-commit/pre-commit-hooks.
[INFO] Initializing environment for https://github.com/jsh9/pydoclint.
[INFO] Installing environment for https://github.com/astral-sh/ruff-pre-commit.
[INFO] Once installed this environment will be reused.
[INFO] This may take a few minutes...
[INFO] Installing environment for https://github.com/pre-commit/pre-commit-hooks.
[INFO] Once installed this environment will be reused.
[INFO] This may take a few minutes...
[INFO] Installing environment for https://github.com/jsh9/pydoclint.
[INFO] Once installed this environment will be reused.
[INFO] This may take a few minutes...
ruff (legacy alias)......................................................Failed

  • hook id: ruff
  • exit code: 1
  • files were modified by this hook

Found 6 errors (6 fixed, 0 remaining).

ruff format..............................................................Failed

  • hook id: ruff-format
  • files were modified by this hook

3 files reformatted, 358 files left unchanged

debug statements (python)................................................Passed
check docstring is first.................................................Passed
mypy.....................................................................Passed
custom repo checks.......................................................Passed
pydoclint................................................................Passed
pre-commit hook(s) made changes.
If you are seeing this message in CI, reproduce locally with: pre-commit run --all-files.
To run pre-commit as part of git workflow, use pre-commit install.
All changes made by hooks:
diff --git a/cognite/client/data_classes/agents/agents.py b/cognite/client/data_classes/agents/agents.py
index faa01d7..094c054 100644
--- a/cognite/client/data_classes/agents/agents.py
+++ b/cognite/client/data_classes/agents/agents.py
@@ -223,9 +223,9 @@ class Agent(AgentCore):
from cognite.client.data_classes.agents.chat import AgentSession
from cognite.client.exceptions import CogniteMissingClientError

  •    if not hasattr(self, '_cognite_client') or self._cognite_client is None:
    
  •    if not hasattr(self, "_cognite_client") or self._cognite_client is None:
           raise CogniteMissingClientError(self)
    
  •    return AgentSession(agent_id=self.external_id, cognite_client=self._cognite_client, cursor=cursor)
    

diff --git a/cognite/client/data_classes/agents/chat.py b/cognite/client/data_classes/agents/chat.py
index 9653b53..3ce820a 100644
--- a/cognite/client/data_classes/agents/chat.py
+++ b/cognite/client/data_classes/agents/chat.py
@@ -1,9 +1,9 @@
from future import annotations

from abc import ABC, abstractmethod
+from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Literal
-from collections.abc import Sequence

from cognite.client.data_classes._base import CogniteObject, CogniteResource, CogniteResourceList
from cognite.client.utils._text import convert_all_keys_to_camel_case
@@ -324,12 +324,8 @@ class AgentSession:
... ])
"""
# Call the underlying agents.chat method with current cursor

  •    response = self._cognite_client.agents.chat(
    
  •        agent_id=self.agent_id,
    
  •        messages=messages,
    
  •        cursor=self._cursor
    
  •    )
    
  •    response = self._cognite_client.agents.chat(agent_id=self.agent_id, messages=messages, cursor=self._cursor)
    
  •    # Update the cursor for the next interaction
       self._cursor = response.cursor
    

diff --git a/tests/tests_unit/test_api/test_agents_session.py b/tests/tests_unit/test_api/test_agents_session.py
index eea4880..2330b0c 100644
--- a/tests/tests_unit/test_api/test_agents_session.py
+++ b/tests/tests_unit/test_api/test_agents_session.py
@@ -8,8 +8,6 @@ from cognite.client import CogniteClient
from cognite.client.data_classes.agents import Agent, AgentSession, Message
from cognite.client.data_classes.agents.chat import (
AgentChatResponse,

  • AgentMessage,
  • TextContent,
    )
    from cognite.client.exceptions import CogniteMissingClientError

@@ -93,13 +91,8 @@ class TestAgentSession:
assert response.text == "Hello! How can I help you today?"
assert session.cursor == "cursor_12345" # Cursor should be updated

  • def test_session_cursor_management(
  •    self,
    
  •    cognite_client: CogniteClient,
    
  •    chat_response_body: dict,
    
  •    followup_response_body: dict
    
  •    self, cognite_client: CogniteClient, chat_response_body: dict, followup_response_body: dict
    

    ) -> None:
    """Test that cursor is automatically managed across multiple interactions."""
    # Set up mock responses for two consecutive calls
    @@ -129,10 +122,7 @@ class TestAgentSession:
    cognite_client.agents._post = MagicMock(return_value=MagicMock(json=lambda: chat_response_body))

       session = cognite_client.agents.start_session("my_agent")
    
  •    messages = [
    
  •        Message("Find temperature sensors"),
    
  •        Message("Show me their recent data")
    
  •    ]
    
  •    messages = [Message("Find temperature sensors"), Message("Show me their recent data")]
       session.chat(messages)
    
       # Verify multiple messages were sent
    

@@ -146,11 +136,7 @@ class TestAgentStartSession:
def test_agent_start_session(self, cognite_client: CogniteClient) -> None:
"""Test starting a session from an Agent instance."""
# Create an agent instance with a cognite client

  •    agent = Agent(
    
  •        external_id="my_agent",
    
  •        name="My Agent",
    
  •        description="Test agent"
    
  •    )
    
  •    agent = Agent(external_id="my_agent", name="My Agent", description="Test agent")
       agent._cognite_client = cognite_client
    
       session = agent.start_session()
    

@@ -162,10 +148,7 @@ class TestAgentStartSession:

 def test_agent_start_session_with_cursor(self, cognite_client: CogniteClient) -> None:
     """Test starting a session from an Agent instance with a cursor."""
  •    agent = Agent(
    
  •        external_id="my_agent",
    
  •        name="My Agent"
    
  •    )
    
  •    agent = Agent(external_id="my_agent", name="My Agent")
       agent._cognite_client = cognite_client
    
       initial_cursor = "existing_cursor_456"
    

@@ -176,10 +159,7 @@ class TestAgentStartSession:

 def test_agent_start_session_without_client_raises_error(self) -> None:
     """Test that starting a session without a cognite client raises an error."""
  •    agent = Agent(
    
  •        external_id="my_agent",
    
  •        name="My Agent"
    
  •    )
    
  •    agent = Agent(external_id="my_agent", name="My Agent")
       # Don't set _cognite_client
    
       with pytest.raises(CogniteMissingClientError):
    

@@ -187,10 +167,7 @@ class TestAgentStartSession:

 def test_agent_start_session_with_none_client_raises_error(self) -> None:
     """Test that starting a session with None cognite client raises an error."""
  •    agent = Agent(
    
  •        external_id="my_agent",
    
  •        name="My Agent"
    
  •    )
    
  •    agent = Agent(external_id="my_agent", name="My Agent")
       agent._cognite_client = None
    
       with pytest.raises(CogniteMissingClientError):
    

@@ -199,10 +176,7 @@ class TestAgentStartSession:

class TestIntegrationWorkflow:
def test_complete_workflow(

  •    self,
    
  •    cognite_client: CogniteClient,
    
  •    chat_response_body: dict,
    
  •    followup_response_body: dict
    
  •    self, cognite_client: CogniteClient, chat_response_body: dict, followup_response_body: dict
    

    ) -> None:
    """Test the complete workflow as described in the user requirements."""
    # Set up mock responses
    @@ -245,4 +219,4 @@ class TestIntegrationWorkflow:
    assert "cursor" not in first_call_args[1]["json"]

       # Second call should have cursor from first response
    
  •    assert second_call_args[1]["json"]["cursor"] == "cursor_12345"
    

\ No newline at end of file

  •    assert second_call_args[1]["json"]["cursor"] == "cursor_12345"
    

Error: Process completed with exit code 1.

Copy link

cursor bot commented Aug 21, 2025

Working on it! I'll start making changes in this branch.

Open in Cursor Open in Web

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants