kiln_ai.datamodel.embedding

 1from typing import TYPE_CHECKING, List, Union
 2
 3from pydantic import BaseModel, Field, PositiveInt
 4from typing_extensions import TypedDict
 5
 6from kiln_ai.datamodel.basemodel import ID_TYPE, FilenameString, KilnParentedModel
 7from kiln_ai.datamodel.datamodel_enums import ModelProviderName
 8
 9if TYPE_CHECKING:
10    from kiln_ai.datamodel.chunk import ChunkedDocument
11    from kiln_ai.datamodel.project import Project
12
13
14class EmbeddingProperties(TypedDict, total=False):
15    dimensions: PositiveInt
16
17
18class EmbeddingConfig(KilnParentedModel):
19    name: FilenameString = Field(
20        description="A name to identify the embedding config.",
21    )
22    description: str | None = Field(
23        default=None,
24        description="A description for your reference, not shared with embedding models.",
25    )
26    model_provider_name: ModelProviderName = Field(
27        description="The provider to use to generate embeddings.",
28    )
29    model_name: str = Field(
30        description="The model to use to generate embeddings.",
31    )
32    properties: EmbeddingProperties = Field(
33        description="Properties to be used to execute the embedding config.",
34    )
35
36    # Workaround to return typed parent without importing Project
37    def parent_project(self) -> Union["Project", None]:
38        if self.parent is None or self.parent.__class__.__name__ != "Project":
39            return None
40        return self.parent  # type: ignore
41
42
43class Embedding(BaseModel):
44    vector: List[float] = Field(description="The vector of the embedding.")
45
46
47class ChunkEmbeddings(KilnParentedModel):
48    embedding_config_id: ID_TYPE = Field(
49        description="The ID of the embedding config used to generate the embeddings.",
50    )
51    embeddings: List[Embedding] = Field(
52        description="The embeddings of the chunks. The embedding at index i corresponds to the chunk at index i in the parent chunked document."
53    )
54
55    def parent_chunked_document(self) -> Union["ChunkedDocument", None]:
56        if self.parent is None or self.parent.__class__.__name__ != "ChunkedDocument":
57            return None
58        return self.parent  # type: ignore
class EmbeddingProperties(typing_extensions.TypedDict):
15class EmbeddingProperties(TypedDict, total=False):
16    dimensions: PositiveInt
dimensions: Annotated[int, Gt(gt=0)]
class EmbeddingConfig(kiln_ai.datamodel.basemodel.KilnParentedModel):
19class EmbeddingConfig(KilnParentedModel):
20    name: FilenameString = Field(
21        description="A name to identify the embedding config.",
22    )
23    description: str | None = Field(
24        default=None,
25        description="A description for your reference, not shared with embedding models.",
26    )
27    model_provider_name: ModelProviderName = Field(
28        description="The provider to use to generate embeddings.",
29    )
30    model_name: str = Field(
31        description="The model to use to generate embeddings.",
32    )
33    properties: EmbeddingProperties = Field(
34        description="Properties to be used to execute the embedding config.",
35    )
36
37    # Workaround to return typed parent without importing Project
38    def parent_project(self) -> Union["Project", None]:
39        if self.parent is None or self.parent.__class__.__name__ != "Project":
40            return None
41        return self.parent  # type: ignore

Base model for Kiln models that have a parent-child relationship. This base class is for child models.

This class provides functionality for managing hierarchical relationships between models, including parent reference handling and file system organization.

Attributes: parent (KilnBaseModel): Reference to the parent model instance. Not persisted, just in memory.

name: Annotated[str, BeforeValidator(func=<function name_validator.<locals>.fn at 0x7f2f1ec0c9a0>, json_schema_input_type=PydanticUndefined)]
description: str | None
model_provider_name: kiln_ai.datamodel.datamodel_enums.ModelProviderName
model_name: str
properties: EmbeddingProperties
def parent_project(self) -> Optional[kiln_ai.datamodel.Project]:
38    def parent_project(self) -> Union["Project", None]:
39        if self.parent is None or self.parent.__class__.__name__ != "Project":
40            return None
41        return self.parent  # type: ignore
def relationship_name() -> str:
713        def relationship_name_method() -> str:
714            return relationship_name

The type of the None singleton.

def parent_type() -> Type[kiln_ai.datamodel.basemodel.KilnParentModel]:
706        def parent_class_method() -> Type[KilnParentModel]:
707            return cls

The type of the None singleton.

model_config = {'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

def model_post_init(self: pydantic.main.BaseModel, context: Any, /) -> None:
337def init_private_attributes(self: BaseModel, context: Any, /) -> None:
338    """This function is meant to behave like a BaseModel method to initialise private attributes.
339
340    It takes context as an argument since that's what pydantic-core passes when calling it.
341
342    Args:
343        self: The BaseModel instance.
344        context: The context.
345    """
346    if getattr(self, '__pydantic_private__', None) is None:
347        pydantic_private = {}
348        for name, private_attr in self.__private_attributes__.items():
349            default = private_attr.get_default()
350            if default is not PydanticUndefined:
351                pydantic_private[name] = default
352        object_setattr(self, '__pydantic_private__', pydantic_private)

This function is meant to behave like a BaseModel method to initialise private attributes.

It takes context as an argument since that's what pydantic-core passes when calling it.

Args: self: The BaseModel instance. context: The context.

class Embedding(pydantic.main.BaseModel):
44class Embedding(BaseModel):
45    vector: List[float] = Field(description="The vector of the embedding.")

!!! abstract "Usage Documentation" Models

A base class for creating Pydantic models.

Attributes: __class_vars__: The names of the class variables defined on the model. __private_attributes__: Metadata about the private attributes of the model. __signature__: The synthesized __init__ [Signature][inspect.Signature] of the model.

__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The core schema of the model.
__pydantic_custom_init__: Whether the model has a custom `__init__` function.
__pydantic_decorators__: Metadata containing the decorators defined on the model.
    This replaces `Model.__validators__` and `Model.__root_validators__` from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to
    __args__, __origin__, __parameters__ in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a [`RootModel`][pydantic.root_model.RootModel].
__pydantic_serializer__: The `pydantic-core` `SchemaSerializer` used to dump instances of the model.
__pydantic_validator__: The `pydantic-core` `SchemaValidator` used to validate instances of the model.

__pydantic_fields__: A dictionary of field names and their corresponding [`FieldInfo`][pydantic.fields.FieldInfo] objects.
__pydantic_computed_fields__: A dictionary of computed field names and their corresponding [`ComputedFieldInfo`][pydantic.fields.ComputedFieldInfo] objects.

__pydantic_extra__: A dictionary containing extra values, if [`extra`][pydantic.config.ConfigDict.extra]
    is set to `'allow'`.
__pydantic_fields_set__: The names of fields explicitly set during instantiation.
__pydantic_private__: Values of private attributes set on the model instance.
vector: List[float]
model_config: ClassVar[pydantic.config.ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class ChunkEmbeddings(kiln_ai.datamodel.basemodel.KilnParentedModel):
48class ChunkEmbeddings(KilnParentedModel):
49    embedding_config_id: ID_TYPE = Field(
50        description="The ID of the embedding config used to generate the embeddings.",
51    )
52    embeddings: List[Embedding] = Field(
53        description="The embeddings of the chunks. The embedding at index i corresponds to the chunk at index i in the parent chunked document."
54    )
55
56    def parent_chunked_document(self) -> Union["ChunkedDocument", None]:
57        if self.parent is None or self.parent.__class__.__name__ != "ChunkedDocument":
58            return None
59        return self.parent  # type: ignore

Base model for Kiln models that have a parent-child relationship. This base class is for child models.

This class provides functionality for managing hierarchical relationships between models, including parent reference handling and file system organization.

Attributes: parent (KilnBaseModel): Reference to the parent model instance. Not persisted, just in memory.

embedding_config_id: Optional[str]
embeddings: List[Embedding]
def parent_chunked_document(self) -> Optional[kiln_ai.datamodel.chunk.ChunkedDocument]:
56    def parent_chunked_document(self) -> Union["ChunkedDocument", None]:
57        if self.parent is None or self.parent.__class__.__name__ != "ChunkedDocument":
58            return None
59        return self.parent  # type: ignore
def relationship_name() -> str:
713        def relationship_name_method() -> str:
714            return relationship_name

The type of the None singleton.

def parent_type() -> Type[kiln_ai.datamodel.basemodel.KilnParentModel]:
706        def parent_class_method() -> Type[KilnParentModel]:
707            return cls

The type of the None singleton.

model_config = {'validate_assignment': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

def model_post_init(self: pydantic.main.BaseModel, context: Any, /) -> None:
337def init_private_attributes(self: BaseModel, context: Any, /) -> None:
338    """This function is meant to behave like a BaseModel method to initialise private attributes.
339
340    It takes context as an argument since that's what pydantic-core passes when calling it.
341
342    Args:
343        self: The BaseModel instance.
344        context: The context.
345    """
346    if getattr(self, '__pydantic_private__', None) is None:
347        pydantic_private = {}
348        for name, private_attr in self.__private_attributes__.items():
349            default = private_attr.get_default()
350            if default is not PydanticUndefined:
351                pydantic_private[name] = default
352        object_setattr(self, '__pydantic_private__', pydantic_private)

This function is meant to behave like a BaseModel method to initialise private attributes.

It takes context as an argument since that's what pydantic-core passes when calling it.

Args: self: The BaseModel instance. context: The context.