Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ravendb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@
)

from ravendb.documents.operations.etl.configuration import EtlConfiguration, RavenEtlConfiguration
from ravendb.documents.operations.etl.olap import OlapEtlConfiguration
from ravendb.documents.operations.etl.olap.connection import OlapEtlConfiguration
from ravendb.documents.operations.etl.sql import SqlEtlConfiguration
from ravendb.documents.operations.executor import MaintenanceOperationExecutor, SessionOperationExecutor
from ravendb.documents.operations.expiration.configuration import ExpirationConfiguration
Expand Down
16 changes: 16 additions & 0 deletions ravendb/documents/operations/ai/abstract_ai_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
from abc import ABC, abstractmethod
from typing import Dict, Any


class AbstractAiSettings(ABC):
def __init__(self):
self.embeddings_max_concurrent_batches = None

@classmethod
@abstractmethod
def from_json(cls, json_dict: Dict[str, Any]) -> Any:
pass

@abstractmethod
def to_json(self) -> Dict[str, Any]:
pass
132 changes: 132 additions & 0 deletions ravendb/documents/operations/ai/ai_connection_string.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
import enum
from typing import Optional, Dict, Any

from ravendb.serverwide.server_operation_executor import ConnectionStringType
from ravendb.documents.operations.ai.azure_open_ai_settings import AzureOpenAiSettings
from ravendb.documents.operations.ai.embedded_settings import EmbeddedSettings
from ravendb.documents.operations.ai.google_settings import GoogleSettings
from ravendb.documents.operations.ai.hugging_face_settings import HuggingFaceSettings
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.connection_strings import ConnectionString


class AiModelType(enum.Enum):
TEXT_EMBEDDINGS = "TextEmbeddings"
CHAT = "Chat"


class AiConnectionString(ConnectionString): # todo kuba
def __init__(
self,
name: str,
identifier: str,
openai_settings: Optional[OpenAiSettings] = None,
azure_openai_settings: Optional[AzureOpenAiSettings] = None,
ollama_settings: Optional[OllamaSettings] = None,
embedded_settings: Optional[EmbeddedSettings] = None,
google_settings: Optional[GoogleSettings] = None,
huggingface_settings: Optional[HuggingFaceSettings] = None,
mistral_ai_settings: Optional[MistralAiSettings] = None,
model_type: AiModelType = None,
):
super().__init__(name)
self.identifier = identifier
self.openai_settings = openai_settings
self.azure_openai_settings = azure_openai_settings
self.ollama_settings = ollama_settings
self.embedded_settings = embedded_settings
self.google_settings = google_settings
self.huggingface_settings = huggingface_settings
self.mistral_ai_settings = mistral_ai_settings
self.model_type = model_type

if not any(
[
openai_settings,
azure_openai_settings,
ollama_settings,
embedded_settings,
google_settings,
huggingface_settings,
mistral_ai_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"
)

if model_type is None:
raise ValueError("Please provide a model type - AiModelType.TEXT_EMBEDDINGS or AiModelType.CHAT")

settings_set_count = 0
for setting in [
openai_settings,
azure_openai_settings,
ollama_settings,
embedded_settings,
google_settings,
huggingface_settings,
mistral_ai_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"
)

@property
def get_type(self):
return ConnectionStringType.AI.value

def to_json(self) -> Dict[str, Any]:
return {
"Name": self.name,
"Identifier": self.identifier,
"OpenAiSettings": self.openai_settings.to_json() if self.openai_settings else None,
"AzureOpenAiSettings": self.azure_openai_settings.to_json() if self.azure_openai_settings else None,
"OllamaSettings": self.ollama_settings.to_json() if self.ollama_settings else None,
"EmbeddedSettings": self.embedded_settings.to_json() if self.embedded_settings else None,
"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,
"ModelType": self.model_type.value if self.model_type else None,
"Type": self.get_type,
}

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "AiConnectionString":
return cls(
name=json_dict["Name"],
identifier=json_dict["Identifier"],
openai_settings=(
OpenAiSettings.from_json(json_dict["OpenAiSettings"]) if json_dict.get("OpenAiSettings") else None
),
azure_openai_settings=(
AzureOpenAiSettings.from_json(json_dict["AzureOpenAiSettings"])
if json_dict.get("AzureOpenAiSettings")
else None
),
ollama_settings=(
OllamaSettings.from_json(json_dict["OllamaSettings"]) if json_dict.get("OllamaSettings") else None
),
embedded_settings=(
EmbeddedSettings.from_json(json_dict["EmbeddedSettings"]) if json_dict.get("EmbeddedSettings") else None
),
google_settings=(
GoogleSettings.from_json(json_dict["GoogleSettings"]) if json_dict.get("GoogleSettings") else None
),
huggingface_settings=(
HuggingFaceSettings.from_json(json_dict["HuggingFaceSettings"])
if json_dict.get("HuggingFaceSettings")
else None
),
mistral_ai_settings=(
MistralAiSettings.from_json(json_dict["MistralAiSettings"])
if json_dict.get("MistralAiSettings")
else None
),
model_type=AiModelType(json_dict["ModelType"]) if json_dict.get("ModelType") else None,
)
38 changes: 38 additions & 0 deletions ravendb/documents/operations/ai/azure_open_ai_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
from typing import Dict, Any

from ravendb.documents.operations.ai.open_ai_base_settings import OpenAiBaseSettings


class AzureOpenAiSettings(OpenAiBaseSettings):
def __init__(
self,
api_key: str = None,
endpoint: str = None,
model: str = None,
deployment_name: str = None,
dimensions: int = None,
temperature: float = None,
):
super().__init__(api_key, endpoint, model, dimensions, temperature)
self.deployment_name = deployment_name

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "AzureOpenAiSettings":
return cls(
api_key=json_dict["ApiKey"] if "ApiKey" in json_dict else None,
endpoint=json_dict["Endpoint"] if "Endpoint" in json_dict else None,
model=json_dict["Model"] if "Model" in json_dict else None,
dimensions=json_dict["Dimensions"] if "Dimensions" in json_dict else None,
temperature=json_dict["Temperature"] if "Temperature" in json_dict else None,
deployment_name=json_dict["DeploymentName"] if "DeploymentName" in json_dict else None,
)

def to_json(self) -> Dict[str, Any]:
return {
"ApiKey": self.api_key,
"Endpoint": self.endpoint,
"Model": self.model,
"Dimensions": self.dimensions,
"Temperature": self.temperature,
"DeploymentName": self.deployment_name,
}
15 changes: 15 additions & 0 deletions ravendb/documents/operations/ai/embedded_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from typing import Dict, Any

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class EmbeddedSettings(AbstractAiSettings):
def __init__(self):
super().__init__()

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "EmbeddedSettings":
return cls()

def to_json(self) -> Dict[str, Any]:
return {"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches}
38 changes: 38 additions & 0 deletions ravendb/documents/operations/ai/google_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
from enum import Enum
from typing import Dict, Any

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class GoogleAiVersion(Enum):
V1 = "V1"
V1_Beta = "V1_Beta"


class GoogleSettings(AbstractAiSettings):
def __init__(
self, model: str = None, api_key: str = None, ai_version: GoogleAiVersion = None, dimensions: int = None
):
super().__init__()
self.model = model
self.api_key = api_key
self.ai_version = ai_version
self.dimensions = dimensions

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "GoogleSettings":
return cls(
model=json_dict["Model"],
api_key=json_dict["ApiKey"],
ai_version=GoogleAiVersion(json_dict["AiVersion"]),
dimensions=json_dict["Dimensions"],
)

def to_json(self) -> Dict[str, Any]:
return {
"Model": self.model,
"ApiKey": self.api_key,
"AiVersion": self.ai_version.value,
"Dimensions": self.dimensions,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
27 changes: 27 additions & 0 deletions ravendb/documents/operations/ai/hugging_face_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from typing import Dict, Any

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class HuggingFaceSettings(AbstractAiSettings):
def __init__(self, api_key: str = None, model: str = None, endpoint: str = None):
super().__init__()
self.api_key = api_key
self.model = model
self.endpoint = endpoint

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "HuggingFaceSettings":
return cls(
api_key=json_dict["ApiKey"],
model=json_dict["Model"],
endpoint=json_dict["Endpoint"],
)

def to_json(self) -> Dict[str, Any]:
return {
"ApiKey": self.api_key,
"Model": self.model,
"Endpoint": self.endpoint,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
27 changes: 27 additions & 0 deletions ravendb/documents/operations/ai/mistral_ai_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from typing import Dict, Any

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class MistralAiSettings(AbstractAiSettings):
def __init__(self, api_key: str = None, model: str = None, endpoint: str = None):
super().__init__()
self.api_key = api_key
self.model = model
self.endpoint = endpoint

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "MistralAiSettings":
return cls(
api_key=json_dict["ApiKey"],
model=json_dict["Model"],
endpoint=json_dict["Endpoint"],
)

def to_json(self) -> Dict[str, Any]:
return {
"ApiKey": self.api_key,
"Model": self.model,
"Endpoint": self.endpoint,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
41 changes: 41 additions & 0 deletions ravendb/documents/operations/ai/ollama_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
from typing import Dict, Any

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class OllamaSettings(AbstractAiSettings):
def __init__(
self,
uri: str = None,
model: str = None,
think: bool = None,
temperature: float = None,
embeddings_max_concurrent_batches: int = None,
):
super().__init__()
self.uri = uri
self.model = model
self.think = think
self.temperature = temperature
self.embeddings_max_concurrent_batches = embeddings_max_concurrent_batches

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "OllamaSettings":
return cls(
uri=json_dict["Uri"] if "Uri" in json_dict else None,
model=json_dict["Model"] if "Model" in json_dict else None,
think=json_dict["Think"] if "Think" in json_dict else None,
temperature=json_dict["Temperature"] if "Temperature" in json_dict else None,
embeddings_max_concurrent_batches=(
json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
),
)

def to_json(self) -> Dict[str, Any]:
return {
"Uri": self.uri,
"Model": self.model,
"Think": self.think,
"Temperature": self.temperature,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
20 changes: 20 additions & 0 deletions ravendb/documents/operations/ai/open_ai_base_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
from abc import ABC

from ravendb.documents.operations.ai.abstract_ai_settings import AbstractAiSettings


class OpenAiBaseSettings(AbstractAiSettings, ABC):
def __init__(
self,
api_key: str = None,
endpoint: str = None,
model: str = None,
dimensions: int = None,
temperature: float = None,
):
super().__init__()
self.api_key = api_key
self.endpoint = endpoint
self.model = model
self.dimensions = dimensions
self.temperature = temperature
Loading