RDBC-934 Add Handle & Receive, clean-up

Author: poissoncorp

ravendb/__init__.py

@@ -80,7 +80,6 @@
 # AI Operations
 from ravendb.documents.ai import (
     AiOperations,
-    IAiConversationOperations,
     AiConversation,
     AiConversationResult,
     AiAgentParametersBuilder,

ravendb/documents/ai/__init__.py

@@ -1,12 +1,10 @@
 from .ai_operations import AiOperations
-from .ai_conversation_operations import IAiConversationOperations
 from .ai_conversation import AiConversation
 from .ai_conversation_result import AiConversationResult
 from .ai_agent_parameters_builder import AiAgentParametersBuilder, IAiAgentParametersBuilder
 
 __all__ = [
     "AiOperations",
-    "IAiConversationOperations",
     "AiConversation",
     "AiConversationResult",
     "AiAgentParametersBuilder",

ravendb/documents/ai/ai_conversation.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
+
 import json
-from typing import List, Dict, Any, Optional, Union, TypeVar, TYPE_CHECKING
+import traceback
+from typing import List, Dict, Any, Optional, TypeVar, TYPE_CHECKING, Callable
 
-from ravendb.documents.ai.ai_conversation_operations import IAiConversationOperations
 from ravendb.documents.ai.ai_conversation_result import AiConversationResult
 
 if TYPE_CHECKING:
@@ -16,7 +17,12 @@
 TResponse = TypeVar("TResponse")
 
 
-class AiConversation(IAiConversationOperations[TResponse]):
+class AiHandleErrorStrategy:
+    SEND_ERRORS_TO_MODEL = "SendErrorsToModel"
+    RAISE_IMMEDIATELY = "RaiseImmediately"
+
+
+class AiConversation:
     """
     Implementation of AI conversation operations for managing conversations with AI agents.
 
@@ -26,6 +32,8 @@ class AiConversation(IAiConversationOperations[TResponse]):
             result = conversation.run()
     """
 
+    _invocations: Dict[str, Callable[[AiAgentActionRequest], None]] = {}
+
     def __init__(
         self,
         store: DocumentStore,
@@ -41,12 +49,20 @@ def __init__(
         self._change_vector = change_vector
         self._user_prompt: Optional[str] = None
         self._action_responses: List[AiAgentActionResponse] = []
-        self._last_result: Optional[ConversationResult[TResponse]] = None
+        self._last_result: Optional[ConversationResult] = None
+
+    def __enter__(self) -> AiConversation:
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Context manager exit - cleanup resources."""
+        pass
 
     @classmethod
     def with_conversation_id(
         cls, store: DocumentStore, conversation_id: str, change_vector: str = None
-    ) -> AiConversation[TResponse]:
+    ) -> AiConversation:
         """
         Creates a conversation instance for continuing an existing conversation.
 
@@ -74,35 +90,27 @@ def required_actions(self) -> List[AiAgentActionRequest]:
             return self._last_result.action_requests
         return []
 
-    def add_action_response(self, action_id: str, action_response: Union[str, TResponse]) -> None:
+    def add_action_response(self, action_id: str, action_response: str) -> None:
         """
         Adds a response for a given action request.
 
         Args:
             action_id: The ID of the action to respond to
-            action_response: The response content (string or typed response object)
+            action_response: The response content
         """
         from ravendb.documents.operations.ai.agents import AiAgentActionResponse
 
         response = AiAgentActionResponse(tool_id=action_id)
 
         if isinstance(action_response, str):
             response.content = action_response
-        else:
-            # More robust JSON serialization
-            try:
-                response.content = json.dumps(
-                    action_response.__dict__ if hasattr(action_response, "__dict__") else action_response, default=str
-                )
-            except (TypeError, ValueError) as e:
-                response.content = str(action_response)
 
         self._action_responses.append(response)
 
-    def run(self) -> AiConversationResult[TResponse]:
+    def run(self) -> AiConversationResult:
         """
         Executes one "turn" of the conversation:
-        sends the current prompt, processes any required actions,
+        sends the current prompt or replies to any required actions,
         and awaits the agent's reply.
         """
         from ravendb.documents.operations.ai.agents import RunConversationOperation
@@ -150,7 +158,7 @@ def run(self) -> AiConversationResult[TResponse]:
         self._action_responses.clear()
 
         # Convert to AiConversationResult
-        conversation_result = AiConversationResult[TResponse]()
+        conversation_result = AiConversationResult()
         conversation_result.conversation_id = result.conversation_id
         conversation_result.change_vector = result.change_vector
         conversation_result.response = result.response
@@ -173,12 +181,92 @@ def set_user_prompt(self, user_prompt: str) -> None:
             raise ValueError("User prompt cannot be empty or whitespace-only")
         self._user_prompt = user_prompt
 
-    def __enter__(self) -> AiConversation[TResponse]:
-        """Context manager entry."""
-        return self
+    def handle(
+        self,
+        action_name: str,
+        action: Callable[[dict], Any],
+        ai_handle_error: AiHandleErrorStrategy,
+    ) -> None:
+        self.handle_ai_agent_action_request(action_name, lambda _, args: action(args), ai_handle_error)
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Context manager exit - cleanup resources."""
-        # Clear any pending data
-        self._user_prompt = None
-        self._action_responses.clear()
+    def handle_ai_agent_action_request(
+        self,
+        action_name: str,
+        action: Callable[[AiAgentActionRequest, dict], Any],
+        ai_handle_error: AiHandleErrorStrategy = AiHandleErrorStrategy.SEND_ERRORS_TO_MODEL,
+    ) -> None:
+        def wrapped_no_return(request: AiAgentActionRequest, args: dict) -> Any:
+            result = action(request, args)
+            self.add_action_response(request.tool_id, result)
+
+        self.receive(action_name, wrapped_no_return, ai_handle_error)
+
+    def receive(
+        self,
+        action_name: str,
+        action: Callable[[AiAgentActionRequest, dict], None],
+        ai_handle_error: AiHandleErrorStrategy = AiHandleErrorStrategy.SEND_ERRORS_TO_MODEL,
+    ):
+        t = self.AiActionContext(self, lambda request, args: action(request, args), ai_handle_error)
+        self._add_action(action_name, t.execute)
+
+    def _add_action(self, action_name: str, action: Callable[[AiAgentActionRequest], Any]):
+        if action_name in self._invocations:
+            raise ValueError(f"Action '{action_name}' already exists")
+
+        self._invocations[action_name] = action
+
+    class AiActionContext:
+        def __init__(
+            self,
+            conversation: AiConversation,
+            action: Callable[[AiAgentActionRequest, dict], Any],
+            ai_handle_error: AiHandleErrorStrategy,
+        ):
+            self._conversation = conversation
+            self._action = action
+            self._ai_handle_error = ai_handle_error
+
+        def execute(self, action_request: AiAgentActionRequest):
+            args = json.loads(action_request.arguments)
+            self.invoke(action_request, args)
+
+        def invoke(self, action_request: AiAgentActionRequest, args: dict):
+            try:
+                self._action(action_request, args)
+            except Exception as e:
+                if self._ai_handle_error == AiHandleErrorStrategy.SEND_ERRORS_TO_MODEL:
+                    self._conversation.add_action_response(action_request.tool_id, self.create_error_message_for_llm(e))
+                else:
+                    raise e
+
+        @staticmethod
+        def create_error_message_for_llm(exc: Exception) -> str:
+            parts = []
+
+            current = exc
+            indent = 0
+
+            while current is not None:
+                prefix = "  " * indent
+                header = f"{prefix}{current.__class__.__name__}: {current}"
+                parts.append(header)
+
+                tb = "".join(traceback.format_exception(type(current), current, current.__traceback__))
+                tb_lines = tb.strip().splitlines()
+
+                # indent the traceback block
+                indented_tb = "\n".join(prefix + "  " + line for line in tb_lines)
+                parts.append(indented_tb)
+
+                # Move to next exception in the chain
+                if current.__cause__:
+                    current = current.__cause__
+                elif current.__context__ and not current.__suppress_context__:
+                    current = current.__context__
+                else:
+                    current = None
+
+                indent += 1
+
+            return "\n".join(parts)

ravendb/documents/ai/ai_conversation_operations.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
-from typing import TYPE_CHECKING, Optional, Dict, Any, Type
+from typing import TYPE_CHECKING, Dict, Any, Type
+from ravendb.documents.ai.ai_conversation import AiConversation
 
 if TYPE_CHECKING:
     from ravendb.documents.store.definition import DocumentStore
@@ -8,7 +9,6 @@
         AiAgentConfigurationResult,
         GetAiAgentsResponse,
     )
-    from ravendb.documents.ai.ai_conversation_operations import IAiConversationOperations
 
 
 class AiOperations:
@@ -67,7 +67,7 @@ def get_agents(self, agent_id: str = None) -> GetAiAgentsResponse:
         operation = GetAiAgentOperation(agent_id)
         return self._store.maintenance.send(operation)
 
-    def conversation(self, agent_id: str, parameters: Dict[str, Any] = None) -> IAiConversationOperations:
+    def conversation(self, agent_id: str, parameters: Dict[str, Any] = None) -> AiConversation:
         """
         Creates a new conversation with the specified AI agent.
 
@@ -78,11 +78,10 @@ def conversation(self, agent_id: str, parameters: Dict[str, Any] = None) -> IAiC
         Returns:
             Conversation operations interface for managing the conversation
         """
-        from ravendb.documents.ai.ai_conversation import AiConversation
 
         return AiConversation(self._store, agent_id, parameters)
 
-    def conversation_with_id(self, conversation_id: str, change_vector: str = None) -> IAiConversationOperations:
+    def conversation_with_id(self, conversation_id: str, change_vector: str = None) -> AiConversation:
         """
         Continues an existing conversation by its ID.
 

ravendb/documents/ai/ai_operations.py

@@ -34,7 +34,7 @@ def __init__(self, configuration: AiAgentConfiguration, schema_type: Type = None
                 "The 'schema_type' parameter is deprecated and will be removed in 8.0 version of Python client."
                 " Use 'sample_object' or 'output_schema' inside AiAgentConfiguration instead.",
                 DeprecationWarning,
-                stacklevel=2
+                stacklevel=2,
             )
         if configuration is None:
             raise ValueError("configuration cannot be None")

ravendb/documents/operations/ai/agents/add_or_update_ai_agent_operation.py

@@ -12,17 +12,21 @@
 TSchema = TypeVar("TSchema")
 
 
-@dataclass
 class AiAgentActionRequest:
     """Represents an action request from an AI agent."""
 
-    name: Optional[str] = None
-    tool_id: Optional[str] = None
-    arguments: Optional[str] = None
+    def __init__(self, name: str = None, tool_id: str = None, arguments: str = None):
+        self.name = name
+        self.tool_id = tool_id
+        self.arguments = arguments
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentActionRequest:
-        return cls(name=json_dict.get("Name"), tool_id=json_dict.get("ToolId"), arguments=json.loads(json_dict.get("Arguments")))
+        return cls(
+            name=json_dict.get("Name"),
+            tool_id=json_dict.get("ToolId"),
+            arguments=json_dict.get("Arguments"),
+        )
 
     def to_json(self) -> Dict[str, Any]:
         return {