From e7384ef4f7113aaa4bbe60d894a05687f1200ce6 Mon Sep 17 00:00:00 2001
From: Bui Dai <ss1280dzzz@gmail.com>
Date: Mon, 26 Jan 2026 21:47:56 +0700
Subject: [PATCH] fix: embedding config not apply during document processing

---
 .../settings/EmbeddingConfigForm.tsx          |  28 +++++
 ragitect/api/v1/chat.py                       |  42 ++-----
 .../services/document_processing_service.py   |   4 +-
 ragitect/services/embedding.py                |  49 +++++++-
 tests/api/v1/test_chat_streaming.py           |  18 +--
 .../test_document_processing_service.py       |   9 +-
 tests/unit/test_embedding_generation.py       | 111 +++++++++++++++++-
 7 files changed, 200 insertions(+), 61 deletions(-)

diff --git a/frontend/src/components/settings/EmbeddingConfigForm.tsx b/frontend/src/components/settings/EmbeddingConfigForm.tsx
index 8f04f8d..6513c3a 100644
--- a/frontend/src/components/settings/EmbeddingConfigForm.tsx
+++ b/frontend/src/components/settings/EmbeddingConfigForm.tsx
@@ -17,6 +17,8 @@ import { Label } from '@/components/ui/label';
 import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card';
 import { Badge } from '@/components/ui/badge';
 import { Alert, AlertDescription } from '@/components/ui/alert';
+import { Switch } from '@/components/ui/switch';
+import { Separator } from '@/components/ui/separator';
 import {
   Select,
   SelectContent,
@@ -153,6 +155,15 @@ export function EmbeddingConfigForm() {
     }));
   }, [currentProvider]);
 
+  // Handle active toggle
+  const handleToggle = useCallback((checked: boolean) => {
+    setFormState(prev => ({
+      ...prev,
+      isEnabled: checked,
+      hasChanges: true,
+    }));
+  }, []);
+
   // Handler for testing connection
   const handleTestConnection = useCallback(async () => {
     setFormState(prev => ({
@@ -433,6 +444,23 @@ export function EmbeddingConfigForm() {
           </div>
         )}
 
+        <Separator />
+
+        {/* Set as Active Toggle */}
+        <div className="flex items-center justify-between">
+          <div className="space-y-0.5">
+            <Label htmlFor="is-active">Set as active</Label>
+            <p className="text-sm text-muted-foreground">
+              Use this embedding provider for document processing
+            </p>
+          </div>
+          <Switch
+            id="is-active"
+            checked={formState.isEnabled}
+            onCheckedChange={handleToggle}
+          />
+        </div>
+
         {/* Action Buttons */}
         <div className="flex gap-2 pt-2">
           <Button
diff --git a/ragitect/api/v1/chat.py b/ragitect/api/v1/chat.py
index f664f9e..fb0de9c 100644
--- a/ragitect/api/v1/chat.py
+++ b/ragitect/api/v1/chat.py
@@ -30,8 +30,11 @@
 from ragitect.services.database.repositories.document_repo import DocumentRepository
 from ragitect.services.database.repositories.vector_repo import VectorRepository
 from ragitect.services.database.repositories.workspace_repo import WorkspaceRepository
-from ragitect.services.embedding import create_embeddings_model, embed_text
-from ragitect.services.llm_config_service import get_active_embedding_config
+from ragitect.services.embedding import (
+    create_embeddings_model,
+    embed_text,
+    get_embedding_model_from_config,
+)
 from ragitect.services.llm_factory import create_llm_with_provider
 
 
@@ -166,21 +169,8 @@ async def retrieve_context_with_graph(
     Returns:
         List of chunks with content and metadata
     """
-    # Get embedding configuration and create model
-    embedding_config = await get_active_embedding_config(session)
-
-    if embedding_config:
-        config = EmbeddingConfig(
-            provider=embedding_config.provider_name,
-            model=embedding_config.model_name or "nomic-embed-text",
-            base_url=embedding_config.config_data.get("base_url"),
-            api_key=embedding_config.config_data.get("api_key"),
-            dimension=embedding_config.config_data.get("dimension", 768),
-        )
-    else:
-        config = EmbeddingConfig()
-
-    embedding_model = create_embeddings_model(config)
+    # Get embedding model from active DB config (or defaults)
+    embedding_model = await get_embedding_model_from_config(session)
 
     # Create async embedding function for graph
     async def embed_fn(text: str) -> list[float]:
@@ -373,22 +363,8 @@ async def chat_stream(
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-    # Get active embedding configuration and model
-    embed_config_dto = await get_active_embedding_config(session)
-    if embed_config_dto is None:
-        # Use default config if no active config found
-        embed_config = EmbeddingConfig()
-    else:
-        # Convert DTO to EmbeddingConfig
-        embed_config = EmbeddingConfig(
-            provider=embed_config_dto.provider_name,
-            model=embed_config_dto.model_name or "all-MiniLM-L6-v2",
-            api_key=embed_config_dto.api_key,
-            base_url=embed_config_dto.base_url,
-            dimension=embed_config_dto.dimension,
-        )
-
-    embed_model = create_embeddings_model(embed_config)
+    # Get embedding model from active DB config (or defaults)
+    embed_model = await get_embedding_model_from_config(session)
 
     async def embed_fn(text: str) -> list[float]:
         """Embedding function for runtime dependency injection."""
diff --git a/ragitect/services/document_processing_service.py b/ragitect/services/document_processing_service.py
index 243029e..10aa836 100644
--- a/ragitect/services/document_processing_service.py
+++ b/ragitect/services/document_processing_service.py
@@ -24,7 +24,7 @@
 from ragitect.services.config import load_document_config
 from ragitect.services.database.repositories.document_repo import DocumentRepository
 from ragitect.services.document_processor import process_file_bytes, split_document
-from ragitect.services.embedding import create_embeddings_model, embed_documents
+from ragitect.services.embedding import embed_documents, get_embedding_model_from_config
 
 logger = logging.getLogger(__name__)
 
@@ -135,7 +135,7 @@ async def process_document(self, document_id: UUID) -> None:
             # Generate embeddings if there are chunks
             if chunks:
                 try:
-                    embedding_model = create_embeddings_model()
+                    embedding_model = await get_embedding_model_from_config(self.session)
                     embeddings = await embed_documents(embedding_model, chunks)
                     logger.info(
                         f"Generated {len(embeddings)} embeddings for document {document_id}"
diff --git a/ragitect/services/embedding.py b/ragitect/services/embedding.py
index 55a9148..0261226 100644
--- a/ragitect/services/embedding.py
+++ b/ragitect/services/embedding.py
@@ -7,7 +7,7 @@
 """
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import httpx
 from langchain_core.embeddings import Embeddings
@@ -15,6 +15,9 @@
 
 from ragitect.services.config import EmbeddingConfig
 
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
 logger = logging.getLogger(__name__)
 
 
@@ -307,3 +310,47 @@ def get_embedding_dimension(config: EmbeddingConfig | None = None) -> int:
     if config is None:
         return 768  # Default for nomic-embed-text
     return config.dimension
+
+
+async def get_embedding_model_from_config(session: "AsyncSession") -> Embeddings:
+    """Get embedding model using active DB config, fallback to env vars, then defaults.
+
+    Priority order:
+    1. Active embedding config from database (UI-configured)
+    2. Environment variables (EMBEDDING_PROVIDER, EMBEDDING_API_KEY, etc.)
+    3. Hardcoded defaults (Ollama)
+
+    This is the DRY helper that encapsulates embedding config lookup, avoiding
+    duplicate patterns across chat.py and document_processing_service.py.
+
+    Args:
+        session: Database session for config lookup
+
+    Returns:
+        Embeddings model configured from DB, env vars, or defaults
+    """
+    # Import inside function to avoid circular imports
+    from ragitect.services.llm_config_service import get_active_embedding_config
+    from ragitect.services.config import load_embedding_config
+
+    embedding_config = await get_active_embedding_config(session)
+    if embedding_config:
+        logger.info(
+            f"Using embedding config from DB: provider={embedding_config.provider_name}, "
+            f"model={embedding_config.model_name}"
+        )
+        config = EmbeddingConfig(
+            provider=embedding_config.provider_name,
+            model=embedding_config.model_name or "qwen3-embedding:0.6b",
+            base_url=embedding_config.base_url,
+            api_key=embedding_config.api_key,
+            dimension=embedding_config.dimension,
+        )
+    else:
+        # Fall back to environment variables
+        config = load_embedding_config()
+        logger.info(
+            f"No active embedding config in DB, using env vars: "
+            f"provider={config.provider}, model={config.model}"
+        )
+    return create_embeddings_model(config)
diff --git a/tests/api/v1/test_chat_streaming.py b/tests/api/v1/test_chat_streaming.py
index 0849824..8f0f38e 100644
--- a/tests/api/v1/test_chat_streaming.py
+++ b/tests/api/v1/test_chat_streaming.py
@@ -32,23 +32,10 @@ def setup_langgraph_streaming_mocks(mocker):
     Returns:
         dict with mock objects for assertions
     """
-    # Mock embedding config - required for both retrieval and streaming
-    mock_embed_config = mocker.MagicMock()
-    mock_embed_config.provider_name = "ollama"
-    mock_embed_config.model_name = "all-MiniLM-L6-v2"
-    mock_embed_config.api_key = None
-    mock_embed_config.base_url = None
-    mock_embed_config.dimension = 768
-
-    mocker.patch(
-        "ragitect.api.v1.chat.get_active_embedding_config",
-        return_value=mock_embed_config,
-    )
-
-    # Mock embedding model and embed function
+    # Mock embedding model - now using the DRY helper
     mock_embed_model = mocker.MagicMock()
     mocker.patch(
-        "ragitect.api.v1.chat.create_embeddings_model",
+        "ragitect.api.v1.chat.get_embedding_model_from_config",
         return_value=mock_embed_model,
     )
 
@@ -96,7 +83,6 @@ async def mock_embed_fn(model, text: str):
     )
 
     return {
-        "embed_config": mock_embed_config,
         "embed_model": mock_embed_model,
         "vector_repo": mock_vector_repo,
         "llm": mock_llm,
diff --git a/tests/services/test_document_processing_service.py b/tests/services/test_document_processing_service.py
index dbea649..22d836f 100644
--- a/tests/services/test_document_processing_service.py
+++ b/tests/services/test_document_processing_service.py
@@ -104,8 +104,8 @@ async def test_process_document_success(
                 mock_model = MagicMock()
                 mock_model.aembed_documents = AsyncMock(return_value=[[0.1] * 768])
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model",
-                    return_value=mock_model,
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(return_value=mock_model),
                 ):
                     with patch(
                         "ragitect.services.document_processing_service.embed_documents",
@@ -247,8 +247,11 @@ async def test_process_document_integration_with_repo_methods(
                 "ragitect.services.document_processing_service.split_document",
                 return_value=["Chunk"],
             ):
+                mock_model = MagicMock()
+                mock_model.aembed_documents = AsyncMock(return_value=[[0.1] * 768])
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model"
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(return_value=mock_model),
                 ):
                     with patch(
                         "ragitect.services.document_processing_service.embed_documents",
diff --git a/tests/unit/test_embedding_generation.py b/tests/unit/test_embedding_generation.py
index 7593332..e4e9a49 100644
--- a/tests/unit/test_embedding_generation.py
+++ b/tests/unit/test_embedding_generation.py
@@ -74,6 +74,99 @@ def test_create_embeddings_model_ollama_requires_base_url(self):
             create_embeddings_model(config)
 
 
+class TestGetEmbeddingModelFromConfig:
+    """Tests for the get_embedding_model_from_config helper function."""
+
+    @pytest.fixture
+    def mock_session(self):
+        """Mock AsyncSession."""
+        return MagicMock()
+
+    @pytest.fixture
+    def mock_openai_config_dto(self):
+        """Mock OpenAI embedding configuration DTO."""
+        dto = MagicMock()
+        dto.provider_name = "openai"
+        dto.model_name = "text-embedding-3-small"
+        dto.base_url = None
+        dto.api_key = "sk-test-key-12345"
+        dto.dimension = 1536
+        return dto
+
+    @pytest.fixture
+    def mock_ollama_config_dto(self):
+        """Mock Ollama embedding configuration DTO."""
+        dto = MagicMock()
+        dto.provider_name = "ollama"
+        dto.model_name = "qwen3-embedding:0.6b"
+        dto.base_url = "http://localhost:11434"
+        dto.api_key = None
+        dto.dimension = 768
+        return dto
+
+    @pytest.mark.asyncio
+    async def test_get_embedding_model_with_openai_config(
+        self, mock_session, mock_openai_config_dto
+    ):
+        """Test helper returns OpenAI model when active config is OpenAI."""
+        from ragitect.services.embedding import get_embedding_model_from_config
+
+        with patch(
+            "ragitect.services.llm_config_service.get_active_embedding_config",
+            new=AsyncMock(return_value=mock_openai_config_dto),
+        ):
+            model = await get_embedding_model_from_config(mock_session)
+
+        # Should return an embeddings model (OpenAI)
+        assert model is not None
+        assert hasattr(model, "aembed_documents")
+
+    @pytest.mark.asyncio
+    async def test_get_embedding_model_with_ollama_config(
+        self, mock_session, mock_ollama_config_dto
+    ):
+        """Test helper returns Ollama model when active config is Ollama."""
+        from ragitect.services.embedding import get_embedding_model_from_config
+
+        with patch(
+            "ragitect.services.llm_config_service.get_active_embedding_config",
+            new=AsyncMock(return_value=mock_ollama_config_dto),
+        ):
+            model = await get_embedding_model_from_config(mock_session)
+
+        # Should return an embeddings model (Ollama)
+        assert model is not None
+        assert hasattr(model, "aembed_documents")
+
+    @pytest.mark.asyncio
+    async def test_get_embedding_model_with_no_active_config(self, mock_session):
+        """Test helper returns env-configured model when no active DB config."""
+        from ragitect.services.embedding import get_embedding_model_from_config
+        from ragitect.services.config import EmbeddingConfig
+
+        # Mock no DB config and mock env config to Ollama defaults
+        env_config = EmbeddingConfig(
+            provider="ollama",
+            model="qwen3-embedding:0.6b",
+            base_url="http://localhost:11434",
+            dimension=768,
+        )
+
+        with patch(
+            "ragitect.services.llm_config_service.get_active_embedding_config",
+            new=AsyncMock(return_value=None),
+        ):
+            with patch(
+                "ragitect.services.config.load_embedding_config",
+                return_value=env_config,
+            ):
+                model = await get_embedding_model_from_config(mock_session)
+
+        # Should return default embeddings model (Ollama)
+        assert model is not None
+        assert hasattr(model, "aembed_documents")
+
+
 class TestBatchEmbedding:
     """Tests for batch embedding generation."""
 
@@ -307,8 +400,8 @@ async def test_process_document_generates_embeddings(
                 )
 
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model",
-                    return_value=mock_model,
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(return_value=mock_model),
                 ):
                     with patch(
                         "ragitect.services.document_processing_service.embed_documents",
@@ -354,8 +447,11 @@ async def test_process_document_status_transition_to_embedding(
                 "ragitect.services.document_processing_service.process_file_bytes",
                 return_value=("Text", {"file_type": ".txt"}),
             ):
+                mock_model = MagicMock()
+                mock_model.aembed_documents = AsyncMock(return_value=[[0.1] * 768])
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model"
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(return_value=mock_model),
                 ):
                     with patch(
                         "ragitect.services.document_processing_service.embed_documents",
@@ -403,8 +499,8 @@ async def test_process_document_embedding_failure_sets_error_status(
                 return_value=("Text", {"file_type": ".txt"}),
             ):
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model",
-                    side_effect=Exception("Embedding model unavailable"),
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(side_effect=Exception("Embedding model unavailable")),
                 ):
                     with pytest.raises(Exception, match="Embedding model unavailable"):
                         await service.process_document(sample_document.id)
@@ -437,8 +533,11 @@ async def test_process_document_stores_chunks_with_embeddings(
                 "ragitect.services.document_processing_service.process_file_bytes",
                 return_value=("Text", {"file_type": ".txt"}),
             ):
+                mock_model = MagicMock()
+                mock_model.aembed_documents = AsyncMock(return_value=embeddings)
                 with patch(
-                    "ragitect.services.document_processing_service.create_embeddings_model"
+                    "ragitect.services.document_processing_service.get_embedding_model_from_config",
+                    new=AsyncMock(return_value=mock_model),
                 ):
                     with patch(
                         "ragitect.services.document_processing_service.embed_documents",