Merge pull request #9 from NadaBhm/feature/nada-rag

Feature/nada rag

Author: NadaBhm

src/lib/rag/__init__.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+from .api import ask_about_repo, get_repo_context
 from .config import RAGConfig, get_rag_config
 from .embeddings import EmbeddingClient, get_embedding
 from .ingestion import ingest_repo, ingest_text
@@ -29,4 +30,6 @@
     "ask_repo",
     "retrieve_context",
     "similarity_search",
+    "ask_about_repo",
+    "get_repo_context",
 ]

src/lib/rag/api.py

@@ -0,0 +1,57 @@
+"""
+RAG Public API — Semantic Retrieval 
+===============================================
+Single entry-point for the Orchestrator and Chat backend.
+
+Usage:
+    from lib.rag.api import ask_about_repo
+    answer = ask_about_repo(job_id="550e-...", question="What framework?")
+"""
+
+from __future__ import annotations
+
+import logging
+
+from .config import RAGConfig, get_rag_config
+from .retrieval import ask_repo, retrieve_context
+
+logger = logging.getLogger(__name__)
+
+
+def ask_about_repo(
+    job_id: str,
+    question: str,
+    top_k: int = 5,
+    config: RAGConfig | None = None,
+) -> str:
+    """
+    Ask a natural-language question about an analyzed repository.
+
+    This is the primary integration point for the Orchestrator Chat node.
+    It retrieves relevant chunks from Qdrant and generates an answer via Gemini.
+
+    Args:
+        job_id: The CodeSec analysis job ID (links to Qdrant collection).
+        question: User question in any language.
+        top_k: Number of document chunks to retrieve.
+
+    Returns:
+        LLM-generated answer, or a fallback message if no context is found.
+    """
+    logger.info("[RAG API] job=%s question=%r", job_id, question)
+    return ask_repo(question, job_id, top_k, config)
+
+
+def get_repo_context(
+    job_id: str,
+    question: str,
+    top_k: int = 5,
+    config: RAGConfig | None = None,
+) -> str:
+    """
+    Retrieve raw context chunks for a question (debug / custom prompting).
+
+    Returns the formatted context string that would be fed to the LLM,
+    without actually calling the LLM.
+    """
+    return retrieve_context(question, job_id, top_k, config)

src/lib/rag/config.py

@@ -14,7 +14,7 @@
 DEFAULT_QDRANT_COLLECTION: Final[str] = os.getenv("QDRANT_COLLECTION", "devguard_repos")
 
 # Hugging Face embeddings (local, free)
-DEFAULT_HF_MODEL: Final[str] = os.getenv("HF_MODEL", "BAAI/bge-base-en-v1.5")
+DEFAULT_HF_MODEL: Final[str] = os.getenv("HF_MODEL", "BAAI/bge-large-en-v1.5")
 DEFAULT_EMBEDDING_DIM: Final[int] = int(os.getenv("EMBEDDING_DIM", "768"))
 
 # Gemini LLM (free tier)

src/lib/rag/embeddings.py

@@ -37,13 +37,20 @@ def _load_model(self) -> SentenceTransformer:
 
         return EmbeddingClient._model_cache[model_name]
 
-    def embed(self, texts: list[str]) -> list[list[float]]:
-        """Generate embeddings for texts."""
+    def embed(self, texts: list[str], is_query: bool = False) -> list[list[float]]:
+        """Generate embeddings for texts.
+
+        Args:
+            texts: List of texts to embed.
+            is_query: If True, apply BGE query prefix for retrieval-optimized
+                     embedding. Documents should use is_query=False (no prefix).
+        """
         if not texts:
             return []
 
-        # BGE models need instruction prefix for retrieval
-        if "bge" in self.config.hf_model.lower():
+        # BGE models: query gets prefix, documents stay raw
+        # Official BGE usage: only query uses instruction
+        if is_query and "bge" in self.config.hf_model.lower():
             texts = [
                 "Represent this sentence for searching relevant passages: " + t
                 for t in texts
@@ -60,12 +67,12 @@ def embed(self, texts: list[str]) -> list[list[float]]:
             logger.error("Embedding generation failed: %s", exc)
             raise
 
-    def embed_single(self, text: str) -> list[float]:
+    def embed_single(self, text: str, is_query: bool = False) -> list[float]:
         """Embed single text."""
-        return self.embed([text])[0]
+        return self.embed([text], is_query=is_query)[0]
 
 
 def get_embedding(text: str, config: RAGConfig | None = None) -> list[float]:
     """Convenience function."""
     client = EmbeddingClient(config)
-    return client.embed_single(text)
+    return client.embed_single(text, is_query=True)

src/lib/rag/retrieval.py

@@ -31,7 +31,8 @@ def similarity_search(
     client = QdrantClient(url=config.qdrant_url)
     embedder = EmbeddingClient(config)
 
-    query_vector = embedder.embed_single(query)
+    # Use query-specific embedding for better retrieval quality
+    query_vector = embedder.embed_single(query, is_query=True)
 
     try:
         result = client.query_points(

src/lib/rag/tests/test_api.py

@@ -0,0 +1,53 @@
+"""Tests for RAG Public API."""
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from lib.rag.api import ask_about_repo, get_repo_context
+
+
+class TestAskAboutRepo:
+    """Primary integration point for Orchestrator Chat."""
+
+    @patch("lib.rag.api.ask_repo")
+    def test_ask_about_repo_delegates(self, mock_ask_repo):
+        mock_ask_repo.return_value = "It uses FastAPI."
+
+        result = ask_about_repo(
+            job_id="550e8400-e29b-41d4-a716-446655440000",
+            question="What framework?",
+            top_k=3,
+        )
+
+        assert result == "It uses FastAPI."
+        mock_ask_repo.assert_called_once_with(
+            "What framework?",
+            "550e8400-e29b-41d4-a716-446655440000",
+            3,
+            None,
+        )
+
+    @patch("lib.rag.api.ask_repo")
+    def test_ask_about_repo_no_context(self, mock_ask_repo):
+        mock_ask_repo.return_value = "No relevant context found for this repository."
+
+        result = ask_about_repo(job_id="test-job", question="random?")
+
+        assert "No relevant context" in result
+
+
+class TestGetRepoContext:
+    """Raw context retrieval for debugging."""
+
+    @patch("lib.rag.api.retrieve_context")
+    def test_get_repo_context_delegates(self, mock_retrieve):
+        mock_retrieve.return_value = "[Source: README.md]\nThis is a FastAPI project."
+
+        result = get_repo_context(
+            job_id="test-job",
+            question="What framework?",
+            top_k=5,
+        )
+
+        assert "FastAPI" in result
+        mock_retrieve.assert_called_once_with("What framework?", "test-job", 5, None)

src/lib/rag/tests/test_embeddings.py

@@ -31,15 +31,15 @@ def test_model_caching(self, mock_st_class):
         assert client1.model is client2.model
 
     @patch("lib.rag.embeddings.SentenceTransformer")
-    def test_bge_prefix_added(self, mock_st_class):
-        """BGE models need instruction prefix."""
+    def test_bge_query_prefix_added(self, mock_st_class):
+        """BGE queries need instruction prefix."""
         mock_model = MagicMock()
         mock_model.encode.return_value = np.array([[0.1] * 768])
         mock_st_class.return_value = mock_model
 
         config = RAGConfig(hf_model="BAAI/bge-base-en-v1.5")
         client = EmbeddingClient(config)
-        client.embed(["hello world"])
+        client.embed(["hello world"], is_query=True)
 
         call_args = mock_model.encode.call_args
         texts = call_args[0][0]
@@ -48,6 +48,23 @@ def test_bge_prefix_added(self, mock_st_class):
             for t in texts
         )
 
+    @patch("lib.rag.embeddings.SentenceTransformer")
+    def test_bge_document_no_prefix(self, mock_st_class):
+        """BGE documents should NOT get prefix — only queries do."""
+        mock_model = MagicMock()
+        mock_model.encode.return_value = np.array([[0.1] * 768])
+        mock_st_class.return_value = mock_model
+
+        config = RAGConfig(hf_model="BAAI/bge-base-en-v1.5")
+        client = EmbeddingClient(config)
+        client.embed(["hello world"], is_query=False)
+
+        call_args = mock_model.encode.call_args
+        texts = call_args[0][0]
+        assert not any(t.startswith("Represent") for t in texts)
+        # Raw text preserved
+        assert texts[0] == "hello world"
+
     @patch("lib.rag.embeddings.SentenceTransformer")
     def test_non_bge_no_prefix(self, mock_st_class):
         """Non-BGE models should not get prefix."""
@@ -57,7 +74,7 @@ def test_non_bge_no_prefix(self, mock_st_class):
 
         config = RAGConfig(hf_model="sentence-transformers/all-MiniLM-L6-v2")
         client = EmbeddingClient(config)
-        client.embed(["hello world"])
+        client.embed(["hello world"], is_query=True)
 
         call_args = mock_model.encode.call_args
         texts = call_args[0][0]
@@ -84,7 +101,6 @@ def test_embed_single(self, mock_st_class):
 
         assert isinstance(result, list)
         assert len(result) == 768
-        # Value is normalized by encode(); just check it's a valid float
         assert isinstance(result[0], float)
 
     @patch("lib.rag.embeddings.SentenceTransformer")