RDBC-1005 Add Embeddings Generation Task operations

Author: poissoncorp

ravendb/__init__.py

@@ -110,6 +110,15 @@
     AddOrUpdateAiAgentOperation,
     DeleteAiAgentOperation,
 )
+from ravendb.documents.operations.ai import (
+    ChunkingOptions,
+    ChunkingMethod,
+    EmbeddingPathConfiguration,
+    EmbeddingsTransformation,
+    EmbeddingsGenerationConfiguration,
+    AddEmbeddingsGenerationOperation,
+    UpdateEmbeddingsGenerationOperation,
+)
 
 from ravendb.documents.operations.etl.configuration import EtlConfiguration, RavenEtlConfiguration
 from ravendb.documents.operations.etl.olap.connection import OlapEtlConfiguration

ravendb/documents/ai/ai_conversation.py

@@ -145,9 +145,7 @@ def add_artificial_action_with_response(self, tool_id: str, action_response) ->
         else:
             content = json.dumps(action_response)
 
-        self._artificial_actions.append(
-            AiAgentArtificialActionResponse(tool_id=tool_id, content=content)
-        )
+        self._artificial_actions.append(AiAgentArtificialActionResponse(tool_id=tool_id, content=content))
 
     def run(self) -> AiAnswer:
         """

ravendb/documents/operations/ai/__init__.py

@@ -14,6 +14,12 @@
 )
 from ravendb.documents.operations.ai.add_gen_ai_operation import AddGenAiOperation
 from ravendb.documents.operations.ai.update_gen_ai_operation import UpdateGenAiOperation
+from ravendb.documents.operations.ai.chunking_options import ChunkingOptions, ChunkingMethod
+from ravendb.documents.operations.ai.embedding_path_configuration import EmbeddingPathConfiguration
+from ravendb.documents.operations.ai.embeddings_transformation import EmbeddingsTransformation
+from ravendb.documents.operations.ai.embeddings_generation_configuration import EmbeddingsGenerationConfiguration
+from ravendb.documents.operations.ai.add_embeddings_generation_operation import AddEmbeddingsGenerationOperation
+from ravendb.documents.operations.ai.update_embeddings_generation_operation import UpdateEmbeddingsGenerationOperation
 
 __all__ = [
     "AiConnectionString",
@@ -28,4 +34,11 @@
     "AddEmbeddingsGenerationOperationResult",
     "AddGenAiOperation",
     "UpdateGenAiOperation",
+    "ChunkingOptions",
+    "ChunkingMethod",
+    "EmbeddingPathConfiguration",
+    "EmbeddingsTransformation",
+    "EmbeddingsGenerationConfiguration",
+    "AddEmbeddingsGenerationOperation",
+    "UpdateEmbeddingsGenerationOperation",
 ]

ravendb/documents/operations/ai/add_embeddings_generation_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 AddEmbeddingsGenerationOperationResult
+import requests
+
+from ravendb.util.util import RaftIdGenerator
+
+if TYPE_CHECKING:
+    from ravendb.documents.operations.ai.embeddings_generation_configuration import EmbeddingsGenerationConfiguration
+
+
+class AddEmbeddingsGenerationOperation(MaintenanceOperation[AddEmbeddingsGenerationOperationResult]):
+    """
+    Operation to add a new Embeddings Generation task to the database.
+    """
+
+    def __init__(
+        self,
+        configuration: EmbeddingsGenerationConfiguration,
+        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[AddEmbeddingsGenerationOperationResult]:
+        return AddEmbeddingsGenerationCommand(self._configuration, self._starting_point, conventions)
+
+
+class AddEmbeddingsGenerationCommand(RavenCommand[AddEmbeddingsGenerationOperationResult]):
+    def __init__(
+        self,
+        configuration: EmbeddingsGenerationConfiguration,
+        starting_point: StartingPointChangeVector,
+        conventions: DocumentConventions,
+    ):
+        super().__init__(AddEmbeddingsGenerationOperationResult)
+        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 = AddEmbeddingsGenerationOperationResult()
+            return
+
+        response_json = json.loads(response)
+        self.result = AddEmbeddingsGenerationOperationResult.from_json(response_json)
+
+    def get_raft_unique_request_id(self) -> str:
+        return RaftIdGenerator.new_id()

ravendb/documents/operations/ai/chunking_options.py

@@ -1,8 +1,7 @@
 from enum import Enum
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, List
 
 
-# todo: EmbeddingsGenerationConfiguration
 class ChunkingMethod(Enum):
     PLAIN_TEXT_SPLIT = "PlainTextSplit"
     PLAIN_TEXT_SPLIT_LINES = "PlainTextSplitLines"
@@ -12,6 +11,14 @@ class ChunkingMethod(Enum):
     HTML_STRIP = "HtmlStrip"
 
 
+# Methods that support overlap tokens
+METHODS_SUPPORTING_OVERLAP_TOKENS = {
+    ChunkingMethod.PLAIN_TEXT_SPLIT,
+    ChunkingMethod.PLAIN_TEXT_SPLIT_LINES,
+    ChunkingMethod.PLAIN_TEXT_SPLIT_PARAGRAPHS,
+}
+
+
 class ChunkingOptions:
     def __init__(
         self,
@@ -26,9 +33,9 @@ def __init__(
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "ChunkingOptions":
         return cls(
-            chunking_method=ChunkingMethod(json_dict["ChunkingMethod"]) if json_dict.get("ChunkingMethod") else None,
-            max_tokens_per_chunk=json_dict.get("MaxTokensPerChunk", 512),
-            overlap_tokens=json_dict.get("OverlapTokens", 0),
+            chunking_method=ChunkingMethod(json_dict["ChunkingMethod"]),
+            max_tokens_per_chunk=json_dict.get("MaxTokensPerChunk", None),
+            overlap_tokens=json_dict.get("OverlapTokens", None),
         )
 
     def to_json(self) -> Dict[str, Any]:
@@ -38,6 +45,37 @@ def to_json(self) -> Dict[str, Any]:
             "OverlapTokens": self.overlap_tokens,
         }
 
+    def validate(self, source: str, errors: List[str]) -> None:
+        """Validates the chunking options.
+
+        Args:
+            source: The source context for error messages (e.g., 'embeddings.generate').
+            errors: List to append validation errors to.
+        """
+        if self.max_tokens_per_chunk <= 0:
+            errors.append(f"{source}: MaxTokensPerChunk must be greater than 0.")
+
+        if self.overlap_tokens < 0:
+            errors.append(f"{source}: OverlapTokens cannot be negative.")
+
+        if self.overlap_tokens > self.max_tokens_per_chunk:
+            errors.append(f"{source}: OverlapTokens cannot be greater than MaxTokensPerChunk.")
+
+        if self.overlap_tokens > 0 and self.chunking_method not in METHODS_SUPPORTING_OVERLAP_TOKENS:
+            errors.append(
+                f"{source}: OverlapTokens is only supported for PlainTextSplit, "
+                f"PlainTextSplitLines, and PlainTextSplitParagraphs chunking methods."
+            )
+
+    @staticmethod
+    def are_equal(left: Optional["ChunkingOptions"], right: Optional["ChunkingOptions"]) -> bool:
+        """Check if two ChunkingOptions instances are equal, handling None values."""
+        if left is None and right is None:
+            return True
+        if left is None or right is None:
+            return False
+        return left == right
+
     def __eq__(self, other: object) -> bool:
         if other is None:
             return False

ravendb/documents/operations/ai/embedding_path_configuration.py

@@ -3,23 +3,42 @@
 from ravendb.documents.operations.ai.chunking_options import ChunkingOptions
 
 
-# todo: EmbeddingsGenerationConfiguration
 class EmbeddingPathConfiguration:
     def __init__(self, path: Optional[str] = None, chunking_options: Optional[ChunkingOptions] = None):
         self.path = path
         self.chunking_options = chunking_options
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "EmbeddingPathConfiguration":
+        chunking_data = json_dict.get("ChunkingOptions")
         return cls(
-            path=json_dict.get("Path"),
-            chunking_options=(
-                ChunkingOptions.from_json(json_dict["ChunkingOptions"]) if json_dict.get("ChunkingOptions") else None
-            ),
+            path=json_dict["Path"],
+            chunking_options=ChunkingOptions.from_json(chunking_data) if chunking_data else None,
         )
 
     def to_json(self) -> Dict[str, Any]:
         return {
             "Path": self.path,
             "ChunkingOptions": self.chunking_options.to_json() if self.chunking_options else None,
         }
+
+    @staticmethod
+    def are_equal(left: Optional["EmbeddingPathConfiguration"], right: Optional["EmbeddingPathConfiguration"]) -> bool:
+        """Check if two EmbeddingPathConfiguration instances are equal, handling None values."""
+        if left is None and right is None:
+            return True
+        if left is None or right is None:
+            return False
+        return left == right
+
+    def __eq__(self, other: object) -> bool:
+        if other is None:
+            return False
+        if self is other:
+            return True
+        if not isinstance(other, EmbeddingPathConfiguration):
+            return False
+        return self.path == other.path and ChunkingOptions.are_equal(self.chunking_options, other.chunking_options)
+
+    def __hash__(self) -> int:
+        return hash((self.path, self.chunking_options))