diff --git a/services/ai-service/main.py b/services/ai-service/main.py index c11249c..bc2fdb7 100644 --- a/services/ai-service/main.py +++ b/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) \ No newline at end of file diff --git a/services/ai-service/src/api/app.py b/services/ai-service/src/api/app.py index df6d238..e4dcd7e 100644 --- a/services/ai-service/src/api/app.py +++ b/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": , + "detail": + } + """ 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": , + "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) \ No newline at end of file diff --git a/services/ai-service/src/api/routers/frameworks.py b/services/ai-service/src/api/routers/frameworks.py index 9e64447..93b968b 100644 --- a/services/ai-service/src/api/routers/frameworks.py +++ b/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"], - ) + ) \ No newline at end of file diff --git a/services/ai-service/src/api/routers/health.py b/services/ai-service/src/api/routers/health.py index d1b4da6..bcf8f6e 100644 --- a/services/ai-service/src/api/routers/health.py +++ b/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"} \ No newline at end of file diff --git a/services/ai-service/src/core/evaluator.py b/services/ai-service/src/core/evaluator.py index f770b63..b80ca41 100644 --- a/services/ai-service/src/core/evaluator.py +++ b/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}" - } + } \ No newline at end of file diff --git a/services/ai-service/src/core/framework_extraction.py b/services/ai-service/src/core/framework_extraction.py index 60cdbad..0484eba 100644 --- a/services/ai-service/src/core/framework_extraction.py +++ b/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 \ No newline at end of file diff --git a/services/ai-service/src/embeddings/gemma_embedder.py b/services/ai-service/src/embeddings/gemma_embedder.py index 8d7b06b..486307e 100644 --- a/services/ai-service/src/embeddings/gemma_embedder.py +++ b/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) \ No newline at end of file diff --git a/services/ai-service/src/embeddings/haystack_retriever.py b/services/ai-service/src/embeddings/haystack_retriever.py index dc7c1b0..13f256e 100644 --- a/services/ai-service/src/embeddings/haystack_retriever.py +++ b/services/ai-service/src/embeddings/haystack_retriever.py @@ -20,13 +20,13 @@ def __init__( top_k: int = 5 ): """ - Initialize Haystack retriever with Qdrant. + Create a Haystack-compatible retriever backed by a Qdrant collection and a GemmaEmbedder. - Args: - qdrant_client: QdrantClient instance - collection_name: Name of Qdrant collection - embedder: GemmaEmbedder instance - top_k: Number of documents to retrieve + Parameters: + qdrant_client (QdrantClient): Active Qdrant client used to query and scroll the specified collection. + collection_name (str): Name of the Qdrant collection to load documents from and search. + embedder (GemmaEmbedder): Embedder used to convert queries into vector embeddings. + top_k (int): Default number of nearest documents to return for retrieval operations. """ self.qdrant_client = qdrant_client self.collection_name = collection_name @@ -38,7 +38,11 @@ def __init__( self._load_documents_from_qdrant() def _load_documents_from_qdrant(self): - """Load all documents from Qdrant into Haystack format.""" + """ + Load all documents from the configured Qdrant collection into the retriever's in-memory Haystack Document list. + + Each loaded Document's content is taken from the payload field "text" (defaults to an empty string). The Document meta includes "chunk_id", "framework_name" (defaults to "unknown"), and any other payload fields except "text", "chunk_id", and "framework_name". If an error occurs while loading, the in-memory documents list is reset to an empty list. + """ try: scroll_result = self.qdrant_client.scroll( collection_name=self.collection_name, @@ -64,19 +68,17 @@ def _load_documents_from_qdrant(self): def retrieve_documents(self, query: str, top_k: Optional[int] = None) -> List[Dict]: """ - Retrieve documents using Haystack with custom Gemma embedder. + Retrieve the top-k nearest documents for a text query using the Gemma embedder and Qdrant. + + Parameters: + query (str): Query text to embed and search. + top_k (Optional[int]): If provided, overrides the instance default number of results to return. - Args: - query: Query text - top_k: Number of documents to retrieve (overrides default) - Returns: - List of dictionaries with retrieved documents: - { - "text": str, - "score": float, - "metadata": dict - } + List[dict]: A list of result dictionaries with keys: + - "text" (str): The document text. + - "score" (float): Similarity score for the result. + - "metadata" (dict): Associated metadata for the document. """ if top_k is None: top_k = self.top_k @@ -104,16 +106,10 @@ def create_retrieval_pipeline( top_k: int = 5 ) -> HaystackQdrantRetriever: """ - Create a retrieval pipeline with Haystack and Qdrant. + Create a HaystackQdrantRetriever configured with the given Qdrant client, collection, embedder, and top_k. - Args: - qdrant_client: QdrantClient instance - collection_name: Name of Qdrant collection - embedder: GemmaEmbedder instance - top_k: Number of documents to retrieve - Returns: - HaystackQdrantRetriever instance + A HaystackQdrantRetriever configured to query the specified Qdrant collection. """ return HaystackQdrantRetriever( qdrant_client=qdrant_client, @@ -129,14 +125,14 @@ def retrieve_documents( top_k: Optional[int] = None ) -> List[Dict]: """ - Retrieve documents for a query. + Retrieve documents matching the query using the provided retriever. + + Parameters: + retriever (HaystackQdrantRetriever): Retriever instance to perform the search. + query (str): The text query to embed and search. + top_k (Optional[int]): Maximum number of results to return; if None, use the retriever's default. - Args: - retriever: HaystackQdrantRetriever instance - query: Query text - top_k: Number of documents to retrieve - Returns: - List of retrieved document dictionaries + List[Dict]: A list of result dictionaries, each containing at least the keys `text`, `score`, and `metadata`. """ - return retriever.retrieve_documents(query, top_k=top_k) + return retriever.retrieve_documents(query, top_k=top_k) \ No newline at end of file diff --git a/services/ai-service/src/embeddings/qdrant_manager.py b/services/ai-service/src/embeddings/qdrant_manager.py index 14435b7..fd2642e 100644 --- a/services/ai-service/src/embeddings/qdrant_manager.py +++ b/services/ai-service/src/embeddings/qdrant_manager.py @@ -25,16 +25,18 @@ def initialize_qdrant( url: Optional[str] = None ) -> QdrantClient: """ - Initialize Qdrant client and create collection if it doesn't exist. + Initialize and return a configured Qdrant client and ensure the specified collection exists. + + If neither `path` nor `url` is provided, a default local directory under the project config (config/vector_db) is used. If `url` is provided the client targets a remote Qdrant server; otherwise a local client is used. If the named collection is missing, it will be created with vectors of size `vector_size` using cosine distance. + + Parameters: + collection_name (str): Name of the Qdrant collection to use or create. + vector_size (int): Dimensionality of vectors stored in the collection. + path (Optional[str]): Local filesystem path for a local Qdrant instance; when omitted and `url` is not provided, a default config/vector_db path is used. + url (Optional[str]): Remote Qdrant server URL; when provided the client will connect remotely. - Args: - collection_name: Name of the collection - vector_size: Size of the embedding vectors - path: Local path for Qdrant (default: config/vector_db) - url: Qdrant server URL (for remote/cloud) - Returns: - QdrantClient instance + QdrantClient: A Qdrant client configured for the requested collection. """ # Default to local path if neither path nor url provided if path is None and url is None: @@ -74,12 +76,13 @@ def create_collection( vector_size: int ) -> None: """ - Create a new collection in Qdrant. + Create a Qdrant collection configured for cosine similarity. + + Creates a collection named `collection_name` with vectors of length `vector_size` and sets the distance metric to COSINE. - Args: - client: QdrantClient instance - collection_name: Name of the collection - vector_size: Size of the embedding vectors + Parameters: + collection_name (str): Name of the collection to create. + vector_size (int): Length of embedding vectors to store in the collection. """ client.create_collection( collection_name=collection_name, @@ -97,13 +100,22 @@ def add_documents( embeddings: List[List[float]] ) -> None: """ - Add documents with embeddings to Qdrant collection. + Upsert document chunks and their embeddings into the specified Qdrant collection. - Args: - client: QdrantClient instance - collection_name: Name of the collection - documents: List of document dictionaries (from chunker) - embeddings: List of embedding vectors + Constructs a payload for each document and inserts points into the collection in batches. Each point ID is derived deterministically from the document's `chunk_id` using a fixed namespace, ensuring stable identifiers across runs. + + Parameters: + client (QdrantClient): Qdrant client instance to use for upserts. + collection_name (str): Target Qdrant collection name. + documents (List[Dict]): List of document chunk dictionaries. Expected keys: + - "text" (str): Chunk text (optional, defaults to empty string). + - "chunk_id" (str): Chunk identifier (optional; a UUID will be generated if missing). + - "framework_name" (str): Origin framework name (optional, defaults to "unknown"). + - "metadata" (dict): Additional payload fields to include (optional). + embeddings (List[List[float]]): Corresponding list of embedding vectors for each document. + + Raises: + ValueError: If the number of documents does not match the number of embeddings. """ if len(documents) != len(embeddings): raise ValueError(f"Number of documents ({len(documents)}) must match number of embeddings ({len(embeddings)})") @@ -146,25 +158,21 @@ def search_similar( score_threshold: Optional[float] = None ) -> List[Dict]: """ - Search for similar documents in Qdrant collection. + Finds nearest documents in the specified Qdrant collection for a given query embedding. - For local Qdrant, uses query_points with Query object. - This is the correct method for local mode. + Filters results by an optional minimum similarity score. + + Parameters: + collection_name (str): Name of the Qdrant collection to query. + query_embedding (List[float]): Embedding vector used as the search query. + top_k (int): Maximum number of results to return. + score_threshold (Optional[float]): Minimum similarity score required for returned results. - Args: - client: QdrantClient instance - collection_name: Name of the collection - query_embedding: Query embedding vector - top_k: Number of results to return - score_threshold: Minimum similarity score threshold - Returns: - List of dictionaries with search results: - { - "text": str, - "score": float, - "metadata": dict - } + List[Dict]: A list of result dictionaries. Each dictionary contains: + - "text" (str): The stored document text. + - "score" (float): The similarity score for the result. + - "metadata" (dict): Payload metadata including "chunk_id", "framework_name", and any additional fields. """ try: # Use query_points - can accept vector directly as list[float] @@ -256,11 +264,17 @@ def fetch_by_filter( limit: int = 1000, ) -> List[Dict[str, Any]]: """ - Fetch points matching a filter (no vector search). Used for exact lookups - e.g. control_id + source=json. - + Retrieve points that match the given payload filter from a Qdrant collection. + + Parameters: + query_filter (Filter): Payload filter to apply when selecting points. + limit (int): Maximum number of points to return (default 1000). + Returns: - List of {"text": str, "metadata": dict} for each matching point. + List[Dict[str, Any]]: A list of dictionaries for each matching point with keys: + - "text" (str): The stored text (empty string if missing). + - "metadata" (dict): Metadata dictionary containing "chunk_id", "framework_name" + (defaults to "unknown" if missing), and any other payload fields. """ results, _ = client.scroll( collection_name=collection_name, @@ -293,10 +307,21 @@ def search_similar_filtered( score_threshold: Optional[float] = None, ) -> List[Dict[str, Any]]: """ - Vector similarity search with optional payload filter. - + Perform a vector similarity search in a Qdrant collection using a query embedding with optional payload filtering. + + Parameters: + client (QdrantClient): Qdrant client used to run the query. + collection_name (str): Name of the collection to search. + query_embedding (List[float]): Query vector used for nearest-neighbor search. + top_k (int): Maximum number of results to return. + query_filter (Optional[Filter]): Optional payload filter to restrict returned points. + score_threshold (Optional[float]): Optional minimum score threshold to include a result. + Returns: - List of {"text": str, "score": float, "metadata": dict}. + List[Dict[str, Any]]: A list of result dictionaries. Each dictionary contains: + - `text` (str): The stored text payload for the point (empty string if missing). + - `score` (float): The similarity score for the match. + - `metadata` (dict): Payload fields with at least `chunk_id` and `framework_name` (defaults to "unknown"), plus any other payload keys. """ kwargs: Dict[str, Any] = { "collection_name": collection_name, @@ -335,10 +360,6 @@ def delete_collection( collection_name: str ) -> None: """ - Delete a collection from Qdrant. - - Args: - client: QdrantClient instance - collection_name: Name of the collection to delete + Delete the specified collection from the Qdrant instance. """ - client.delete_collection(collection_name=collection_name) + client.delete_collection(collection_name=collection_name) \ No newline at end of file diff --git a/services/ai-service/src/processing/pdf_parser.py b/services/ai-service/src/processing/pdf_parser.py index 8109931..f6f092f 100644 --- a/services/ai-service/src/processing/pdf_parser.py +++ b/services/ai-service/src/processing/pdf_parser.py @@ -9,17 +9,17 @@ def extract_text_from_pdf(pdf_path: str) -> str: """ - Extract text from a PDF file with support for multilingual content. + Extracts and returns text from the PDF at the given path, preserving page breaks as two newline separators. + + Parameters: + pdf_path (str): Path to the PDF file to extract. - Args: - pdf_path: Path to the PDF file - Returns: - Extracted text as a string - + str: Concatenated text extracted from all pages, with pages separated by two newline characters. + Raises: - FileNotFoundError: If PDF file doesn't exist - ValueError: If PDF is corrupted or cannot be read + FileNotFoundError: If the file at `pdf_path` does not exist. + ValueError: If no text could be extracted from the PDF. """ pdf_path_obj = Path(pdf_path) @@ -40,4 +40,4 @@ def extract_text_from_pdf(pdf_path: str) -> str: if not text_content: raise ValueError(f"No text could be extracted from PDF: {pdf_path}") - return "\n\n".join(text_content) + return "\n\n".join(text_content) \ No newline at end of file diff --git a/services/ai-service/src/processing/text_chunker.py b/services/ai-service/src/processing/text_chunker.py index d23ac15..cf38d86 100644 --- a/services/ai-service/src/processing/text_chunker.py +++ b/services/ai-service/src/processing/text_chunker.py @@ -15,30 +15,24 @@ def chunk_text( preserve_sentences: bool = True ) -> List[Dict]: """ - Split text into fixed-size chunks with overlap, preserving semantic boundaries. - - Improved version that: - - Uses larger default chunk sizes (1500 chars) for better context - - Preserves paragraph boundaries when possible - - Filters out chunks that are mostly formatting artifacts - - Better sentence boundary detection with larger lookback window - - Args: - text: Text to chunk - chunk_size: Size of each chunk in characters (default: 1500) - overlap: Number of characters to overlap between chunks (default: 200) - framework_name: Optional framework name for metadata - preserve_sentences: If True, try to break at sentence boundaries - + Split long text into overlapping chunks that prefer paragraph or sentence boundaries. + + The function produces sequential text segments of up to `chunk_size` characters with `overlap` between adjacent segments. When `preserve_sentences` is True the function attempts to end chunks at paragraph breaks (double newlines) or sentence boundaries within a lookback window; resulting chunks that are mostly formatting artifacts are discarded. + + Parameters: + text (str): Input text to split. + chunk_size (int): Maximum number of characters per chunk (default 1500). + overlap (int): Number of characters to overlap between consecutive chunks (default 200). + framework_name (Optional[str]): Optional label included in each chunk's metadata; "unknown" if omitted. + preserve_sentences (bool): If True, prefer paragraph or sentence boundaries when choosing chunk end positions. + Returns: - List of dictionaries, each containing: - { - "text": str, - "chunk_id": int, - "chunk_index": int, - "framework_name": str, - "metadata": dict - } + List[dict]: A list of chunk dictionaries. Each dictionary contains: + - "text" (str): Chunk content. + - "chunk_id" (int): Sequential chunk identifier. + - "chunk_index" (int): Alias of chunk_id. + - "framework_name" (str): Provided framework name or "unknown". + - "metadata" (dict): Includes "start_char", "end_char", "length", "chunk_size", and "overlap". """ if not text or len(text.strip()) == 0: return [] @@ -133,14 +127,16 @@ def chunk_text( def _is_valid_chunk(text: str, min_meaningful_chars: int = 100) -> bool: """ - Check if a chunk is valid (not mostly formatting artifacts). + Determine whether a text chunk contains sufficient meaningful content versus formatting artifacts. + + Accepts a chunk only if it has at least one meaningful line, the ratio of meaningful characters to total characters is at least 0.3, and the total meaningful characters are at least min_meaningful_chars. Chunks shorter than 20 characters are rejected. + + Parameters: + text: Chunk text to evaluate. + min_meaningful_chars: Minimum count of meaningful characters required for acceptance. - Args: - text: Chunk text to validate - min_meaningful_chars: Minimum number of meaningful characters required - Returns: - True if chunk is valid, False if it's mostly formatting + True if the chunk is considered meaningful, False otherwise. """ if not text or len(text) < 20: return False @@ -195,15 +191,22 @@ def chunk_text_by_sentences( framework_name: Optional[str] = None ) -> List[Dict]: """ - Chunk text by sentences instead of fixed character size. + Split text into chunks where each chunk contains a fixed number of sentences. + + Sentences are detected by runs of sentence-ending punctuation followed by whitespace (e.g., ". ", "! ", "? ") or by double newlines; empty fragments are ignored. + + Parameters: + text (str): Input text to split. + sentences_per_chunk (int): Number of sentences to include in each chunk. + framework_name (Optional[str]): Optional label stored in each chunk's `framework_name` field; defaults to `"unknown"` when not provided. - Args: - text: Text to chunk - sentences_per_chunk: Number of sentences per chunk - framework_name: Optional framework name for metadata - Returns: - List of chunk dictionaries + List[Dict]: A list of chunk dictionaries. Each dictionary contains: + - text: chunk text (joined sentences). + - chunk_id: numeric chunk identifier. + - chunk_index: numeric chunk index (same as `chunk_id`). + - framework_name: provided framework name or `"unknown"`. + - metadata: dict with `sentence_start`, `sentence_end`, `sentences_per_chunk`, and `length` (character count of the chunk text). """ # Split into sentences sentence_pattern = r'[.!?]+\s+|[\n]{2,}' @@ -232,4 +235,4 @@ def chunk_text_by_sentences( chunks.append(chunk_dict) chunk_id += 1 - return chunks + return chunks \ No newline at end of file diff --git a/services/ai-service/src/rag/_shared.py b/services/ai-service/src/rag/_shared.py index 3f3ee09..aa5df5d 100644 --- a/services/ai-service/src/rag/_shared.py +++ b/services/ai-service/src/rag/_shared.py @@ -11,8 +11,13 @@ def get_shared_embedder(): - """Return a single embedder instance, creating and caching on first use.""" + """ + Return the shared, per-process embedder instance, creating and caching it on first use. + + Returns: + The embedder instance cached for the process; created and stored on the first call. + """ global _cached_embedder if _cached_embedder is None: _cached_embedder = load_gemma_embedder() - return _cached_embedder + return _cached_embedder \ No newline at end of file diff --git a/services/ai-service/src/rag/ingestion.py b/services/ai-service/src/rag/ingestion.py index 9a69fde..8857439 100644 --- a/services/ai-service/src/rag/ingestion.py +++ b/services/ai-service/src/rag/ingestion.py @@ -17,13 +17,12 @@ def index_framework( pdf_paths: Optional[List[str]] = None, ) -> None: """ - Index a framework for RAG: JSON control cards + PDF chunks from vector_db input. - - - JSON cards: one per control from config/frameworks/{framework_name}/*.json. - - PDF chunks: from data/inputs/vector_db/{framework_name}/*.pdf, or vector_db/*.pdf - if no subdir. Use pdf_paths when provided instead. - - Collection name: {framework_name}_rag. + Index a framework's JSON control cards and PDF text chunks into a Qdrant collection for RAG. + + Scans JSON control cards under config/frameworks/{framework_name} and converts each control into a textual document; extracts and chunks text from PDFs found under data/inputs/vector_db/{framework_name} (or the global vector_db directory) unless overridden by pdf_paths. Computes embeddings with the shared embedder and writes documents and embeddings to the collection named "{framework_name}_rag". If no documents are collected the function returns without modifying the vector store. + + Parameters: + pdf_paths (Optional[List[str]]): Optional explicit list of PDF file paths to index; when provided, these paths override automatic PDF discovery. """ embedder = get_shared_embedder() dim = embedder.get_embedding_dim() @@ -87,4 +86,4 @@ def index_framework( texts = [d["text"] for d in all_docs] embeddings = embedder.embed_batch(texts) - add_documents(client, collection, all_docs, embeddings) + add_documents(client, collection, all_docs, embeddings) \ No newline at end of file diff --git a/services/ai-service/src/rag/retrieval.py b/services/ai-service/src/rag/retrieval.py index 4a51e23..7a45b83 100644 --- a/services/ai-service/src/rag/retrieval.py +++ b/services/ai-service/src/rag/retrieval.py @@ -56,18 +56,17 @@ def retrieve_control_details( top_k_pdf: int = 5, ) -> Dict[str, Any]: """ - Retrieve all details for a control: JSON cards (filtered by control_id) + top-k - PDF chunks (semantic search). Safe to use as an agent tool. - - Args: - control_id: Control identifier (e.g. DSI.OE.01, DG.1). - framework_name: Framework name (e.g. NDI) matching the indexed collection. - top_k_pdf: Max PDF chunks to return. Default 5. - + Retrieve JSON control cards and relevant PDF passages for a given control identifier from a framework-specific RAG collection. + + Parameters: + control_id (str): Control identifier (e.g., "DSI.OE.01", "DG.1"). + framework_name (str): Name of the indexed framework/collection (e.g., "NDI"). + top_k_pdf (int): Maximum number of PDF chunks to return. + Returns: - { - "json_cards": [{"text": str, "metadata": dict}, ...], - "pdf_chunks": [{"text": str, "score": float, "metadata": dict}, ...], + dict: { + "json_cards": List[dict] — each dict contains `text` (str) and `metadata` (dict) for JSON-based control cards; + "pdf_chunks": List[dict] — each dict contains `text` (str), `score` (float), and `metadata` (dict) for retrieved PDF passages. } """ embedder = get_shared_embedder() @@ -106,4 +105,4 @@ def retrieve_control_details( return { "json_cards": json_cards, "pdf_chunks": pdf_chunks, - } + } \ No newline at end of file diff --git a/services/ai-service/src/services/framework_service.py b/services/ai-service/src/services/framework_service.py index 5a00395..f23e7bb 100644 --- a/services/ai-service/src/services/framework_service.py +++ b/services/ai-service/src/services/framework_service.py @@ -15,6 +15,14 @@ def _project_root() -> Path: + """ + Get the repository's project root directory. + + Resolves this file's location and returns the Path three levels above it. + + Returns: + Path: Path object pointing to the repository root directory. + """ return Path(__file__).resolve().parent.parent.parent @@ -27,21 +35,18 @@ def setup_framework( pdf_sections: list[tuple[str, bytes]], ) -> dict[str, Any]: """ - Process multiple PDFs with section names: extract controls, save JSON per section, and index into vector DB. - - Steps: - 1. Save uploaded PDFs to a temp directory with section names as filenames. - 2. Extract controls from each PDF in parallel (existing core logic). - 3. Save one JSON per section under config/frameworks/{framework_name}/{section_name}.json. - 4. Persist PDFs to data/inputs/vector_db/{framework_name}/ and call index_framework to populate vector DB. - 5. Return summary (framework_name, total_controls, sections, created_at). - - Args: - framework_name: Framework identifier. - pdf_sections: List of (section_name, pdf_bytes). Section name is used as the JSON filename. - + Orchestrates extraction of controls from labeled PDF sections, saves per-section JSON, indexes PDFs into the vector DB, and returns a summary. + + Parameters: + framework_name (str): Identifier used for saved files and vector DB organization. + pdf_sections (list[tuple[str, bytes]]): List of (section_name, pdf_bytes); each section_name is used as the JSON and PDF filename. + Returns: - Dict with keys: framework_name, total_controls, sections (list of {section_name, controls_count, json_path}), created_at. + dict: Summary containing: + - framework_name: the provided framework_name. + - total_controls: total number of extracted controls across all sections. + - sections: list of objects with keys `section_name`, `controls_count`, and `json_path`. + - created_at: UTC timestamp when the summary was created. """ if not pdf_sections: raise ValueError("At least one PDF section is required") @@ -98,4 +103,4 @@ def setup_framework( "created_at": datetime.utcnow(), } finally: - shutil.rmtree(temp_dir, ignore_errors=True) + shutil.rmtree(temp_dir, ignore_errors=True) \ No newline at end of file diff --git a/services/ai-service/src/utils/framework_utils.py b/services/ai-service/src/utils/framework_utils.py index ba87114..365207c 100644 --- a/services/ai-service/src/utils/framework_utils.py +++ b/services/ai-service/src/utils/framework_utils.py @@ -9,6 +9,12 @@ def _project_root() -> Path: + """ + Resolve the project root directory for this module. + + Returns: + Path: The filesystem path to the project root (three parent directories above this file). + """ return Path(__file__).parent.parent.parent @@ -46,13 +52,13 @@ def save_extraction_json( def get_input_paths() -> Dict[str, Path]: """ - Get standard input directory paths. - + Return standard project input directories. + Returns: - Dictionary with paths: - - frameworks: Path to framework PDFs directory - - applicants: Path to applicant PDFs directory - - vector_db: Path to vector DB input PDFs directory + Mapping of input directory names to their project-relative Paths: + - `frameworks`: Path to data/inputs/frameworks + - `applicants`: Path to data/inputs/applicants + - `vector_db`: Path to data/inputs/vector_db """ project_root = _project_root() return { @@ -65,9 +71,11 @@ def get_input_paths() -> Dict[str, Path]: def list_framework_jsons(framework_name: str) -> List[Tuple[str, Dict[str, Any]]]: """ Load all JSON files for a framework from config/frameworks/{framework_name}/. - + + If the framework directory does not exist, returns an empty list. Files that cannot be opened or parsed are skipped. + Returns: - List of (stem, parsed_json) tuples. Stem = filename without .json (e.g. PoliciesEn-1). + List[Tuple[str, Dict[str, Any]]]: A list of (stem, parsed_json) tuples where `stem` is the filename without the `.json` extension. """ base = _project_root() / "config" / "frameworks" / framework_name if not base.is_dir(): @@ -85,10 +93,15 @@ def list_framework_jsons(framework_name: str) -> List[Tuple[str, Dict[str, Any]] def get_vector_db_pdf_paths(framework_name: str) -> List[Path]: """ - PDF paths used as input for the vector DB. Looks in data/inputs/vector_db/. - - Prefers vector_db/{framework_name}/*.pdf. If that dir has no PDFs, falls back to - vector_db/*.pdf (flat). + Locate PDF files to use as input for the vector database for a given framework. + + Searches data/inputs/vector_db/{framework_name} for files with extensions `.pdf` or `.PDF` and returns those if any are present; otherwise returns PDFs from data/inputs/vector_db (flat). Returned paths are sorted and may be empty. + + Parameters: + framework_name (str): Framework subdirectory name to prefer under data/inputs/vector_db. + + Returns: + List[Path]: A list of Path objects pointing to found PDF files; may be empty. """ paths = get_input_paths() vdb = paths["vector_db"] @@ -98,4 +111,4 @@ def get_vector_db_pdf_paths(framework_name: str) -> List[Path]: if pdfs: return pdfs pdfs = sorted(vdb.glob("*.pdf")) + sorted(vdb.glob("*.PDF")) - return pdfs + return pdfs \ No newline at end of file