Merge pull request #638 from tisnik/lcore-642-byok-rag-configuration

LCORE-642: BYOK RAG configuration

Author: tisnik

docs/config.png

@@ -19,6 +19,14 @@ class "AuthenticationConfiguration" as src.models.config.AuthenticationConfigura
 class "AuthorizationConfiguration" as src.models.config.AuthorizationConfiguration {
   access_rules : Optional[list[AccessRule]]
 }
+class "ByokRag" as src.models.config.ByokRag {
+  db_path : Annotated
+  embedding_dimension : Annotated
+  embedding_model : Annotated
+  rag_id : Annotated
+  rag_type : Annotated
+  vector_db_id : Annotated
+}
 class "CORSConfiguration" as src.models.config.CORSConfiguration {
   allow_credentials : bool
   allow_headers : list[str]
@@ -29,6 +37,7 @@ class "CORSConfiguration" as src.models.config.CORSConfiguration {
 class "Configuration" as src.models.config.Configuration {
   authentication : Optional[AuthenticationConfiguration]
   authorization : Optional[AuthorizationConfiguration]
+  byok_rag : Optional[list[ByokRag]]
   conversation_cache : Optional[ConversationCacheConfiguration]
   customization : Optional[Customization]
   database : Optional[DatabaseConfiguration]
@@ -155,6 +164,7 @@ class "UserDataCollection" as src.models.config.UserDataCollection {
 src.models.config.AccessRule --|> src.models.config.ConfigurationBase
 src.models.config.AuthenticationConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.AuthorizationConfiguration --|> src.models.config.ConfigurationBase
+src.models.config.ByokRag --|> src.models.config.ConfigurationBase
 src.models.config.CORSConfiguration --|> src.models.config.ConfigurationBase
 src.models.config.Configuration --|> src.models.config.ConfigurationBase
 src.models.config.ConversationCacheConfiguration --|> src.models.config.ConfigurationBase

docs/config.puml

@@ -127,3 +127,16 @@
 CACHE_TYPE_SQLITE = "sqlite"
 CACHE_TYPE_POSTGRES = "postgres"
 CACHE_TYPE_NOOP = "noop"
+
+# BYOK RAG
+# Default RAG type for bring-your-own-knowledge RAG configurations, that type
+# needs to be supported by Llama Stack
+DEFAULT_RAG_TYPE = "inline::faiss"
+
+# Default sentence transformer model for embedding generation, that type needs
+# to be supported by Llama Stack and configured properly in providers and
+# models sections
+DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
+
+# Default embedding vector dimension for the sentence transformer model
+DEFAULT_EMBEDDING_DIMENSION = 768

docs/config.svg

@@ -13,6 +13,7 @@
     ConfigDict,
     Field,
     model_validator,
+    constr,
     FilePath,
     AnyHttpUrl,
     PositiveInt,
@@ -545,6 +546,19 @@ def check_cache_configuration(self) -> Self:
         return self
 
 
+class ByokRag(ConfigurationBase):
+    """BYOK RAG configuration."""
+
+    rag_id: constr(min_length=1)  # type:ignore
+    rag_type: constr(min_length=1) = constants.DEFAULT_RAG_TYPE  # type:ignore
+    embedding_model: constr(min_length=1) = (  # type:ignore
+        constants.DEFAULT_EMBEDDING_MODEL
+    )
+    embedding_dimension: PositiveInt = constants.DEFAULT_EMBEDDING_DIMENSION
+    vector_db_id: constr(min_length=1)  # type:ignore
+    db_path: FilePath
+
+
 class Configuration(ConfigurationBase):
     """Global service configuration."""
 
@@ -563,6 +577,7 @@ class Configuration(ConfigurationBase):
     conversation_cache: ConversationCacheConfiguration = Field(
         default_factory=ConversationCacheConfiguration
     )
+    byok_rag: list[ByokRag] = Field(default_factory=list)
 
     def dump(self, filename: str = "configuration.json") -> None:
         """Dump actual configuration into JSON file."""

src/constants.py

@@ -0,0 +1,130 @@
+"""Unit tests for ByokRag model."""
+
+from pathlib import Path
+
+import pytest
+
+from pydantic import ValidationError
+
+from models.config import ByokRag
+
+from constants import (
+    DEFAULT_RAG_TYPE,
+    DEFAULT_EMBEDDING_MODEL,
+    DEFAULT_EMBEDDING_DIMENSION,
+)
+
+
+def test_byok_rag_configuration_default_values() -> None:
+    """Test the ByokRag constructor."""
+
+    byok_rag = ByokRag(
+        rag_id="rag_id",
+        vector_db_id="vector_db_id",
+        db_path="tests/configuration/rag.txt",
+    )
+    assert byok_rag is not None
+    assert byok_rag.rag_id == "rag_id"
+    assert byok_rag.rag_type == DEFAULT_RAG_TYPE
+    assert byok_rag.embedding_model == DEFAULT_EMBEDDING_MODEL
+    assert byok_rag.embedding_dimension == DEFAULT_EMBEDDING_DIMENSION
+    assert byok_rag.vector_db_id == "vector_db_id"
+    assert byok_rag.db_path == Path("tests/configuration/rag.txt")
+
+
+def test_byok_rag_configuration_nondefault_values() -> None:
+    """Test the ByokRag constructor."""
+
+    byok_rag = ByokRag(
+        rag_id="rag_id",
+        rag_type="rag_type",
+        embedding_model="embedding_model",
+        embedding_dimension=1024,
+        vector_db_id="vector_db_id",
+        db_path="tests/configuration/rag.txt",
+    )
+    assert byok_rag is not None
+    assert byok_rag.rag_id == "rag_id"
+    assert byok_rag.rag_type == "rag_type"
+    assert byok_rag.embedding_model == "embedding_model"
+    assert byok_rag.embedding_dimension == 1024
+    assert byok_rag.vector_db_id == "vector_db_id"
+    assert byok_rag.db_path == Path("tests/configuration/rag.txt")
+
+
+def test_byok_rag_configuration_wrong_dimension() -> None:
+    """Test the ByokRag constructor."""
+
+    with pytest.raises(ValidationError, match="should be greater than 0"):
+        _ = ByokRag(
+            rag_id="rag_id",
+            rag_type="rag_type",
+            embedding_model="embedding_model",
+            embedding_dimension=-1024,
+            vector_db_id="vector_db_id",
+            db_path="tests/configuration/rag.txt",
+        )
+
+
+def test_byok_rag_configuration_empty_rag_id() -> None:
+    """Test the ByokRag constructor."""
+
+    with pytest.raises(
+        ValidationError, match="String should have at least 1 character"
+    ):
+        _ = ByokRag(
+            rag_id="",
+            rag_type="rag_type",
+            embedding_model="embedding_model",
+            embedding_dimension=1024,
+            vector_db_id="vector_db_id",
+            db_path="tests/configuration/rag.txt",
+        )
+
+
+def test_byok_rag_configuration_empty_rag_type() -> None:
+    """Test the ByokRag constructor."""
+
+    with pytest.raises(
+        ValidationError, match="String should have at least 1 character"
+    ):
+        _ = ByokRag(
+            rag_id="rag_id",
+            rag_type="",
+            embedding_model="embedding_model",
+            embedding_dimension=1024,
+            vector_db_id="vector_db_id",
+            db_path="tests/configuration/rag.txt",
+        )
+
+
+def test_byok_rag_configuration_empty_embedding_model() -> None:
+    """Test the ByokRag constructor."""
+
+    with pytest.raises(
+        ValidationError, match="String should have at least 1 character"
+    ):
+        _ = ByokRag(
+            rag_id="rag_id",
+            rag_type="rag_type",
+            embedding_model="",
+            embedding_dimension=1024,
+            vector_db_id="vector_db_id",
+            db_path="tests/configuration/rag.txt",
+        )
+
+
+def test_byok_rag_configuration_empty_vector_db_id() -> None:
+    """Test the ByokRag constructor."""
+
+    with pytest.raises(
+        ValidationError, match="String should have at least 1 character"
+    ):
+        _ = ByokRag(
+            rag_id="rag_id",
+            rag_type="rag_type",
+            embedding_model="embedding_model",
+            embedding_dimension=1024,
+            vector_db_id="",
+            db_path="tests/configuration/rag.txt",
+        )

src/models/config.py

@@ -88,6 +88,7 @@ def test_dump_configuration(tmp_path) -> None:
         assert "customization" in content
         assert "inference" in content
         assert "database" in content
+        assert "byok_rag" in content
 
         # check the whole deserialized JSON file content
         assert content == {
@@ -169,6 +170,7 @@ def test_dump_configuration(tmp_path) -> None:
                 "sqlite": None,
                 "type": None,
             },
+            "byok_rag": [],
         }