kiln_ai.datamodel.rag

 1from typing import TYPE_CHECKING, Union
 2
 3from pydantic import Field, model_validator
 4
 5from kiln_ai.datamodel.basemodel import ID_TYPE, FilenameString, KilnParentedModel
 6from kiln_ai.utils.validation import ToolNameString
 7
 8if TYPE_CHECKING:
 9    from kiln_ai.datamodel.project import Project
10
11
12class RagConfig(KilnParentedModel):
13    """Configuration for Retrieval-Augmented Generation (RAG) on a project's documents."""
14
15    name: FilenameString = Field(
16        description="A name to identify this RAG configuration for your own reference.",
17    )
18
19    is_archived: bool = Field(
20        default=False,
21        description="Whether the RAG configuration is archived. Archived RAG configurations are not shown in the UI and are not available for use.",
22    )
23
24    description: str | None = Field(
25        default=None,
26        description="A description of the RAG configuration for you and your team. Will not be used in prompts/training/validation.",
27    )
28
29    tool_name: ToolNameString = Field(
30        description="A name for the model to identify the Search Tool in conversations.",
31    )
32
33    tool_description: str = Field(
34        description="A description of the purpose of the tool. The model will use this description to understand the tool's capabilities.",
35        max_length=128,
36    )
37
38    extractor_config_id: ID_TYPE = Field(
39        description="The ID of the extractor config used to extract the documents.",
40    )
41
42    chunker_config_id: ID_TYPE = Field(
43        description="The ID of the chunker config used to chunk the documents.",
44    )
45
46    embedding_config_id: ID_TYPE = Field(
47        description="The ID of the embedding config used to embed the documents.",
48    )
49
50    vector_store_config_id: ID_TYPE = Field(
51        description="The ID of the vector store config used to store the documents.",
52    )
53
54    reranker_config_id: ID_TYPE | None = Field(
55        default=None,
56        description="The ID of the reranker config used to rerank the documents. If None, no reranking will be performed.",
57    )
58
59    tags: list[str] | None = Field(
60        default=None,
61        description="List of document tags to filter by. If None, all documents in the project are used.",
62    )
63
64    # Workaround to return typed parent without importing Project
65    def parent_project(self) -> Union["Project", None]:
66        if self.parent is None or self.parent.__class__.__name__ != "Project":
67            return None
68        return self.parent  # type: ignore
69
70    @model_validator(mode="after")
71    def validate_tags(self):
72        if self.tags is not None:
73            if len(self.tags) == 0:
74                raise ValueError("Tags cannot be an empty list.")
75            for tag in self.tags:
76                if not tag:
77                    raise ValueError("Tags cannot be empty.")
78                if " " in tag:
79                    raise ValueError("Tags cannot contain spaces. Try underscores.")
80
81        if self.tool_name.strip() == "":
82            raise ValueError("Tool name cannot be empty.")
83        if self.tool_description.strip() == "":
84            raise ValueError("Tool description cannot be empty.")
85
86        return self