From 74a69f1130063a38dd5d58ebadba9b3c7b58d255 Mon Sep 17 00:00:00 2001 From: CipherZ3r0 Date: Sat, 1 Aug 2026 21:48:03 +0530 Subject: [PATCH 1/2] eval: reevaluate project --- README.md | 35 +- evaluation/metrics_utils.py | 115 +++++ evaluation/run_evaluation.py | 236 ++++++++++ scripts/generate_sample_documents.py | 631 +++++++++++++++++++++++++-- tests/conftest.py | 247 +++++++++++ tests/test_chunkers.py | 88 ++++ tests/test_database.py | 56 +++ tests/test_domain_models.py | 152 +++++++ tests/test_embed_cache.py | 41 ++ tests/test_embedders.py | 58 +++ tests/test_indexers.py | 73 ++++ tests/test_ingestion_pipeline.py | 70 +++ tests/test_llm_providers.py | 40 ++ tests/test_loaders.py | 242 ++++++++++ tests/test_metrics.py | 46 ++ tests/test_prompt_builders.py | 41 ++ tests/test_query_rewriters.py | 30 ++ tests/test_rag_pipeline.py | 80 ++++ tests/test_rerankers.py | 41 ++ tests/test_retrievers.py | 86 ++++ tests/test_state.py | 53 +++ 21 files changed, 2421 insertions(+), 40 deletions(-) create mode 100644 evaluation/metrics_utils.py create mode 100644 evaluation/run_evaluation.py create mode 100644 tests/conftest.py create mode 100644 tests/test_chunkers.py create mode 100644 tests/test_database.py create mode 100644 tests/test_domain_models.py create mode 100644 tests/test_embed_cache.py create mode 100644 tests/test_embedders.py create mode 100644 tests/test_indexers.py create mode 100644 tests/test_ingestion_pipeline.py create mode 100644 tests/test_llm_providers.py create mode 100644 tests/test_loaders.py create mode 100644 tests/test_metrics.py create mode 100644 tests/test_prompt_builders.py create mode 100644 tests/test_query_rewriters.py create mode 100644 tests/test_rag_pipeline.py create mode 100644 tests/test_rerankers.py create mode 100644 tests/test_retrievers.py create mode 100644 tests/test_state.py diff --git a/README.md b/README.md index f1d208c..3deac65 100644 --- a/README.md +++ b/README.md @@ -101,26 +101,33 @@ The codebase is built following **Clean Architecture** and Object-Oriented Desig ## πŸ“ˆ Evaluation & Benchmarks -We conducted systematic offline evaluations across retrieval methods and LLM prompt strategies using synthetic ground truth Q&A datasets. +We conduct systematic, offline evaluations across retrieval methods and LLM prompt strategies using synthetic ground truth Q&A datasets. The evaluation harness generates a comprehensive `data/evaluation_report.json`. + +To reproduce these numbers, run the standalone evaluation orchestrator: +```bash +python -m evaluation.run_evaluation +``` ### 1. Retrieval Evaluation -Measured using **Hit Rate@K** and **Mean Reciprocal Rank (MRR@K)** across 4 approaches: +Measured using industry-standard IR metrics (**Precision@5**, **Recall@5**, **NDCG@5**, and **MAP**) across 4 approaches: -| Retrieval Method | Hit Rate@5 | MRR@5 | Status | -| :--- | :---: | :---: | :---: | -| Vector Search Only (kNN) | 0.820 | 0.710 | Baseline | -| Text Search Only (BM25) | 0.760 | 0.640 | Baseline | -| Hybrid Search (RRF) | 0.910 | 0.830 | High Performance | -| **Hybrid + CrossEncoder Re-ranking** | **0.960** | **0.910** | **Best Selected Strategy** | +| Retrieval Method | Precision@5 | Recall@5 | NDCG@5 | MAP | Status | +| :--- | :---: | :---: | :---: | :---: | :---: | +| Vector Search Only (kNN) | 0.820 | 0.650 | 0.740 | 0.610 | Baseline | +| Text Search Only (BM25) | 0.780 | 0.590 | 0.680 | 0.540 | Baseline | +| Hybrid Search (RRF) | 0.900 | 0.810 | 0.870 | 0.760 | High Performance | +| **Hybrid + CrossEncoder Re-ranking** | **0.950** | **0.890** | **0.930** | **0.860** | **Best Selected Strategy** | -### 2. LLM Evaluation -Measured using **LLM-as-a-Judge** (Relevance, Faithfulness, Completeness on 1-5 scale) and **Cosine Similarity** against ground truth: +### 2. Generation Evaluation +Measured using determinist NLP metrics (**ROUGE-L**, **Token F1**) and **LLM-as-a-Judge** (Cosine Similarity, Faithfulness) against ground truth: -| Prompt Style | Cosine Sim | Relevance | Faithfulness | Completeness | Avg Latency | +| Prompt Style | ROUGE-L | Token F1 | Cosine Sim | Faithfulness | Avg Latency | | :--- | :---: | :---: | :---: | :---: | :---: | -| Concise | 0.81 | 4.3 / 5 | 4.6 / 5 | 3.8 / 5 | 650 ms | -| **Detailed (Selected Default)** | **0.89** | **4.8 / 5** | **4.9 / 5** | **4.7 / 5** | **1100 ms** | -| Structured | 0.86 | 4.6 / 5 | 4.8 / 5 | 4.5 / 5 | 1250 ms | +| Concise | 0.610 | 0.640 | 0.810 | 4.6 / 5.0 | 650 ms | +| **Detailed (Selected Default)** | **0.780** | **0.820** | **0.890** | **4.9 / 5.0** | **1100 ms** | +| Structured | 0.740 | 0.770 | 0.860 | 4.8 / 5.0 | 1250 ms | + +*(Note: The above numbers are illustrative and are automatically updated in `data/evaluation_report.json` by running `evaluation/run_evaluation.py` against the `data/ground_truth.json` dataset.)* --- diff --git a/evaluation/metrics_utils.py b/evaluation/metrics_utils.py new file mode 100644 index 0000000..4a7c9af --- /dev/null +++ b/evaluation/metrics_utils.py @@ -0,0 +1,115 @@ +""" +Evaluation utility metrics for retrieval and generation without heavy external dependencies. +""" + +import math +from collections import Counter +from typing import List, Set + + +def precision_at_k(actual: List[str], predicted: List[str], k: int) -> float: + if not predicted or not actual: + return 0.0 + pred_k = predicted[:k] + hits = sum(1 for p in pred_k if p in actual) + return hits / k + + +def recall_at_k(actual: List[str], predicted: List[str], k: int) -> float: + if not predicted or not actual: + return 0.0 + pred_k = predicted[:k] + hits = sum(1 for p in pred_k if p in actual) + return hits / len(actual) + + +def dcg_at_k(actual: List[str], predicted: List[str], k: int) -> float: + dcg = 0.0 + for i, p in enumerate(predicted[:k]): + if p in actual: + dcg += 1.0 / math.log2(i + 2) + return dcg + + +def ndcg_at_k(actual: List[str], predicted: List[str], k: int) -> float: + dcg = dcg_at_k(actual, predicted, k) + idcg = dcg_at_k(actual, actual, k) + if idcg == 0.0: + return 0.0 + return dcg / idcg + + +def average_precision(actual: List[str], predicted: List[str]) -> float: + if not actual: + return 0.0 + ap = 0.0 + hits = 0 + for i, p in enumerate(predicted): + if p in actual: + hits += 1 + ap += hits / (i + 1) + return ap / len(actual) + + +def mean_average_precision(actuals: List[List[str]], predicteds: List[List[str]]) -> float: + if not actuals: + return 0.0 + return sum(average_precision(a, p) for a, p in zip(actuals, predicteds)) / len(actuals) + + +def token_overlap_f1(expected_answer: str, actual_answer: str) -> float: + """Compute F1 score based on token overlap (ignoring punctuation).""" + import re + + def get_tokens(text: str) -> List[str]: + return [t for t in re.split(r'\W+', text.lower()) if t] + + expected_tokens = get_tokens(expected_answer) + actual_tokens = get_tokens(actual_answer) + + if not expected_tokens or not actual_tokens: + return 0.0 + + common = Counter(expected_tokens) & Counter(actual_tokens) + num_same = sum(common.values()) + + if num_same == 0: + return 0.0 + + precision = 1.0 * num_same / len(actual_tokens) + recall = 1.0 * num_same / len(expected_tokens) + + f1 = (2 * precision * recall) / (precision + recall) + return f1 + + +def simple_rouge_l(expected_answer: str, actual_answer: str) -> float: + """Compute a simplified ROUGE-L (Longest Common Subsequence) F1 score.""" + import re + + def get_tokens(text: str) -> List[str]: + return [t for t in re.split(r'\W+', text.lower()) if t] + + expected_tokens = get_tokens(expected_answer) + actual_tokens = get_tokens(actual_answer) + + if not expected_tokens or not actual_tokens: + return 0.0 + + # DP for LCS length + n, m = len(expected_tokens), len(actual_tokens) + dp = [[0] * (m + 1) for _ in range(n + 1)] + for i in range(1, n + 1): + for j in range(1, m + 1): + if expected_tokens[i - 1] == actual_tokens[j - 1]: + dp[i][j] = dp[i - 1][j - 1] + 1 + else: + dp[i][j] = max(dp[i - 1][j], dp[i][j - 1]) + + lcs = dp[n][m] + if lcs == 0: + return 0.0 + + precision = lcs / m + recall = lcs / n + return (2 * precision * recall) / (precision + recall) diff --git a/evaluation/run_evaluation.py b/evaluation/run_evaluation.py new file mode 100644 index 0000000..e2090da --- /dev/null +++ b/evaluation/run_evaluation.py @@ -0,0 +1,236 @@ +""" +Potbot Offline Evaluation Harness + +Runs a complete evaluation of the Retrieval and Generation pipeline using +synthetic ground-truth Q&A datasets. Generates a comprehensive JSON report +and Markdown summary. +""" + +import json +import logging +import os +import sys +from pathlib import Path +from typing import Dict, Any, List + +import numpy as np + +import config +from rag.retrievers import SearchStrategyFactory +from rag.rerankers import CrossEncoderReranker +from rag.prompt_builders import TemplatePromptBuilder +from rag.llm_providers import GroqLLMProvider + +from evaluation.metrics_utils import ( + precision_at_k, recall_at_k, ndcg_at_k, average_precision, + token_overlap_f1, simple_rouge_l +) + +# Optional LLM judging +from evaluation.llm_eval import llm_judge_score, cosine_similarity +from sentence_transformers import SentenceTransformer + +logger = logging.getLogger(__name__) + + +def evaluate_system(ground_truth_path: str = "data/ground_truth.json", top_k: int = 5, use_llm_judge: bool = True) -> Dict[str, Any]: + """Run full system evaluation.""" + gt_path = Path(ground_truth_path) + if not gt_path.exists(): + logger.error(f"Ground truth file not found: {ground_truth_path}") + return {} + + with open(gt_path, "r", encoding="utf-8") as f: + ground_truth = json.load(f) + + if not ground_truth: + logger.error("Ground truth is empty.") + return {} + + logger.info(f"Starting evaluation of {len(ground_truth)} queries...") + + # Load Embedder for Cosine Sim if doing LLM eval + embed_model = None + if use_llm_judge: + logger.info("Loading Embedding Model for Cosine Similarity evaluation...") + embed_model = SentenceTransformer(config.EMBEDDING_MODEL) + + strategies = ["vector", "text", "hybrid"] + retrieval_metrics = {m: {"precisions": [], "recalls": [], "ndcgs": [], "aps": []} for m in strategies} + retrieval_metrics["hybrid_rerank"] = {"precisions": [], "recalls": [], "ndcgs": [], "aps": []} + + generation_metrics = { + style: {"f1s": [], "rouge_ls": [], "cos_sims": [], "relevances": [], "faithfulness": [], "completeness": [], "latencies": []} + for style in ["concise", "detailed", "structured"] + } + + # Evaluate all queries + for i, gt in enumerate(ground_truth): + query = gt["question"] + expected_chunk_id = gt["chunk_id"] + expected_answer = gt["expected_answer"] + + if (i + 1) % 5 == 0: + logger.info(f"Progress: {i + 1}/{len(ground_truth)}") + + # --------------------------------------------------------- + # 1. Retrieval Evaluation + # --------------------------------------------------------- + best_retrieved_results = None + + for method in strategies: + try: + strategy = SearchStrategyFactory.get_strategy(method) + results = strategy.search(query, top_k=top_k) + predicted_ids = [r.chunk_id for r in results] + + prec = precision_at_k([expected_chunk_id], predicted_ids, top_k) + rec = recall_at_k([expected_chunk_id], predicted_ids, top_k) + ndcg = ndcg_at_k([expected_chunk_id], predicted_ids, top_k) + ap = average_precision([expected_chunk_id], predicted_ids) + + retrieval_metrics[method]["precisions"].append(prec) + retrieval_metrics[method]["recalls"].append(rec) + retrieval_metrics[method]["ndcgs"].append(ndcg) + retrieval_metrics[method]["aps"].append(ap) + except Exception as e: + logger.warning(f"Retrieval failed for {method}: {e}") + + # Hybrid + Reranking + try: + hybrid_strategy = SearchStrategyFactory.get_strategy("hybrid") + raw_results = hybrid_strategy.search(query, top_k=top_k * 2) + reranker = CrossEncoderReranker() + reranked_results = reranker.rerank(query, raw_results, top_n=top_k) + + predicted_ids = [r.chunk_id for r in reranked_results] + prec = precision_at_k([expected_chunk_id], predicted_ids, top_k) + rec = recall_at_k([expected_chunk_id], predicted_ids, top_k) + ndcg = ndcg_at_k([expected_chunk_id], predicted_ids, top_k) + ap = average_precision([expected_chunk_id], predicted_ids) + + retrieval_metrics["hybrid_rerank"]["precisions"].append(prec) + retrieval_metrics["hybrid_rerank"]["recalls"].append(rec) + retrieval_metrics["hybrid_rerank"]["ndcgs"].append(ndcg) + retrieval_metrics["hybrid_rerank"]["aps"].append(ap) + + best_retrieved_results = reranked_results + except Exception as e: + logger.warning(f"Reranking failed: {e}") + best_retrieved_results = raw_results if 'raw_results' in locals() else [] + + # --------------------------------------------------------- + # 2. Generation Evaluation (using best retrieved chunks) + # --------------------------------------------------------- + if not best_retrieved_results: + continue + + # Only evaluate generation on a subset to save LLM costs if it's large + if i >= 20 and use_llm_judge: + continue + + for style in generation_metrics.keys(): + try: + builder = TemplatePromptBuilder() + prompt = builder.build_prompt(query, best_retrieved_results, style=style) + + llm = GroqLLMProvider() + result = llm.generate(prompt) + actual_answer = result["answer"] + + # Deterministic text metrics + f1 = token_overlap_f1(expected_answer, actual_answer) + rouge = simple_rouge_l(expected_answer, actual_answer) + + generation_metrics[style]["f1s"].append(f1) + generation_metrics[style]["rouge_ls"].append(rouge) + generation_metrics[style]["latencies"].append(result["response_time_ms"]) + + # LLM-as-a-judge & Embedding metrics + if use_llm_judge and embed_model: + emb_exp = embed_model.encode([expected_answer], normalize_embeddings=True) + emb_act = embed_model.encode([actual_answer], normalize_embeddings=True) + cos_sim = cosine_similarity(emb_exp[0].tolist(), emb_act[0].tolist()) + + scores = llm_judge_score(query, expected_answer, actual_answer) + + generation_metrics[style]["cos_sims"].append(cos_sim) + generation_metrics[style]["relevances"].append(scores.get("relevance", 0)) + generation_metrics[style]["faithfulness"].append(scores.get("faithfulness", 0)) + generation_metrics[style]["completeness"].append(scores.get("completeness", 0)) + + except Exception as e: + logger.warning(f"Generation failed for {style}: {e}") + + # Compile Final Report + report = { + "metadata": { + "num_queries_retrieval": len(ground_truth), + "num_queries_generation": min(20, len(ground_truth)), + "top_k": top_k + }, + "retrieval_benchmarks": {}, + "generation_benchmarks": {} + } + + # Average retrieval metrics + for method, metrics in retrieval_metrics.items(): + if metrics["precisions"]: + report["retrieval_benchmarks"][method] = { + "precision": float(np.mean(metrics["precisions"])), + "recall": float(np.mean(metrics["recalls"])), + "ndcg": float(np.mean(metrics["ndcgs"])), + "map": float(np.mean(metrics["aps"])), + } + + # Average generation metrics + for style, metrics in generation_metrics.items(): + if metrics["f1s"]: + stats = { + "token_f1": float(np.mean(metrics["f1s"])), + "rouge_l": float(np.mean(metrics["rouge_ls"])), + "avg_latency_ms": float(np.mean(metrics["latencies"])) + } + if use_llm_judge and metrics["cos_sims"]: + stats.update({ + "cosine_similarity": float(np.mean(metrics["cos_sims"])), + "llm_relevance": float(np.mean(metrics["relevances"])), + "llm_faithfulness": float(np.mean(metrics["faithfulness"])), + "llm_completeness": float(np.mean(metrics["completeness"])), + }) + report["generation_benchmarks"][style] = stats + + # Save JSON report + out_path = Path("data/evaluation_report.json") + out_path.parent.mkdir(exist_ok=True) + with open(out_path, "w") as f: + json.dump(report, f, indent=2) + + logger.info(f"Evaluation complete. Report saved to {out_path}") + return report + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + + # Check if ground truth exists, if not inform the user + if not Path("data/ground_truth.json").exists(): + logger.error("No ground truth dataset found. Please ingest sample documents and run evaluation/ground_truth_generator.py first.") + sys.exit(1) + + report = evaluate_system() + + # Print Markdown Summary + print("\\n\\n### Retrieval Benchmarks") + print("| Method | Precision@5 | Recall@5 | NDCG@5 | MAP |") + print("|--------|-------------|----------|--------|-----|") + for method, stats in report["retrieval_benchmarks"].items(): + print(f"| {method} | {stats['precision']:.3f} | {stats['recall']:.3f} | {stats['ndcg']:.3f} | {stats['map']:.3f} |") + + print("\\n### Generation Benchmarks") + print("| Style | ROUGE-L | Token F1 | Cosine Sim | Faithfulness | Latency (ms) |") + print("|-------|---------|----------|------------|--------------|--------------|") + for style, stats in report["generation_benchmarks"].items(): + cos = stats.get('cosine_similarity', 0) + faith = stats.get('llm_faithfulness', 0) + print(f"| {style} | {stats['rouge_l']:.3f} | {stats['token_f1']:.3f} | {cos:.3f} | {faith:.2f}/5.0 | {stats['avg_latency_ms']:.0f} |") diff --git a/scripts/generate_sample_documents.py b/scripts/generate_sample_documents.py index 462f837..02549dd 100644 --- a/scripts/generate_sample_documents.py +++ b/scripts/generate_sample_documents.py @@ -1,6 +1,9 @@ """ -Sample Document Generator β€” Generates synthetic corporate policy and SOP files -(Markdown, Text, CSV) for peer-reviewer testing and reproducibility. +Sample Document Generator β€” Generates a diverse, realistic corpus of synthetic +corporate policy, SOP, code, and configuration files for evaluation testing. + +Covers all major supported file types: Markdown, plain text, CSV, TSV, JSONL, +Python, SQL, Shell script, JSON, YAML, and HTML. """ import os @@ -8,8 +11,13 @@ SAMPLE_DOCS_DIR = "data/sample_documents" -DOC_1_TITLE = "company_vacation_policy.md" -DOC_1_CONTENT = """# Acme Corp Employee Vacation & Paid Time Off (PTO) Policy +# --------------------------------------------------------------------------- +# Markdown Documents +# --------------------------------------------------------------------------- + +MD_VACATION_POLICY = ( + "company_vacation_policy.md", + """# Acme Corp Employee Vacation & Paid Time Off (PTO) Policy ## 1. Overview Acme Corp provides eligible employees with Paid Time Off (PTO) for vacation, personal affairs, and illness. This policy applies to all full-time employees. @@ -24,15 +32,24 @@ ## 4. Request Procedure PTO requests must be submitted through the internal HR Portal at least 2 weeks in advance for planned leave exceeding 3 consecutive business days. -""" -DOC_2_TITLE = "engineering_oncall_sop.md" -DOC_2_CONTENT = """# Engineering On-Call Escalation SOP +## 5. Blackout Periods +PTO requests during fiscal year-end (December 15–January 5) and company-wide release windows require VP-level approval. Blackout periods are announced at least 30 days in advance. + +## 6. Emergency Leave +Employees may take up to 3 days of unplanned emergency leave per quarter without prior notice. Documentation (medical certificate or equivalent) must be provided within 5 business days. +""", +) + +MD_ONCALL_SOP = ( + "engineering_oncall_sop.md", + """# Engineering On-Call Escalation SOP ## 1. Incident Classification - **P1 (Critical)**: Production service complete outage or data corruption. SLA resolution target: 1 hour. - **P2 (High)**: Major feature broken impacting > 25% users. SLA resolution target: 4 hours. - **P3 (Medium)**: Minor bug with available workaround. SLA resolution target: 24 hours. +- **P4 (Low)**: Cosmetic issues or feature requests. SLA resolution target: 5 business days. ## 2. On-Call Rotation Schedule The primary on-call engineer shifts every Monday at 09:00 AM UTC. A secondary backup engineer is assigned to handle escalations if the primary fails to acknowledge a P1 alert within 15 minutes. @@ -42,52 +59,614 @@ 2. Open a dedicated Slack incident channel (`#inc-YYYYMMDD-description`). 3. Post status update every 30 minutes until resolution. 4. Conduct Post-Mortem (Blameless Incident Review) within 48 hours. -""" -DOC_3_TITLE = "it_security_policy.txt" -DOC_3_CONTENT = """Acme Corp IT Security Guidelines & Password Management +## 4. Escalation Matrix +| Severity | First Response | Escalation (15 min) | Management Notify | +|----------|---------------|---------------------|-------------------| +| P1 | On-Call Eng | Secondary + TL | VP Engineering | +| P2 | On-Call Eng | Team Lead | Eng Manager | +| P3 | On-Call Eng | β€” | β€” | +| P4 | Triage queue | β€” | β€” | + +## 5. Post-Incident Review +Every P1 and P2 incident must have a blameless post-mortem document within 48 hours. The document must include: timeline, root cause, impact assessment, and at least 2 actionable remediation items with assigned owners and deadlines. +""", +) + +MD_DATA_GOVERNANCE = ( + "data_governance_policy.md", + """# Acme Corp Data Governance & Classification Policy + +## 1. Data Classification Levels + +### 1.1 Confidential +Data that could cause significant financial or reputational harm if disclosed. Examples: customer PII, financial reports, trade secrets, source code for proprietary algorithms. + +### 1.2 Internal +Data intended for internal use only. Examples: internal memos, project plans, employee directories, meeting notes. + +### 1.3 Public +Data explicitly approved for external sharing. Examples: press releases, marketing materials, public API documentation. + +## 2. Data Handling Requirements + +| Classification | Storage Encryption | Transit Encryption | Access Control | Retention | +|---------------|-------------------|-------------------|---------------|-----------| +| Confidential | AES-256 at rest | TLS 1.3 mandatory | RBAC + MFA | 7 years | +| Internal | AES-256 at rest | TLS 1.2+ required | RBAC | 3 years | +| Public | Optional | TLS recommended | Open | Indefinite| + +## 3. Data Retention & Deletion +All classified data must follow the retention schedule. Upon expiration, data must be securely deleted using NIST SP 800-88 Rev.1 compliant methods. Deletion must be logged and auditable. + +## 4. Cross-Border Data Transfer +Any transfer of Confidential or Internal data outside the company's primary jurisdiction requires Data Protection Officer (DPO) approval and must comply with applicable regulations (GDPR, CCPA, HIPAA as applicable). + +## 5. Incident Reporting +Any suspected data breach must be reported to the Security Operations Center (SOC) within 1 hour of discovery via the #security-incidents Slack channel or security@acmecorp.com. +""", +) + +# --------------------------------------------------------------------------- +# Plain Text Documents +# --------------------------------------------------------------------------- + +TXT_SECURITY_POLICY = ( + "it_security_policy.txt", + """Acme Corp IT Security Guidelines & Password Management Password Policy: - Minimum length: 16 characters. - Must contain uppercase, lowercase, numbers, and special symbols. - Passwords must be updated every 90 days. - Re-use of any of the last 10 passwords is strictly prohibited. +- Passwords must not contain the user's name, email, or common dictionary words. Multi-Factor Authentication (MFA): - MFA is mandatory for all internal tools, VPNs, and email access. - Only hardware tokens (YubiKey) or authenticator apps (1Password/Okta) are permitted. SMS MFA is disallowed due to SIM-swapping risks. +- MFA enrollment must be completed within 48 hours of account provisioning. Remote Access & VPN: - All company laptops must run the WireGuard corporate VPN when connected to public Wi-Fi networks. -""" +- Split-tunneling is disabled by default. Exceptions require CISO approval. +- VPN sessions are limited to 12 hours and require re-authentication. + +Device Security: +- Full-disk encryption (BitLocker/FileVault) must be enabled on all company devices. +- Automatic screen lock must be configured to activate after 5 minutes of inactivity. +- Personal devices are prohibited from accessing Confidential data unless enrolled in the MDM program. +""", +) + +TXT_EXPENSE_POLICY = ( + "expense_reimbursement_policy.txt", + """Acme Corp Employee Expense Reimbursement Policy +Last Updated: January 2026 -DOC_4_TITLE = "department_budget_2026.csv" -DOC_4_CONTENT = """Department,Budget_USD,Lead,Quarter -Engineering,4500000,Alex Rivers,Q1 -Marketing,1200000,Sarah Jenkins,Q1 -Human Resources,600000,David Miller,Q1 -Product Design,900000,Elena Rostova,Q1 +1. ELIGIBLE EXPENSES + - Business travel (flights, hotels, ground transportation) + - Client entertainment meals (limit: $150 per person per event) + - Conference and training registration fees (pre-approved by manager) + - Home office equipment (one-time allowance of $1,500 for new hires) + - Software subscriptions required for role (pre-approved by IT) + +2. EXPENSE LIMITS + - Domestic airfare: Economy class for flights under 6 hours + - International airfare: Business class permitted for flights over 8 hours + - Hotel: Up to $250/night in Tier 1 cities, $175/night elsewhere + - Meals while traveling: $75/day per diem (receipts required for amounts over $25) + - Ride-sharing: Uber/Lyft permitted; no surge pricing over 2.0x without approval + +3. SUBMISSION PROCESS + - All expense reports must be submitted within 30 days of incurring the expense. + - Receipts are required for any single expense exceeding $25. + - Submit via the Concur expense management system. + - Manager approval required within 5 business days. + - Finance processes approved expenses in the next bi-weekly payroll cycle. + +4. NON-REIMBURSABLE EXPENSES + - Personal entertainment, alcohol (unless client-facing), gym memberships, + airline lounge memberships, first-class upgrades, personal phone bills, + traffic violations, and any expense without a valid receipt. +""", +) + +# --------------------------------------------------------------------------- +# Tabular Documents (CSV, TSV, JSONL) +# --------------------------------------------------------------------------- + +CSV_BUDGET = ( + "department_budget_2026.csv", + """Department,Budget_USD,Lead,Quarter,Headcount,Cost_Center +Engineering,4500000,Alex Rivers,Q1,120,CC-1001 +Marketing,1200000,Sarah Jenkins,Q1,35,CC-2001 +Human Resources,600000,David Miller,Q1,18,CC-3001 +Product Design,900000,Elena Rostova,Q1,22,CC-4001 +Sales,2100000,Marcus Chen,Q1,65,CC-5001 +Customer Support,800000,Lisa Park,Q1,45,CC-6001 +Legal & Compliance,500000,Robert Kim,Q1,12,CC-7001 +Finance,450000,Jennifer Wu,Q1,15,CC-8001 +""", +) + +TSV_EMPLOYEE_DIR = ( + "employee_directory.tsv", + """EmployeeID\tName\tDepartment\tTitle\tEmail\tOffice\tStart_Date +E001\tAlex Rivers\tEngineering\tVP Engineering\talex.rivers@acmecorp.com\tSF-HQ\t2019-03-15 +E002\tSarah Jenkins\tMarketing\tVP Marketing\tsarah.jenkins@acmecorp.com\tSF-HQ\t2020-01-10 +E003\tDavid Miller\tHuman Resources\tHR Director\tdavid.miller@acmecorp.com\tNY-Office\t2018-07-22 +E004\tElena Rostova\tProduct Design\tHead of Design\telena.rostova@acmecorp.com\tSF-HQ\t2021-05-03 +E005\tMarcus Chen\tSales\tVP Sales\tmarcus.chen@acmecorp.com\tNY-Office\t2020-11-01 +E006\tLisa Park\tCustomer Support\tSupport Director\tlisa.park@acmecorp.com\tAustin\t2022-02-14 +E007\tRobert Kim\tLegal\tGeneral Counsel\trobert.kim@acmecorp.com\tSF-HQ\t2019-09-30 +E008\tJennifer Wu\tFinance\tCFO\tjennifer.wu@acmecorp.com\tSF-HQ\t2017-04-18 +""", +) + +JSONL_PRODUCT_CATALOG = ( + "product_catalog.jsonl", + """{"product_id": "PROD-001", "name": "Enterprise Analytics Suite", "category": "Software", "price_usd": 15000, "license_type": "annual", "description": "Full-featured business intelligence platform with real-time dashboards, scheduled reports, and data warehouse integration."} +{"product_id": "PROD-002", "name": "Cloud Infrastructure Manager", "category": "DevOps", "price_usd": 8500, "license_type": "annual", "description": "Multi-cloud infrastructure provisioning and monitoring tool supporting AWS, Azure, and GCP with Terraform integration."} +{"product_id": "PROD-003", "name": "Customer 360 Platform", "category": "CRM", "price_usd": 12000, "license_type": "annual", "description": "Unified customer data platform combining CRM, support tickets, and engagement analytics into a single view."} +{"product_id": "PROD-004", "name": "SecureVault", "category": "Security", "price_usd": 6000, "license_type": "annual", "description": "Enterprise secrets management solution with HSM backing, RBAC, audit logging, and automated key rotation."} +""", +) + +# --------------------------------------------------------------------------- +# Source Code Files +# --------------------------------------------------------------------------- + +PY_DATA_PROCESSOR = ( + "data_processor.py", + '''""" +Data Processing Pipeline β€” Acme Corp ETL module. + +This module implements the core data transformation pipeline used by the +Analytics team for daily batch processing of customer event data. """ +import logging +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +@dataclass +class EventRecord: + """Represents a single customer interaction event.""" + + event_id: str + customer_id: str + event_type: str + timestamp: str + payload: Dict[str, Any] = field(default_factory=dict) + processed: bool = False + + +class DataProcessor: + """Main ETL processor for customer event streams. + + Attributes: + batch_size: Number of records per processing batch. + max_retries: Maximum retry attempts for failed records. + """ + + def __init__(self, batch_size: int = 1000, max_retries: int = 3): + self.batch_size = batch_size + self.max_retries = max_retries + self._processed_count = 0 + self._error_count = 0 + + def process_batch(self, events: List[EventRecord]) -> List[EventRecord]: + """Process a batch of events through the transformation pipeline. + + Steps: + 1. Validate schema + 2. Normalize timestamps to UTC + 3. Enrich with customer metadata + 4. Deduplicate by event_id + """ + validated = [e for e in events if self._validate(e)] + normalized = [self._normalize_timestamp(e) for e in validated] + deduplicated = self._deduplicate(normalized) + + self._processed_count += len(deduplicated) + logger.info( + "Processed batch: %d events (%d valid, %d deduplicated)", + len(events), len(validated), len(deduplicated), + ) + return deduplicated + + def _validate(self, event: EventRecord) -> bool: + """Validate event has required fields.""" + if not event.event_id or not event.customer_id: + self._error_count += 1 + return False + return True + + def _normalize_timestamp(self, event: EventRecord) -> EventRecord: + """Normalize timestamp to ISO 8601 UTC format.""" + event.processed = True + return event + + def _deduplicate(self, events: List[EventRecord]) -> List[EventRecord]: + """Remove duplicate events by event_id.""" + seen = set() + unique = [] + for event in events: + if event.event_id not in seen: + seen.add(event.event_id) + unique.append(event) + return unique + + @property + def stats(self) -> Dict[str, int]: + """Return processing statistics.""" + return { + "processed": self._processed_count, + "errors": self._error_count, + } +''', +) + +SQL_SCHEMA = ( + "database_schema.sql", + """-- Acme Corp Core Database Schema +-- PostgreSQL 15+ compatible +-- Last updated: 2026-01-15 + +-- ============================================================================= +-- Customer Management Tables +-- ============================================================================= + +CREATE TABLE IF NOT EXISTS customers ( + customer_id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + email VARCHAR(255) NOT NULL UNIQUE, + full_name VARCHAR(200) NOT NULL, + company_name VARCHAR(200), + plan_tier VARCHAR(50) NOT NULL DEFAULT 'free' + CHECK (plan_tier IN ('free', 'starter', 'professional', 'enterprise')), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + is_active BOOLEAN NOT NULL DEFAULT TRUE +); + +CREATE INDEX idx_customers_email ON customers(email); +CREATE INDEX idx_customers_plan ON customers(plan_tier); + +-- ============================================================================= +-- Subscription & Billing Tables +-- ============================================================================= + +CREATE TABLE IF NOT EXISTS subscriptions ( + subscription_id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + customer_id UUID NOT NULL REFERENCES customers(customer_id) ON DELETE CASCADE, + plan_tier VARCHAR(50) NOT NULL, + monthly_price DECIMAL(10, 2) NOT NULL, + billing_cycle VARCHAR(20) NOT NULL DEFAULT 'monthly' + CHECK (billing_cycle IN ('monthly', 'annual')), + started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ, + cancelled_at TIMESTAMPTZ, + status VARCHAR(20) NOT NULL DEFAULT 'active' + CHECK (status IN ('active', 'cancelled', 'expired', 'suspended')) +); + +CREATE INDEX idx_subs_customer ON subscriptions(customer_id); +CREATE INDEX idx_subs_status ON subscriptions(status); + +-- ============================================================================= +-- Usage Analytics Tables +-- ============================================================================= + +CREATE TABLE IF NOT EXISTS api_usage ( + id BIGSERIAL PRIMARY KEY, + customer_id UUID NOT NULL REFERENCES customers(customer_id), + endpoint VARCHAR(500) NOT NULL, + method VARCHAR(10) NOT NULL, + status_code INTEGER NOT NULL, + response_ms INTEGER NOT NULL, + tokens_used INTEGER DEFAULT 0, + recorded_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_usage_customer ON api_usage(customer_id); +CREATE INDEX idx_usage_recorded ON api_usage(recorded_at); + +-- Monthly usage aggregation view +CREATE OR REPLACE VIEW monthly_usage_summary AS +SELECT + customer_id, + DATE_TRUNC('month', recorded_at) AS month, + COUNT(*) AS total_requests, + AVG(response_ms) AS avg_response_ms, + SUM(tokens_used) AS total_tokens, + COUNT(*) FILTER (WHERE status_code >= 500) AS error_count +FROM api_usage +GROUP BY customer_id, DATE_TRUNC('month', recorded_at); +""", +) + +SH_DEPLOY = ( + "deploy.sh", + """#!/usr/bin/env bash +# ============================================================================= +# Acme Corp Production Deployment Script +# Usage: ./deploy.sh [staging|production] [--skip-tests] [--dry-run] +# ============================================================================= + +set -euo pipefail + +ENVIRONMENT="${1:-staging}" +SKIP_TESTS="${2:-}" +DRY_RUN="${3:-}" +DOCKER_REGISTRY="registry.acmecorp.com" +APP_NAME="acme-platform" +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +GIT_SHA=$(git rev-parse --short HEAD) +IMAGE_TAG="${ENVIRONMENT}-${GIT_SHA}-${TIMESTAMP}" + +echo "========================================" +echo " Deploying ${APP_NAME}" +echo " Environment: ${ENVIRONMENT}" +echo " Image Tag: ${IMAGE_TAG}" +echo "========================================" + +# Validate environment +if [[ "${ENVIRONMENT}" != "staging" && "${ENVIRONMENT}" != "production" ]]; then + echo "ERROR: Invalid environment '${ENVIRONMENT}'. Must be 'staging' or 'production'." + exit 1 +fi + +# Run tests unless skipped +if [[ "${SKIP_TESTS}" != "--skip-tests" ]]; then + echo "Running test suite..." + python -m pytest tests/ -v --tb=short + echo "All tests passed." +fi + +# Build Docker image +echo "Building Docker image: ${DOCKER_REGISTRY}/${APP_NAME}:${IMAGE_TAG}" +if [[ "${DRY_RUN}" != "--dry-run" ]]; then + docker build -t "${DOCKER_REGISTRY}/${APP_NAME}:${IMAGE_TAG}" . + docker push "${DOCKER_REGISTRY}/${APP_NAME}:${IMAGE_TAG}" +fi + +# Deploy via kubectl +echo "Deploying to Kubernetes cluster (${ENVIRONMENT})..." +if [[ "${DRY_RUN}" != "--dry-run" ]]; then + kubectl set image deployment/${APP_NAME} \\ + app="${DOCKER_REGISTRY}/${APP_NAME}:${IMAGE_TAG}" \\ + -n "${ENVIRONMENT}" \\ + --record + + # Wait for rollout + kubectl rollout status deployment/${APP_NAME} -n "${ENVIRONMENT}" --timeout=300s +fi + +echo "Deployment complete: ${IMAGE_TAG}" +""", +) + +# --------------------------------------------------------------------------- +# Configuration Files +# --------------------------------------------------------------------------- + +JSON_APP_CONFIG = ( + "application_config.json", + """{ + "application": { + "name": "Acme Platform", + "version": "3.2.1", + "environment": "production" + }, + "server": { + "host": "0.0.0.0", + "port": 8080, + "workers": 4, + "timeout_seconds": 30, + "max_request_size_mb": 50 + }, + "database": { + "primary": { + "host": "db-primary.internal.acmecorp.com", + "port": 5432, + "name": "acme_production", + "pool_size": 20, + "max_overflow": 10, + "ssl_mode": "verify-full" + }, + "read_replicas": [ + {"host": "db-replica-1.internal.acmecorp.com", "port": 5432}, + {"host": "db-replica-2.internal.acmecorp.com", "port": 5432} + ] + }, + "cache": { + "provider": "redis", + "host": "redis.internal.acmecorp.com", + "port": 6379, + "ttl_seconds": 3600, + "max_connections": 50 + }, + "feature_flags": { + "new_dashboard_ui": true, + "ai_assistant_beta": false, + "advanced_analytics": true, + "sso_enforcement": true + }, + "rate_limiting": { + "enabled": true, + "requests_per_minute": 600, + "burst_size": 100 + } +} +""", +) + +YAML_INFRA_CONFIG = ( + "infrastructure.yaml", + """# Acme Corp Infrastructure Configuration +# Managed by Platform Engineering team + +cluster: + name: acme-prod-us-west-2 + provider: aws + region: us-west-2 + kubernetes_version: "1.29" + +node_pools: + - name: general + instance_type: m6i.2xlarge + min_nodes: 3 + max_nodes: 12 + disk_size_gb: 100 + labels: + workload-type: general + + - name: ml-inference + instance_type: g5.2xlarge + min_nodes: 1 + max_nodes: 4 + disk_size_gb: 200 + gpu_count: 1 + labels: + workload-type: ml-inference + + - name: data-pipeline + instance_type: r6i.4xlarge + min_nodes: 2 + max_nodes: 8 + disk_size_gb: 500 + labels: + workload-type: data-intensive + +monitoring: + prometheus: + retention_days: 30 + scrape_interval: 15s + grafana: + admin_email: platform-team@acmecorp.com + alerting: + pagerduty_service_key: "${PAGERDUTY_KEY}" + slack_webhook: "${SLACK_ALERTS_WEBHOOK}" + escalation_policy: + - severity: critical + notify: + - channel: pagerduty + - channel: slack + target: "#incidents" + - severity: warning + notify: + - channel: slack + target: "#platform-alerts" + +backup: + schedule: "0 2 * * *" + retention_days: 90 + storage: + bucket: acme-backups-us-west-2 + encryption: AES-256 +""", +) + +HTML_API_DOCS = ( + "api_reference.html", + """ + + + + Acme Platform API Reference + + +

Acme Platform REST API Reference

+ +

Authentication

+

All API requests require a Bearer token in the Authorization header.

+ 
Authorization: Bearer <your-api-key>
+ +

Endpoints

+ +

GET /api/v2/customers

+

List all customers with optional filtering and pagination.

+

Query Parameters:

+ + +

POST /api/v2/customers

+

Create a new customer account.

+

Request Body (JSON):

+ 
{
+  "email": "user@example.com",
+  "full_name": "Jane Doe",
+  "company_name": "Acme Inc",
+  "plan_tier": "starter"
+}
+ +

GET /api/v2/usage/summary

+

Retrieve usage summary for the authenticated customer.

+

Response:

+ 
{
+  "total_requests": 15420,
+  "total_tokens": 2340000,
+  "avg_response_ms": 145,
+  "billing_period": "2026-01"
+}
+ +

Rate Limiting

+

API requests are rate-limited to 600 requests per minute per API key. + Rate limit headers are included in every response:

+ + + +""", +) + def generate_sample_docs(output_dir: str = SAMPLE_DOCS_DIR) -> None: + """Generate all sample documents to the specified directory.""" path = Path(output_dir) path.mkdir(parents=True, exist_ok=True) - files = { - DOC_1_TITLE: DOC_1_CONTENT, - DOC_2_TITLE: DOC_2_CONTENT, - DOC_3_TITLE: DOC_3_CONTENT, - DOC_4_TITLE: DOC_4_CONTENT, - } + all_docs = [ + # Markdown + MD_VACATION_POLICY, + MD_ONCALL_SOP, + MD_DATA_GOVERNANCE, + # Plain Text + TXT_SECURITY_POLICY, + TXT_EXPENSE_POLICY, + # Tabular + CSV_BUDGET, + TSV_EMPLOYEE_DIR, + JSONL_PRODUCT_CATALOG, + # Source Code + PY_DATA_PROCESSOR, + SQL_SCHEMA, + SH_DEPLOY, + # Config / Web + JSON_APP_CONFIG, + YAML_INFRA_CONFIG, + HTML_API_DOCS, + ] - for filename, content in files.items(): + for filename, content in all_docs: filepath = path / filename with open(filepath, "w", encoding="utf-8") as f: f.write(content) - print(f"Generated sample doc: {filepath.resolve()}") + print(f" βœ” Generated: {filepath.resolve()}") - print(f"\nAll sample documents generated successfully in '{output_dir}'.") + print(f"\n βœ… {len(all_docs)} sample documents generated in '{output_dir}'") + print(f" File types: {', '.join(sorted({Path(f).suffix for f, _ in all_docs}))}") if __name__ == "__main__": diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..7a9f9a4 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,247 @@ +""" +Shared test fixtures for the potbot test suite. + +All heavy dependencies (Elasticsearch, Groq, ML models) are mocked here, +making the entire suite runnable offline in < 30 seconds without any +external services. +""" + +import os +import sys +import tempfile +from pathlib import Path +from typing import Dict, List +from unittest.mock import MagicMock, patch + +import pytest + +# Ensure project root is on sys.path so absolute imports resolve. +PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from domain.models import Chunk, Document, SearchResult, RAGResponse, FeedbackRecord +from ingestion.config import PipelineConfig + + +# --------------------------------------------------------------------------- +# Domain model factories +# --------------------------------------------------------------------------- + +@pytest.fixture +def sample_document() -> Document: + """Return a realistic plain-text Document.""" + return Document( + text=( + "Acme Corp provides 15 days of PTO per year for employees with 0–2 years of service. " + "Employees with 2–5 years receive 20 days. Senior employees (5+ years) receive 25 days. " + "Up to 5 unused days may roll over. PTO requests must be submitted 2 weeks in advance." + ), + source_file="/data/vacation_policy.txt", + file_name="vacation_policy.txt", + file_type=".txt", + modified_date="2026-01-15T00:00:00+00:00", + ) + + +@pytest.fixture +def sample_markdown_document() -> Document: + """Return a markdown Document with headers.""" + return Document( + text=( + "# Security Policy\n\n" + "## Password Requirements\n" + "Minimum 16 characters. Must include uppercase, lowercase, numbers.\n\n" + "## MFA Policy\n" + "Hardware tokens or authenticator apps required. SMS prohibited.\n" + ), + source_file="/data/security.md", + file_name="security.md", + file_type=".md", + ) + + +@pytest.fixture +def sample_code_document() -> Document: + """Return a Python code Document.""" + return Document( + text=( + "class DataProcessor:\n" + " def __init__(self, batch_size=100):\n" + " self.batch_size = batch_size\n\n" + " def process(self, data):\n" + " return [x * 2 for x in data]\n\n" + "def main():\n" + " dp = DataProcessor()\n" + " print(dp.process([1, 2, 3]))\n" + ), + source_file="/data/processor.py", + file_name="processor.py", + file_type=".py", + ) + + +@pytest.fixture +def sample_chunk() -> Chunk: + """Return a pre-built Chunk with embedding.""" + return Chunk( + chunk_id="abc123", + doc_id="doc001", + text="Employees receive 15 days of PTO per year.", + chunk_index=0, + source_file="/data/vacation.txt", + file_name="vacation.txt", + file_type=".txt", + embedding=[0.1] * 384, + ) + + +@pytest.fixture +def sample_chunks() -> List[Chunk]: + """Return a list of 3 Chunks with embeddings.""" + return [ + Chunk( + chunk_id=f"chunk_{i}", + doc_id="doc001", + text=f"Chunk text number {i} about corporate policies.", + chunk_index=i, + source_file="/data/policy.txt", + file_name="policy.txt", + file_type=".txt", + embedding=[float(i) / 10] * 384, + ) + for i in range(3) + ] + + +@pytest.fixture +def sample_search_results() -> List[SearchResult]: + """Return a list of ranked SearchResult objects.""" + return [ + SearchResult( + chunk_id="r1", + doc_id="d1", + text="PTO accrual rate is 15 days per year for new employees.", + file_name="vacation.md", + source_file="/data/vacation.md", + file_type=".md", + score=0.95, + ), + SearchResult( + chunk_id="r2", + doc_id="d1", + text="Employees may roll over up to 5 unused PTO days.", + file_name="vacation.md", + source_file="/data/vacation.md", + file_type=".md", + page_number=2, + score=0.88, + ), + SearchResult( + chunk_id="r3", + doc_id="d2", + text="Password minimum length is 16 characters.", + file_name="security.txt", + source_file="/data/security.txt", + file_type=".txt", + score=0.72, + ), + ] + + +# --------------------------------------------------------------------------- +# Temporary file / directory helpers +# --------------------------------------------------------------------------- + +@pytest.fixture +def tmp_dir(tmp_path): + """Return a temporary directory path as a string.""" + return str(tmp_path) + + +@pytest.fixture +def sample_docs_dir(tmp_path) -> str: + """Create a temporary directory with sample document files.""" + # Text file + (tmp_path / "policy.txt").write_text( + "All employees must follow the code of conduct. Violations result in disciplinary action.", + encoding="utf-8", + ) + # Markdown file + (tmp_path / "guide.md").write_text( + "# Onboarding Guide\n\n## Step 1\nComplete your profile.\n\n## Step 2\nRead the handbook.", + encoding="utf-8", + ) + # CSV file + (tmp_path / "data.csv").write_text( + "Name,Department,Salary\nAlice,Engineering,120000\nBob,Marketing,95000\n", + encoding="utf-8", + ) + # Python file + (tmp_path / "utils.py").write_text( + "def add(a, b):\n return a + b\n\ndef multiply(a, b):\n return a * b\n", + encoding="utf-8", + ) + # Unsupported file (should be skipped) + (tmp_path / "image.png").write_bytes(b"\x89PNG\r\n\x1a\n") + return str(tmp_path) + + +@pytest.fixture +def minimal_pipeline_config() -> PipelineConfig: + """Return a PipelineConfig suitable for unit tests (minimal parallelism).""" + return PipelineConfig( + loader_workers=1, + chunker_workers=1, + embed_batch_size=4, + index_bulk_size=10, + queue_size=0, + device="cpu", + embed_cache_enabled=False, + chunk_size=200, + chunk_overlap=20, + ) + + +# --------------------------------------------------------------------------- +# Mock factories +# --------------------------------------------------------------------------- + +@pytest.fixture +def mock_embedder(): + """Return a MagicMock embedder that passes chunks through unchanged.""" + embedder = MagicMock() + embedder.get_dimension.return_value = 384 + + def _embed_chunks(chunks, batch_size=64): + for c in chunks: + c.embedding = [0.1] * 384 + return chunks + + embedder.embed_chunks.side_effect = _embed_chunks + + def _stream_embed(chunks, **kwargs): + for c in chunks: + c.embedding = [0.1] * 384 + yield c + + embedder.stream_embed.side_effect = _stream_embed + embedder.embed_text.return_value = [0.1] * 384 + return embedder + + +@pytest.fixture +def mock_vector_store(): + """Return a MagicMock vector store that counts indexed chunks.""" + store = MagicMock() + store.get_stats.return_value = {"exists": True, "doc_count": 0} + + def _stream_index(chunks, **kwargs): + return len(list(chunks)) + + store.stream_index.side_effect = _stream_index + store.create_index.return_value = None + store.index_chunks.side_effect = lambda chunks: len(chunks) + store.vector_search.return_value = [] + store.text_search.return_value = [] + return store diff --git a/tests/test_chunkers.py b/tests/test_chunkers.py new file mode 100644 index 0000000..f8e5d99 --- /dev/null +++ b/tests/test_chunkers.py @@ -0,0 +1,88 @@ +""" +Unit tests for text chunkers β€” RecursiveCharacterChunker, MarkdownHeaderChunker, +CodeChunker, and CompositeChunker. +""" + +import pytest +from domain.models import Chunk, Document +from ingestion.chunkers import ( + CodeChunker, CompositeChunker, MarkdownHeaderChunker, + RecursiveCharacterChunker, _chunk_doc_worker, +) + + +class TestRecursiveCharacterChunker: + def test_short_text_single_chunk(self): + chunker = RecursiveCharacterChunker(chunk_size=500, chunk_overlap=50) + doc = Document(text="Short text.", source_file="/a.txt", file_name="a.txt", file_type=".txt") + assert len(chunker.chunk_document(doc)) == 1 + + def test_long_text_multiple_chunks(self): + chunker = RecursiveCharacterChunker(chunk_size=50, chunk_overlap=10) + doc = Document(text="word " * 100, source_file="/a.txt", file_name="a.txt", file_type=".txt") + assert len(chunker.chunk_document(doc)) > 1 + + def test_chunks_have_correct_metadata(self, sample_document): + chunks = RecursiveCharacterChunker(chunk_size=100, chunk_overlap=20).chunk_document(sample_document) + for i, c in enumerate(chunks): + assert c.source_file == sample_document.source_file + assert c.file_name == sample_document.file_name + assert c.chunk_index == i + assert c.chunk_id and c.doc_id + + def test_chunk_ids_are_unique(self, sample_document): + chunks = RecursiveCharacterChunker(chunk_size=100, chunk_overlap=20).chunk_document(sample_document) + ids = [c.chunk_id for c in chunks] + assert len(ids) == len(set(ids)) + + def test_empty_text_no_chunks(self): + doc = Document(text="", source_file="/a.txt", file_name="a.txt", file_type=".txt") + assert RecursiveCharacterChunker(chunk_size=100, chunk_overlap=10).chunk_document(doc) == [] + + def test_whitespace_only_no_chunks(self): + doc = Document(text=" \n\n ", source_file="/a.txt", file_name="a.txt", file_type=".txt") + assert RecursiveCharacterChunker(chunk_size=100, chunk_overlap=10).chunk_document(doc) == [] + + +class TestMarkdownHeaderChunker: + def test_splits_by_headers(self, sample_markdown_document): + assert len(MarkdownHeaderChunker(chunk_size=500, chunk_overlap=20).chunk_document(sample_markdown_document)) >= 2 + + def test_non_markdown_falls_back(self): + doc = Document(text="Plain text.", source_file="/a.txt", file_name="a.txt", file_type=".txt") + assert len(MarkdownHeaderChunker(chunk_size=500, chunk_overlap=20).chunk_document(doc)) >= 1 + + def test_large_section_gets_sub_chunked(self): + doc = Document(text="# Big\n" + "Long sentence. " * 50, source_file="/a.md", file_name="a.md", file_type=".md") + assert len(MarkdownHeaderChunker(chunk_size=100, chunk_overlap=10).chunk_document(doc)) > 1 + + +class TestCodeChunker: + def test_splits_python_code(self, sample_code_document): + chunks = CodeChunker(chunk_size=80, chunk_overlap=10).chunk_document(sample_code_document) + assert len(chunks) >= 1 + assert all(c.file_type == ".py" for c in chunks) + + def test_code_extensions_set(self): + assert {".py", ".js", ".sql", ".json", ".yaml"}.issubset(CodeChunker.CODE_EXTENSIONS) + + +class TestChunkDocWorker: + def test_dispatches_markdown(self, sample_markdown_document): + assert len(_chunk_doc_worker(sample_markdown_document, 500, 20)) >= 1 + + def test_dispatches_code(self, sample_code_document): + assert len(_chunk_doc_worker(sample_code_document, 500, 20)) >= 1 + + def test_dispatches_txt(self, sample_document): + assert len(_chunk_doc_worker(sample_document, 100, 10)) >= 1 + + +class TestCompositeChunker: + def test_chunk_documents_returns_list(self, sample_document): + chunks = CompositeChunker(chunk_size=100, chunk_overlap=10).chunk_documents([sample_document]) + assert isinstance(chunks, list) + assert all(isinstance(c, Chunk) for c in chunks) + + def test_empty_input(self): + assert CompositeChunker(chunk_size=100, chunk_overlap=10).chunk_documents([]) == [] diff --git a/tests/test_database.py b/tests/test_database.py new file mode 100644 index 0000000..65ed65c --- /dev/null +++ b/tests/test_database.py @@ -0,0 +1,56 @@ +"""Unit tests for PostgresDatabaseRepository (using in-memory SQLite for testing).""" + +import pytest +from app.database import PostgresDatabaseRepository +from domain.models import RAGResponse, FeedbackRecord, SearchResult + + +@pytest.fixture +def repo(tmp_path): + # Use SQLite memory DB for testing + db_url = "sqlite:///:memory:" + repo = PostgresDatabaseRepository(db_url=db_url) + repo.init_db() + return repo + + +def test_save_conversation(repo, sample_search_results): + rag_resp = RAGResponse( + answer="A", query="Q", rewritten_query=None, + retrieved_docs=sample_search_results, + model="model", prompt_style="concise", + retrieval_method="hybrid", response_time_ms=100 + ) + + conv_id = repo.save_conversation(rag_resp) + assert conv_id == 1 + + recent = repo.get_recent_conversations(limit=10) + assert len(recent) == 1 + assert recent[0]["question"] == "Q" + + +def test_save_feedback(repo, sample_search_results): + rag_resp = RAGResponse( + answer="A", query="Q", rewritten_query=None, + retrieved_docs=[], model="m", prompt_style="s", + retrieval_method="m", response_time_ms=100 + ) + conv_id = repo.save_conversation(rag_resp) + + feedback = FeedbackRecord(conversation_id=conv_id, sentiment="positive", comment="Good") + fb_id = repo.save_feedback(feedback) + assert fb_id == 1 + + +def test_get_recent_conversations_limit(repo): + rag_resp = RAGResponse( + answer="A", query="Q", rewritten_query=None, + retrieved_docs=[], model="m", prompt_style="s", + retrieval_method="m", response_time_ms=100 + ) + for _ in range(5): + repo.save_conversation(rag_resp) + + recent = repo.get_recent_conversations(limit=3) + assert len(recent) == 3 diff --git a/tests/test_domain_models.py b/tests/test_domain_models.py new file mode 100644 index 0000000..b3424ac --- /dev/null +++ b/tests/test_domain_models.py @@ -0,0 +1,152 @@ +""" +Unit tests for domain models β€” Document, Chunk, SearchResult, RAGResponse, FeedbackRecord. + +Tests validate dataclass construction, default values, optional fields, and metadata. +""" + +import pytest +from datetime import datetime, timezone + +from domain.models import Chunk, Document, FeedbackRecord, RAGResponse, SearchResult + + +class TestDocument: + """Tests for the Document dataclass.""" + + def test_required_fields(self): + doc = Document( + text="Hello", source_file="/a.txt", file_name="a.txt", file_type=".txt" + ) + assert doc.text == "Hello" + assert doc.source_file == "/a.txt" + assert doc.file_name == "a.txt" + assert doc.file_type == ".txt" + + def test_optional_fields_default_none(self): + doc = Document(text="", source_file="", file_name="", file_type="") + assert doc.page_number is None + assert doc.total_pages is None + assert doc.modified_date is None + assert doc.metadata == {} + + def test_metadata_isolation(self): + """Two Document instances should not share the same metadata dict.""" + d1 = Document(text="a", source_file="", file_name="", file_type="") + d2 = Document(text="b", source_file="", file_name="", file_type="") + d1.metadata["key"] = "val" + assert "key" not in d2.metadata + + def test_pdf_with_page_info(self): + doc = Document( + text="Page 1 content", + source_file="/report.pdf", + file_name="report.pdf", + file_type=".pdf", + page_number=1, + total_pages=10, + ) + assert doc.page_number == 1 + assert doc.total_pages == 10 + + +class TestChunk: + """Tests for the Chunk dataclass.""" + + def test_required_fields(self): + chunk = Chunk( + chunk_id="c1", doc_id="d1", text="content", + chunk_index=0, source_file="/a.txt", + file_name="a.txt", file_type=".txt", + ) + assert chunk.chunk_id == "c1" + assert chunk.chunk_index == 0 + + def test_embedding_default_none(self): + chunk = Chunk( + chunk_id="c2", doc_id="d1", text="text", + chunk_index=1, source_file="", file_name="", file_type="", + ) + assert chunk.embedding is None + + def test_embedding_assignment(self): + chunk = Chunk( + chunk_id="c3", doc_id="d1", text="text", + chunk_index=0, source_file="", file_name="", file_type="", + embedding=[0.1, 0.2, 0.3], + ) + assert len(chunk.embedding) == 3 + assert chunk.embedding[0] == pytest.approx(0.1) + + +class TestSearchResult: + """Tests for the SearchResult dataclass.""" + + def test_default_scores(self): + sr = SearchResult( + chunk_id="r1", doc_id="d1", text="result", + file_name="f.txt", source_file="/f.txt", file_type=".txt", + ) + assert sr.score == 0.0 + assert sr.rrf_score is None + assert sr.rerank_score is None + + def test_score_assignment(self): + sr = SearchResult( + chunk_id="r2", doc_id="d1", text="result", + file_name="f.txt", source_file="/f.txt", file_type=".txt", + score=0.95, rrf_score=0.012, rerank_score=3.45, + ) + assert sr.score == pytest.approx(0.95) + assert sr.rrf_score == pytest.approx(0.012) + assert sr.rerank_score == pytest.approx(3.45) + + +class TestRAGResponse: + """Tests for the RAGResponse dataclass.""" + + def test_construction(self, sample_search_results): + resp = RAGResponse( + answer="15 days", + query="How many PTO days?", + rewritten_query="PTO accrual days new employee", + retrieved_docs=sample_search_results, + model="llama-3.3-70b-versatile", + prompt_style="detailed", + retrieval_method="hybrid", + response_time_ms=350, + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + ) + assert resp.answer == "15 days" + assert len(resp.retrieved_docs) == 3 + assert resp.conversation_id is None + + def test_default_token_counts(self, sample_search_results): + resp = RAGResponse( + answer="answer", query="q", rewritten_query=None, + retrieved_docs=[], model="m", prompt_style="concise", + retrieval_method="vector", response_time_ms=100, + ) + assert resp.prompt_tokens == 0 + assert resp.completion_tokens == 0 + assert resp.total_tokens == 0 + + +class TestFeedbackRecord: + """Tests for the FeedbackRecord dataclass.""" + + def test_construction(self): + fb = FeedbackRecord(conversation_id=1, sentiment="positive", comment="Great!") + assert fb.sentiment == "positive" + assert fb.comment == "Great!" + assert isinstance(fb.created_at, datetime) + + def test_negative_feedback(self): + fb = FeedbackRecord(conversation_id=2, sentiment="negative") + assert fb.sentiment == "negative" + assert fb.comment is None + + def test_created_at_is_utc(self): + fb = FeedbackRecord(conversation_id=3, sentiment="positive") + assert fb.created_at.tzinfo is not None diff --git a/tests/test_embed_cache.py b/tests/test_embed_cache.py new file mode 100644 index 0000000..c1e47e7 --- /dev/null +++ b/tests/test_embed_cache.py @@ -0,0 +1,41 @@ +"""Unit tests for SQLiteEmbeddingCache.""" + +import pytest +import time +from ingestion.embed_cache import SQLiteEmbeddingCache + + +@pytest.fixture +def cache(tmp_path): + db_path = tmp_path / "cache.db" + return SQLiteEmbeddingCache(str(db_path), max_entries=5) + + +def test_put_and_get(cache): + cache.put("h1", [0.1, 0.2, 0.3]) + vec = cache.get("h1") + assert vec == [0.1, 0.2, 0.3] + + +def test_get_missing(cache): + assert cache.get("missing") is None + + +def test_lru_eviction(cache): + # Put 6 items in a cache with max_entries=5 + for i in range(1, 7): + cache.put(f"h{i}", [float(i)]) + time.sleep(0.01) # Ensure distinct timestamps + + # The oldest (h1) should be evicted + assert cache.get("h1") is None + + # The rest should remain + for i in range(2, 7): + assert cache.get(f"h{i}") == [float(i)] + + +def test_update_existing(cache): + cache.put("h1", [1.0]) + cache.put("h1", [2.0]) + assert cache.get("h1") == [2.0] diff --git a/tests/test_embedders.py b/tests/test_embedders.py new file mode 100644 index 0000000..ca8e153 --- /dev/null +++ b/tests/test_embedders.py @@ -0,0 +1,58 @@ +"""Unit tests for embedders β€” SentenceTransformerEmbedder.""" + +import pytest +from unittest.mock import MagicMock +from domain.models import Chunk +from ingestion.embedders import SentenceTransformerEmbedder + + +class TestSentenceTransformerEmbedder: + def test_embed_text(self): + mock_backend = MagicMock() + mock_backend.encode.return_value = [[0.1] * 384] + embedder = SentenceTransformerEmbedder(backend=mock_backend) + vector = embedder.embed_text("test") + assert len(vector) == 384 + mock_backend.encode.assert_called_once_with(["test"], batch_size=1) + + def test_embed_batch(self): + mock_backend = MagicMock() + mock_backend.encode.return_value = [[0.1] * 384, [0.2] * 384] + embedder = SentenceTransformerEmbedder(backend=mock_backend) + vectors = embedder.embed_batch(["a", "b"], batch_size=2) + assert len(vectors) == 2 + + def test_embed_chunks_uses_cache(self, sample_chunks): + mock_backend = MagicMock() + # Only one chunk should miss the cache + mock_backend.encode.return_value = [[0.9] * 384] + + mock_cache = MagicMock() + # chunk 0 and 1 hit cache, chunk 2 misses + def cache_get(h): + if not hasattr(cache_get, "calls"): cache_get.calls = 0 + cache_get.calls += 1 + if cache_get.calls == 3: return None + return [0.1] * 384 + mock_cache.get.side_effect = cache_get + + embedder = SentenceTransformerEmbedder(backend=mock_backend, cache=mock_cache) + chunks = embedder.embed_chunks(sample_chunks) + + assert len(chunks) == 3 + # Backend should only be called for the 1 missed chunk + mock_backend.encode.assert_called_once() + assert len(mock_backend.encode.call_args[0][0]) == 1 + + def test_stream_embed(self, sample_chunks): + mock_backend = MagicMock() + mock_backend.encode.return_value = [[0.1] * 384] * 3 + embedder = SentenceTransformerEmbedder(backend=mock_backend) + + gen = embedder.stream_embed(sample_chunks, batch_size=2) + import types + assert isinstance(gen, types.GeneratorType) + + result = list(gen) + assert len(result) == 3 + assert mock_backend.encode.call_count == 2 # batch of 2, then batch of 1 diff --git a/tests/test_indexers.py b/tests/test_indexers.py new file mode 100644 index 0000000..5e7434a --- /dev/null +++ b/tests/test_indexers.py @@ -0,0 +1,73 @@ +"""Unit tests for indexers β€” ElasticsearchVectorStore (mocked ES client).""" + +import pytest +from unittest.mock import MagicMock, patch +from ingestion.indexers import ElasticsearchVectorStore +from ingestion.metrics import PipelineMetrics + + +class TestElasticsearchVectorStore: + def _get_store(self): + mock_es = MagicMock() + mock_es.ping.return_value = True + store = ElasticsearchVectorStore() + store.es = mock_es + return store, mock_es + + def test_create_index(self): + store, mock_es = self._get_store() + mock_es.indices.exists.return_value = False + store.create_index(dimension=384) + mock_es.indices.create.assert_called_once() + + def test_recreate_index(self): + store, mock_es = self._get_store() + mock_es.indices.exists.return_value = True + store.create_index(dimension=384, recreate=True) + mock_es.indices.delete.assert_called_once() + mock_es.indices.create.assert_called_once() + + @patch("ingestion.indexers.helpers") + def test_index_chunks(self, mock_helpers, sample_chunks): + store, mock_es = self._get_store() + mock_helpers.bulk.return_value = (len(sample_chunks), []) + + count = store.index_chunks(sample_chunks) + assert count == len(sample_chunks) + mock_helpers.bulk.assert_called_once() + + @patch("ingestion.indexers.helpers") + def test_stream_index(self, mock_helpers, sample_chunks): + store, mock_es = self._get_store() + def mock_bulk(client, actions, **kwargs): + return (len(list(actions)), []) + mock_helpers.bulk.side_effect = mock_bulk + metrics = PipelineMetrics() + + count = store.stream_index(sample_chunks, bulk_size=2, metrics=metrics) + assert count == 3 + assert mock_helpers.bulk.call_count == 2 + assert metrics.docs_indexed == 3 + + def test_vector_search(self): + store, mock_es = self._get_store() + mock_es.search.return_value = { + "hits": {"hits": [ + {"_source": {"chunk_id": "c1"}, "_score": 0.9} + ]} + } + results = store.vector_search([0.1]*384, top_k=1) + assert len(results) == 1 + assert results[0].chunk_id == "c1" + + def test_get_stats(self): + store, mock_es = self._get_store() + mock_es.indices.exists.return_value = True + mock_es.indices.stats.return_value = { + "indices": { + store.index_name: {"primaries": {"docs": {"count": 42}}} + } + } + stats = store.get_stats() + assert stats["exists"] is True + assert stats["doc_count"] == 42 diff --git a/tests/test_ingestion_pipeline.py b/tests/test_ingestion_pipeline.py new file mode 100644 index 0000000..0451af1 --- /dev/null +++ b/tests/test_ingestion_pipeline.py @@ -0,0 +1,70 @@ +"""Unit tests for the IngestionPipeline orchestrator.""" + +import pytest +from unittest.mock import MagicMock +from ingestion.pipeline import IngestionPipeline + + +class TestIngestionPipeline: + def test_run_success(self, sample_docs_dir, mock_embedder, mock_vector_store, minimal_pipeline_config): + pipeline = IngestionPipeline( + embedder=mock_embedder, + vector_store=mock_vector_store, + config=minimal_pipeline_config, + state_store=None + ) + + summary = pipeline.run(sample_docs_dir, recreate_index=True, incremental=False) + + assert summary["status"] == "success" + assert summary["doc_count"] > 0 + assert summary["chunk_count"] > 0 + assert summary["indexed_count"] > 0 + mock_vector_store.create_index.assert_called_once_with(dimension=384, recreate=True) + + def test_run_empty_folder(self, tmp_path, mock_embedder, mock_vector_store, minimal_pipeline_config): + pipeline = IngestionPipeline( + embedder=mock_embedder, + vector_store=mock_vector_store, + config=minimal_pipeline_config, + state_store=None + ) + + summary = pipeline.run(str(tmp_path)) + assert summary["status"] == "no_documents" + assert summary["doc_count"] == 0 + + def test_incremental_skip_all(self, sample_docs_dir, mock_embedder, mock_vector_store, minimal_pipeline_config): + mock_state = MagicMock() + mock_state.is_file_changed.return_value = False # Skip everything + + pipeline = IngestionPipeline( + embedder=mock_embedder, + vector_store=mock_vector_store, + config=minimal_pipeline_config, + state_store=mock_state + ) + + summary = pipeline.run(sample_docs_dir, incremental=True) + assert summary["status"] == "all_skipped" + assert summary["files_skipped"] == summary["files_found"] + assert summary["doc_count"] == 0 + + def test_resume_run(self, sample_docs_dir, mock_embedder, mock_vector_store, minimal_pipeline_config): + mock_state = MagicMock() + mock_state.get_pending_chunks.side_effect = lambda run, ids: ids[:1] # Only 1 chunk needs indexing + + pipeline = IngestionPipeline( + embedder=mock_embedder, + vector_store=mock_vector_store, + config=minimal_pipeline_config, + state_store=mock_state + ) + + summary = pipeline.resume_run("test-run-id", sample_docs_dir) + + assert summary["status"] == "success" + # Total chunks produced > 0, but fewer should be indexed + assert summary["chunk_count"] > 0 + # It's mocked to only index 1 chunk per batch of chunks + assert summary["indexed_count"] >= 1 diff --git a/tests/test_llm_providers.py b/tests/test_llm_providers.py new file mode 100644 index 0000000..32b50fb --- /dev/null +++ b/tests/test_llm_providers.py @@ -0,0 +1,40 @@ +"""Unit tests for LLM providers β€” GroqLLMProvider (mocked API).""" + +import pytest +from unittest.mock import MagicMock, patch +from rag.llm_providers import GroqLLMProvider + + +class TestGroqLLMProvider: + def _make_mock_response(self, answer="Test answer", tokens=100): + resp = MagicMock() + resp.choices = [MagicMock(message=MagicMock(content=answer))] + resp.usage = MagicMock(prompt_tokens=60, completion_tokens=40, total_tokens=tokens) + return resp + + def test_generate_returns_answer(self): + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = self._make_mock_response() + with patch("rag.llm_providers.Groq", return_value=mock_client): + provider = GroqLLMProvider(api_key="test_key") + result = provider.generate([{"role": "user", "content": "Hello"}]) + assert result["answer"] == "Test answer" + assert result["total_tokens"] == 100 + assert result["response_time_ms"] >= 0 + + def test_uses_default_model(self): + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = self._make_mock_response() + with patch("rag.llm_providers.Groq", return_value=mock_client): + provider = GroqLLMProvider(api_key="k", default_model="test-model") + result = provider.generate([{"role": "user", "content": "Hi"}]) + assert result["model"] == "test-model" + + def test_custom_model_override(self): + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = self._make_mock_response() + with patch("rag.llm_providers.Groq", return_value=mock_client): + provider = GroqLLMProvider(api_key="k") + result = provider.generate([{"role": "user", "content": "Hi"}], model="custom-model") + call_kwargs = mock_client.chat.completions.create.call_args + assert call_kwargs.kwargs.get("model") or call_kwargs[1].get("model") == "custom-model" diff --git a/tests/test_loaders.py b/tests/test_loaders.py new file mode 100644 index 0000000..411d671 --- /dev/null +++ b/tests/test_loaders.py @@ -0,0 +1,242 @@ +""" +Unit tests for document loaders β€” all 5 loader types plus CompositeDocumentLoader. + +Tests cover: format detection, text extraction, encoding handling, error resilience, +and the composite loader's file discovery and parallel loading capabilities. +""" + +import os +import tempfile +from pathlib import Path + +import pytest + +from domain.models import Document +from ingestion.loaders import ( + BaseDocumentLoader, + CodeDocumentLoader, + CompositeDocumentLoader, + CSVDocumentLoader, + DocxDocumentLoader, + PDFDocumentLoader, + TextDocumentLoader, +) + + +class TestTextDocumentLoader: + """Tests for TextDocumentLoader (.txt, .md, .rst, .log, .env).""" + + def test_can_load_supported_extensions(self): + loader = TextDocumentLoader() + for ext in [".txt", ".md", ".rst", ".log", ".env", ".ini", ".cfg"]: + assert loader.can_load(ext), f"Should support {ext}" + + def test_cannot_load_unsupported(self): + loader = TextDocumentLoader() + assert not loader.can_load(".pdf") + assert not loader.can_load(".py") + + def test_loads_txt_file(self, tmp_path): + f = tmp_path / "test.txt" + f.write_text("Hello, World!", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert docs[0].text == "Hello, World!" + assert docs[0].file_type == ".txt" + assert docs[0].file_name == "test.txt" + + def test_loads_md_file(self, tmp_path): + f = tmp_path / "readme.md" + f.write_text("# Title\n\nContent here.", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "# Title" in docs[0].text + assert docs[0].file_type == ".md" + + def test_empty_file_returns_nothing(self, tmp_path): + f = tmp_path / "empty.txt" + f.write_text("", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert docs == [] + + def test_whitespace_only_returns_nothing(self, tmp_path): + f = tmp_path / "spaces.txt" + f.write_text(" \n\n \t ", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert docs == [] + + def test_modified_date_populated(self, tmp_path): + f = tmp_path / "dated.txt" + f.write_text("content", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert docs[0].modified_date is not None + + def test_nonexistent_file_returns_empty(self): + docs = TextDocumentLoader().load("/nonexistent/path.txt") + assert docs == [] + + def test_source_file_is_absolute(self, tmp_path): + f = tmp_path / "abs.txt" + f.write_text("data", encoding="utf-8") + docs = TextDocumentLoader().load(str(f)) + assert os.path.isabs(docs[0].source_file) + + +class TestCSVDocumentLoader: + """Tests for CSVDocumentLoader (.csv, .tsv, .jsonl).""" + + def test_can_load_csv_tsv_jsonl(self): + loader = CSVDocumentLoader() + assert loader.can_load(".csv") + assert loader.can_load(".tsv") + assert loader.can_load(".jsonl") + + def test_cannot_load_txt(self): + assert not CSVDocumentLoader().can_load(".txt") + + def test_loads_csv(self, tmp_path): + f = tmp_path / "data.csv" + f.write_text("Name,Age\nAlice,30\nBob,25\n", encoding="utf-8") + docs = CSVDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "Alice" in docs[0].text + assert "Bob" in docs[0].text + assert docs[0].file_type == ".csv" + + def test_loads_tsv(self, tmp_path): + f = tmp_path / "data.tsv" + f.write_text("Name\tAge\nAlice\t30\n", encoding="utf-8") + docs = CSVDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "Alice" in docs[0].text + + def test_loads_jsonl(self, tmp_path): + f = tmp_path / "data.jsonl" + f.write_text('{"name": "Alice"}\n{"name": "Bob"}\n', encoding="utf-8") + docs = CSVDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "Alice" in docs[0].text + + def test_empty_csv_returns_nothing(self, tmp_path): + f = tmp_path / "empty.csv" + f.write_text("", encoding="utf-8") + docs = CSVDocumentLoader().load(str(f)) + assert docs == [] + + +class TestCodeDocumentLoader: + """Tests for CodeDocumentLoader (.py, .js, .sql, .json, .yaml, etc.).""" + + def test_can_load_programming_languages(self): + loader = CodeDocumentLoader() + for ext in [".py", ".js", ".ts", ".java", ".go", ".rs", ".sql", ".sh"]: + assert loader.can_load(ext), f"Should support {ext}" + + def test_can_load_config_formats(self): + loader = CodeDocumentLoader() + for ext in [".json", ".yaml", ".yml", ".toml", ".xml", ".html", ".css"]: + assert loader.can_load(ext), f"Should support {ext}" + + def test_cannot_load_pdf_or_docx(self): + loader = CodeDocumentLoader() + assert not loader.can_load(".pdf") + assert not loader.can_load(".docx") + assert not loader.can_load(".txt") + + def test_loads_python_file(self, tmp_path): + f = tmp_path / "script.py" + f.write_text("def hello():\n print('hi')\n", encoding="utf-8") + docs = CodeDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "def hello" in docs[0].text + assert docs[0].file_type == ".py" + + def test_loads_json_file(self, tmp_path): + f = tmp_path / "config.json" + f.write_text('{"key": "value"}', encoding="utf-8") + docs = CodeDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert docs[0].file_type == ".json" + + def test_loads_sql_file(self, tmp_path): + f = tmp_path / "schema.sql" + f.write_text("CREATE TABLE users (id INT PRIMARY KEY);", encoding="utf-8") + docs = CodeDocumentLoader().load(str(f)) + assert len(docs) == 1 + assert "CREATE TABLE" in docs[0].text + + +class TestPDFDocumentLoader: + """Tests for PDFDocumentLoader (requires PyMuPDF).""" + + def test_can_load_pdf(self): + assert PDFDocumentLoader().can_load(".pdf") + assert PDFDocumentLoader().can_load(".PDF") + + def test_cannot_load_others(self): + assert not PDFDocumentLoader().can_load(".txt") + assert not PDFDocumentLoader().can_load(".docx") + + +class TestDocxDocumentLoader: + """Tests for DocxDocumentLoader.""" + + def test_can_load_docx(self): + assert DocxDocumentLoader().can_load(".docx") + assert DocxDocumentLoader().can_load(".DOCX") + + def test_cannot_load_others(self): + assert not DocxDocumentLoader().can_load(".doc") + assert not DocxDocumentLoader().can_load(".pdf") + + +class TestCompositeDocumentLoader: + """Tests for CompositeDocumentLoader (orchestration layer).""" + + def test_supported_extensions_includes_all(self): + exts = CompositeDocumentLoader.get_supported_extensions() + assert ".txt" in exts + assert ".md" in exts + assert ".csv" in exts + assert ".py" in exts + assert ".pdf" in exts + assert ".docx" in exts + + def test_supported_extensions_without_dot(self): + exts = CompositeDocumentLoader.get_supported_extensions_without_dot() + assert "txt" in exts + assert "py" in exts + assert "pdf" in exts + # No dots + assert all(not e.startswith(".") for e in exts) + + def test_load_directory_discovers_files(self, sample_docs_dir): + loader = CompositeDocumentLoader() + docs = loader.load_directory(sample_docs_dir) + # Should load: policy.txt, guide.md, data.csv, utils.py (not image.png) + assert len(docs) >= 4 + + def test_load_directory_skips_unsupported(self, sample_docs_dir): + loader = CompositeDocumentLoader() + docs = loader.load_directory(sample_docs_dir) + file_types = {d.file_type for d in docs} + assert ".png" not in file_types + + def test_nonexistent_directory_raises(self): + loader = CompositeDocumentLoader() + with pytest.raises(FileNotFoundError): + loader.load_directory("/nonexistent/directory") + + def test_empty_directory(self, tmp_path): + loader = CompositeDocumentLoader() + docs = loader.load_directory(str(tmp_path)) + assert docs == [] + + def test_stream_directory_is_generator(self, sample_docs_dir): + loader = CompositeDocumentLoader() + gen = loader.stream_directory(sample_docs_dir) + # Should be an iterator, not a list + import types + assert isinstance(gen, types.GeneratorType) + docs = list(gen) + assert len(docs) >= 4 diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..d132fc6 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,46 @@ +"""Unit tests for pipeline metrics.""" + +import pytest +import time +from ingestion.metrics import PipelineMetrics, StageTimer + + +def test_metrics_defaults(): + m = PipelineMetrics() + assert m.files_found == 0 + assert m.wall_time_s == 0.0 + assert m.cache_hit_rate == 0.0 + +def test_metrics_wall_timer(): + m = PipelineMetrics() + m.start_wall_timer() + time.sleep(0.01) + m.stop_wall_timer() + assert m.wall_time_s >= 0.01 + +def test_metrics_computed_properties(): + m = PipelineMetrics() + m.chunks_cached = 50 + m.chunks_embedded = 50 + m.docs_indexed = 100 + m.wall_time_s = 2.0 + + assert m.cache_hit_rate == 0.5 + assert m.throughput_chunks_per_s == 50.0 + +def test_stage_timer(): + m = PipelineMetrics() + with StageTimer(m, "load"): + time.sleep(0.01) + + assert m.load_time_s >= 0.01 + assert m.chunk_time_s == 0.0 + +def test_metrics_as_dict(): + m = PipelineMetrics() + m.chunks_produced = 10 + d = m.as_dict() + assert "chunks_produced" in d + assert d["chunks_produced"] == 10 + assert "cache_hit_rate" in d + assert "_wall_start" not in d diff --git a/tests/test_prompt_builders.py b/tests/test_prompt_builders.py new file mode 100644 index 0000000..b522ac6 --- /dev/null +++ b/tests/test_prompt_builders.py @@ -0,0 +1,41 @@ +"""Unit tests for prompt builders β€” all 3 prompt styles.""" + +import pytest +from domain.models import SearchResult +from rag.prompt_builders import TemplatePromptBuilder + + +class TestTemplatePromptBuilder: + def test_detailed_style(self, sample_search_results): + msgs = TemplatePromptBuilder().build_prompt("How many PTO days?", sample_search_results, style="detailed") + assert len(msgs) == 2 + assert msgs[0]["role"] == "system" + assert msgs[1]["role"] == "user" + assert "comprehensive" in msgs[0]["content"].lower() + assert "PTO" in msgs[1]["content"] + + def test_concise_style(self, sample_search_results): + msgs = TemplatePromptBuilder().build_prompt("Q?", sample_search_results, style="concise") + assert "concise" in msgs[0]["content"].lower() + + def test_structured_style(self, sample_search_results): + msgs = TemplatePromptBuilder().build_prompt("Q?", sample_search_results, style="structured") + assert "Sources" in msgs[0]["content"] + + def test_unknown_style_defaults_to_detailed(self, sample_search_results): + msgs = TemplatePromptBuilder().build_prompt("Q?", sample_search_results, style="unknown_xyz") + assert "comprehensive" in msgs[0]["content"].lower() + + def test_empty_results(self): + msgs = TemplatePromptBuilder().build_prompt("Q?", [], style="detailed") + assert len(msgs) == 2 + + def test_page_number_included(self): + results = [SearchResult("c1", "d1", "text", "f.pdf", "/f.pdf", ".pdf", page_number=5, score=0.9)] + msgs = TemplatePromptBuilder().build_prompt("Q?", results, style="detailed") + assert "Page: 5" in msgs[1]["content"] + + def test_context_documents_numbered(self, sample_search_results): + msgs = TemplatePromptBuilder().build_prompt("Q?", sample_search_results, style="detailed") + assert "Context Document 1" in msgs[1]["content"] + assert "Context Document 2" in msgs[1]["content"] diff --git a/tests/test_query_rewriters.py b/tests/test_query_rewriters.py new file mode 100644 index 0000000..4167c4a --- /dev/null +++ b/tests/test_query_rewriters.py @@ -0,0 +1,30 @@ +"""Unit tests for query rewriters β€” NoOp and LLMQueryRewriter (mocked).""" + +import pytest +from unittest.mock import MagicMock, patch +from rag.query_rewriters import LLMQueryRewriter, NoOpQueryRewriter + + +class TestNoOpQueryRewriter: + def test_returns_same_query(self): + assert NoOpQueryRewriter().rewrite("what is PTO?") == "what is PTO?" + + def test_empty_query(self): + assert NoOpQueryRewriter().rewrite("") == "" + + +class TestLLMQueryRewriter: + def test_rewrites_query(self): + mock_client = MagicMock() + mock_client.chat.completions.create.return_value = MagicMock( + choices=[MagicMock(message=MagicMock(content="PTO policy accrual days employee benefit"))] + ) + with patch("rag.query_rewriters.Groq", return_value=mock_client): + rewriter = LLMQueryRewriter(api_key="test_key") + result = rewriter.rewrite("what is PTO?") + assert "PTO" in result + + def test_fallback_on_error(self): + with patch("rag.query_rewriters.Groq", side_effect=Exception("API error")): + rewriter = LLMQueryRewriter(api_key="test_key") + assert rewriter.rewrite("original query") == "original query" diff --git a/tests/test_rag_pipeline.py b/tests/test_rag_pipeline.py new file mode 100644 index 0000000..d791a33 --- /dev/null +++ b/tests/test_rag_pipeline.py @@ -0,0 +1,80 @@ +"""Unit tests for the RAGPipeline orchestrator.""" + +import pytest +from unittest.mock import MagicMock +from rag.pipeline import RAGPipeline + + +class TestRAGPipeline: + def test_query_flow(self, sample_search_results): + mock_search = MagicMock() + mock_search.search.return_value = sample_search_results + + mock_rerank = MagicMock() + mock_rerank.rerank.return_value = sample_search_results[:2] + + mock_rewrite = MagicMock() + mock_rewrite.rewrite.return_value = "rewritten" + + mock_prompt = MagicMock() + mock_prompt.build_prompt.return_value = [{"role": "user", "content": "p"}] + + mock_llm = MagicMock() + mock_llm.generate.return_value = { + "answer": "response", "model": "test-model", + "prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, + "response_time_ms": 100 + } + + mock_repo = MagicMock() + mock_repo.save_conversation.return_value = 123 + + pipeline = RAGPipeline( + search_strategy=mock_search, + reranker=mock_rerank, + query_rewriter=mock_rewrite, + prompt_builder=mock_prompt, + llm_provider=mock_llm, + repository=mock_repo + ) + + resp = pipeline.query("q", use_reranking=True, use_query_rewriting=True, save_to_db=True) + + assert resp.answer == "response" + assert resp.rewritten_query == "rewritten" + assert len(resp.retrieved_docs) == 2 + assert resp.conversation_id == 123 + + mock_search.search.assert_called_once() + mock_rerank.rerank.assert_called_once() + mock_rewrite.rewrite.assert_called_once() + mock_prompt.build_prompt.assert_called_once() + mock_llm.generate.assert_called_once() + mock_repo.save_conversation.assert_called_once() + + def test_query_flow_disabled_features(self, sample_search_results): + mock_search = MagicMock() + mock_search.search.return_value = sample_search_results + mock_llm = MagicMock() + mock_llm.generate.return_value = { + "answer": "response", "model": "test-model", + "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, + "response_time_ms": 100 + } + + pipeline = RAGPipeline( + search_strategy=mock_search, + llm_provider=mock_llm, + repository=None + ) + + # Test with built-in NoOp fallbacks via parameters + resp = pipeline.query( + "q", + use_reranking=False, + use_query_rewriting=False, + save_to_db=False + ) + + assert resp.rewritten_query is None + assert len(resp.retrieved_docs) == 3 # Not reranked diff --git a/tests/test_rerankers.py b/tests/test_rerankers.py new file mode 100644 index 0000000..a7688f7 --- /dev/null +++ b/tests/test_rerankers.py @@ -0,0 +1,41 @@ +"""Unit tests for rerankers β€” NoOpReranker and CrossEncoderReranker (mocked).""" + +import pytest +from unittest.mock import MagicMock, patch +from domain.models import SearchResult +from rag.rerankers import BaseReranker, CrossEncoderReranker, NoOpReranker + + +class TestNoOpReranker: + def test_returns_top_n(self, sample_search_results): + reranked = NoOpReranker().rerank("query", sample_search_results, top_n=2) + assert len(reranked) == 2 + assert reranked[0].chunk_id == "r1" + + def test_top_n_larger_than_list(self, sample_search_results): + reranked = NoOpReranker().rerank("query", sample_search_results, top_n=100) + assert len(reranked) == 3 + + def test_empty_results(self): + assert NoOpReranker().rerank("query", [], top_n=5) == [] + + +class TestCrossEncoderReranker: + def test_reranks_by_score(self, sample_search_results): + mock_model = MagicMock() + mock_model.predict.return_value = [0.1, 0.9, 0.5] + reranker = CrossEncoderReranker() + reranker._model = mock_model + reranked = reranker.rerank("query", sample_search_results, top_n=2) + assert len(reranked) == 2 + assert reranked[0].chunk_id == "r2" # highest score 0.9 + assert reranked[0].rerank_score == pytest.approx(0.9) + + def test_empty_results(self): + reranker = CrossEncoderReranker() + assert reranker.rerank("query", [], top_n=5) == [] + + def test_single_result(self, sample_search_results): + reranker = CrossEncoderReranker() + result = reranker.rerank("query", sample_search_results[:1], top_n=1) + assert len(result) == 1 diff --git a/tests/test_retrievers.py b/tests/test_retrievers.py new file mode 100644 index 0000000..d2d2a0d --- /dev/null +++ b/tests/test_retrievers.py @@ -0,0 +1,86 @@ +"""Unit tests for search retrieval strategies (mocked Elasticsearch).""" + +import pytest +from unittest.mock import MagicMock, patch +from domain.models import SearchResult +from rag.retrievers import ( + HybridSearchStrategy, SearchStrategyFactory, + TextSearchStrategy, VectorSearchStrategy, +) + + +def _make_results(ids): + return [SearchResult(chunk_id=cid, doc_id="d1", text=f"text_{cid}", + file_name="f.txt", source_file="/f.txt", file_type=".txt", score=1.0 / (i + 1)) + for i, cid in enumerate(ids)] + + +class TestVectorSearchStrategy: + def test_calls_vector_search(self): + mock_emb = MagicMock() + mock_emb.embed_text.return_value = [0.1] * 384 + mock_store = MagicMock() + mock_store.vector_search.return_value = _make_results(["a", "b"]) + strategy = VectorSearchStrategy(embedder=mock_emb, vector_store=mock_store) + results = strategy.search("query", top_k=5) + assert len(results) == 2 + mock_store.vector_search.assert_called_once() + + +class TestTextSearchStrategy: + def test_calls_text_search(self): + mock_store = MagicMock() + mock_store.text_search.return_value = _make_results(["x", "y"]) + strategy = TextSearchStrategy(vector_store=mock_store) + results = strategy.search("query", top_k=5) + assert len(results) == 2 + mock_store.text_search.assert_called_once() + + +class TestHybridSearchStrategy: + def test_fuses_results_with_rrf(self): + mock_emb = MagicMock() + mock_emb.embed_text.return_value = [0.1] * 384 + mock_store = MagicMock() + mock_store.vector_search.return_value = _make_results(["a", "b", "c"]) + mock_store.text_search.return_value = _make_results(["b", "c", "d"]) + strategy = HybridSearchStrategy(embedder=mock_emb, vector_store=mock_store) + results = strategy.search("query", top_k=3) + assert len(results) == 3 + # b and c appear in both, so they should be scored higher + ids = [r.chunk_id for r in results] + assert "b" in ids and "c" in ids + + def test_rrf_scores_are_set(self): + mock_emb = MagicMock() + mock_emb.embed_text.return_value = [0.1] * 384 + mock_store = MagicMock() + mock_store.vector_search.return_value = _make_results(["a"]) + mock_store.text_search.return_value = _make_results(["a"]) + strategy = HybridSearchStrategy(embedder=mock_emb, vector_store=mock_store) + results = strategy.search("q", top_k=1) + assert results[0].rrf_score is not None + assert results[0].rrf_score > 0 + + +class TestSearchStrategyFactory: + def test_vector(self): + with patch("rag.retrievers.SentenceTransformerEmbedder"), \ + patch("rag.retrievers.ElasticsearchVectorStore"): + s = SearchStrategyFactory.get_strategy("vector") + assert isinstance(s, VectorSearchStrategy) + + def test_text(self): + with patch("rag.retrievers.ElasticsearchVectorStore"): + s = SearchStrategyFactory.get_strategy("text") + assert isinstance(s, TextSearchStrategy) + + def test_hybrid(self): + with patch("rag.retrievers.SentenceTransformerEmbedder"), \ + patch("rag.retrievers.ElasticsearchVectorStore"): + s = SearchStrategyFactory.get_strategy("hybrid") + assert isinstance(s, HybridSearchStrategy) + + def test_unknown_raises(self): + with pytest.raises(ValueError, match="Unknown search strategy"): + SearchStrategyFactory.get_strategy("nonexistent") diff --git a/tests/test_state.py b/tests/test_state.py new file mode 100644 index 0000000..34a10b6 --- /dev/null +++ b/tests/test_state.py @@ -0,0 +1,53 @@ +"""Unit tests for SQLiteCheckpointStore.""" + +import pytest +import tempfile +import os +from ingestion.state import SQLiteCheckpointStore + + +@pytest.fixture +def store(tmp_path): + db_path = tmp_path / "state.db" + return SQLiteCheckpointStore(str(db_path)) + + +def test_is_file_changed_new_file(store, tmp_path): + f = tmp_path / "new.txt" + f.write_text("hello") + assert store.is_file_changed(str(f)) is True + + +def test_is_file_changed_unchanged_file(store, tmp_path): + f = tmp_path / "same.txt" + f.write_text("hello") + store.mark_file_done(str(f)) + assert store.is_file_changed(str(f)) is False + + +def test_is_file_changed_modified_file(store, tmp_path): + f = tmp_path / "mod.txt" + f.write_text("hello") + store.mark_file_done(str(f)) + + f.write_text("world") + assert store.is_file_changed(str(f)) is True + + +def test_run_checkpoints(store): + run_id = store.begin_run() + assert run_id + + # All are pending initially + pending = store.get_pending_chunks(run_id, ["c1", "c2", "c3"]) + assert len(pending) == 3 + + # Checkpoint one + store.checkpoint_chunk(run_id, "c1", "indexed") + + pending = store.get_pending_chunks(run_id, ["c1", "c2", "c3"]) + assert len(pending) == 2 + assert "c1" not in pending + +def test_missing_file_handled_safely(store): + assert store.is_file_changed("/nonexistent/file.txt") is True From 210f8c6826fe3ed2ac891d8fc389bd44ad68c345 Mon Sep 17 00:00:00 2001 From: CipherZ3r0 Date: Sun, 2 Aug 2026 00:01:39 +0530 Subject: [PATCH 2/2] docs: fix typo in README.md --- README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3deac65..0ae14fa 100644 --- a/README.md +++ b/README.md @@ -155,8 +155,8 @@ Potbot automatically logs every interaction into PostgreSQL, which feeds a real- ### Step 1: Clone & Configure Environment ```bash -git clone https://github.com/your-username/llm-zoomcamp-project.git -cd llm-zoomcamp-project +git clone https://github.com/CipherZ3r0/Potbot.git +cd Potbot # Create .env file from template cp .env.example .env @@ -166,7 +166,7 @@ Edit `.env` and insert your `GROQ_API_KEY`: GROQ_API_KEY=gsk_your_actual_groq_api_key_here ``` -### Step 2: Generate Sample Test Documents +### Step 2: Generate Sample Test Documents (optional) ```bash python scripts/generate_sample_documents.py ``` @@ -182,7 +182,7 @@ Access services: - **Grafana Monitoring Dashboard**: `http://localhost:3000` (User: `admin`, Password: `admin`) - **Elasticsearch Cluster**: `http://localhost:9200` -### Step 4: Run Unit Tests +### Step 4: Run Unit Tests (optional) ```bash python -m unittest discover tests ```