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
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.
The type of the None singleton.
Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].
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.
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.
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.
The type of the None singleton.
Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].
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.