RDBC-934 Sync-up Python AI Agents .run() implementation with C#

Author: poissoncorp

ravendb/__init__.py

@@ -100,6 +100,7 @@
     AiAgentActionRequest,
     AiAgentActionResponse,
     AiUsage,
+    AiConversationCreationOptions,
     GetAiAgentOperation,
     GetAiAgentsResponse,
     AddOrUpdateAiAgentOperation,

ravendb/documents/ai/__init__.py

@@ -2,11 +2,14 @@
 from .ai_conversation import AiConversation
 from .ai_conversation_result import AiConversationResult
 from .ai_agent_parameters_builder import AiAgentParametersBuilder, IAiAgentParametersBuilder
+from .ai_answer import AiAnswer, AiConversationStatus
 
 __all__ = [
     "AiOperations",
     "AiConversation",
     "AiConversationResult",
     "AiAgentParametersBuilder",
     "IAiAgentParametersBuilder",
+    "AiAnswer",
+    "AiConversationStatus",
 ]

ravendb/documents/ai/ai_answer.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+from typing import Optional, TypeVar, Generic, TYPE_CHECKING
+from datetime import timedelta
+import enum
+
+if TYPE_CHECKING:
+    from ravendb.documents.operations.ai.agents import AiUsage
+
+TAnswer = TypeVar("TAnswer")
+
+
+class AiConversationStatus(enum.Enum):
+    """
+    Represents the status of an AI conversation.
+    """
+
+    DONE = "Done"
+    ACTION_REQUIRED = "ActionRequired"
+
+    def __str__(self):
+        return self.value
+
+
+class AiAnswer(Generic[TAnswer]):
+    """
+    Represents the answer from an AI conversation turn.
+    
+    This class contains the AI's response, the conversation status,
+    token usage statistics, and timing information.
+    """
+
+    def __init__(
+        self,
+        answer: Optional[TAnswer] = None,
+        status: AiConversationStatus = AiConversationStatus.DONE,
+        usage: Optional[AiUsage] = None,
+        elapsed: Optional[timedelta] = None,
+    ):
+        """
+        Initialize an AiAnswer instance.
+
+        Args:
+            answer: The answer content produced by the AI
+            status: The status of the conversation (Done or ActionRequired)
+            usage: Token usage reported by the model
+            elapsed: The total time elapsed to produce the answer
+        """
+        self.answer = answer
+        self.status = status
+        self.usage = usage
+        self.elapsed = elapsed
+
+    def __str__(self) -> str:
+        """String representation for debugging."""
+        return (
+            f"AiAnswer(status={self.status.value}, "
+            f"has_answer={self.answer is not None}, "
+            f"elapsed={self.elapsed})"
+        )
+
+    def __repr__(self) -> str:
+        """Detailed representation for debugging."""
+        return (
+            f"AiAnswer(answer={self.answer!r}, "
+            f"status={self.status!r}, "
+            f"usage={self.usage!r}, "
+            f"elapsed={self.elapsed!r})"
+        )
+

ravendb/documents/ai/ai_conversation.py

@@ -3,16 +3,17 @@
 import json
 import traceback
 from typing import List, Dict, Any, Optional, TypeVar, TYPE_CHECKING, Callable
+from datetime import timedelta
 
-from ravendb.documents.ai.ai_conversation_result import AiConversationResult
+from ravendb.documents.ai.ai_answer import AiAnswer, AiConversationStatus
+from ravendb.documents.operations.ai.agents import (
+    AiAgentActionRequest,
+    AiAgentActionResponse,
+    AiConversationCreationOptions,
+)
 
 if TYPE_CHECKING:
     from ravendb.documents.store.definition import DocumentStore
-    from ravendb.documents.operations.ai.agents import (
-        AiAgentActionRequest,
-        AiAgentActionResponse,
-        ConversationResult,
-    )
 
 TResponse = TypeVar("TResponse")
 
@@ -32,24 +33,26 @@ class AiConversation:
             result = conversation.run()
     """
 
-    _invocations: Dict[str, Callable[[AiAgentActionRequest], None]] = {}
-
     def __init__(
         self,
         store: DocumentStore,
         agent_id: str = None,
-        parameters: Dict[str, Any] = None,
+        options: AiConversationCreationOptions = None,
         conversation_id: str = None,
         change_vector: str = None,
     ):
         self._store = store
         self._agent_id = agent_id
-        self._parameters = parameters or {}
+        self._options = options or AiConversationCreationOptions()
         self._conversation_id = conversation_id
         self._change_vector = change_vector
-        self._user_prompt: Optional[str] = None
+
+        self._prompt_parts: List[str] = []
         self._action_responses: List[AiAgentActionResponse] = []
-        self._last_result: Optional[ConversationResult] = None
+        self._action_requests: Optional[List[AiAgentActionRequest]] = None
+
+        # Action handlers
+        self._invocations: Dict[str, Callable[[AiAgentActionRequest], None]] = {}
 
     def __enter__(self) -> AiConversation:
         """Context manager entry."""
@@ -85,10 +88,13 @@ def required_actions(self) -> List[AiAgentActionRequest]:
         """
         Gets the list of action requests that need to be fulfilled before
         the conversation can continue.
+
+        Raises:
+            RuntimeError: If run() hasn't been called yet
         """
-        if self._last_result and self._last_result.action_requests:
-            return self._last_result.action_requests
-        return []
+        if self._action_requests is None:
+            raise RuntimeError("You have to call run() first")
+        return self._action_requests
 
     def add_action_response(self, action_id: str, action_response: str) -> None:
         """
@@ -107,69 +113,133 @@ def add_action_response(self, action_id: str, action_response: str) -> None:
 
         self._action_responses.append(response)
 
-    def run(self) -> AiConversationResult:
+    def run(self, answer_type: type = dict) -> AiAnswer:
+        """
+        Executes the conversation loop, automatically handling action requests
+        until the conversation is complete or no handlers are available.
+
+        Args:
+            answer_type: The expected type of the answer (default: dict)
+
+        Returns:
+            AiAnswer with the final response, status, usage, and elapsed time
         """
-        Executes one "turn" of the conversation:
-        sends the current prompt or replies to any required actions,
-        and awaits the agent's reply.
+        while True:
+            r = self._run_internal(answer_type)
+            if self._handle_server_reply(r):
+                return r
+
+    def _run_internal(self, answer_type: type = dict) -> AiAnswer:
         """
-        from ravendb.documents.operations.ai.agents import RunConversationOperation
+        Internal method that executes a single server call.
 
-        if self._conversation_id:
-            # Continue existing conversation
-            if not self._agent_id:
-                raise ValueError("Agent ID is required for conversation continuation")
+        Args:
+            answer_type: The expected type of the answer
 
-            operation = RunConversationOperation(
-                self._conversation_id,
-                self._user_prompt,
-                self._action_responses,
-                self._change_vector,
-            )
-            # Set agent ID for conversation continuation
-            operation._agent_id = self._agent_id
-        else:
-            # Start new conversation
-            if not self._agent_id:
-                raise ValueError("Agent ID is required for new conversations")
-
-            operation = RunConversationOperation(
-                self._agent_id,
-                self._user_prompt,
-                self._parameters,
+        Returns:
+            AiAnswer from this single turn
+        """
+        from ravendb.documents.operations.ai.agents import RunConversationOperation
+        import time
+
+        # If we already went to the server and have nothing new to tell it, we're done
+        if (
+            self._action_requests is not None
+            and len(self._prompt_parts) == 0
+            and len(self._action_responses) == 0
+        ):
+            return AiAnswer(
+                answer=None,
+                status=AiConversationStatus.DONE,
+                usage=None,
+                elapsed=None,
             )
 
-        # Execute the operation
-        result = self._store.maintenance.send(operation)
-        self._last_result = result
+        # Build the operation
+        if not self._agent_id:
+            raise ValueError("Agent ID is required")
+
+        # If we don't have a conversation ID yet, generate one with the prefix
+        # The server will complete it with a unique ID
+        if not self._conversation_id:
+            self._conversation_id = "conversations/"
+
+        # Create operation with all required parameters
+        operation = RunConversationOperation(
+            agent_id=self._agent_id,
+            conversation_id=self._conversation_id,
+            prompt_parts=self._prompt_parts,  # Always send list, even if empty
+            action_responses=self._action_responses,  # Always send list, even if empty
+            options=self._options,
+            change_vector=self._change_vector,
+        )
 
-        # Update conversation state for future calls
-        if result.conversation_id:
-            self._conversation_id = result.conversation_id
-        if result.change_vector:
+        try:
+            # Track elapsed time
+            start_time = time.time()
+            result = self._store.maintenance.send(operation)
+            elapsed = timedelta(seconds=time.time() - start_time)
+
+            # Update conversation state
             self._change_vector = result.change_vector
+            self._conversation_id = result.conversation_id
+            self._action_requests = result.action_requests or []
+
+            # Build AiAnswer
+            return AiAnswer(
+                answer=result.response,
+                status=AiConversationStatus.ACTION_REQUIRED if len(self._action_requests) > 0 else AiConversationStatus.DONE,
+                usage=result.usage,
+                elapsed=elapsed,
+            )
+        # except ConcurrencyException as e:
+        #     self._change_vector = e.actual_change_vector
+        #     raise
+        finally:
+            # Clear the user prompt and tool responses after running the conversation
+            self._prompt_parts.clear()
+            self._action_responses.clear()
+
+    def _handle_server_reply(self, answer: AiAnswer) -> bool:
+        """
+        Handles the server reply by invoking registered action handlers.
 
-        # Preserve agent ID for future conversation turns
-        if not self._agent_id and hasattr(operation, "_agent_id"):
-            self._agent_id = operation._agent_id
+        Args:
+            answer: The answer from the server
 
-        # Clear processed data for next turn
-        self._user_prompt = None
-        self._action_responses.clear()
+        Returns:
+            True if the conversation is done, False if it should continue
+        """
+        if answer.status == AiConversationStatus.DONE:
+            return True
 
-        # Convert to AiConversationResult
-        conversation_result = AiConversationResult()
-        conversation_result.conversation_id = result.conversation_id
-        conversation_result.change_vector = result.change_vector
-        conversation_result.response = result.response
-        conversation_result.usage = result.usage
-        conversation_result.action_requests = result.action_requests or []
+        if len(self._action_requests) == 0:
+            raise RuntimeError(
+                f"There are no action requests to process, but Status was {answer.status}, should not be possible."
+            )
 
-        return conversation_result
+        # Process each action request
+        for action in self._action_requests:
+            if action.name in self._invocations:
+                # Invoke the registered handler
+                # Error handling is done by the invocation based on the error strategy
+                self._invocations[action.name](action)
+            else:
+                # No handler registered for this action
+                raise RuntimeError(
+                    f"There is no action defined for action '{action.name}' on agent '{self._agent_id}' "
+                    f"({self._conversation_id}), but it was invoked by the model with: {action.arguments}. "
+                    f"Did you forget to call receive() or handle()?"
+                )
+
+        # If we have nothing to tell the server (no action responses), we're done
+        # Otherwise, continue the loop to send the responses
+        return len(self._action_responses) == 0
 
     def set_user_prompt(self, user_prompt: str) -> None:
         """
-        Sets the next user prompt to send to the AI agent.
+        Sets the user prompt to send to the AI agent.
+        Clears any existing prompt parts and adds the new prompt.
 
         Args:
             user_prompt: The prompt text to send to the agent
@@ -179,7 +249,23 @@ def set_user_prompt(self, user_prompt: str) -> None:
         """
         if not user_prompt or user_prompt.isspace():
             raise ValueError("User prompt cannot be empty or whitespace-only")
-        self._user_prompt = user_prompt
+        self._prompt_parts.clear()
+        self._prompt_parts.append(user_prompt)
+
+    def add_user_prompt(self, *prompts: str) -> None:
+        """
+        Adds one or more user prompts to the conversation.
+
+        Args:
+            *prompts: One or more prompt strings to add
+
+        Raises:
+            ValueError: If any prompt is empty or whitespace-only
+        """
+        for prompt in prompts:
+            if not prompt or prompt.isspace():
+                raise ValueError("User prompt cannot be empty or whitespace-only")
+            self._prompt_parts.append(prompt)
 
     def handle(
         self,

ravendb/documents/ai/ai_conversation_result.py

@@ -13,12 +13,19 @@ class AiConversationResult(Generic[TResponse]):
     usage statistics, and any action requests that need to be fulfilled.
     """
 
-    def __init__(self):
-        self.conversation_id: Optional[str] = None
-        self.change_vector: Optional[str] = None
-        self.response: Optional[TResponse] = None
-        self.usage: Optional[AiUsage] = None
-        self.action_requests: List[AiAgentActionRequest] = []
+    def __init__(
+        self,
+        conversation_id: Optional[str] = None,
+        change_vector: Optional[str] = None,
+        response: Optional[TResponse] = None,
+        usage: Optional[AiUsage] = None,
+        action_requests: Optional[List[AiAgentActionRequest]] = None,
+    ):
+        self.conversation_id: Optional[str] = conversation_id
+        self.change_vector: Optional[str] = change_vector
+        self.response: Optional[TResponse] = response
+        self.usage: Optional[AiUsage] = usage
+        self.action_requests: List[AiAgentActionRequest] = action_requests or []
 
     def __str__(self) -> str:
         """String representation for debugging."""