RDBC-934 Vertex & Overlap tokens

Author: poissoncorp

ravendb/documents/operations/ai/ai_connection_string.py

@@ -9,6 +9,8 @@
 from ravendb.documents.operations.ai.mistral_ai_settings import MistralAiSettings
 from ravendb.documents.operations.ai.ollama_settings import OllamaSettings
 from ravendb.documents.operations.ai.open_ai_settings import OpenAiSettings
+from ravendb.documents.operations.ai.vertex_settings import VertexSettings
+
 from ravendb.documents.operations.connection_strings import ConnectionString
 
 
@@ -29,6 +31,7 @@ def __init__(
         google_settings: Optional[GoogleSettings] = None,
         huggingface_settings: Optional[HuggingFaceSettings] = None,
         mistral_ai_settings: Optional[MistralAiSettings] = None,
+        vertex_settings: Optional[VertexSettings] = None,
         model_type: AiModelType = None,
     ):
         super().__init__(name)
@@ -40,6 +43,7 @@ def __init__(
         self.google_settings = google_settings
         self.huggingface_settings = huggingface_settings
         self.mistral_ai_settings = mistral_ai_settings
+        self.vertex_settings = vertex_settings
         self.model_type = model_type
 
         if not any(
@@ -51,10 +55,11 @@ def __init__(
                 google_settings,
                 huggingface_settings,
                 mistral_ai_settings,
+                vertex_settings,
             ]
         ):
             raise ValueError(
-                "Please provide at least one of the following settings: openai_settings, azure_openai_settings, ollama_settings, embedded_settings, google_settings, huggingface_settings, mistral_ai_settings"
+                "Please provide at least one of the following settings: openai_settings, azure_openai_settings, ollama_settings, embedded_settings, google_settings, huggingface_settings, mistral_ai_settings, vertex_settings"
             )
 
         if model_type is None:
@@ -69,12 +74,13 @@ def __init__(
             google_settings,
             huggingface_settings,
             mistral_ai_settings,
+            vertex_settings,
         ]:
             if setting:
                 settings_set_count += 1 if setting else 0
             if settings_set_count > 1:
                 raise ValueError(
-                    "Please provide only one of the following settings: openai_settings, azure_openai_settings, ollama_settings, embedded_settings, google_settings, huggingface_settings, mistral_ai_settings"
+                    "Please provide only one of the following settings: openai_settings, azure_openai_settings, ollama_settings, embedded_settings, google_settings, huggingface_settings, mistral_ai_settings, vertex_settings"
                 )
 
     @property
@@ -92,6 +98,7 @@ def to_json(self) -> Dict[str, Any]:
             "GoogleSettings": self.google_settings.to_json() if self.google_settings else None,
             "HuggingFaceSettings": self.huggingface_settings.to_json() if self.huggingface_settings else None,
             "MistralAiSettings": self.mistral_ai_settings.to_json() if self.mistral_ai_settings else None,
+            "VertexSettings": self.vertex_settings.to_json() if self.vertex_settings else None,
             "ModelType": self.model_type.value if self.model_type else None,
             "Type": self.get_type,
         }
@@ -128,5 +135,8 @@ def from_json(cls, json_dict: Dict[str, Any]) -> "AiConnectionString":
                 if json_dict.get("MistralAiSettings")
                 else None
             ),
+            vertex_settings=(
+                VertexSettings.from_json(json_dict["VertexSettings"]) if json_dict.get("VertexSettings") else None
+            ),
             model_type=AiModelType(json_dict["ModelType"]) if json_dict.get("ModelType") else None,
         )

ravendb/documents/operations/ai/chunking_options.py

@@ -0,0 +1,39 @@
+from enum import Enum
+from typing import Dict, Any, Optional
+
+
+# todo: EmbeddingsGenerationConfiguration
+class ChunkingMethod(Enum):
+    PLAIN_TEXT_SPLIT = "PlainTextSplit"
+    PLAIN_TEXT_SPLIT_LINES = "PlainTextSplitLines"
+    PLAIN_TEXT_SPLIT_PARAGRAPHS = "PlainTextSplitParagraphs"
+    MARK_DOWN_SPLIT_LINES = "MarkDownSplitLines"
+    MARK_DOWN_SPLIT_PARAGRAPHS = "MarkDownSplitParagraphs"
+    HTML_STRIP = "HtmlStrip"
+
+
+class ChunkingOptions:
+    def __init__(
+        self,
+        chunking_method: Optional[ChunkingMethod] = None,
+        max_tokens_per_chunk: int = 512,
+        overlap_tokens: int = 0,
+    ):
+        self.chunking_method = chunking_method
+        self.max_tokens_per_chunk = max_tokens_per_chunk
+        self.overlap_tokens = overlap_tokens
+
+    @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),
+        )
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "ChunkingMethod": self.chunking_method.value if self.chunking_method else None,
+            "MaxTokensPerChunk": self.max_tokens_per_chunk,
+            "OverlapTokens": self.overlap_tokens,
+        }

ravendb/documents/operations/ai/embedding_path_configuration.py

@@ -0,0 +1,25 @@
+from typing import Dict, Any, Optional
+
+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":
+        return cls(
+            path=json_dict.get("Path"),
+            chunking_options=(
+                ChunkingOptions.from_json(json_dict["ChunkingOptions"]) if json_dict.get("ChunkingOptions") 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,
+        }

ravendb/documents/operations/ai/embeddings_transformation.py

@@ -0,0 +1,33 @@
+from typing import Dict, Any, Optional
+
+from ravendb.documents.operations.ai.chunking_options import ChunkingOptions, ChunkingMethod
+
+
+# todo: EmbeddingsGenerationConfiguration
+class EmbeddingsTransformation:
+    GENERATE_EMBEDDINGS_FUNCTION_NAME = "embeddings.generate"
+
+    def __init__(
+        self,
+        script: Optional[str] = None,
+        chunking_options: Optional[ChunkingOptions] = None,
+    ):
+        self.script = script
+        self.chunking_options = chunking_options or ChunkingOptions(
+            chunking_method=ChunkingMethod.PLAIN_TEXT_SPLIT, max_tokens_per_chunk=256
+        )
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "EmbeddingsTransformation":
+        return cls(
+            script=json_dict.get("Script"),
+            chunking_options=(
+                ChunkingOptions.from_json(json_dict["ChunkingOptions"]) if json_dict.get("ChunkingOptions") else None
+            ),
+        )
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "Script": self.script,
+            "ChunkingOptions": self.chunking_options.to_json() if self.chunking_options else None,
+        }

ravendb/documents/operations/ai/vertex_settings.py

@@ -0,0 +1,42 @@
+from enum import Enum
+from typing import Dict, Any, Optional
+
+from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings
+
+
+class VertexAIVersion(Enum):
+    V1 = "V1"
+    V1_BETA = "V1_Beta"
+
+
+class VertexSettings(AbstractAiSettings):
+    def __init__(
+        self,
+        model: Optional[str] = None,
+        google_credentials_json: Optional[str] = None,
+        location: Optional[str] = None,
+        ai_version: Optional[VertexAIVersion] = None,
+    ):
+        super().__init__()
+        self.model = model
+        self.google_credentials_json = google_credentials_json
+        self.location = location
+        self.ai_version = ai_version
+
+    @classmethod
+    def from_json(cls, json_dict: Dict[str, Any]) -> "VertexSettings":
+        return cls(
+            model=json_dict.get("Model"),
+            google_credentials_json=json_dict.get("GoogleCredentialsJson"),
+            location=json_dict.get("Location"),
+            ai_version=VertexAIVersion(json_dict["AiVersion"]) if json_dict.get("AiVersion") else None,
+        )
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "Model": self.model,
+            "GoogleCredentialsJson": self.google_credentials_json,
+            "AiVersion": self.ai_version.value if self.ai_version else None,
+            "Location": self.location,
+            "EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
+        }