📝 Add docstrings to feature/ai/rag_&_Prompt_extraction

services/ai-service/main.py

@@ -49,7 +49,11 @@ def setup_framework(pdf_paths: str | list[str], framework_name: str) -> list[Pat
 
 
 def main():
-    """Main entry point."""
+    """
+    Prints a CLI header, shows configured input/output directories, and displays a usage example for the framework extraction tool.
+
+    The function retrieves input paths via get_input_paths(), prints a banner and the locations for framework and applicant PDFs and framework outputs, and prints a sample call for setup_framework.
+    """
     print("\n" + "=" * 60)
     print("Compliance Framework Extraction System")
     print("=" * 60)
@@ -68,4 +72,4 @@ def main():
 if __name__ == "__main__":
     framework_path = "data/inputs/frameworks/NDI"
     framework_name = "NDI"
-    setup_framework(framework_path, framework_name)
+    setup_framework(framework_path, framework_name)

services/ai-service/src/api/app.py

@@ -29,6 +29,21 @@
 
 @app.exception_handler(ExtractionError)
 async def extraction_error_handler(request, exc: ExtractionError):
+    """
+    Handle ExtractionError exceptions by returning a standardized JSON error response.
+
+    Parameters:
+        request: The incoming HTTP request that triggered the exception.
+        exc (ExtractionError): The extraction error containing `message` and `framework_name`.
+
+    Returns:
+        JSONResponse: Response with status code 500 and JSON body:
+            {
+                "error": "extraction_failed",
+                "message": <exc.message>,
+                "detail": <exc.framework_name>
+            }
+    """
     return JSONResponse(
         status_code=500,
         content={
@@ -41,6 +56,21 @@ async def extraction_error_handler(request, exc: ExtractionError):
 
 @app.exception_handler(ValueError)
 async def value_error_handler(request, exc: ValueError):
+    """
+    Convert a ValueError into a standardized HTTP 400 JSON response.
+
+    Parameters:
+        request: The incoming HTTP request that triggered the error.
+        exc (ValueError): The exception whose string representation is used as the response `message`.
+
+    Returns:
+        JSONResponse: A response with status code 400 and JSON body:
+            {
+                "error": "bad_request",
+                "message": <str(exc)>,
+                "detail": None
+            }
+    """
     return JSONResponse(
         status_code=400,
         content={
@@ -53,4 +83,4 @@ async def value_error_handler(request, exc: ValueError):
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)

services/ai-service/src/api/routers/frameworks.py

@@ -18,9 +18,20 @@ async def setup_framework(
     files: list[UploadFile] = File(..., description="PDF files (same order as section_names)"),
 ) -> SetupFrameworkResponse:
     """
-    Setup a new framework by extracting controls from multiple PDFs.
-    Each PDF has a section name; controls are saved as config/frameworks/{framework_name}/{section_name}.json.
-    Send section_names as one string: comma-separated names, one per file, in the same order as files.
+    Initialize a framework by extracting controls from multiple uploaded PDF sections.
+
+    Section names must be provided as a single comma-separated string with one name per file, in the same order as the uploaded files.
+
+    Parameters:
+        framework_name: Framework identifier.
+        section_names: Comma-separated section names (one name per PDF, same order as `files`).
+        files: Uploaded PDF files corresponding to `section_names`.
+
+    Returns:
+        SetupFrameworkResponse: Summary of the created framework, including `framework_name`, `total_controls`, `sections`, and `created_at`.
+
+    Raises:
+        HTTPException: Status 400 for input validation errors (mismatched counts between section names and files, empty or non-PDF uploads, empty file bodies, or empty section names).
     """
     section_names_list = [s.strip() for s in section_names.split(",") if s.strip()]
     if len(section_names_list) != len(files):
@@ -60,4 +71,4 @@ async def setup_framework(
         total_controls=result["total_controls"],
         sections=[ControlSummary(**s) for s in result["sections"]],
         created_at=result["created_at"],
-    )
+    )

services/ai-service/src/api/routers/health.py

@@ -7,5 +7,10 @@
 
 @router.get("/health")
 def health() -> dict[str, str]:
-    """Basic health check."""
-    return {"status": "ok"}
+    """
+    Return a basic service health status.
+
+    Returns:
+        dict[str, str]: A dictionary containing {"status": "ok"} when the service is healthy.
+    """
+    return {"status": "ok"}

services/ai-service/src/core/evaluator.py

@@ -18,15 +18,18 @@ def evaluate_applicant(
     controls_json: Dict[str, Any]
 ) -> Dict[str, Any]:
     """
-    Evaluate applicant documents using the saved evaluation prompt and controls.
+    Evaluate applicant documents against a saved evaluation prompt and controls to produce a structured evaluation report.
+
+    Parameters:
+        applicant_docs (List[str]): Text contents of applicant documents to evaluate; documents are combined in order and may be truncated if excessively long.
+        evaluation_prompt (str): The saved evaluation prompt that defines evaluation criteria and instructions for the model.
+        controls_json (Dict[str, Any]): Reference controls/framework as a JSON-serializable dictionary used for context during evaluation.
 
-    Args:
-        applicant_docs: List of text content from applicant documents
-        evaluation_prompt: Saved evaluation prompt from setup phase
-        controls_json: Saved controls JSON from setup phase
-
     Returns:
-        Dictionary containing evaluation report with scores/compliance status
+        Dict[str, Any]: The evaluation report parsed from the model's response. On successful JSON parsing this is the structured report (scores, compliance status, etc.). If the model returns non-JSON text that cannot be parsed, returns a fallback dictionary with keys `report_type` set to `"text"`, `content` containing the raw response, and `error` describing the JSON parse failure.
+
+    Raises:
+        ValueError: If the OPENROUTER_API_KEY environment variable is missing, if the API response is missing or empty, or if the model returns empty content.
     """
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
@@ -107,4 +110,4 @@ def evaluate_applicant(
                 "report_type": "text",
                 "content": result_text,
                 "error": f"Could not parse as JSON: {e}"
-            }
+            }

services/ai-service/src/core/framework_extraction.py

@@ -20,15 +20,20 @@ def extract_controls_from_framework(
     pdf_text: str, framework_name: str, use_fallback_prompt: bool = False
 ) -> Dict[str, Any]:
     """
-    Extract compliance controls from PDF text using LLM.
-
-    Args:
-        pdf_text: Extracted text from framework PDF
-        framework_name: Name of the framework
-        use_fallback_prompt: If True, use a stricter prompt that forbids empty controls (for retries).
-
+    Extract compliance controls from text extracted from a framework PDF.
+
+    Parameters:
+        pdf_text (str): Full text extracted from the PDF to be analyzed.
+        framework_name (str): Name of the framework being processed; included in the returned result.
+        use_fallback_prompt (bool): When True, uses a stricter retry prompt that enforces extracting at least one control
+            (used for retry attempts when initial extraction yields no controls).
+
     Returns:
-        Dictionary with controls array
+        result (Dict[str, Any]): Dictionary with keys:
+            - "framework_name" (str): The provided framework_name.
+            - "controls" (List[Dict[str, Any]]): List of extracted control objects; each control is a dict expected to
+              contain fields such as `id`, `description`, `calculation`, `threshold`, and `scale`. The list is empty
+              when no extractable controls are found or parsing fails.
     """
     api_key = os.getenv("OPENROUTER_API_KEY")
     if not api_key:
@@ -187,20 +192,25 @@ def extract_controls_from_pdfs(
     pdf_paths_list: List[str], framework_name: str
 ) -> List[List[Dict[str, Any]]]:
     """
-    Extract controls from multiple PDFs in parallel.
-    Each PDF gets one LLM call.
-
-    Args:
-        pdf_paths_list: List of PDF file paths
-        framework_name: Name of the framework
-
+    Extract controls from each PDF path in pdf_paths_list using the LLM in parallel.
+
     Returns:
-        List of controls arrays (one per PDF)
+        List[List[Dict[str, Any]]]: A list where each element is the extracted controls list for the corresponding PDF in pdf_paths_list. Empty list for PDFs that failed or produced no controls.
     """
     MAX_RETRIES = 3
 
     def extract_pdf(pdf_path: str) -> List[Dict[str, Any]]:
-        """Extract controls from a single PDF. Retries with fallback prompt if empty."""
+        """
+        Extracts compliance controls from a single PDF file.
+
+        Attempts to extract text from the PDF and parse controls for the current framework. If the initial extraction yields no controls, retries up to MAX_RETRIES using a stricter fallback prompt to force extraction. Returns an empty list if no text is found, if extraction fails after retries, or if an exception occurs.
+
+        Parameters:
+            pdf_path (str): Path to the PDF file to process.
+
+        Returns:
+            List[Dict[str, Any]]: A list of extracted control objects (possibly empty).
+        """
         try:
             pdf_text = extract_text_from_pdf(pdf_path)
             if not (pdf_text and pdf_text.strip()):
@@ -225,4 +235,4 @@ def extract_pdf(pdf_path: str) -> List[Dict[str, Any]]:
     with ThreadPoolExecutor(max_workers=len(pdf_paths_list)) as executor:
         controls_arrays = list(executor.map(extract_pdf, pdf_paths_list))
 
-    return controls_arrays
+    return controls_arrays

services/ai-service/src/embeddings/gemma_embedder.py

@@ -28,14 +28,12 @@ def __init__(
         token: Optional[str] = None
     ):
         """
-        Initialize EmbeddingGemma embedder.
+        Create a GemmaEmbedder configured to load the specified HuggingFace embedding model.
 
-        Args:
-            model_name: HuggingFace model name (default: google/embeddinggemma-300m)
-                       This is the 300M parameter embedding model designed for embeddings.
-                       Model page: https://huggingface.co/google/embeddinggemma-300m
-            device: Device to use ('cpu', 'cuda', or None for auto-detection)
-            token: HuggingFace token for gated models (or set HF_TOKEN env var)
+        Parameters:
+            model_name (str): HuggingFace model identifier to load (default: "google/embeddinggemma-300m").
+            device (Optional[str]): Target device ("cpu", "cuda"), or None to auto-detect CUDA if available then "cpu".
+            token (Optional[str]): HuggingFace access token for gated models; if not provided, the `HF_TOKEN` or `HUGGINGFACE_TOKEN` environment variable is used.
         """
         self.model_name = model_name
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -44,7 +42,17 @@ def __init__(
         self._load_model()
 
     def _load_model(self):
-        """Load the EmbeddingGemma model using sentence-transformers."""
+        """
+        Load and initialize the SentenceTransformer embedding model, handling offline mode and optional HuggingFace authentication.
+
+        This method:
+        - Respects HF_HUB_OFFLINE (environment) to load local files only.
+        - If a token is provided and login has not yet been performed, performs a one-time HuggingFace login.
+        - Builds model loading kwargs (including token when appropriate) and instantiates SentenceTransformer for self.model.
+
+        Raises:
+            RuntimeError: If the model cannot be loaded. If the failure appears related to gated access or HTTP 401/403, the exception message will include actionable steps to obtain a HuggingFace token and accept the model license.
+        """
         global _LOGIN_DONE
         offline = os.getenv("HF_HUB_OFFLINE", "").strip().lower() == "1"
         local_files_only = offline
@@ -91,26 +99,27 @@ def _load_model(self):
 
     def embed_text(self, text: str) -> List[float]:
         """
-        Generate embedding for a single text.
+        Produce an embedding for a single input string.
+
+        Parameters:
+            text (str): The input text to embed.
 
-        Args:
-            text: Text to embed
-
         Returns:
-            Embedding vector as list of floats
+            List[float]: Embedding vector for the provided text.
         """
         return self.embed_batch([text])[0]
 
     def embed_batch(self, texts: List[str], batch_size: int = 32) -> List[List[float]]:
         """
-        Generate embeddings for a batch of texts.
+        Produce embeddings for a list of texts.
+
+        Parameters:
+            texts (List[str]): Input texts to convert into embeddings.
+            batch_size (int): Maximum number of texts processed at once.
 
-        Args:
-            texts: List of texts to embed
-            batch_size: Batch size for processing
-
         Returns:
-            List of embedding vectors
+            List[List[float]]: A list where each element is the embedding vector (list of floats)
+            corresponding to the input text at the same index.
         """
         if not texts:
             return []
@@ -128,10 +137,13 @@ def embed_batch(self, texts: List[str], batch_size: int = 32) -> List[List[float
 
     def get_embedding_dim(self) -> int:
         """
-        Get the dimension of embeddings produced by this model.
+        Get the dimensionality of embedding vectors produced by the loaded model.
 
         Returns:
-            Embedding dimension (768 for EmbeddingGemma-300M)
+            int: Number of dimensions in each embedding vector (for example, 768 for google/embeddinggemma-300m).
+
+        Raises:
+            RuntimeError: If the underlying model is not loaded.
         """
         if self.model is None:
             raise RuntimeError("Model not loaded")
@@ -145,14 +157,14 @@ def load_gemma_embedder(
     token: Optional[str] = None
 ) -> GemmaEmbedder:
     """
-    Load and return a Gemma embedder instance.
+    Create a GemmaEmbedder configured for the specified model, device, and HuggingFace token.
+
+    Parameters:
+        model_name (str): HuggingFace model identifier to load (default: "google/embeddinggemma-300m").
+        device (Optional[str]): Device to run the model on; use "cpu", "cuda", or None to auto-detect.
+        token (Optional[str]): HuggingFace access token for gated models; if omitted, the function will read HF_TOKEN or HUGGINGFACE_TOKEN from the environment.
 
-    Args:
-        model_name: HuggingFace model name (default: google/embeddinggemma-300m)
-        device: Device to use ('cpu', 'cuda', or None for auto-detection)
-        token: HuggingFace token for gated models (or set HF_TOKEN env var)
-
     Returns:
-        GemmaEmbedder instance
+        GemmaEmbedder: An initialized GemmaEmbedder instance ready to produce embeddings.
     """
-    return GemmaEmbedder(model_name=model_name, device=device, token=token)
+    return GemmaEmbedder(model_name=model_name, device=device, token=token)