Manage agent conversational sessions and cursors

cognite/client/_api/agents/agents.py

@@ -5,7 +5,7 @@
 
 from cognite.client._api_client import APIClient
 from cognite.client.data_classes.agents import Agent, AgentList, AgentUpsert
-from cognite.client.data_classes.agents.chat import AgentChatResponse, Message, MessageList
+from cognite.client.data_classes.agents.chat import AgentChatResponse, AgentSession, Message, MessageList
 from cognite.client.utils._experimental import FeaturePreviewWarning
 from cognite.client.utils._identifier import IdentifierSequence
 from cognite.client.utils.useful_types import SequenceNotStr
@@ -313,3 +313,39 @@ def chat(
         )
 
         return AgentChatResponse._load(response.json(), cognite_client=self._cognite_client)
+
+    def start_session(self, agent_id: str, cursor: str | None = None) -> AgentSession:
+        """Start a new chat session with an agent.
+
+        This creates a session object that automatically manages cursor state across
+        multiple chat interactions, providing a more convenient interface for
+        multi-turn conversations.
+
+        Args:
+            agent_id (str): External ID that uniquely identifies the agent.
+            cursor (str | None): Optional cursor to continue an existing conversation.
+                If None, starts a fresh session.
+
+        Returns:
+            AgentSession: A session object for chatting with the agent.
+
+        Examples:
+
+            Start a new session and have a conversation:
+
+                >>> from cognite.client import CogniteClient
+                >>> from cognite.client.data_classes.agents import Message
+                >>> client = CogniteClient()
+                >>> session = client.agents.start_session("my_agent")
+                >>> response = session.chat(Message("Hello, how can you help me?"))
+                >>> print(response.text)
+                >>> followup = session.chat(Message("Tell me more about that"))
+                >>> print(followup.text)
+
+            Continue an existing session with a cursor:
+
+                >>> session = client.agents.start_session("my_agent", cursor="existing_cursor_123")
+                >>> response = session.chat(Message("Continue our previous conversation"))
+        """
+        self._warnings.warn()
+        return AgentSession(agent_id=agent_id, cognite_client=self._cognite_client, cursor=cursor)

cognite/client/data_classes/agents/__init__.py

@@ -26,6 +26,7 @@
     AgentMessage,
     AgentMessageList,
     AgentReasoningItem,
+    AgentSession,
     Message,
     MessageContent,
     MessageList,
@@ -41,6 +42,7 @@
     "AgentMessage",
     "AgentMessageList",
     "AgentReasoningItem",
+    "AgentSession",
     "AgentTool",
     "AgentToolList",
     "AgentToolUpsert",

cognite/client/data_classes/agents/agents.py

@@ -6,6 +6,7 @@
 
 if TYPE_CHECKING:
     from cognite.client import CogniteClient
+    from cognite.client.data_classes.agents.chat import AgentSession
 
 from cognite.client.data_classes._base import (
     CogniteResourceList,
@@ -192,6 +193,41 @@ def _load(cls, resource: dict[str, Any], cognite_client: CogniteClient | None =
         instance._unknown_properties = {key: value for key, value in resource.items() if key not in existing}
         return instance
 
+    def start_session(self, cursor: str | None = None) -> AgentSession:
+        """Start a new chat session with this agent.
+
+        This creates a session object that automatically manages cursor state across
+        multiple chat interactions, providing a more convenient interface for
+        multi-turn conversations.
+
+        Args:
+            cursor (str | None): Optional cursor to continue an existing conversation.
+                If None, starts a fresh session.
+
+        Returns:
+            AgentSession: A session object for chatting with this agent.
+
+        Examples:
+
+            Start a session from an agent instance:
+
+                >>> from cognite.client import CogniteClient
+                >>> from cognite.client.data_classes.agents import Message
+                >>> client = CogniteClient()
+                >>> agent = client.agents.retrieve("my_agent")
+                >>> session = agent.start_session()
+                >>> response = session.chat(Message("Hello!"))
+                >>> print(response.text)
+        """
+        # Import here to avoid circular imports
+        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:
+            raise CogniteMissingClientError(self)
+
+        return AgentSession(agent_id=self.external_id, cognite_client=self._cognite_client, cursor=cursor)
+
 
 class AgentUpsertList(CogniteResourceList[AgentUpsert], ExternalIDTransformerMixin):
     _RESOURCE = AgentUpsert

cognite/client/data_classes/agents/chat.py

@@ -1,6 +1,7 @@
 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
 
@@ -271,3 +272,70 @@ def _load(cls, data: dict[str, Any], cognite_client: CogniteClient | None = None
         )
 
         return instance
+
+
+class AgentSession:
+    """A session for chatting with an agent that manages cursor state automatically.
+
+    This class provides a higher-level interface for chatting with agents by automatically
+    managing the conversation cursor state across multiple chat interactions.
+
+    Args:
+        agent_id (str): The external ID of the agent to chat with.
+        cognite_client (CogniteClient): The Cognite client instance.
+        cursor (str | None): Optional initial cursor for continuing an existing session.
+    """
+
+    def __init__(self, agent_id: str, cognite_client: CogniteClient, cursor: str | None = None) -> None:
+        self.agent_id = agent_id
+        self._cognite_client = cognite_client
+        self._cursor = cursor
+
+    def chat(self, messages: Message | Sequence[Message]) -> AgentChatResponse:
+        """Send a message or messages to the agent and return the response.
+
+        The cursor state is automatically managed and updated after each interaction.
+
+        Args:
+            messages (Message | Sequence[Message]): The message(s) to send to the agent.
+
+        Returns:
+            AgentChatResponse: The response from the agent.
+
+        Examples:
+
+            Simple message:
+
+                >>> from cognite.client.data_classes.agents import Message
+                >>> session = client.agents.start_session("my_agent")
+                >>> response = session.chat(Message("Hello, how can you help me?"))
+                >>> print(response.text)
+
+            Follow-up message (cursor automatically managed):
+
+                >>> followup = session.chat(Message("Tell me more about that"))
+                >>> print(followup.text)
+
+            Multiple messages at once:
+
+                >>> response = session.chat([
+                ...     Message("Find the temperature sensors"),
+                ...     Message("Show me their recent data")
+                ... ])
+        """
+        # 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)
+
+        # Update the cursor for the next interaction
+        self._cursor = response.cursor
+
+        return response
+
+    @property
+    def cursor(self) -> str | None:
+        """Get the current cursor state for this session.
+
+        Returns:
+            str | None: The current cursor, or None if no conversation has taken place yet.
+        """
+        return self._cursor