From 0c506a4595500e62e18c95a8154ae8eed5ad02be Mon Sep 17 00:00:00 2001 From: Gracjan Sadowicz Date: Mon, 8 Dec 2025 16:52:18 +0100 Subject: [PATCH 1/3] RDBC-967 7.1.4 AI changes --- ravendb/__init__.py | 5 ++ ravendb/documents/ai/__init__.py | 5 ++ ravendb/documents/ai/ai_conversation.py | 7 +- ravendb/documents/ai/content_part.py | 61 +++++++++++++++ .../operations/ai/agents/__init__.py | 2 + .../ai/agents/ai_agent_configuration.py | 75 +++++++++++++++---- .../ai/agents/run_conversation_operation.py | 49 ++++++++++-- .../operations/ai/azure_open_ai_settings.py | 2 + 8 files changed, 181 insertions(+), 25 deletions(-) create mode 100644 ravendb/documents/ai/content_part.py diff --git a/ravendb/__init__.py b/ravendb/__init__.py index 0c1a7bea..85e9782c 100644 --- a/ravendb/__init__.py +++ b/ravendb/__init__.py @@ -82,10 +82,15 @@ AiOperations, AiConversation, AiConversationResult, + ContentPart, + TextPart, + AiMessagePromptFields, + AiMessagePromptTypes, ) from ravendb.documents.operations.ai.agents import ( AiAgentConfiguration, AiAgentConfigurationResult, + AiAgentParameter, AiAgentToolAction, AiAgentToolQuery, AiAgentPersistenceConfiguration, diff --git a/ravendb/documents/ai/__init__.py b/ravendb/documents/ai/__init__.py index de693a7f..9c2c750a 100644 --- a/ravendb/documents/ai/__init__.py +++ b/ravendb/documents/ai/__init__.py @@ -2,6 +2,7 @@ from .ai_conversation import AiConversation from .ai_conversation_result import AiConversationResult from .ai_answer import AiAnswer, AiConversationStatus +from .content_part import ContentPart, TextPart, AiMessagePromptFields, AiMessagePromptTypes __all__ = [ "AiOperations", @@ -9,4 +10,8 @@ "AiConversationResult", "AiAnswer", "AiConversationStatus", + "ContentPart", + "TextPart", + "AiMessagePromptFields", + "AiMessagePromptTypes", ] diff --git a/ravendb/documents/ai/ai_conversation.py b/ravendb/documents/ai/ai_conversation.py index d08f8850..9e4fd649 100644 --- a/ravendb/documents/ai/ai_conversation.py +++ b/ravendb/documents/ai/ai_conversation.py @@ -6,6 +6,7 @@ from datetime import timedelta from ravendb.documents.ai.ai_answer import AiAnswer, AiConversationStatus +from ravendb.documents.ai.content_part import ContentPart, TextPart from ravendb.documents.operations.ai.agents import ( AiAgentActionRequest, AiAgentActionResponse, @@ -53,7 +54,7 @@ def __init__( self._conversation_id = conversation_id self._change_vector = change_vector - self._prompt_parts: List[str] = [] + self._prompt_parts: List[ContentPart] = [] self._action_responses: List[AiAgentActionResponse] = [] self._action_requests: Optional[List[AiAgentActionRequest]] = None @@ -269,7 +270,7 @@ 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._prompt_parts.clear() - self._prompt_parts.append(user_prompt) + self.add_user_prompt(user_prompt) def add_user_prompt(self, *prompts: str) -> None: """ @@ -284,7 +285,7 @@ def add_user_prompt(self, *prompts: str) -> None: 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) + self._prompt_parts.append(TextPart(prompt)) def handle( self, diff --git a/ravendb/documents/ai/content_part.py b/ravendb/documents/ai/content_part.py new file mode 100644 index 00000000..e46ee2fd --- /dev/null +++ b/ravendb/documents/ai/content_part.py @@ -0,0 +1,61 @@ +from __future__ import annotations +from typing import Dict, Any + + +class AiMessagePromptFields: + """Constants for AI message prompt field names.""" + + TEXT = "text" + TYPE = "type" + + +class AiMessagePromptTypes: + """Constants for AI message prompt types.""" + + TEXT = "text" + + +class ContentPart: + """ + Base class for content parts in AI prompts. + Content parts allow structured prompt content with different types (text, etc.). + """ + + def __init__(self, content_type: str): + self._type = content_type + + @property + def type(self) -> str: + return self._type + + def to_json(self) -> Dict[str, Any]: + """ + Converts the content part to a JSON-serializable dictionary. + Subclasses should override this method to include their specific fields. + """ + return {AiMessagePromptFields.TYPE: self._type} + + +class TextPart(ContentPart): + """ + Represents a text content part in AI prompts. + """ + + def __init__(self, text: str): + super().__init__(AiMessagePromptTypes.TEXT) + self._text = text + + @property + def text(self) -> str: + return self._text + + @text.setter + def text(self, value: str): + self._text = value + + def to_json(self) -> Dict[str, Any]: + return { + AiMessagePromptFields.TYPE: self._type, + AiMessagePromptFields.TEXT: self._text, + } + diff --git a/ravendb/documents/operations/ai/agents/__init__.py b/ravendb/documents/operations/ai/agents/__init__.py index a88ce169..87fbd158 100644 --- a/ravendb/documents/operations/ai/agents/__init__.py +++ b/ravendb/documents/operations/ai/agents/__init__.py @@ -1,5 +1,6 @@ from .ai_agent_configuration import ( AiAgentConfiguration, + AiAgentParameter, AiAgentToolAction, AiAgentToolQuery, AiAgentPersistenceConfiguration, @@ -33,6 +34,7 @@ __all__ = [ "AiAgentConfiguration", "AiAgentConfigurationResult", + "AiAgentParameter", "AiAgentToolAction", "AiAgentToolQuery", "AiAgentPersistenceConfiguration", diff --git a/ravendb/documents/operations/ai/agents/ai_agent_configuration.py b/ravendb/documents/operations/ai/agents/ai_agent_configuration.py index f609e2ff..b460987b 100644 --- a/ravendb/documents/operations/ai/agents/ai_agent_configuration.py +++ b/ravendb/documents/operations/ai/agents/ai_agent_configuration.py @@ -1,5 +1,48 @@ from __future__ import annotations -from typing import List, Set, Optional, Dict, Any +from typing import List, Optional, Dict, Any, Union + + +class AiAgentParameter: + """ + Represents a parameter for an AI agent configuration. + Parameters can be used to pass values to the agent's system prompt. + """ + + def __init__( + self, + name: str = None, + description: str = None, + send_to_model: bool = None, + ): + """ + Initialize an agent parameter. + + Args: + name: The parameter name. Cannot be null or empty. + description: A human-readable description. May be null or empty. + send_to_model: When False, the parameter is hidden from the model + (it will not be included in prompts/echo messages). + When True, the parameter is exposed to the model. + If None (default), treated as exposed. + """ + self.name = name + self.description: Optional[str] = description + self.send_to_model: Optional[bool] = send_to_model + + def to_json(self) -> Dict[str, Any]: + return { + "Name": self.name, + "Description": self.description, + "SendToModel": self.send_to_model, + } + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentParameter: + return cls( + name=json_dict.get("name") or json_dict.get("Name"), + description=json_dict.get("description") or json_dict.get("Description"), + send_to_model=json_dict.get("sendToModel") if "sendToModel" in json_dict else json_dict.get("SendToModel"), + ) class AiAgentToolQuery: @@ -260,7 +303,7 @@ def __init__( queries: List[AiAgentToolQuery] = None, actions: List[AiAgentToolAction] = None, persistence: AiAgentPersistenceConfiguration = None, - parameters: Set[str] = None, + parameters: List[Union[str, AiAgentParameter]] = None, chat_trimming: AiAgentChatTrimmingConfiguration = None, max_model_iterations_per_call: int = None, ): @@ -273,14 +316,24 @@ def __init__( self.queries: List[AiAgentToolQuery] = queries or [] self.actions: List[AiAgentToolAction] = actions or [] self.persistence: Optional[AiAgentPersistenceConfiguration] = persistence - self.parameters: Set[str] = parameters or set() + self.parameters: List[AiAgentParameter] = self._normalize_parameters(parameters) self.chat_trimming: Optional[AiAgentChatTrimmingConfiguration] = chat_trimming self.max_model_iterations_per_call: Optional[int] = max_model_iterations_per_call - def to_json(self) -> Dict[str, Any]: - # Convert parameters set to list of parameter objects using list comprehension - parameters_list = [{"Name": param_name, "Description": None} for param_name in self.parameters] + @staticmethod + def _normalize_parameters(parameters: List[Union[str, AiAgentParameter]]) -> List[AiAgentParameter]: + """Convert a list of strings or AiAgentParameter objects to a list of AiAgentParameter objects.""" + if not parameters: + return [] + result = [] + for param in parameters: + if isinstance(param, str): + result.append(AiAgentParameter(name=param)) + else: + result.append(param) + return result + def to_json(self) -> Dict[str, Any]: return { "Identifier": self.identifier, "Name": self.name, @@ -291,7 +344,7 @@ def to_json(self) -> Dict[str, Any]: "Queries": [q.to_json() for q in self.queries], "Actions": [a.to_json() for a in self.actions], "Persistence": self.persistence.to_json() if self.persistence else None, - "Parameters": parameters_list, + "Parameters": [p.to_json() for p in self.parameters], "ChatTrimming": self.chat_trimming.to_json() if self.chat_trimming else None, "MaxModelIterationsPerCall": self.max_model_iterations_per_call, } @@ -321,13 +374,7 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiAgentConfiguration: params_data = json_dict.get("parameters") or json_dict.get("Parameters") if params_data: - # Handle both string list and object list formats - if params_data and isinstance(params_data[0], dict): - # New format: list of objects with name property - instance.parameters = set(param.get("name") or param.get("Name") for param in params_data) - else: - # Old format: list of strings - instance.parameters = set(params_data) + instance.parameters = [AiAgentParameter.from_json(param) for param in params_data] trimming_data = json_dict.get("chatTrimming") or json_dict.get("ChatTrimming") if trimming_data: diff --git a/ravendb/documents/operations/ai/agents/run_conversation_operation.py b/ravendb/documents/operations/ai/agents/run_conversation_operation.py index 75ed4526..63d8b3d8 100644 --- a/ravendb/documents/operations/ai/agents/run_conversation_operation.py +++ b/ravendb/documents/operations/ai/agents/run_conversation_operation.py @@ -9,6 +9,7 @@ from ravendb.http.server_node import ServerNode import requests from ravendb.http.misc import ResponseDisposeHandling +from ravendb.documents.ai.content_part import ContentPart TSchema = TypeVar("TSchema") @@ -64,6 +65,7 @@ class AiUsage: completion_tokens: int = 0 total_tokens: int = 0 cached_tokens: int = 0 + reasoning_tokens: int = 0 @classmethod def from_json(cls, json_dict: Dict[str, Any]) -> AiUsage: @@ -72,6 +74,7 @@ def from_json(cls, json_dict: Dict[str, Any]) -> AiUsage: completion_tokens=json_dict.get("CompletionTokens", 0), total_tokens=json_dict.get("TotalTokens", 0), cached_tokens=json_dict.get("CachedTokens", 0), + reasoning_tokens=json_dict.get("ReasoningTokens", 0), ) def to_json(self) -> Dict[str, Any]: @@ -80,8 +83,36 @@ def to_json(self) -> Dict[str, Any]: "CompletionTokens": self.completion_tokens, "TotalTokens": self.total_tokens, "CachedTokens": self.cached_tokens, + "ReasoningTokens": self.reasoning_tokens, } + @staticmethod + def get_usage_difference(current: AiUsage, previous: AiUsage) -> AiUsage: + """ + Calculate the usage difference between current and previous usage. + + Args: + current: The current usage statistics + previous: The previous usage statistics + + Returns: + An AiUsage object representing the difference + """ + previous_total_without_reasoning = ( + previous.completion_tokens - previous.reasoning_tokens + previous.prompt_tokens + ) + return AiUsage( + # in case the model gives us crappy results and current.prompt_tokens - previous_total_without_reasoning < 0 + prompt_tokens=max(current.prompt_tokens - previous_total_without_reasoning, 0), + # in case the model gives us crappy results and current.total_tokens - previous_total_without_reasoning < 0 + total_tokens=max(current.total_tokens - previous_total_without_reasoning, 0), + # we don't want to subtract cached tokens, as they are only for the last response + cached_tokens=current.cached_tokens, + # we don't want to subtract completion tokens, as they are only for the last response + completion_tokens=current.completion_tokens, + reasoning_tokens=current.reasoning_tokens, + ) + class ConversationResult(Generic[TSchema]): def __init__( @@ -161,11 +192,11 @@ class ConversationRequestBody: def __init__( self, action_responses: Optional[List[AiAgentActionResponse]] = None, - user_prompt: Optional[List[str]] = None, + user_prompt: Optional[List[ContentPart]] = None, creation_options: Optional[AiConversationCreationOptions] = None, ): self.action_responses: Optional[List[AiAgentActionResponse]] = action_responses - self.user_prompt: Optional[List[str]] = user_prompt # List of prompt parts + self.user_prompt: Optional[List[ContentPart]] = user_prompt # List of ContentPart objects self.creation_options: Optional[AiConversationCreationOptions] = creation_options def to_json(self) -> Dict[str, Any]: @@ -182,8 +213,10 @@ def to_json(self) -> Dict[str, Any]: None if self.action_responses is None else [resp.to_json() for resp in self.action_responses] ) - # UserPrompt: null if None, otherwise array (even if empty) - result["UserPrompt"] = self.user_prompt + # UserPrompt: null if None, otherwise array of ContentPart JSON objects + result["UserPrompt"] = ( + None if self.user_prompt is None else [part.to_json() for part in self.user_prompt] + ) # CreationOptions: always present (create empty if None, matching C# behavior) result["CreationOptions"] = (self.creation_options or AiConversationCreationOptions()).to_json() @@ -203,7 +236,7 @@ def __init__( self, agent_id: str, conversation_id: str, - prompt_parts: Optional[List[str]] = None, + prompt_parts: Optional[List[ContentPart]] = None, action_responses: Optional[List[AiAgentActionResponse]] = None, options: Optional[AiConversationCreationOptions] = None, change_vector: Optional[str] = None, @@ -216,7 +249,7 @@ def __init__( Args: agent_id: The ID of the AI agent (required) conversation_id: The ID of the conversation (required) - prompt_parts: List of prompt strings to send to the agent + prompt_parts: List of ContentPart objects to send to the agent action_responses: List of action responses from previous turn options: Creation options including parameters and expiration change_vector: Change vector for optimistic concurrency @@ -258,7 +291,7 @@ def __init__( self, agent_id: str, conversation_id: str, - prompt_parts: Optional[List[str]] = None, + prompt_parts: Optional[List[ContentPart]] = None, action_responses: Optional[List[AiAgentActionResponse]] = None, options: Optional[AiConversationCreationOptions] = None, change_vector: Optional[str] = None, @@ -309,7 +342,7 @@ def create_request(self, node: ServerNode) -> requests.Request: # Build request body with correct structure to match .NET client request_body = ConversationRequestBody( action_responses=self._action_responses, - user_prompt="".join(self._prompt_parts), + user_prompt=self._prompt_parts, creation_options=self._options, ) diff --git a/ravendb/documents/operations/ai/azure_open_ai_settings.py b/ravendb/documents/operations/ai/azure_open_ai_settings.py index b7e512f7..b636b0d8 100644 --- a/ravendb/documents/operations/ai/azure_open_ai_settings.py +++ b/ravendb/documents/operations/ai/azure_open_ai_settings.py @@ -14,6 +14,8 @@ def __init__( temperature: float = None, ): super().__init__(api_key, endpoint, model, dimensions, temperature) + if deployment_name is None: + raise ValueError("deployment_name cannot be None") self.deployment_name = deployment_name @classmethod From 19fedd93e4f403cb3161f577a3433d6e36d90d43 Mon Sep 17 00:00:00 2001 From: Gracjan Sadowicz Date: Tue, 9 Dec 2025 23:12:37 +0100 Subject: [PATCH 2/3] RDBC-967 GenAI Tasks Operations --- ravendb/documents/__init__.py | 5 + ravendb/documents/ai/content_part.py | 1 - ravendb/documents/operations/ai/__init__.py | 31 + .../abstract_ai_integration_configuration.py | 46 ++ .../operations/ai/add_gen_ai_operation.py | 75 +++ .../ai/agents/run_conversation_operation.py | 4 +- .../operations/ai/ai_connection_string.py | 56 ++ .../ai/ai_task_identifier_helper.py | 94 +++ .../ai/ai_task_operation_results.py | 49 ++ .../operations/ai/gen_ai_configuration.py | 208 ++++++ .../operations/ai/gen_ai_transformation.py | 33 + .../operations/ai/update_gen_ai_operation.py | 99 +++ ravendb/documents/operations/etl/__init__.py | 18 + .../documents/operations/etl/configuration.py | 122 +++- .../operations/etl/etl_operation_results.py | 32 + ravendb/documents/operations/etl/etl_type.py | 12 + .../operations/etl/transformation.py | 426 ++++++++++++ ravendb/documents/operations/ongoing_tasks.py | 361 ++++++++++ .../documents/starting_point_change_vector.py | 22 + ravendb/tests/gen_ai_tests/__init__.py | 0 .../gen_ai_tests/test_gen_ai_operations.py | 617 ++++++++++++++++++ 21 files changed, 2301 insertions(+), 10 deletions(-) create mode 100644 ravendb/documents/operations/ai/abstract_ai_integration_configuration.py create mode 100644 ravendb/documents/operations/ai/add_gen_ai_operation.py create mode 100644 ravendb/documents/operations/ai/ai_task_identifier_helper.py create mode 100644 ravendb/documents/operations/ai/ai_task_operation_results.py create mode 100644 ravendb/documents/operations/ai/gen_ai_configuration.py create mode 100644 ravendb/documents/operations/ai/gen_ai_transformation.py create mode 100644 ravendb/documents/operations/ai/update_gen_ai_operation.py create mode 100644 ravendb/documents/operations/etl/etl_operation_results.py create mode 100644 ravendb/documents/operations/etl/etl_type.py create mode 100644 ravendb/documents/operations/etl/transformation.py create mode 100644 ravendb/documents/starting_point_change_vector.py create mode 100644 ravendb/tests/gen_ai_tests/__init__.py create mode 100644 ravendb/tests/gen_ai_tests/test_gen_ai_operations.py diff --git a/ravendb/documents/__init__.py b/ravendb/documents/__init__.py index e69de29b..28cd3c0b 100644 --- a/ravendb/documents/__init__.py +++ b/ravendb/documents/__init__.py @@ -0,0 +1,5 @@ +from ravendb.documents.starting_point_change_vector import StartingPointChangeVector + +__all__ = [ + "StartingPointChangeVector", +] diff --git a/ravendb/documents/ai/content_part.py b/ravendb/documents/ai/content_part.py index e46ee2fd..754023b8 100644 --- a/ravendb/documents/ai/content_part.py +++ b/ravendb/documents/ai/content_part.py @@ -58,4 +58,3 @@ def to_json(self) -> Dict[str, Any]: AiMessagePromptFields.TYPE: self._type, AiMessagePromptFields.TEXT: self._text, } - diff --git a/ravendb/documents/operations/ai/__init__.py b/ravendb/documents/operations/ai/__init__.py index e69de29b..f43bd77f 100644 --- a/ravendb/documents/operations/ai/__init__.py +++ b/ravendb/documents/operations/ai/__init__.py @@ -0,0 +1,31 @@ +from ravendb.documents.operations.ai.ai_connection_string import ( + AiConnectionString, + AiModelType, + AiConnectorType, +) +from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper +from ravendb.documents.operations.ai.gen_ai_transformation import GenAiTransformation +from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration +from ravendb.documents.operations.ai.abstract_ai_integration_configuration import AbstractAiIntegrationConfiguration +from ravendb.documents.operations.ai.ai_task_operation_results import ( + AddAiTaskOperationResult, + AddGenAiOperationResult, + AddEmbeddingsGenerationOperationResult, +) +from ravendb.documents.operations.ai.add_gen_ai_operation import AddGenAiOperation +from ravendb.documents.operations.ai.update_gen_ai_operation import UpdateGenAiOperation + +__all__ = [ + "AiConnectionString", + "AiModelType", + "AiConnectorType", + "AiTaskIdentifierHelper", + "GenAiTransformation", + "GenAiConfiguration", + "AbstractAiIntegrationConfiguration", + "AddAiTaskOperationResult", + "AddGenAiOperationResult", + "AddEmbeddingsGenerationOperationResult", + "AddGenAiOperation", + "UpdateGenAiOperation", +] diff --git a/ravendb/documents/operations/ai/abstract_ai_integration_configuration.py b/ravendb/documents/operations/ai/abstract_ai_integration_configuration.py new file mode 100644 index 00000000..fadea02f --- /dev/null +++ b/ravendb/documents/operations/ai/abstract_ai_integration_configuration.py @@ -0,0 +1,46 @@ +from __future__ import annotations +from abc import ABC +from typing import TYPE_CHECKING, Optional, List + +from ravendb.documents.operations.etl.configuration import EtlConfiguration +from ravendb.documents.operations.ai.ai_connection_string import AiConnectionString, AiConnectorType +from ravendb.documents.operations.etl.transformation import Transformation + +if TYPE_CHECKING: + pass + + +class AbstractAiIntegrationConfiguration(EtlConfiguration[AiConnectionString], ABC): + """ + Base class for AI integration configurations. + Extends EtlConfiguration with AiConnectionString as the connection type. + """ + + def __init__( + self, + name: Optional[str] = None, + task_id: int = 0, + connection_string_name: Optional[str] = None, + mentor_node: Optional[str] = None, + pin_to_mentor_node: bool = False, + transforms: Optional[List[Transformation]] = None, + disabled: bool = False, + allow_etl_on_non_encrypted_channel: bool = False, + ): + super().__init__( + name=name, + task_id=task_id, + connection_string_name=connection_string_name, + mentor_node=mentor_node, + pin_to_mentor_node=pin_to_mentor_node, + transforms=transforms, + disabled=disabled, + allow_etl_on_non_encrypted_channel=allow_etl_on_non_encrypted_channel, + ) + + @property + def ai_connector_type(self) -> AiConnectorType: + """Returns the AI connector type based on the active provider in the connection.""" + if self.connection: + return self.connection.get_active_provider() + return AiConnectorType.NONE diff --git a/ravendb/documents/operations/ai/add_gen_ai_operation.py b/ravendb/documents/operations/ai/add_gen_ai_operation.py new file mode 100644 index 00000000..71b094c1 --- /dev/null +++ b/ravendb/documents/operations/ai/add_gen_ai_operation.py @@ -0,0 +1,75 @@ +from __future__ import annotations +import json +from typing import Optional, TYPE_CHECKING +from urllib.parse import quote + +from ravendb.documents.operations.definitions import MaintenanceOperation +from ravendb.documents.conventions import DocumentConventions +from ravendb.http.raven_command import RavenCommand +from ravendb.http.server_node import ServerNode +from ravendb.documents.starting_point_change_vector import StartingPointChangeVector +from ravendb.documents.operations.ai.ai_task_operation_results import AddGenAiOperationResult +import requests + +from ravendb.util.util import RaftIdGenerator + +if TYPE_CHECKING: + from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration + + +class AddGenAiOperation(MaintenanceOperation[AddGenAiOperationResult]): + """ + Operation to add a new GenAI task to the database. + """ + + def __init__( + self, + configuration: GenAiConfiguration, + starting_point: Optional[StartingPointChangeVector] = StartingPointChangeVector.LAST_DOCUMENT, + ): + if configuration is None: + raise ValueError("configuration cannot be None") + + self._configuration = configuration + self._starting_point = starting_point + + def get_command(self, conventions: DocumentConventions) -> RavenCommand[AddGenAiOperationResult]: + return AddGenAiCommand(self._configuration, self._starting_point, conventions) + + +class AddGenAiCommand(RavenCommand[AddGenAiOperationResult]): + def __init__( + self, + configuration: GenAiConfiguration, + starting_point: StartingPointChangeVector, + conventions: DocumentConventions, + ): + super().__init__(AddGenAiOperationResult) + self._configuration = configuration + self._starting_point = starting_point + self._conventions = conventions + + def is_read_request(self) -> bool: + return False + + def create_request(self, node: ServerNode) -> requests.Request: + url = f"{node.url}/databases/{node.database}/admin/etl?changeVector={quote(self._starting_point.value)}" + + body_json = self._configuration.to_json() + body = json.dumps(body_json) + + request = requests.Request("PUT", url) + request.headers = {"Content-Type": "application/json"} + request.data = body + return request + + def set_response(self, response: str, from_cache: bool) -> None: + if response is None: + self.result = AddGenAiOperationResult() + return + + response_json = json.loads(response) + self.result = AddGenAiOperationResult.from_json(response_json) + + def get_raft_unique_request_id(self) -> str: + return RaftIdGenerator.new_id() diff --git a/ravendb/documents/operations/ai/agents/run_conversation_operation.py b/ravendb/documents/operations/ai/agents/run_conversation_operation.py index 63d8b3d8..54dc293e 100644 --- a/ravendb/documents/operations/ai/agents/run_conversation_operation.py +++ b/ravendb/documents/operations/ai/agents/run_conversation_operation.py @@ -214,9 +214,7 @@ def to_json(self) -> Dict[str, Any]: ) # UserPrompt: null if None, otherwise array of ContentPart JSON objects - result["UserPrompt"] = ( - None if self.user_prompt is None else [part.to_json() for part in self.user_prompt] - ) + result["UserPrompt"] = None if self.user_prompt is None else [part.to_json() for part in self.user_prompt] # CreationOptions: always present (create empty if None, matching C# behavior) result["CreationOptions"] = (self.creation_options or AiConversationCreationOptions()).to_json() diff --git a/ravendb/documents/operations/ai/ai_connection_string.py b/ravendb/documents/operations/ai/ai_connection_string.py index dda821e3..9e39e520 100644 --- a/ravendb/documents/operations/ai/ai_connection_string.py +++ b/ravendb/documents/operations/ai/ai_connection_string.py @@ -19,6 +19,18 @@ class AiModelType(enum.Enum): CHAT = "Chat" +class AiConnectorType(enum.Enum): + NONE = "None" + OPEN_AI = "OpenAi" + AZURE_OPEN_AI = "AzureOpenAi" + OLLAMA = "Ollama" + EMBEDDED = "Embedded" + GOOGLE = "Google" + HUGGING_FACE = "HuggingFace" + MISTRAL_AI = "MistralAi" + VERTEX = "Vertex" + + class AiConnectionString(ConnectionString): def __init__( self, @@ -87,6 +99,50 @@ def __init__( def get_type(self): return ConnectionStringType.AI.value + def get_active_provider(self) -> AiConnectorType: + """Returns the active AI connector type based on which settings are configured.""" + if self.openai_settings: + return AiConnectorType.OPEN_AI + if self.azure_openai_settings: + return AiConnectorType.AZURE_OPEN_AI + if self.ollama_settings: + return AiConnectorType.OLLAMA + if self.embedded_settings: + return AiConnectorType.EMBEDDED + if self.google_settings: + return AiConnectorType.GOOGLE + if self.huggingface_settings: + return AiConnectorType.HUGGING_FACE + if self.mistral_ai_settings: + return AiConnectorType.MISTRAL_AI + if self.vertex_settings: + return AiConnectorType.VERTEX + return AiConnectorType.NONE + + def using_encrypted_communication_channel(self) -> bool: + """Returns True if the connection uses HTTPS (encrypted communication).""" + active_settings = None + if self.openai_settings: + active_settings = self.openai_settings + elif self.azure_openai_settings: + active_settings = self.azure_openai_settings + elif self.ollama_settings: + active_settings = self.ollama_settings + elif self.google_settings: + active_settings = self.google_settings + elif self.huggingface_settings: + active_settings = self.huggingface_settings + elif self.mistral_ai_settings: + active_settings = self.mistral_ai_settings + elif self.vertex_settings: + active_settings = self.vertex_settings + + if active_settings and hasattr(active_settings, "endpoint") and active_settings.endpoint: + return active_settings.endpoint.lower().startswith("https://") + + # Embedded settings don't have an endpoint + return False + def to_json(self) -> Dict[str, Any]: return { "Name": self.name, diff --git a/ravendb/documents/operations/ai/ai_task_identifier_helper.py b/ravendb/documents/operations/ai/ai_task_identifier_helper.py new file mode 100644 index 00000000..1107659f --- /dev/null +++ b/ravendb/documents/operations/ai/ai_task_identifier_helper.py @@ -0,0 +1,94 @@ +import unicodedata +from typing import List, Tuple + + +class AiTaskIdentifierHelper: + """Helper class for validating and generating AI task identifiers.""" + + @staticmethod + def validate_identifier(identifier: str) -> Tuple[bool, List[str]]: + """ + Validates an AI task identifier. + + Args: + identifier: The identifier to validate. + + Returns: + A tuple of (is_valid, errors). If valid, errors is an empty list. + """ + errors: List[str] = [] + + if not identifier or identifier.isspace(): + errors.append("Identifier cannot be empty or contain only whitespace;") + return False, errors + + # Check that the string is already normalized (contains only a-z, 0-9 and hyphens) + normalized = unicodedata.normalize("NFD", identifier) + if identifier != normalized: + errors.append("Identifier contains diacritical marks or non-ASCII characters;") + + # Check that there are no uppercase letters + if any(c.isupper() for c in identifier): + errors.append("Identifier contains uppercase letters;") + + # Check for invalid characters and collect them + invalid_chars = set() + for c in identifier: + if not (("a" <= c <= "z") or ("0" <= c <= "9") or c == "-"): + invalid_chars.add(c) + + if invalid_chars: + chars_str = ", ".join(f"'{c}'" for c in sorted(invalid_chars)) + errors.append( + f"Identifier contains invalid characters: {chars_str}. " + f"Only lowercase letters (a-z), numbers (0-9) and hyphens (-) are allowed." + ) + + # Check that there are no consecutive hyphens + if "--" in identifier: + errors.append("Identifier contains consecutive hyphens;") + + # Check that the string does not end with a hyphen + if identifier.endswith("-"): + errors.append("Identifier ends with a hyphen;") + + return len(errors) == 0, errors + + @staticmethod + def generate_identifier(input_str: str) -> str: + """ + Generates a valid identifier from an input string. + + Args: + input_str: The input string to convert to an identifier. + + Returns: + A valid identifier string, or a default identifier if input is empty. + """ + if not input_str or input_str.isspace(): + return None + + result = [] + last_was_hyphen = False + + # First normalize to FormD to separate letters from their diacritics + normalized = unicodedata.normalize("NFD", input_str) + + for c in normalized: + # Check if this is a letter that needs to be preserved + if "a" <= c <= "z" or "0" <= c <= "9": + result.append(c) + last_was_hyphen = False + elif "A" <= c <= "Z": + result.append(c.lower()) + last_was_hyphen = False + elif not last_was_hyphen and len(result) > 0: + # Add hyphen for any other character + result.append("-") + last_was_hyphen = True + + # Trim any trailing hyphens + final_result = "".join(result).rstrip("-") + + # Ensure we have at least one character + return final_result if final_result else "AiConnectionStringIdentifier" diff --git a/ravendb/documents/operations/ai/ai_task_operation_results.py b/ravendb/documents/operations/ai/ai_task_operation_results.py new file mode 100644 index 00000000..f4ea57e9 --- /dev/null +++ b/ravendb/documents/operations/ai/ai_task_operation_results.py @@ -0,0 +1,49 @@ +from __future__ import annotations +from typing import Dict, Any, Optional + +from ravendb.documents.operations.etl.etl_operation_results import AddEtlOperationResult + + +class AddAiTaskOperationResult(AddEtlOperationResult): + """Base result class for AI task add operations.""" + + def __init__( + self, + raft_command_index: int = 0, + task_id: int = 0, + identifier: Optional[str] = None, + ): + super().__init__(raft_command_index=raft_command_index, task_id=task_id) + self.identifier = identifier + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "AddAiTaskOperationResult": + return cls( + raft_command_index=json_dict.get("RaftCommandIndex", 0), + task_id=json_dict.get("TaskId", 0), + identifier=json_dict.get("Identifier"), + ) + + +class AddGenAiOperationResult(AddAiTaskOperationResult): + """Result of adding a GenAI task.""" + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "AddGenAiOperationResult": + return cls( + raft_command_index=json_dict.get("RaftCommandIndex", 0), + task_id=json_dict.get("TaskId", 0), + identifier=json_dict.get("Identifier"), + ) + + +class AddEmbeddingsGenerationOperationResult(AddAiTaskOperationResult): + """Result of adding an Embeddings Generation task.""" + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "AddEmbeddingsGenerationOperationResult": + return cls( + raft_command_index=json_dict.get("RaftCommandIndex", 0), + task_id=json_dict.get("TaskId", 0), + identifier=json_dict.get("Identifier"), + ) diff --git a/ravendb/documents/operations/ai/gen_ai_configuration.py b/ravendb/documents/operations/ai/gen_ai_configuration.py new file mode 100644 index 00000000..bff5dac7 --- /dev/null +++ b/ravendb/documents/operations/ai/gen_ai_configuration.py @@ -0,0 +1,208 @@ +from __future__ import annotations +from typing import Dict, Any, Optional, List + +from ravendb.documents.operations.ai.abstract_ai_integration_configuration import AbstractAiIntegrationConfiguration +from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper +from ravendb.documents.operations.ai.gen_ai_transformation import GenAiTransformation +from ravendb.documents.operations.ai.agents.ai_agent_configuration import AiAgentToolQuery +from ravendb.documents.operations.etl.etl_type import EtlType +from ravendb.documents.operations.etl.transformation import Transformation + + +class GenAiConfiguration(AbstractAiIntegrationConfiguration): + """ + Configuration for GenAI document processing in RavenDB. + This configuration defines how documents from a collection are processed using AI. + """ + + DEFAULT_MAX_CONCURRENCY = 4 + TRANSFORMATION_NAME = "GenAi-transform-script" + + def __init__( + self, + name: str = None, + identifier: str = None, + collection: str = None, + connection_string_name: str = None, + prompt: str = None, + json_schema: str = None, + sample_object: str = None, + update_script: str = None, + gen_ai_transformation: GenAiTransformation = None, + max_concurrency: int = None, + queries: List[AiAgentToolQuery] = None, + enable_tracing: bool = False, + expiration_in_sec: int = None, + disabled: bool = False, + mentor_node: str = None, + pin_to_mentor_node: bool = False, + task_id: int = 0, + allow_etl_on_non_encrypted_channel: bool = False, + ): + super().__init__( + name=name, + task_id=task_id, + connection_string_name=connection_string_name, + mentor_node=mentor_node, + pin_to_mentor_node=pin_to_mentor_node, + disabled=disabled, + allow_etl_on_non_encrypted_channel=allow_etl_on_non_encrypted_channel, + ) + + self.identifier = identifier + self.collection = collection + self.prompt = prompt + self.json_schema = json_schema + self.sample_object = sample_object + self.update_script = update_script + self.gen_ai_transformation = gen_ai_transformation + self.max_concurrency = max_concurrency if max_concurrency is not None else self.DEFAULT_MAX_CONCURRENCY + self.queries: List[AiAgentToolQuery] = queries or [] + self.enable_tracing = enable_tracing + self.expiration_in_sec: Optional[int] = expiration_in_sec + + self._transforms: Optional[List[Transformation]] = None + + @property + def etl_type(self) -> EtlType: + return EtlType.GEN_AI + + def get_destination(self) -> str: + """Returns the destination identifier for this configuration.""" + return self.identifier + + def get_default_task_name(self) -> str: + """Returns the default task name for this configuration.""" + return self.identifier + + def using_encrypted_communication_channel(self) -> bool: + """Returns True if the connection uses encrypted communication.""" + if self.connection: + return self.connection.using_encrypted_communication_channel() + return False + + # todo: use validate in __init__ or in some more suitable place + def validate( + self, + validate_name: bool = True, + # todo: validate_connection: bool = True, + validate_identifier: bool = True, + ) -> List[str]: + """ + Validates the GenAI configuration. + + Args: + validate_name: Whether to validate the name field. + validate_identifier: Whether to validate the identifier format. + + Returns: + A list of validation error messages. Empty list if valid. + """ + errors: List[str] = [] + + # Validate identifier using AiTaskIdentifierHelper + if validate_identifier: + is_valid, id_errors = AiTaskIdentifierHelper.validate_identifier(self.identifier) + if not is_valid: + errors.extend(id_errors) + + if validate_name and not self.name: + errors.append("Name of GenAi configuration cannot be empty") + + if not self._test_mode and not self.connection_string_name: + errors.append("ConnectionStringName cannot be empty") + + if not self.collection: + errors.append("Collection must be provided") + + if self.gen_ai_transformation is None: + errors.append("GenAiTransformation must be specified") + else: + is_valid, error = self.gen_ai_transformation.validate_script() + if not is_valid: + errors.append(error) + + if not self._test_mode: + if not self.prompt: + errors.append("Prompt must be provided") + + if not self.json_schema and not self.sample_object: + errors.append("You must provide either a JSON schema or a sample object") + + if not self.update_script: + errors.append("You must provide an update function") + + return errors + + def generate_identifier(self) -> str: + """Generates a valid identifier based on the configuration name.""" + return AiTaskIdentifierHelper.generate_identifier(self.name) + + @property + def transforms(self) -> List[Transformation]: + """ + GenAiConfiguration uses a single transformation based on GenAiTransformation. + This property is provided for compatibility with the base EtlConfiguration. + """ + if self._transforms is None: + self._transforms = [ + Transformation( + name=self.TRANSFORMATION_NAME, + collections=[self.collection] if self.collection else [], + script=self.gen_ai_transformation.script if self.gen_ai_transformation else None, + ) + ] + return self._transforms + + @transforms.setter + def transforms(self, value: List[Transformation]): + raise NotImplementedError( + "GenAiConfiguration doesn't support multiple transformations. " + "Please use gen_ai_transformation property instead." + ) + + def to_json(self) -> Dict[str, Any]: + result = super().to_json() + result.update( + { + "Identifier": self.identifier, + "AiConnectorType": self.ai_connector_type.value if self.ai_connector_type else None, + "Collection": self.collection, + "Prompt": self.prompt, + "JsonSchema": self.json_schema, + "SampleObject": self.sample_object, + "UpdateScript": self.update_script, + "GenAiTransformation": self.gen_ai_transformation.to_json() if self.gen_ai_transformation else None, + "MaxConcurrency": self.max_concurrency, + "Queries": [q.to_json() for q in self.queries] if self.queries else None, + "EnableTracing": self.enable_tracing, + "ExpirationInSec": self.expiration_in_sec, + } + ) + return result + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "GenAiConfiguration": + transformation_data = json_dict.get("GenAiTransformation") + queries_data = json_dict.get("Queries") + + return cls( + name=json_dict.get("Name"), + task_id=json_dict.get("TaskId", 0), + identifier=json_dict.get("Identifier"), + collection=json_dict.get("Collection"), + connection_string_name=json_dict.get("ConnectionStringName"), + prompt=json_dict.get("Prompt"), + json_schema=json_dict.get("JsonSchema"), + sample_object=json_dict.get("SampleObject"), + update_script=json_dict.get("UpdateScript"), + gen_ai_transformation=GenAiTransformation.from_json(transformation_data) if transformation_data else None, + max_concurrency=json_dict.get("MaxConcurrency", cls.DEFAULT_MAX_CONCURRENCY), + queries=[AiAgentToolQuery.from_json(q) for q in queries_data] if queries_data else None, + enable_tracing=json_dict.get("EnableTracing", False), + expiration_in_sec=json_dict.get("ExpirationInSec"), + disabled=json_dict.get("Disabled", False), + mentor_node=json_dict.get("MentorNode"), + pin_to_mentor_node=json_dict.get("PinToMentorNode", False), + allow_etl_on_non_encrypted_channel=json_dict.get("AllowEtlOnNonEncryptedChannel", False), + ) diff --git a/ravendb/documents/operations/ai/gen_ai_transformation.py b/ravendb/documents/operations/ai/gen_ai_transformation.py new file mode 100644 index 00000000..7c4984ea --- /dev/null +++ b/ravendb/documents/operations/ai/gen_ai_transformation.py @@ -0,0 +1,33 @@ +from typing import Dict, Any, Tuple + + +class GenAiTransformation: + """ + Represents the transformation script configuration for GenAI processing. + The script must call ai.genContext(ctx) function. + """ + + def __init__(self, script: str = None): + self.script = script + + def validate_script(self) -> Tuple[bool, str]: + """ + Validates that the script contains the required ai.genContext call. + + Returns: + A tuple of (is_valid, error_message). If valid, error_message is empty. + """ + if self.script and "ai.genContext" in self.script: + return True, "" + return False, "You must call the ai.genContext(ctx) function in your script" + + def to_json(self) -> Dict[str, Any]: + return { + "Script": self.script, + } + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "GenAiTransformation": + return cls( + script=json_dict.get("Script"), + ) diff --git a/ravendb/documents/operations/ai/update_gen_ai_operation.py b/ravendb/documents/operations/ai/update_gen_ai_operation.py new file mode 100644 index 00000000..86f32cb4 --- /dev/null +++ b/ravendb/documents/operations/ai/update_gen_ai_operation.py @@ -0,0 +1,99 @@ +from __future__ import annotations +import json +from typing import Optional, List, TYPE_CHECKING +from urllib.parse import quote + +from ravendb.documents.operations.definitions import MaintenanceOperation +from ravendb.documents.conventions import DocumentConventions +from ravendb.http.raven_command import RavenCommand +from ravendb.http.server_node import ServerNode +from ravendb.documents.starting_point_change_vector import StartingPointChangeVector +from ravendb.documents.operations.etl.etl_operation_results import UpdateEtlOperationResult +import requests + +from ravendb.util.util import RaftIdGenerator + +if TYPE_CHECKING: + from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration + + +class UpdateGenAiOperation(MaintenanceOperation[UpdateEtlOperationResult]): + """ + Operation to update an existing GenAI task in the database. + """ + + def __init__( + self, + task_id: int, + configuration: GenAiConfiguration, + starting_point: Optional[StartingPointChangeVector] = None, + reset: bool = False, + ): + if configuration is None: + raise ValueError("configuration cannot be None") + + self._task_id = task_id + self._configuration = configuration + self._starting_point = starting_point or StartingPointChangeVector.DO_NOT_CHANGE + self._reset = reset + + def get_command(self, conventions: DocumentConventions) -> RavenCommand[UpdateEtlOperationResult]: + transformations_to_reset: Optional[List[str]] = None + if self._reset: + transformations_to_reset = [self._configuration.TRANSFORMATION_NAME] + + return UpdateGenAiCommand( + self._task_id, + self._configuration, + self._starting_point, + transformations_to_reset, + conventions, + ) + + +class UpdateGenAiCommand(RavenCommand[UpdateEtlOperationResult]): + def __init__( + self, + task_id: int, + configuration: GenAiConfiguration, + starting_point: StartingPointChangeVector, + transformations_to_reset: Optional[List[str]], + conventions: DocumentConventions, + ): + super().__init__(UpdateEtlOperationResult) + self._task_id = task_id + self._configuration = configuration + self._starting_point = starting_point + self._transformations_to_reset = transformations_to_reset + self._conventions = conventions + + def is_read_request(self) -> bool: + return False + + def create_request(self, node: ServerNode) -> requests.Request: + url = f"{node.url}/databases/{node.database}/admin/etl?id={self._task_id}" + + if self._transformations_to_reset: + for transformation in self._transformations_to_reset: + url += f"&reset={quote(transformation)}" + + url += f"&changeVector={quote(self._starting_point.value)}" + + body_json = self._configuration.to_json() + body = json.dumps(body_json) + + request = requests.Request("PUT", url) + request.headers = {"Content-Type": "application/json"} + request.data = body + return request + + def set_response(self, response: str, from_cache: bool) -> None: + if response is None: + self.result = UpdateEtlOperationResult() + return + + response_json = json.loads(response) + self.result = UpdateEtlOperationResult.from_json(response_json) + + def get_raft_unique_request_id(self) -> str: + return RaftIdGenerator.new_id() diff --git a/ravendb/documents/operations/etl/__init__.py b/ravendb/documents/operations/etl/__init__.py index e69de29b..79e4eee4 100644 --- a/ravendb/documents/operations/etl/__init__.py +++ b/ravendb/documents/operations/etl/__init__.py @@ -0,0 +1,18 @@ +from ravendb.documents.operations.etl.etl_type import EtlType +from ravendb.documents.operations.etl.transformation import Transformation +from ravendb.documents.operations.etl.configuration import ( + EtlConfiguration, + RavenConnectionString, + RavenEtlConfiguration, +) +from ravendb.documents.operations.etl.etl_operation_results import AddEtlOperationResult, UpdateEtlOperationResult + +__all__ = [ + "EtlType", + "Transformation", + "EtlConfiguration", + "RavenConnectionString", + "RavenEtlConfiguration", + "AddEtlOperationResult", + "UpdateEtlOperationResult", +] diff --git a/ravendb/documents/operations/etl/configuration.py b/ravendb/documents/operations/etl/configuration.py index f01b1c0e..f503a62a 100644 --- a/ravendb/documents/operations/etl/configuration.py +++ b/ravendb/documents/operations/etl/configuration.py @@ -1,9 +1,14 @@ -from typing import Optional, Generic, TypeVar, List, Dict +from __future__ import annotations +from abc import ABC, abstractmethod +from typing import Optional, Generic, TypeVar, List, Dict, Any from ravendb.documents.operations.connection_strings import ConnectionString +from ravendb.documents.operations.etl.etl_type import EtlType +from ravendb.documents.operations.etl.transformation import Transformation import ravendb.serverwide.server_operation_executor -_T = TypeVar("_T") + +_T = TypeVar("_T", bound=ConnectionString) class RavenConnectionString(ConnectionString): @@ -33,10 +38,115 @@ def from_json(cls, json_dict: Dict) -> "RavenConnectionString": ) -# todo: implement -class EtlConfiguration(ConnectionString, Generic[_T]): - pass +class EtlConfiguration(ABC, Generic[_T]): + """ + Base class for ETL (Extract, Transform, Load) configurations. + """ + + def __init__( + self, + name: Optional[str] = None, + task_id: int = 0, + connection_string_name: Optional[str] = None, + mentor_node: Optional[str] = None, + pin_to_mentor_node: bool = False, + transforms: Optional[List[Transformation]] = None, + disabled: bool = False, + allow_etl_on_non_encrypted_channel: bool = False, + ): + self._initialized: bool = False + self._test_mode: bool = False + self._connection: Optional[_T] = None + self._transforms: List[Transformation] = transforms if transforms is not None else [] + + self.task_id = task_id + self.name = name + self.mentor_node = mentor_node + self.pin_to_mentor_node = pin_to_mentor_node + self.connection_string_name = connection_string_name + self.disabled = disabled + self.allow_etl_on_non_encrypted_channel = allow_etl_on_non_encrypted_channel + + @property + def initialized(self) -> bool: + return self._initialized + + @property + def test_mode(self) -> bool: + return self._test_mode + + @test_mode.setter + def test_mode(self, value: bool): + self._test_mode = value + + @property + def connection(self) -> Optional[_T]: + return self._connection + + @property + def transforms(self) -> List[Transformation]: + return self._transforms + + @transforms.setter + def transforms(self, value: List[Transformation]): + self._transforms = value + + def initialize(self, connection_string: _T) -> None: + """Initialize the configuration with a connection string.""" + if self._initialized: + return + self._connection = connection_string + self._initialized = True + + @abstractmethod + def get_destination(self) -> str: + """Returns the destination for this ETL configuration.""" + pass + + @abstractmethod + def get_default_task_name(self) -> str: + """Returns the default task name for this configuration.""" + pass + + @property + @abstractmethod + def etl_type(self) -> EtlType: + """Returns the ETL type for this configuration.""" + pass + + @abstractmethod + def using_encrypted_communication_channel(self) -> bool: + """Returns True if the connection uses encrypted communication.""" + pass + + def to_json(self) -> Dict[str, Any]: + return { + "EtlType": self.etl_type.value if self.etl_type else None, + "Name": self.name, + "TaskId": self.task_id, + "ConnectionStringName": self.connection_string_name, + "MentorNode": self.mentor_node, + "PinToMentorNode": self.pin_to_mentor_node, + "AllowEtlOnNonEncryptedChannel": self.allow_etl_on_non_encrypted_channel, + "Transforms": [t.to_json() for t in self.transforms] if self.transforms else [], + "Disabled": self.disabled, + } + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "EtlConfiguration": + raise NotImplementedError("Subclasses must implement from_json") class RavenEtlConfiguration(EtlConfiguration[RavenConnectionString]): - pass + @property + def etl_type(self) -> EtlType: + return EtlType.RAVEN + + def get_destination(self) -> str: + return self.connection.database if self.connection else "" + + def get_default_task_name(self) -> str: + return f"Raven ETL to {self.get_destination()}" + + def using_encrypted_communication_channel(self) -> bool: + return False diff --git a/ravendb/documents/operations/etl/etl_operation_results.py b/ravendb/documents/operations/etl/etl_operation_results.py new file mode 100644 index 00000000..42738fc8 --- /dev/null +++ b/ravendb/documents/operations/etl/etl_operation_results.py @@ -0,0 +1,32 @@ +from __future__ import annotations +from typing import Dict, Any + + +class AddEtlOperationResult: + """Result of adding an ETL task.""" + + def __init__(self, raft_command_index: int = 0, task_id: int = 0): + self.raft_command_index = raft_command_index + self.task_id = task_id + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "AddEtlOperationResult": + return cls( + raft_command_index=json_dict.get("RaftCommandIndex", 0), + task_id=json_dict.get("TaskId", 0), + ) + + +class UpdateEtlOperationResult: + """Result of updating an ETL task.""" + + def __init__(self, raft_command_index: int = 0, task_id: int = 0): + self.raft_command_index = raft_command_index + self.task_id = task_id + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "UpdateEtlOperationResult": + return cls( + raft_command_index=json_dict.get("RaftCommandIndex", 0), + task_id=json_dict.get("TaskId", 0), + ) diff --git a/ravendb/documents/operations/etl/etl_type.py b/ravendb/documents/operations/etl/etl_type.py new file mode 100644 index 00000000..a01ce560 --- /dev/null +++ b/ravendb/documents/operations/etl/etl_type.py @@ -0,0 +1,12 @@ +import enum + + +class EtlType(enum.Enum): + RAVEN = "Raven" + SQL = "Sql" + OLAP = "Olap" + ELASTIC_SEARCH = "ElasticSearch" + QUEUE = "Queue" + SNOWFLAKE = "Snowflake" + EMBEDDINGS_GENERATION = "EmbeddingsGeneration" + GEN_AI = "GenAi" diff --git a/ravendb/documents/operations/etl/transformation.py b/ravendb/documents/operations/etl/transformation.py new file mode 100644 index 00000000..190f6ad0 --- /dev/null +++ b/ravendb/documents/operations/etl/transformation.py @@ -0,0 +1,426 @@ +from __future__ import annotations +import re +from typing import TYPE_CHECKING, Dict, Any, List, Optional + +if TYPE_CHECKING: + from ravendb.documents.operations.etl.etl_type import EtlType + + +class CountersTransformation: + """ + Internal class for counters transformation validation. + """ + + LOAD = "loadCounter" + ADD = "addCounter" + MARKER = "$counter/" + + _ADD_METHOD_REGEX = re.compile(ADD) + _LOAD_BEHAVIOR_METHOD_REGEX = re.compile(r"function\s+loadCountersOf(\w+)Behavior\s*\(.+\)") + _LOAD_BEHAVIOR_METHOD_NAME_REGEX = re.compile(r"loadCountersOf(\w+)Behavior") + + def __init__(self, parent: "Transformation"): + self._parent = parent + self.is_adding_counters: bool = False + self.collection_to_load_behavior_function: Optional[Dict[str, str]] = None + + def validate(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self._parent.script or "" + self.is_adding_counters = len(self._ADD_METHOD_REGEX.findall(script)) > 0 + + if self.is_adding_counters and etl_type == EtlType.SQL: + errors.append("Adding counters isn't supported by SQL ETL") + + self._fill_collection_to_load_counter_behavior_function(errors, etl_type) + + def _fill_collection_to_load_counter_behavior_function(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self._parent.script or "" + counter_behaviors = list(self._LOAD_BEHAVIOR_METHOD_REGEX.finditer(script)) + if not counter_behaviors: + return + + if etl_type == EtlType.SQL: + errors.append("Load counter behavior functions aren't supported by SQL ETL") + return + + self.collection_to_load_behavior_function = {} + for match in counter_behaviors: + if len(match.groups()) != 1: + errors.append( + "Invalid load counters behavior function. It is expected to have the following signature: " + "loadCountersOfBehavior(docId, counterName) and return 'true' if counter should be loaded to a destination" + ) + continue + + function_signature = match.group(0) + collection = match.group(1) + + function_name_match = self._LOAD_BEHAVIOR_METHOD_NAME_REGEX.search(function_signature) + function_name = function_name_match.group(0) if function_name_match else "" + + if collection not in self._parent.collections: + script_collections = ", ".join(f"'{c}'" for c in self._parent.collections) + errors.append( + f"There is '{function_name}' function defined in '{self._parent.name}' script while the processed collections " + f"({script_collections}) doesn't include '{collection}'. " + "loadCountersOfBehavior() function is meant to be defined only for counters of docs from collections that " + "are loaded to the same collection on a destination side" + ) + else: + collections_from_script = self._parent.get_collections_from_script() + if collections_from_script and collection not in collections_from_script: + errors.append( + f"`{function_name}` function where Defined while there is not load to {collection}. " + "Load behavior function apply only if load to default collection" + ) + + if collection in self.collection_to_load_behavior_function: + errors.append(f"There are multiple '{function_name}' functions defined") + + self.collection_to_load_behavior_function[collection] = function_name + + +class TimeSeriesTransformation: + """ + Internal class for time series transformation validation. + """ + + MARKER = "$timeSeries/" + ADD_TIME_SERIES_NAME = "addTimeSeries" + ADD_TIME_SERIES_SIGNATURE = "addTimeSeries(timeSeriesReference)" + LOAD_TIME_SERIES_NAME = "loadTimeSeries" + LOAD_TIME_SERIES_SIGNATURE = "loadTimeSeries(timeSeriesName, from, to)" + HAS_TIME_SERIES_NAME = "hasTimeSeries" + GET_TIME_SERIES_NAME = "getTimeSeries" + LOAD_TIME_SERIES_BEHAVIOR_SIGNATURE = "loadTimeSeriesOfBehavior(docId, timeSeriesName)" + LOAD_TIME_SERIES_BEHAVIOR_PARAMS_COUNT = 2 + + _ADD_TIME_SERIES_REGEX = re.compile(ADD_TIME_SERIES_NAME) + _LOAD_TIME_SERIES_BEHAVIOR_REGEX = re.compile( + r"function\s+(?PloadTimeSeriesOf(?P[A-Za-z]\w*)Behavior)\s*\(\s*" + r"((?P[a-zA-Z]\w*)\s*(?:,\s*(?P[a-zA-Z]\w*)\s*)*)?\s*\)" + ) + + def __init__(self, parent: "Transformation"): + self._parent = parent + self.is_adding_time_series: bool = False + self.collection_to_load_behavior_function: Optional[Dict[str, str]] = None + + def validate(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self._parent.script or "" + self.is_adding_time_series = len(self._ADD_TIME_SERIES_REGEX.findall(script)) > 0 + + if self.is_adding_time_series and etl_type == EtlType.SQL: + errors.append("Adding time series isn't supported by SQL ETL") + + self._fill_collection_to_load_time_series_behavior_function(errors, etl_type) + + def _fill_collection_to_load_time_series_behavior_function(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self._parent.script or "" + time_series_behaviors = list(self._LOAD_TIME_SERIES_BEHAVIOR_REGEX.finditer(script)) + if not time_series_behaviors: + return + + if etl_type == EtlType.SQL: + errors.append("Load time series behavior functions aren't supported by SQL ETL") + return + + self.collection_to_load_behavior_function = {} + for match in time_series_behaviors: + function_name = match.group("func_name") + collection = match.group("collection") + + # Count params by checking captured groups + params = [g for g in [match.group("param"), match.group("param2")] if g] + if len(params) > self.LOAD_TIME_SERIES_BEHAVIOR_PARAMS_COUNT: + errors.append( + f"'{function_name} function defined with {len(params)}. " + f"The signature should be {self.LOAD_TIME_SERIES_BEHAVIOR_SIGNATURE}" + ) + + if collection not in self._parent.collections: + script_collections = ", ".join(f"'{c}'" for c in self._parent.collections) + errors.append( + f"There is '{function_name}' function defined in '{self._parent.name}' script while the processed collections " + f"({script_collections}) doesn't include '{collection}'. " + f"{self.LOAD_TIME_SERIES_BEHAVIOR_SIGNATURE} function is meant to be defined only for time series of docs from collections that " + "are loaded to the same collection on a destination side" + ) + else: + collections_from_script = self._parent.get_collections_from_script() + if collections_from_script and collection not in collections_from_script: + errors.append( + f"`{function_name}` function where Defined while there is not load to {collection}. " + "Load behavior function apply only if load to default collection" + ) + + if collection in self.collection_to_load_behavior_function: + errors.append(f"There are multiple '{function_name}' functions defined") + + self.collection_to_load_behavior_function[collection] = function_name + + +class Transformation: + """ + Represents an ETL transformation script configuration. + """ + + LOAD_TO = "loadTo" + LOAD_ATTACHMENT = "loadAttachment" + ADD_ATTACHMENT = "addAttachment" + ATTACHMENT_MARKER = "$attachment/" + GENERIC_DELETE_DOCUMENTS_BEHAVIOR_FUNCTION_KEY = "$deleteDocumentsBehavior<>" + GENERIC_DELETE_DOCUMENTS_BEHAVIOR_FUNCTION_NAME = "deleteDocumentsBehavior" + + # Regex patterns + _LOAD_TO_METHOD_REGEX = re.compile(r"loadTo(\w+)") + _LOAD_TO_METHOD_REGEX_ALT = re.compile(r"loadTo\('([\w.]*)'\)|loadTo\(\"([\w.]*)\"\)") + _LOAD_ATTACHMENT_METHOD_REGEX = re.compile(LOAD_ATTACHMENT) + _ADD_ATTACHMENT_METHOD_REGEX = re.compile(ADD_ATTACHMENT) + _LEGACY_REPLICATE_TO_METHOD_REGEX = re.compile(r"replicateTo(\w+)") + + # Parameters and function body regex for matching JS functions + _PARAMETERS_AND_FUNCTION_BODY_REGEX = ( + r"\s*\((?:[^)(]+|\((?:[^)(]+|\([^)(]*\))*\))*\)\s*\{(?:[^}{]+|\{(?:[^}{]+|\{[^}{]*\})*\})*\}" + ) + _DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX = re.compile( + r"function\s+deleteDocumentsOf(\w+)Behavior" + _PARAMETERS_AND_FUNCTION_BODY_REGEX, re.DOTALL + ) + _DELETE_DOCUMENTS_BEHAVIOR_METHOD_NAME_REGEX = re.compile(r"deleteDocumentsOf(\w+)Behavior") + _GENERIC_DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX = re.compile( + r"function\s+" + GENERIC_DELETE_DOCUMENTS_BEHAVIOR_FUNCTION_NAME + _PARAMETERS_AND_FUNCTION_BODY_REGEX, + re.DOTALL, + ) + + def __init__( + self, + name: str = None, + disabled: bool = False, + collections: List[str] = None, + apply_to_all_documents: bool = False, + script: str = None, + document_id_postfix: str = None, + ): + self.name = name + self.disabled = disabled + self.collections: List[str] = collections or [] + self.apply_to_all_documents = apply_to_all_documents + self.script = script + self.document_id_postfix = document_id_postfix + + # Internal state set during validation + self._cached_collections: Optional[List[str]] = None + self.is_empty_script: bool = False + self.collection_to_delete_documents_behavior_function: Optional[Dict[str, str]] = None + self.is_adding_attachments: bool = False + self.is_loading_attachments: bool = False + + # Nested transformation validators + self.counters = CountersTransformation(self) + self.time_series = TimeSeriesTransformation(self) + + def validate(self, errors: List[str], etl_type: "EtlType") -> bool: + """ + Validates the transformation configuration. + + Args: + errors: List to append validation errors to. + etl_type: The type of ETL this transformation is for. + + Returns: + True if validation passed (no errors), False otherwise. + """ + from ravendb.documents.operations.etl.etl_type import EtlType + + if errors is None: + raise ValueError("errors cannot be None") + + if not self.name or not self.name.strip(): + errors.append("Script name cannot be empty") + + if self.apply_to_all_documents: + if self.collections and len(self.collections) > 0: + errors.append( + f"Collections cannot be specified when ApplyToAllDocuments is set. Script name: '{self.name}'" + ) + else: + if not self.collections or len(self.collections) == 0: + errors.append( + f"Collections need be specified or ApplyToAllDocuments has to be set. Script name: '{self.name}'" + ) + + if self.script and self.script.strip(): + # Check for legacy replicateTo method + if self._LEGACY_REPLICATE_TO_METHOD_REGEX.search(self.script): + errors.append( + f"Found `replicateTo()` method in '{self.name}' script which is not supported. " + "If you are using the SQL replication script from RavenDB 3.x version then please use `loadTo()` instead." + ) + + self.is_adding_attachments = bool(self._ADD_ATTACHMENT_METHOD_REGEX.search(self.script)) + self.is_loading_attachments = bool(self._LOAD_ATTACHMENT_METHOD_REGEX.search(self.script)) + + # Validate counters and time series + self.counters.validate(errors, etl_type) + self.time_series.validate(errors, etl_type) + + # Validate delete behaviors + self._validate_delete_behaviors(errors, etl_type) + + # Check for loadTo calls + collections_from_script = self.get_collections_from_script() + if not collections_from_script: + self._check_empty_script(errors, etl_type) + else: + self.is_empty_script = True + + if self.is_empty_script: + if etl_type not in (EtlType.RAVEN, EtlType.EMBEDDINGS_GENERATION, EtlType.GEN_AI): + errors.append(f"Script '{self.name}' must not be empty") + + return len(errors) == 0 + + def _validate_delete_behaviors(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self.script or "" + delete_behaviors = list(self._DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX.finditer(script)) + + if delete_behaviors: + if etl_type == EtlType.SQL: + errors.append("Delete documents behavior functions aren't supported by SQL ETL") + else: + self.collection_to_delete_documents_behavior_function = {} + + for match in delete_behaviors: + if len(match.groups()) != 1: + errors.append( + "Invalid delete documents behavior function. It is expected to have the following signature: " + "deleteDocumentsOfBehavior(docId) and return 'true' if document deletion should be sent to a destination" + ) + continue + + function = match.group(0) + collection = match.group(1) + + function_name_match = self._DELETE_DOCUMENTS_BEHAVIOR_METHOD_NAME_REGEX.search(function) + function_name = function_name_match.group(0) if function_name_match else "" + + if collection not in self.collections: + script_collections = ", ".join(f"'{c}'" for c in self.collections) + errors.append( + f"There is '{function_name}' function defined in '{self.name}' script while the processed collections " + f"({script_collections}) doesn't include '{collection}'. " + "deleteDocumentsOfBehavior() function is meant to be defined only for documents from collections that " + "are loaded to the same collection on a destination side" + ) + + self.collection_to_delete_documents_behavior_function[collection] = function_name + + # Check for generic delete behavior + generic_delete_behaviors = list(self._GENERIC_DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX.finditer(script)) + + if generic_delete_behaviors: + if etl_type == EtlType.SQL: + errors.append("Delete documents behavior functions aren't supported by SQL ETL") + else: + if len(generic_delete_behaviors) > 1: + errors.append("Generic delete behavior function can be defined just once in the script") + else: + if self.collection_to_delete_documents_behavior_function is None: + self.collection_to_delete_documents_behavior_function = {} + self.collection_to_delete_documents_behavior_function[ + self.GENERIC_DELETE_DOCUMENTS_BEHAVIOR_FUNCTION_KEY + ] = self.GENERIC_DELETE_DOCUMENTS_BEHAVIOR_FUNCTION_NAME + + def _check_empty_script(self, errors: List[str], etl_type: "EtlType") -> None: + from ravendb.documents.operations.etl.etl_type import EtlType + + script = self.script or "" + actual_script = script + + # Skip all delete behavior functions to check if we have empty transformation + delete_behaviors = list(self._DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX.finditer(script)) + for match in delete_behaviors: + actual_script = actual_script.replace(match.group(0), "") + + generic_delete_behaviors = list(self._GENERIC_DELETE_DOCUMENTS_BEHAVIOR_METHOD_REGEX.finditer(script)) + if len(generic_delete_behaviors) == 1: + actual_script = actual_script.replace(generic_delete_behaviors[0].group(0), "") + + if actual_script.strip(): + target_name_map = { + EtlType.RAVEN: "Collection", + EtlType.SQL: "Table", + EtlType.OLAP: "Table", + EtlType.ELASTIC_SEARCH: "Index", + EtlType.QUEUE: "Queue", + } + target_name = target_name_map.get(etl_type) + if target_name is None: + raise ValueError(f"Unknown ETL type: {etl_type}") + + errors.append(f"No `loadTo<{target_name}Name>()` method call found in '{self.name}' script") + else: + self.is_empty_script = True + + def get_collections_from_script(self) -> Optional[List[str]]: + """ + Extracts collection names from loadTo calls in the script. + + Returns: + List of collection names found in the script, or None if no loadTo calls found. + """ + if self._cached_collections is not None: + return self._cached_collections + + script = self.script or "" + if not script: + return None + + matches = self._LOAD_TO_METHOD_REGEX.findall(script) + matches_alt = self._LOAD_TO_METHOD_REGEX_ALT.findall(script) + + if not matches and not matches_alt: + return None + + collections = list(matches) + + # Handle alternative syntax loadTo('CollectionName') or loadTo("CollectionName") + for match in matches_alt: + # match is a tuple of (group1, group2) - one will be empty + collection = match[0] if match[0] else match[1] + if collection: + collections.append(collection) + + self._cached_collections = collections + return self._cached_collections + + def to_json(self) -> Dict[str, Any]: + return { + "Name": self.name, + "Disabled": self.disabled, + "Collections": self.collections, + "ApplyToAllDocuments": self.apply_to_all_documents, + "Script": self.script, + "DocumentIdPostfix": self.document_id_postfix, + } + + @classmethod + def from_json(cls, json_dict: Dict[str, Any]) -> "Transformation": + return cls( + name=json_dict.get("Name"), + disabled=json_dict.get("Disabled", False), + collections=json_dict.get("Collections", []), + apply_to_all_documents=json_dict.get("ApplyToAllDocuments", False), + script=json_dict.get("Script"), + document_id_postfix=json_dict.get("DocumentIdPostfix"), + ) diff --git a/ravendb/documents/operations/ongoing_tasks.py b/ravendb/documents/operations/ongoing_tasks.py index 6318353f..98fbed52 100644 --- a/ravendb/documents/operations/ongoing_tasks.py +++ b/ravendb/documents/operations/ongoing_tasks.py @@ -1,3 +1,4 @@ +from __future__ import annotations import json from enum import Enum from typing import Optional, TYPE_CHECKING, Union @@ -14,6 +15,7 @@ if TYPE_CHECKING: from ravendb.documents.conventions import DocumentConventions + from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration class OngoingTaskType(Enum): @@ -21,10 +23,250 @@ class OngoingTaskType(Enum): RAVEN_ETL = "RavenEtl" SQL_ETL = "SqlEtl" OLAP_ETL = "OlapEtl" + ELASTIC_SEARCH_ETL = "ElasticSearchEtl" + QUEUE_ETL = "QueueEtl" + SNOWFLAKE_ETL = "SnowflakeEtl" BACKUP = "Backup" SUBSCRIPTION = "Subscription" PULL_REPLICATION_AS_HUB = "PullReplicationAsHub" PULL_REPLICATION_AS_SINK = "PullReplicationAsSink" + QUEUE_SINK = "QueueSink" + EMBEDDINGS_GENERATION = "EmbeddingsGeneration" + GEN_AI = "GenAi" + + +class OngoingTaskState(Enum): + NONE = "None" + ENABLED = "Enabled" + DISABLED = "Disabled" + PARTIALLY_ENABLED = "PartiallyEnabled" + + +class OngoingTaskConnectionStatus(Enum): + NONE = "None" + ACTIVE = "Active" + NOT_ACTIVE = "NotActive" + RECONNECT = "Reconnect" + NOT_ON_THIS_NODE = "NotOnThisNode" + + +class NodeId: + """Represents a node identifier in the cluster.""" + + def __init__( + self, + node_tag: Optional[str] = None, + node_url: Optional[str] = None, + responsible_node: Optional[str] = None, + ): + self.node_tag = node_tag + self.node_url = node_url + self.responsible_node = responsible_node + + def to_json(self) -> dict: + return { + "NodeTag": self.node_tag, + "NodeUrl": self.node_url, + "ResponsibleNode": self.responsible_node, + } + + @classmethod + def from_json(cls, json_dict: dict) -> "NodeId": + if json_dict is None: + return None + return cls( + node_tag=json_dict.get("NodeTag"), + node_url=json_dict.get("NodeUrl"), + responsible_node=json_dict.get("ResponsibleNode"), + ) + + +class OngoingTask: + """Base class for ongoing task information.""" + + def __init__( + self, + task_id: Optional[int] = None, + task_type: Optional[OngoingTaskType] = None, + responsible_node: Optional[NodeId] = None, + task_state: Optional[OngoingTaskState] = None, + task_connection_status: Optional[OngoingTaskConnectionStatus] = None, + task_name: Optional[str] = None, + error: Optional[str] = None, + mentor_node: Optional[str] = None, + pin_to_mentor_node: Optional[bool] = None, + ): + self.task_id = task_id + self.task_type = task_type + self.responsible_node = responsible_node + self.task_state = task_state + self.task_connection_status = task_connection_status + self.task_name = task_name + self.error = error + self.mentor_node = mentor_node + self.pin_to_mentor_node = pin_to_mentor_node + + def to_json(self) -> dict: + return { + "TaskId": self.task_id, + "TaskType": self.task_type.value if self.task_type else None, + "ResponsibleNode": self.responsible_node.to_json() if self.responsible_node else None, + "TaskState": self.task_state.value if self.task_state else None, + "TaskConnectionStatus": self.task_connection_status.value if self.task_connection_status else None, + "TaskName": self.task_name, + "Error": self.error, + "MentorNode": self.mentor_node, + "PinToMentorNode": self.pin_to_mentor_node, + } + + @classmethod + def from_json(cls, json_dict: dict) -> "OngoingTask": + if json_dict is None: + return None + task_type_str = json_dict.get("TaskType") + task_state_str = json_dict.get("TaskState") + task_connection_status_str = json_dict.get("TaskConnectionStatus") + + return cls( + task_id=json_dict.get("TaskId"), + task_type=OngoingTaskType(task_type_str) if task_type_str else None, + responsible_node=NodeId.from_json(json_dict.get("ResponsibleNode")), + task_state=OngoingTaskState(task_state_str) if task_state_str else None, + task_connection_status=( + OngoingTaskConnectionStatus(task_connection_status_str) if task_connection_status_str else None + ), + task_name=json_dict.get("TaskName"), + error=json_dict.get("Error"), + mentor_node=json_dict.get("MentorNode"), + pin_to_mentor_node=json_dict.get("PinToMentorNode"), + ) + + +class OngoingTaskGenAi(OngoingTask): + """Ongoing task information for GenAI tasks.""" + + def __init__( + self, + task_id: Optional[int] = None, + responsible_node: Optional[NodeId] = None, + task_state: Optional[OngoingTaskState] = None, + task_connection_status: Optional[OngoingTaskConnectionStatus] = None, + task_name: Optional[str] = None, + error: Optional[str] = None, + mentor_node: Optional[str] = None, + pin_to_mentor_node: Optional[bool] = None, + connection_string_name: Optional[str] = None, + configuration: Optional["GenAiConfiguration"] = None, + change_vector: Optional[str] = None, + ): + super().__init__( + task_id=task_id, + task_type=OngoingTaskType.GEN_AI, + responsible_node=responsible_node, + task_state=task_state, + task_connection_status=task_connection_status, + task_name=task_name, + error=error, + mentor_node=mentor_node, + pin_to_mentor_node=pin_to_mentor_node, + ) + self.connection_string_name = connection_string_name + self.configuration = configuration + self.change_vector = change_vector + + def to_json(self) -> dict: + result = super().to_json() + result["ConnectionStringName"] = self.connection_string_name + result["Configuration"] = self.configuration.to_json() if self.configuration else None + result["ChangeVector"] = self.change_vector + return result + + @classmethod + def from_json(cls, json_dict: dict) -> "OngoingTaskGenAi": + from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration + + if json_dict is None: + return None + + task_state_str = json_dict.get("TaskState") + task_connection_status_str = json_dict.get("TaskConnectionStatus") + config_dict = json_dict.get("Configuration") + + return cls( + task_id=json_dict.get("TaskId"), + responsible_node=NodeId.from_json(json_dict.get("ResponsibleNode")), + task_state=OngoingTaskState(task_state_str) if task_state_str else None, + task_connection_status=( + OngoingTaskConnectionStatus(task_connection_status_str) if task_connection_status_str else None + ), + task_name=json_dict.get("TaskName"), + error=json_dict.get("Error"), + mentor_node=json_dict.get("MentorNode"), + pin_to_mentor_node=json_dict.get("PinToMentorNode"), + connection_string_name=json_dict.get("ConnectionStringName"), + configuration=GenAiConfiguration.from_json(config_dict) if config_dict else None, + change_vector=json_dict.get("ChangeVector"), + ) + + +class OngoingTaskEmbeddingsGeneration(OngoingTask): + """Ongoing task information for Embeddings Generation tasks.""" + + def __init__( + self, + task_id: Optional[int] = None, + responsible_node: Optional[NodeId] = None, + task_state: Optional[OngoingTaskState] = None, + task_connection_status: Optional[OngoingTaskConnectionStatus] = None, + task_name: Optional[str] = None, + error: Optional[str] = None, + mentor_node: Optional[str] = None, + pin_to_mentor_node: Optional[bool] = None, + connection_string_name: Optional[str] = None, + configuration: Optional[dict] = None, # EmbeddingsGenerationConfiguration - not yet implemented + ): + super().__init__( + task_id=task_id, + task_type=OngoingTaskType.EMBEDDINGS_GENERATION, + responsible_node=responsible_node, + task_state=task_state, + task_connection_status=task_connection_status, + task_name=task_name, + error=error, + mentor_node=mentor_node, + pin_to_mentor_node=pin_to_mentor_node, + ) + self.connection_string_name = connection_string_name + self.configuration = configuration + + def to_json(self) -> dict: + result = super().to_json() + result["ConnectionStringName"] = self.connection_string_name + result["Configuration"] = self.configuration + return result + + @classmethod + def from_json(cls, json_dict: dict) -> "OngoingTaskEmbeddingsGeneration": + if json_dict is None: + return None + + task_state_str = json_dict.get("TaskState") + task_connection_status_str = json_dict.get("TaskConnectionStatus") + + return cls( + task_id=json_dict.get("TaskId"), + responsible_node=NodeId.from_json(json_dict.get("ResponsibleNode")), + task_state=OngoingTaskState(task_state_str) if task_state_str else None, + task_connection_status=( + OngoingTaskConnectionStatus(task_connection_status_str) if task_connection_status_str else None + ), + task_name=json_dict.get("TaskName"), + error=json_dict.get("Error"), + mentor_node=json_dict.get("MentorNode"), + pin_to_mentor_node=json_dict.get("PinToMentorNode"), + connection_string_name=json_dict.get("ConnectionStringName"), + configuration=json_dict.get("Configuration"), + ) class ToggleOngoingTaskStateOperation(MaintenanceOperation[ModifyOngoingTaskResult]): @@ -83,3 +325,122 @@ def is_read_request(self) -> bool: def get_raft_unique_request_id(self) -> str: return RaftIdGenerator.new_id() + + +class DeleteOngoingTaskOperation(MaintenanceOperation[ModifyOngoingTaskResult]): + """Operation to delete an ongoing task.""" + + def __init__(self, task_id: int, task_type: OngoingTaskType): + """ + Initialize the delete operation. + + Args: + task_id: The unique identifier of the ongoing task to be deleted. + task_type: The type of the ongoing task. + """ + self._task_id = task_id + self._task_type = task_type + + def get_command(self, conventions: "DocumentConventions") -> RavenCommand[ModifyOngoingTaskResult]: + return DeleteOngoingTaskOperation._DeleteOngoingTaskCommand(self._task_id, self._task_type) + + class _DeleteOngoingTaskCommand(RavenCommand[ModifyOngoingTaskResult], RaftCommand): + def __init__(self, task_id: int, task_type: OngoingTaskType): + super().__init__(ModifyOngoingTaskResult) + self._task_id = task_id + self._task_type = task_type + + def create_request(self, node: ServerNode) -> requests.Request: + url = f"{node.url}/databases/{node.database}/admin/tasks?id={self._task_id}&type={self._task_type.value}" + return requests.Request("DELETE", url) + + def set_response(self, response: Optional[str], from_cache: bool) -> None: + if response is not None: + self.result = ModifyOngoingTaskResult.from_json(json.loads(response)) + + def is_read_request(self) -> bool: + return False + + def get_raft_unique_request_id(self) -> str: + return RaftIdGenerator.new_id() + + +class GetOngoingTaskInfoOperation(MaintenanceOperation[OngoingTask]): + """ + Operation to retrieve detailed information about a specific ongoing task. + Ongoing tasks include various types of tasks such as replication, ETL, backup, and subscriptions. + """ + + def __init__(self, task_id_or_name: Union[int, str], task_type: OngoingTaskType): + """ + Initialize the get operation. + + Args: + task_id_or_name: The unique identifier or name of the ongoing task to retrieve information for. + task_type: The type of the ongoing task, such as replication, ETL, or backup. + + Raises: + ValueError: If the specified task type is PullReplicationAsHub, which is not supported. + """ + if task_type == OngoingTaskType.PULL_REPLICATION_AS_HUB: + raise ValueError( + "PullReplicationAsHub type is not supported. Please use GetPullReplicationTasksInfoOperation instead." + ) + + if isinstance(task_id_or_name, str): + if not task_id_or_name or task_id_or_name.isspace(): + raise ValueError("Task name cannot be null or whitespace.") + self._task_name = task_id_or_name + self._task_id = None + else: + self._task_id = task_id_or_name + self._task_name = None + + self._task_type = task_type + + def get_command(self, conventions: "DocumentConventions") -> RavenCommand[OngoingTask]: + if self._task_name is not None: + return GetOngoingTaskInfoOperation._GetOngoingTaskInfoCommand( + task_name=self._task_name, task_type=self._task_type + ) + return GetOngoingTaskInfoOperation._GetOngoingTaskInfoCommand(task_id=self._task_id, task_type=self._task_type) + + class _GetOngoingTaskInfoCommand(RavenCommand[OngoingTask]): + def __init__( + self, + task_type: OngoingTaskType, + task_id: Optional[int] = None, + task_name: Optional[str] = None, + ): + super().__init__(OngoingTask) + self._task_id = task_id + self._task_name = task_name + self._task_type = task_type + + def create_request(self, node: ServerNode) -> requests.Request: + if self._task_name is not None: + url = ( + f"{node.url}/databases/{node.database}/task" + f"?taskName={Utils.quote_key(self._task_name)}&type={self._task_type.value}" + ) + else: + url = f"{node.url}/databases/{node.database}/task?key={self._task_id}&type={self._task_type.value}" + return requests.Request("GET", url) + + def set_response(self, response: Optional[str], from_cache: bool) -> None: + if response is not None: + json_dict = json.loads(response) + self.result = self._deserialize_task(json_dict) + + def _deserialize_task(self, json_dict: dict) -> OngoingTask: + """Deserialize the task based on its type.""" + if self._task_type == OngoingTaskType.GEN_AI: + return OngoingTaskGenAi.from_json(json_dict) + elif self._task_type == OngoingTaskType.EMBEDDINGS_GENERATION: + return OngoingTaskEmbeddingsGeneration.from_json(json_dict) + else: + # todo: handle more types of tasks + return OngoingTask.from_json(json_dict) + + def is_read_request(self) -> bool: + return False diff --git a/ravendb/documents/starting_point_change_vector.py b/ravendb/documents/starting_point_change_vector.py new file mode 100644 index 00000000..234deb9c --- /dev/null +++ b/ravendb/documents/starting_point_change_vector.py @@ -0,0 +1,22 @@ +from __future__ import annotations +from enum import Enum + + +class StartingPointChangeVector(str, Enum): + """ + Represents a starting point for ETL operations based on change vectors. + """ + + DO_NOT_CHANGE = "DoNotChange" + LAST_DOCUMENT = "LastDocument" + BEGINNING_OF_TIME = "BeginningOfTime" + + @classmethod + def from_value(cls, change_vector: str) -> StartingPointChangeVector: + """Creates a StartingPointChangeVector from a custom change vector value or returns matching enum.""" + # Try to match existing enum values + for member in cls: + if member.value == change_vector: + return member + # For custom change vectors, return as string (will work since we inherit from str) + raise ValueError(f"Unknown change vector: {change_vector}. Use one of: {[m.value for m in cls]}") diff --git a/ravendb/tests/gen_ai_tests/__init__.py b/ravendb/tests/gen_ai_tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py b/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py new file mode 100644 index 00000000..e672ad6c --- /dev/null +++ b/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py @@ -0,0 +1,617 @@ +"""Tests for GenAI task operations.""" + +import unittest + +from ravendb.documents.operations.ai.ai_connection_string import AiConnectionString, AiModelType +from ravendb.documents.operations.ai.add_gen_ai_operation import AddGenAiOperation +from ravendb.documents.operations.ai.gen_ai_configuration import GenAiConfiguration +from ravendb.documents.operations.ai.gen_ai_transformation import GenAiTransformation +from ravendb.documents.operations.ai.update_gen_ai_operation import UpdateGenAiOperation +from ravendb.documents.operations.ai.open_ai_settings import OpenAiSettings +from ravendb.documents.operations.connection_string.put_connection_string_operation import PutConnectionStringOperation +from ravendb.documents.operations.connection_string.remove_connection_string_operation import ( + RemoveConnectionStringOperation, +) +from ravendb.documents.operations.ongoing_tasks import ( + DeleteOngoingTaskOperation, + GetOngoingTaskInfoOperation, + OngoingTaskType, + OngoingTaskGenAi, +) +from ravendb.documents.starting_point_change_vector import StartingPointChangeVector +from ravendb.tests.test_base import TestBase + + +class TestGenAiConfigurationValidation(unittest.TestCase): + """Tests for GenAiConfiguration.validate() method. No server required.""" + + def test_validate_valid_configuration_with_sample_object(self): + """Validates that a complete configuration with sample_object passes validation.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertEqual(0, len(errors), f"Expected no errors, got: {errors}") + + def test_validate_valid_configuration_with_json_schema(self): + """Validates that a complete configuration with json_schema passes validation.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + json_schema='{"type": "object", "properties": {"summary": {"type": "string"}}}', + ) + errors = config.validate() + self.assertEqual(0, len(errors), f"Expected no errors, got: {errors}") + + def test_validate_missing_sample_object_and_json_schema(self): + """Validates that missing both sample_object and json_schema returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + ) + errors = config.validate() + self.assertIn("You must provide either a JSON schema or a sample object", errors) + + def test_validate_missing_update_script(self): + """Validates that missing update_script returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertIn("You must provide an update function", errors) + + def test_validate_missing_prompt(self): + """Validates that missing prompt returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertIn("Prompt must be provided", errors) + + def test_validate_missing_collection(self): + """Validates that missing collection returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertIn("Collection must be provided", errors) + + def test_validate_missing_name(self): + """Validates that missing name returns error.""" + config = GenAiConfiguration( + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertIn("Name of GenAi configuration cannot be empty", errors) + + def test_validate_missing_gen_ai_transformation(self): + """Validates that missing gen_ai_transformation returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertIn("GenAiTransformation must be specified", errors) + + def test_validate_invalid_identifier_uppercase(self): + """Validates that uppercase letters in identifier returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="TestGenAi", # Invalid: uppercase + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertTrue(any("lowercase" in e for e in errors), f"Expected lowercase error, got: {errors}") + + def test_validate_invalid_identifier_special_chars(self): + """Validates that special characters (except hyphen) in identifier returns error.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test/gen_ai", # Invalid: slash and underscore + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + ) + errors = config.validate() + self.assertTrue(any("invalid characters" in e for e in errors), f"Expected identifier error, got: {errors}") + + +class TestGenAiConfigurationSerialization(unittest.TestCase): + """Tests for GenAiConfiguration serialization. No server required.""" + + def test_gen_ai_configuration_to_json(self): + """Tests that to_json() produces correct JSON structure.""" + config = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "test"}', + max_concurrency=4, + enable_tracing=True, + ) + json_data = config.to_json() + + self.assertEqual("TestGenAi", json_data["Name"]) + self.assertEqual("test-gen-ai-1", json_data["Identifier"]) + self.assertEqual("Documents", json_data["Collection"]) + self.assertEqual("my-connection", json_data["ConnectionStringName"]) + self.assertEqual("Test prompt", json_data["Prompt"]) + self.assertEqual("this.Summary = $result;", json_data["UpdateScript"]) + self.assertEqual('{"summary": "test"}', json_data["SampleObject"]) + self.assertEqual(4, json_data["MaxConcurrency"]) + self.assertTrue(json_data["EnableTracing"]) + self.assertEqual("GenAi", json_data["EtlType"]) + + def test_gen_ai_configuration_from_json(self): + """Tests that from_json() correctly deserializes configuration.""" + json_data = { + "Name": "TestGenAi", + "Identifier": "test-gen-ai-1", + "Collection": "Documents", + "ConnectionStringName": "my-connection", + "Prompt": "Test prompt", + "UpdateScript": "this.Summary = $result;", + "SampleObject": '{"summary": "test"}', + "MaxConcurrency": 4, + "EnableTracing": True, + "GenAiTransformation": {"Script": "ai.genContext(ctx);"}, + } + config = GenAiConfiguration.from_json(json_data) + + self.assertEqual("TestGenAi", config.name) + self.assertEqual("test-gen-ai-1", config.identifier) + self.assertEqual("Documents", config.collection) + self.assertEqual("my-connection", config.connection_string_name) + self.assertEqual("Test prompt", config.prompt) + self.assertEqual("this.Summary = $result;", config.update_script) + self.assertEqual('{"summary": "test"}', config.sample_object) + self.assertEqual(4, config.max_concurrency) + self.assertTrue(config.enable_tracing) + + def test_gen_ai_configuration_round_trip(self): + """Tests serialization -> deserialization produces equivalent object.""" + original = GenAiConfiguration( + name="TestGenAi", + identifier="test-gen-ai-1", + collection="Documents", + connection_string_name="my-connection", + prompt="Test prompt", + gen_ai_transformation=GenAiTransformation(script="ai.genContext(ctx);"), + update_script="this.Summary = $result;", + json_schema='{"type": "object"}', + max_concurrency=8, + enable_tracing=False, + ) + json_data = original.to_json() + restored = GenAiConfiguration.from_json(json_data) + + self.assertEqual(original.name, restored.name) + self.assertEqual(original.identifier, restored.identifier) + self.assertEqual(original.collection, restored.collection) + self.assertEqual(original.prompt, restored.prompt) + self.assertEqual(original.update_script, restored.update_script) + self.assertEqual(original.json_schema, restored.json_schema) + self.assertEqual(original.max_concurrency, restored.max_concurrency) + self.assertEqual(original.enable_tracing, restored.enable_tracing) + + +class TestGenAiTransformation(unittest.TestCase): + """Tests for GenAiTransformation. No server required.""" + + def test_gen_ai_transformation_validate_script_valid(self): + """Tests that valid script passes validation.""" + transformation = GenAiTransformation(script="var ctx = {}; ai.genContext(ctx);") + is_valid, error = transformation.validate_script() + self.assertTrue(is_valid) + self.assertEqual("", error) + + def test_gen_ai_transformation_validate_script_missing_ai_gen_context(self): + """Tests that script without ai.genContext fails validation.""" + transformation = GenAiTransformation(script="var ctx = {};") + is_valid, error = transformation.validate_script() + self.assertFalse(is_valid) + self.assertIn("ai.genContext", error) + + def test_gen_ai_transformation_to_json(self): + """Tests transformation serialization.""" + transformation = GenAiTransformation(script="ai.genContext(ctx);") + json_data = transformation.to_json() + self.assertEqual("ai.genContext(ctx);", json_data["Script"]) + + +class TestStartingPointChangeVector(unittest.TestCase): + """Tests for StartingPointChangeVector enum. No server required.""" + + def test_starting_point_change_vector_values(self): + """Tests enum values are correct strings.""" + self.assertEqual("DoNotChange", StartingPointChangeVector.DO_NOT_CHANGE.value) + self.assertEqual("LastDocument", StartingPointChangeVector.LAST_DOCUMENT.value) + self.assertEqual("BeginningOfTime", StartingPointChangeVector.BEGINNING_OF_TIME.value) + + def test_starting_point_change_vector_from_value(self): + """Tests from_value() method.""" + result = StartingPointChangeVector.from_value("LastDocument") + self.assertEqual(StartingPointChangeVector.LAST_DOCUMENT, result) + + result = StartingPointChangeVector.from_value("BeginningOfTime") + self.assertEqual(StartingPointChangeVector.BEGINNING_OF_TIME, result) + + result = StartingPointChangeVector.from_value("DoNotChange") + self.assertEqual(StartingPointChangeVector.DO_NOT_CHANGE, result) + + +class TestAiTaskIdentifierHelper(unittest.TestCase): + """Tests for AiTaskIdentifierHelper. No server required.""" + + def test_validate_identifier_valid(self): + """Tests that valid identifiers pass validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("test-gen-ai-1") + self.assertTrue(is_valid) + self.assertEqual(0, len(errors)) + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("my-task-123") + self.assertTrue(is_valid) + self.assertEqual(0, len(errors)) + + def test_validate_identifier_empty(self): + """Tests that empty identifier fails validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("") + self.assertFalse(is_valid) + self.assertTrue(any("empty" in e for e in errors)) + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier(" ") + self.assertFalse(is_valid) + self.assertTrue(any("empty" in e or "whitespace" in e for e in errors)) + + def test_validate_identifier_uppercase(self): + """Tests that uppercase letters fail validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("TestGenAi") + self.assertFalse(is_valid) + self.assertTrue(any("uppercase" in e for e in errors)) + + def test_validate_identifier_invalid_chars(self): + """Tests that invalid characters fail validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("test/gen_ai") + self.assertFalse(is_valid) + self.assertTrue(any("invalid characters" in e for e in errors)) + + def test_validate_identifier_consecutive_hyphens(self): + """Tests that consecutive hyphens fail validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("test--gen-ai") + self.assertFalse(is_valid) + self.assertTrue(any("consecutive hyphens" in e for e in errors)) + + def test_validate_identifier_ends_with_hyphen(self): + """Tests that identifier ending with hyphen fails validation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + is_valid, errors = AiTaskIdentifierHelper.validate_identifier("test-gen-ai-") + self.assertFalse(is_valid) + self.assertTrue(any("ends with a hyphen" in e for e in errors)) + + def test_generate_identifier_basic(self): + """Tests basic identifier generation.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + result = AiTaskIdentifierHelper.generate_identifier("TestGenAi") + self.assertEqual("testgenai", result) + + def test_generate_identifier_with_spaces(self): + """Tests identifier generation with spaces.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + result = AiTaskIdentifierHelper.generate_identifier("Test Gen Ai") + self.assertEqual("test-gen-ai", result) + + def test_generate_identifier_with_special_chars(self): + """Tests identifier generation with special characters.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + result = AiTaskIdentifierHelper.generate_identifier("Test/Gen_Ai") + self.assertEqual("test-gen-ai", result) + + def test_generate_identifier_empty(self): + """Tests identifier generation with empty input.""" + from ravendb.documents.operations.ai.ai_task_identifier_helper import AiTaskIdentifierHelper + + result = AiTaskIdentifierHelper.generate_identifier("") + self.assertIsNone(result) + + result = AiTaskIdentifierHelper.generate_identifier(" ") + self.assertIsNone(result) + + +class TestGenAiCrudOperations(TestBase): + """Tests for GenAI CRUD operations (require server connection).""" + + CONNECTION_STRING_NAME = "test-ai-connection" + + def setUp(self): + super().setUp() + # Create AI connection string for tests + ai_connection_string = AiConnectionString( + name=self.CONNECTION_STRING_NAME, + identifier="test-ai-identifier", + openai_settings=OpenAiSettings( + api_key="test-api-key", + endpoint="https://api.openai.com", + model="gpt-4", + ), + model_type=AiModelType.CHAT, + ) + self.store.maintenance.send(PutConnectionStringOperation(ai_connection_string)) + self._ai_connection_string = ai_connection_string + self._created_task_ids = [] + + def tearDown(self): + # Clean up created tasks + for task_id in self._created_task_ids: + try: + self.store.maintenance.send(DeleteOngoingTaskOperation(task_id, OngoingTaskType.GEN_AI)) + except Exception: + pass + + # Clean up connection string + try: + self.store.maintenance.send(RemoveConnectionStringOperation(self._ai_connection_string)) + except Exception: + pass + + super().tearDown() + + def _create_valid_config(self, name: str, identifier: str) -> GenAiConfiguration: + """Helper to create a valid GenAiConfiguration.""" + return GenAiConfiguration( + name=name, + identifier=identifier, + collection="Documents", + connection_string_name=self.CONNECTION_STRING_NAME, + prompt="Summarize the document: {{content}}", + gen_ai_transformation=GenAiTransformation(script="var ctx = {}; ai.genContext(ctx);"), + update_script="this.Summary = $result;", + sample_object='{"summary": "A brief summary"}', + max_concurrency=2, + enable_tracing=True, + ) + + def test_add_gen_ai_operation_with_default_starting_point(self): + """Tests adding GenAI task with default LAST_DOCUMENT starting point.""" + config = self._create_valid_config("TestAddDefault", "test-add-default") + + result = self.store.maintenance.send(AddGenAiOperation(config)) + + self.assertIsNotNone(result) + self.assertGreater(result.task_id, 0) + self._created_task_ids.append(result.task_id) + + def test_add_gen_ai_operation_with_beginning_of_time(self): + """Tests adding GenAI task with BEGINNING_OF_TIME starting point.""" + config = self._create_valid_config("TestAddBeginning", "test-add-beginning") + + result = self.store.maintenance.send(AddGenAiOperation(config, StartingPointChangeVector.BEGINNING_OF_TIME)) + + self.assertIsNotNone(result) + self.assertGreater(result.task_id, 0) + self._created_task_ids.append(result.task_id) + + def test_add_gen_ai_operation_returns_task_id(self): + """Tests that add operation returns valid task ID.""" + config = self._create_valid_config("TestReturnsId", "test-returns-id") + + result = self.store.maintenance.send(AddGenAiOperation(config)) + + self.assertIsNotNone(result.task_id) + self.assertIsInstance(result.task_id, int) + self.assertGreater(result.task_id, 0) + self._created_task_ids.append(result.task_id) + + def test_add_gen_ai_operation_returns_identifier(self): + """Tests that add operation returns the identifier.""" + config = self._create_valid_config("TestReturnsIdentifier", "test-returns-identifier") + + result = self.store.maintenance.send(AddGenAiOperation(config)) + + self.assertEqual("test-returns-identifier", result.identifier) + self._created_task_ids.append(result.task_id) + + def test_get_ongoing_task_info_gen_ai(self): + """Tests retrieving GenAI task info by task ID.""" + config = self._create_valid_config("TestGetInfo", "test-get-info") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + task_info = self.store.maintenance.send(GetOngoingTaskInfoOperation(add_result.task_id, OngoingTaskType.GEN_AI)) + + self.assertIsNotNone(task_info) + self.assertIsInstance(task_info, OngoingTaskGenAi) + self.assertEqual(add_result.task_id, task_info.task_id) + + def test_get_ongoing_task_info_gen_ai_by_name(self): + """Tests retrieving GenAI task info by task name.""" + config = self._create_valid_config("TestGetByName", "test-get-by-name") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + task_info = self.store.maintenance.send(GetOngoingTaskInfoOperation("TestGetByName", OngoingTaskType.GEN_AI)) + + self.assertIsNotNone(task_info) + self.assertEqual("TestGetByName", task_info.task_name) + + def test_get_ongoing_task_info_returns_correct_type(self): + """Tests that returned task has correct OngoingTaskType.GEN_AI.""" + config = self._create_valid_config("TestCorrectType", "test-correct-type") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + task_info = self.store.maintenance.send(GetOngoingTaskInfoOperation(add_result.task_id, OngoingTaskType.GEN_AI)) + + self.assertEqual(OngoingTaskType.GEN_AI, task_info.task_type) + + def test_update_gen_ai_operation_basic(self): + """Tests basic update of GenAI task.""" + config = self._create_valid_config("TestUpdateBasic", "test-update-basic") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + # Update configuration + config.prompt = "Updated prompt: Analyze the document" + config.task_id = add_result.task_id + + update_result = self.store.maintenance.send(UpdateGenAiOperation(add_result.task_id, config)) + self._created_task_ids.append(update_result.task_id) + self.assertIsNotNone(update_result) + self.assertNotEqual(add_result.task_id, update_result.task_id) + + def test_update_gen_ai_operation_with_reset(self): + """Tests update with reset=True to reprocess documents.""" + config = self._create_valid_config("TestUpdateReset", "test-update-reset") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + # Update with reset + config.prompt = "Reset prompt: Re-analyze all documents" + config.task_id = add_result.task_id + + update_result = self.store.maintenance.send(UpdateGenAiOperation(add_result.task_id, config, reset=True)) + + self._created_task_ids.append(update_result.task_id) + self.assertIsNotNone(update_result) + self.assertNotEqual(add_result.task_id, update_result.task_id) + + def test_update_gen_ai_operation_with_starting_point(self): + """Tests update with custom starting point.""" + config = self._create_valid_config("TestUpdateStartPoint", "test-update-start-point") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self._created_task_ids.append(add_result.task_id) + + # Update with BEGINNING_OF_TIME starting point + config.prompt = "Updated prompt with new starting point" + config.task_id = add_result.task_id + + update_result = self.store.maintenance.send( + UpdateGenAiOperation( + add_result.task_id, + config, + starting_point=StartingPointChangeVector.BEGINNING_OF_TIME, + ) + ) + + self._created_task_ids.append(update_result.task_id) + self.assertIsNotNone(update_result) + + def test_delete_ongoing_task_gen_ai(self): + """Tests deleting GenAI task.""" + config = self._create_valid_config("TestDelete", "test-delete") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + + # Delete the task + self.store.maintenance.send(DeleteOngoingTaskOperation(add_result.task_id, OngoingTaskType.GEN_AI)) + + # Verify task is deleted by trying to get it + task_info = self.store.maintenance.send(GetOngoingTaskInfoOperation(add_result.task_id, OngoingTaskType.GEN_AI)) + self.assertIsNone(task_info) + + def test_gen_ai_full_lifecycle(self): + """Tests complete lifecycle: add -> get -> update -> get -> delete.""" + # 1. Add task + config = self._create_valid_config("TestLifecycle", "test-lifecycle") + add_result = self.store.maintenance.send(AddGenAiOperation(config)) + self.assertGreater(add_result.task_id, 0) + self.assertEqual("test-lifecycle", add_result.identifier) + + # 2. Get task + task_info = self.store.maintenance.send(GetOngoingTaskInfoOperation(add_result.task_id, OngoingTaskType.GEN_AI)) + self.assertIsNotNone(task_info) + self.assertEqual("TestLifecycle", task_info.task_name) + + # 3. Update task + config.prompt = "Updated lifecycle prompt" + config.max_concurrency = 8 + config.task_id = add_result.task_id + + update_result = self.store.maintenance.send(UpdateGenAiOperation(add_result.task_id, config)) + self.assertNotEqual(add_result.task_id, update_result.task_id) + + self._created_task_ids.append(update_result.task_id) + # 4. Get updated task + updated_task_info = self.store.maintenance.send( + GetOngoingTaskInfoOperation(update_result.task_id, OngoingTaskType.GEN_AI) + ) + self.assertIsNotNone(updated_task_info) + + # 5. Delete task + self.store.maintenance.send(DeleteOngoingTaskOperation(updated_task_info.task_id, OngoingTaskType.GEN_AI)) + + # 6. Verify deleted + deleted_task_info = self.store.maintenance.send( + GetOngoingTaskInfoOperation(updated_task_info.task_id, OngoingTaskType.GEN_AI) + ) + self.assertIsNone(deleted_task_info) From 8e58c9c0f2ff5835db8d18e90f33046b8cb03efa Mon Sep 17 00:00:00 2001 From: Gracjan Sadowicz Date: Wed, 10 Dec 2025 15:00:31 +0100 Subject: [PATCH 3/3] RDBC-967 Skip GenAI tests on CI/CD --- ravendb/tests/gen_ai_tests/test_gen_ai_operations.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py b/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py index e672ad6c..fcea7c06 100644 --- a/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py +++ b/ravendb/tests/gen_ai_tests/test_gen_ai_operations.py @@ -1,5 +1,6 @@ """Tests for GenAI task operations.""" +import os import unittest from ravendb.documents.operations.ai.ai_connection_string import AiConnectionString, AiModelType @@ -389,6 +390,7 @@ def test_generate_identifier_empty(self): self.assertIsNone(result) +@unittest.skipIf(os.environ.get("RAVENDB_LICENSE") is None, "Insufficient license permissions. Skipping on CI/CD.") class TestGenAiCrudOperations(TestBase): """Tests for GenAI CRUD operations (require server connection)."""