Skip to content
Open
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
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
77 changes: 77 additions & 0 deletions ravendb/documents/operations/ai/ai_connection_string.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
import enum
from typing import Optional, Dict, Any

import ravendb.serverwide.server_operation_executor
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

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

def to_json(self) -> Dict[str, Any]:
return {
"Name": self.name,
"Identifier": self.identifier,
"OpenaiSettings": self.openai_settings,
"AzureOpenaiSettings": self.azure_openai_settings,
"ollama_settings": self.ollama_settings,
"EmbeddedSettings": self.embedded_settings,
"GoogleSettings": self.google_settings,
"HuggingfaceSettings": self.huggingface_settings,
"MistralAiSettings": self.mistral_ai_settings,
"ModelType": self.model_type,
"Type": ravendb.serverwide.server_operation_executor.ConnectionStringType.OLAP,
}

@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"]),
azure_openai_settings=AzureOpenAiSettings.from_json(json_dict["AzureOpenaiSettings"]),
ollama_settings=OllamaSettings.from_json(json_dict["ollama_settings"]),
embedded_settings=EmbeddedSettings.from_json(json_dict["EmbeddedSettings"]),
google_settings=GoogleSettings.from_json(json_dict["GoogleSettings"]),
huggingface_settings=HuggingFaceSettings.from_json(json_dict["HuggingfaceSettings"]),
mistral_ai_settings=MistralAiSettings.from_json(json_dict["MistralAiSettings"]),
model_type=AiModelType(json_dict["ModelType"]) if json_dict["ModelType"] else None,
)
39 changes: 39 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,39 @@
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"],
endpoint=json_dict["Endpoint"],
model=json_dict["Model"],
dimensions=json_dict["Dimensions"],
temperature=json_dict["Temperature"],
deployment_name=json_dict["DeploymentName"],
)

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,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
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,
}
28 changes: 28 additions & 0 deletions ravendb/documents/operations/ai/ollama_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
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):
super().__init__()
self.uri = uri
self.model = model
self.think: bool = None
self.temperature: float = None

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "OllamaSettings":
return cls(
uri=json_dict["Uri"],
model=json_dict["Model"],
)

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
43 changes: 43 additions & 0 deletions ravendb/documents/operations/ai/open_ai_settings.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
from typing import Dict, Any

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


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

@classmethod
def from_json(cls, json_dict: Dict[str, Any]) -> "OpenAiSettings":
return cls(
api_key=json_dict["ApiKey"],
endpoint=json_dict["Endpoint"],
model=json_dict["Model"],
dimensions=json_dict["Dimensions"],
temperature=json_dict["Temperature"],
organization_id=json_dict["OrganizationId"],
project_id=json_dict["ProjectId"],
)

def to_json(self) -> Dict[str, Any]:
return {
"ApiKey": self.api_key,
"Endpoint": self.endpoint,
"Model": self.model,
"Dimensions": self.dimensions,
"Temperature": self.temperature,
"OrganizationId": self.organization_id,
"ProjectId": self.project_id,
"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
}
Empty file.
Loading