diff --git a/README.md b/README.md index 4a1251eb..be20d5fe 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,22 @@ Server starts on `http://localhost:8000` . - Config managed via the web UI Config page (stored in DuckDB) +### AirGap / Offline Mode + +Set `HF_HUB_OFFLINE=1` to prevent HuggingFace API requests (useful for fully +disconnected deployments): + +```bash +HF_HUB_OFFLINE=1 uv run python main.py +``` + +By default, the app sets `HF_HUB_OFFLINE=1` (via `os.environ.setdefault` in +`src/main.py`) to avoid accidental network calls in an AirGap context. + +Functions that **explicitly need online access** (model search, model info, +GGUF file listing) temporarily remove the env var during the API call, then +restore it — so searching and downloading models works even in AirGap mode. + ## Testing ```bash diff --git a/pyproject.toml b/pyproject.toml index 7672fc7e..00cf744a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,27 +4,28 @@ version = "0.1.0" description = "Local RAG system for Sigma rules" requires-python = ">=3.12" dependencies = [ - "fastapi>=0.136.3", - "llama-index>=0.14.22", - "llama-index-vector-stores-qdrant>=0.10.1", - "llama-index-embeddings-huggingface>=0.7.0", + "docx2txt>=0.9", + "duckdb>=1.5.4", + "fastapi>=0.139.0", "gitpython>=3.1.50", - "huggingface-hub>=1.17.0", - "tomli-w>=1.2.0", - "uvicorn>=0.48.0", + "hf-xet>=1.5.1", "httpx>=0.28.1", - "python-multipart>=0.0.30", + "huggingface-hub>=1.22.0", "jinja2>=3.1.6", - "puremagic>=2.2.0", - "pyyaml>=6.0.3", - "duckdb>=1.5.3", - "hf-xet>=1.5.0", + "llama-index>=0.14.23", + "llama-index-embeddings-huggingface>=0.7.0", + "llama-index-llms-openai-like>=0.7.2", "llama-index-readers-file>=0.6.0", - "pymupdf>=1.27.2.3", - "docx2txt>=0.9", + "llama-index-vector-stores-qdrant>=0.10.2", + "qdrant-client>=1.18.0", + "puremagic>=2.2.0", + "pymupdf>=1.28.0", + "python-multipart>=0.0.32", "python-pptx>=1.0.2", - "llama-index-llms-openai-like>=0.7.2", - "rich>=14.0.0", + "pyyaml>=6.0.3", + "rich>=15.0.0", + "tomli-w>=1.2.0", + "uvicorn>=0.49.0", ] [tool.ruff] @@ -80,12 +81,12 @@ pythonpath = ["."] [dependency-groups] dev = [ + "mypy", + "pre-commit>=4.6.0", "pytest>=9.0.3", "pytest-asyncio>=1.4.0", - "ruff", - "mypy", "pytest-cov>=7.1.0", - "pre-commit>=4.6.0", + "ruff", "types-pyyaml>=6.0.12.20260518", ] diff --git a/scripts/benchmark_sparse_encoder.py b/scripts/benchmark_sparse_encoder.py new file mode 100644 index 00000000..284b69c7 --- /dev/null +++ b/scripts/benchmark_sparse_encoder.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +"""Benchmark: compare TF-only vs full BM25 (IDF + length norm) sparse encoders. + +Usage: + uv run python scripts/benchmark_sparse_encoder.py +""" + +from __future__ import annotations + +import statistics +import time +from pathlib import Path + +from src.core.search.sparse_encoder import ( + IDFCalculator, + _encode_single, + _token_id, + _tokenize, + bm25_idf_sparse_encoder, + bm25_sparse_encoder, +) + +DATA_DIR = Path("data/sigmaref") + + +def load_corpus() -> list[str]: + """Load text files from the sigmaref directory.""" + texts: list[str] = [] + if not DATA_DIR.exists(): + print(f"[warn] {DATA_DIR} not found — using synthetic corpus") + return _synthetic_corpus(1_000) + for f in DATA_DIR.rglob("*.txt"): + if f.is_file(): + texts.append(f.read_text(encoding="utf-8", errors="replace")) + texts += [p.read_text(encoding="utf-8", errors="replace") for p in DATA_DIR.rglob("*.md")] + print(f"Loaded {len(texts)} documents from {DATA_DIR}") + return texts + + +def _synthetic_corpus(n: int) -> list[str]: + """Generate realistic Sigma rule-like documents.""" + import random + + random.seed(42) + + titles = [ + "Suspicious Process Creation", + "Registry Persistence via Run Keys", + "PowerShell Encoded Command Execution", + "Scheduled Task Creating Remote Access", + "WMI Persistence Script Execution", + "LSASS Memory Access via Mimikatz", + "DNS Query to Dynamic DNS Domain", + "Service Path Without Quotes", + "Office Application Creating Suspicious Files", + "Net.exe User Account Creation", + "BITSAdmin Download to Temp", + "Certutil URL Download", + "CMSTP Execution", + "DLL Side-Loading via AppInit", + "Event Log Cleared by wevtutil", + "Kerberoasting with RC4 Encryption", + "Mailslot Creation for Pipe Communication", + "Netsh Port Forwarding", + "ODBC Driver Registration", + "Print Spooler Adding Printer Driver", + ] + fields = [ + "Image", + "CommandLine", + "ParentImage", + "TargetObject", + "Details", + "PipeName", + "ServiceFileName", + "RegistryKey", + "Payload", + "QueryName", + "DstIp", + "SrcIp", + ] + values = [ + r"*.exe", + r"*.dll", + r"*.ps1", + r"*.vbs", + r"*.js", + r"powershell.exe -enc *", + r"cmd.exe /c *", + r"reg.exe add *", + r"schtasks.exe /create *", + r"wmic.exe *", + r"net.exe user *", + r"certutil.exe -urlcache *", + r"%temp%\\*", + r"%appdata%\\*", + "SYSTEM\\CurrentControlSet\\Services\\*", + "HKLM\\SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Run\\*", + ] + logsource_categories = [ + "process_creation", + "registry_set", + "file_event", + "network_connection", + "wmi_event", + "dns_query", + "windows_sysmon", + "security_audit", + ] + products = ["windows", "linux", "macos"] + statuses = ["stable", "test", "experimental", "deprecated"] + levels = ["high", "medium", "low", "critical"] + tags_pool = [ + "attack.execution", + "attack.persistence", + "attack.defense_evasion", + "attack.credential_access", + "attack.discovery", + "attack.lateral_movement", + "attack.collection", + "attack.command_and_control", + "attack.t1059.001", + "attack.t1547.001", + "attack.t1003.001", + "attack.t1053.005", + "attack.t1047", + "attack.t1087.002", + ] + + docs: list[str] = [] + for i in range(n): + title = random.choice(titles) + status = random.choice(statuses) + level = random.choice(levels) + product = random.choice(products) + category = random.choice(logsource_categories) + field = random.choice(fields) + value = random.choice(values) + tags = " ".join(random.sample(tags_pool, random.randint(2, 4))) + detection = f"{field} contains '{value}'" + + doc = ( + f"title: {title}\n" + f"id: synth-{i:06d}\n" + f"status: {status}\n" + f"description: Detects {title.lower()} technique used by threat actors for persistence " + f"and privilege escalation on {product} systems. Related to MITRE ATT&CK techniques.\n" + f"author: Benchmark Generator\n" + f"date: 2024/01/01\n" + f"tags: {tags}\n" + f"logsource:\n" + f" category: {category}\n" + f" product: {product}\n" + f"detection:\n" + f" selection:\n" + f" {detection}\n" + f" condition: selection\n" + f"falsepositives:\n" + f" - Legitimate administrative activity\n" + f" - Software installation\n" + f"level: {level}\n" + f"references:\n" + f" - https://attack.mitre.org/techniques/T{1000 + (i % 9000):04d}/\n" + ) + docs.append(doc) + + return docs + + +def _overlap(j: list[int], k: list[int]) -> float: + s = set(j) + return sum(1 for x in k if x in s) + + +def benchmark_speed(encoder, texts: list[str], name: str) -> float: + start = time.perf_counter() + for _ in range(3): + encoder(texts) + elapsed = time.perf_counter() - start + avg = elapsed / 3 + print(f" {name}: {avg:.4f}s avg (3 runs, {len(texts)} docs)") + return avg + + +def benchmark_quality( + query_texts: list[str], + corpus: list[str], + idf_map: dict[str, float], + avg_doc_len: float, +) -> None: + """Compare top-10 overlap between TF-only and full BM25.""" + print("\n=== Quality comparison (top-k term overlap) ===\n") + + for q in query_texts: + tf_indices, _ = _encode_single(q) + bm25_indices, _ = _encode_single(q, idf_map=idf_map, avg_doc_len=avg_doc_len) + overlap_frac = _overlap(tf_indices, bm25_indices) / max(len(tf_indices), 1) + print( + f" query: {q[:60]:<60s} " + f"TF terms: {len(tf_indices):>3d} " + f"BM25 terms: {len(bm25_indices):>3d} " + f"overlap: {overlap_frac:.0%}" + ) + + +def compute_corpus_stats( + corpus: list[str], +) -> tuple[dict[str, float], float]: + """Compute IDF map and average document length from the corpus.""" + calc = IDFCalculator() + total_tokens = 0 + for doc in corpus: + total_tokens += calc.add_document(doc) + avg_doc_len = total_tokens / max(len(corpus), 1) + idf_map = calc.idf() + print(f"\nCorpus: {len(corpus)} docs, {len(idf_map)} unique terms, avg len={avg_doc_len:.1f}") + return idf_map, avg_doc_len + + +def main() -> None: + corpus = load_corpus() + idf_map, avg_doc_len = compute_corpus_stats(corpus) + + queries = [ + "process creation with image endswith exe", + "registry modification run keys persistence", + "network connection suspicious ip address", + "powershell encoded command execution", + "scheduled task lateral movement", + ] + + # Speed benchmark + print("\n=== Speed benchmark ===\n") + benchmark_speed(bm25_sparse_encoder, corpus, "TF-only (1+log(tf))") + benchmark_speed( + lambda t: bm25_idf_sparse_encoder(t, idf_map=idf_map, avg_doc_len=avg_doc_len), + corpus, + "Full BM25 (IDF + length norm)", + ) + + # Quality benchmark + benchmark_quality(queries, corpus, idf_map, avg_doc_len) + + # IDF distribution + if idf_map: + values = list(idf_map.values()) + print("\n=== IDF distribution ===\n") + print(f" min: {min(values):.3f}") + print(f" max: {max(values):.3f}") + print(f" mean: {statistics.mean(values):.3f}") + print( + f" <1.0: {sum(1 for v in values if v < 1.0)} terms ({sum(1 for v in values if v < 1.0) / len(values):.1%})" + ) + + # Term ID collision rate + tokens_seen = {t for doc in corpus for t in _tokenize(doc.lower()) if len(t) >= 3} + ids = {_token_id(t) for t in tokens_seen} + print("\n=== Token-ID collision ===\n") + print(f" Unique tokens: {len(tokens_seen)}") + print(f" Unique IDs: {len(ids)}") + print(f" Collision rate: {1 - len(ids) / max(len(tokens_seen), 1):.4%}") + + # Qdrant native BM25 recommendation + print("\n=== Recommendation ===\n") + print(" Qdrant modifier=IDF: available via SparseVectorParams(modifier=Modifier.IDF)") + print(" Custom IDF encoder: implemented (use bm25_idf_sparse_encoder)") + print(" Recommendation: enable modifier=IDF in collection creation") + print(" + keep custom encoder as fallback for offline/embedding") + + +if __name__ == "__main__": + main() diff --git a/src/api/v1/base/repo_router.py b/src/api/v1/base/repo_router.py index 45936e87..4559b4b8 100644 --- a/src/api/v1/base/repo_router.py +++ b/src/api/v1/base/repo_router.py @@ -128,6 +128,9 @@ def create_repo_router( include_scan_endpoint: bool = True, on_delete_cleanup: Callable[[str, str], None] | None = None, use_get_for_sync: bool = False, + select_dirs_worker: WorkerName = WorkerName.GITHUB_DISCOVERY, + scan_workers: list[WorkerName] | None = None, + source_type: str = "", ) -> APIRouter: """Create a configured repository management router. @@ -151,6 +154,13 @@ def create_repo_router( Extra cleanup callback on repo deletion (receives ``org, name``). use_get_for_sync : bool Use GET instead of POST for the single-repo sync endpoint. + select_dirs_worker : WorkerName + Worker type to dispatch after saving selected dirs (default: GITHUB_DISCOVERY). + scan_workers : list[WorkerName] | None + Worker types to dispatch after a scan sync. When ``None``, fires all three discovery + workers (GITHUB, LOCAL, SPEC) for backward compatibility. + source_type : str + Source type identifier stored with selected dirs (e.g. 'github', 'spec'). """ router = APIRouter(prefix=prefix, tags=tags) # type: ignore[arg-type] _sync_lock = Lock() @@ -254,7 +264,7 @@ def clone_with_status() -> None: }, repos_dir=rd, ) - save_selected_dirs(org, name, [], repos_dir=rd) + save_selected_dirs(org, name, [], repos_dir=rd, source_type=source_type) else: save_metadata( org, @@ -424,9 +434,11 @@ async def get_repo_tree( async def select_dirs( org: str, name: str, - request: SelectDirsRequest, + request_body: SelectDirsRequest, + background_tasks: BackgroundTasks, + request: Request, ) -> SelectDirsResponse: - """Save selected directories for a repository.""" + """Save selected directories for a repository and trigger discovery.""" try: _validate_org_name(org, name) except HTTPException as e: @@ -440,11 +452,21 @@ async def select_dirs( ) except Exception as e: return SelectDirsResponse(success=False, error=str(e)) - result = save_selected_dirs(org, name, request.selected, repos_dir=repos_dir) + result = save_selected_dirs( + org, name, request_body.selected, repos_dir=repos_dir, source_type=source_type + ) if result.get("success"): + dispatcher = request.app.state.dispatcher + background_tasks.add_task( + dispatcher.ask_for_worker, + select_dirs_worker, + task_type=select_dirs_worker.value, + collection_name="all", + repo_key=f"{org}/{name}", + ) return SelectDirsResponse( success=True, - message=f"Saved {len(request.selected)} selected directories for {org}/{name}", + message=f"Saved {len(request_body.selected)} selected directories for {org}/{name}", ) return SelectDirsResponse( success=False, @@ -481,11 +503,18 @@ async def scan_repo( def _run_scan() -> None: with _sync_lock: _sync_single_repo(org, name, None) - for wc, coll in [ - (WorkerName.GITHUB_DISCOVERY, "all"), - (WorkerName.LOCAL_DISCOVERY, "local"), - (WorkerName.SPEC_DISCOVERY, "spec"), - ]: + _workers = scan_workers or [ + WorkerName.GITHUB_DISCOVERY, + WorkerName.LOCAL_DISCOVERY, + WorkerName.SPEC_DISCOVERY, + ] + _collection_map = { + WorkerName.GITHUB_DISCOVERY: "all", + WorkerName.LOCAL_DISCOVERY: "local", + WorkerName.SPEC_DISCOVERY: "spec", + } + for wc in _workers: + coll = _collection_map.get(wc, "all") dispatcher.ask_for_worker(wc, task_type=wc.value, collection_name=coll) background_tasks.add_task(_run_scan) diff --git a/src/api/v1/chat/chat.py b/src/api/v1/chat/chat.py index a702a654..ed0d192f 100644 --- a/src/api/v1/chat/chat.py +++ b/src/api/v1/chat/chat.py @@ -5,7 +5,7 @@ import logging from datetime import UTC, datetime -from fastapi import APIRouter, HTTPException, UploadFile, status +from fastapi import APIRouter, Depends, Header, HTTPException, UploadFile, status from fastapi.responses import StreamingResponse from src.application.chat.service import ChatService @@ -18,6 +18,11 @@ _chat_service: ChatService | None = None +def _get_session_id(x_session_id: str | None = Header(None)) -> str | None: + """Extract session ID from ``X-Session-ID`` header.""" + return x_session_id + + def _get_chat_service() -> ChatService: global _chat_service if _chat_service is None: @@ -26,19 +31,26 @@ def _get_chat_service() -> ChatService: @router.get("/history") -async def get_chat_history() -> list[dict]: +async def get_chat_history( + session_id: str | None = Depends(_get_session_id), +) -> list[dict]: """Get chat message history.""" - return _get_chat_service().get_history() + return _get_chat_service().get_history(session_id) @router.delete("/history", status_code=status.HTTP_204_NO_CONTENT) -async def clear_chat_history() -> None: +async def clear_chat_history( + session_id: str | None = Depends(_get_session_id), +) -> None: """Clear chat history and llama.cpp KV cache.""" - await _get_chat_service().clear_history() + await _get_chat_service().clear_history(session_id) @router.post("/message", response_model=ChatMessageResponse) -async def send_chat_message(req: ChatMessageRequest) -> ChatMessageResponse: +async def send_chat_message( + req: ChatMessageRequest, + session_id: str | None = Depends(_get_session_id), +) -> ChatMessageResponse: """Process a chat message and return AI response.""" if not req.message or not req.message.strip(): raise HTTPException( @@ -49,9 +61,13 @@ async def send_chat_message(req: ChatMessageRequest) -> ChatMessageResponse: try: svc = _get_chat_service() response_text = await svc.process_message( - req.message, req.mode, req.model, prompt_id=req.prompt_id + req.message, + req.mode, + req.model, + prompt_id=req.prompt_id, + session_id=session_id, ) - citations = svc.get_last_citations() + citations = svc.get_last_citations(session_id) # Format citations as [sigma:rule_id] in response if citations and response_text: @@ -73,8 +89,18 @@ async def send_chat_message(req: ChatMessageRequest) -> ChatMessageResponse: @router.post("/upload") -async def upload_sigma_rule(file: UploadFile) -> dict: - """Upload and validate a Sigma rule YAML file.""" +async def upload_sigma_rule( + file: UploadFile, + session_id: str | None = Depends(_get_session_id), + index_after_upload: bool = False, +) -> dict: + """Upload and validate a Sigma rule YAML file. + + Args: + file: The YAML file to upload + session_id: Optional session ID + index_after_upload: If True, index the rule to Qdrant immediately + """ if not file.filename: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -102,13 +128,22 @@ async def upload_sigma_rule(file: UploadFile) -> dict: content.decode("utf-8") except UnicodeDecodeError: raise HTTPException(status_code=400, detail="File is not valid UTF-8 text") - rule_data = await _get_chat_service().validate_and_store_yaml(content) + rule = await _get_chat_service().validate_and_store_yaml(content, session_id) - return { - "rule_name": rule_data.get("name", "Unknown"), - "rule_id": rule_data.get("id", ""), + result: dict = { + "rule_name": rule.name, + "rule_id": rule.id, "validated": True, } + + if index_after_upload: + from src.application.documents.indexing import index_sigma_rules + + indexed = await index_sigma_rules([rule], mode="flat", index_after_upload=True) + result["indexed"] = indexed.get("indexed", 0) + result["collection"] = indexed.get("collection") + + return result except ValidationError as e: raise HTTPException( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, @@ -123,7 +158,10 @@ async def upload_sigma_rule(file: UploadFile) -> dict: @router.post("/message/stream") -async def send_chat_message_stream(req: ChatMessageRequest): +async def send_chat_message_stream( + req: ChatMessageRequest, + session_id: str | None = Depends(_get_session_id), +): """Process a chat message and stream the LLM response.""" if not req.message or not req.message.strip(): raise HTTPException( @@ -135,7 +173,11 @@ async def generate(): """Generate SSE events from LLM stream.""" try: async for token in _get_chat_service().process_message_stream( - req.message, req.mode, req.model, prompt_id=req.prompt_id + req.message, + req.mode, + req.model, + prompt_id=req.prompt_id, + session_id=session_id, ): yield f"data: {token}\n\n" yield "data: [DONE]\n\n" diff --git a/src/api/v1/chat/search.py b/src/api/v1/chat/search.py index b2d6b484..93b52080 100644 --- a/src/api/v1/chat/search.py +++ b/src/api/v1/chat/search.py @@ -6,7 +6,7 @@ from fastapi import APIRouter, HTTPException, status -from src.core.search.engine import SearchEngine +from src.core.search.engine import SearchEngine, format_result_by_collection from src.api.v1.chat.schemas import SearchRequest, SearchResponse logger = logging.getLogger(__name__) @@ -29,8 +29,9 @@ async def search_rules_endpoint( clean_query = request.query.replace("`", "").rstrip("?").strip() engine = SearchEngine(use_router=request.use_router) results = await engine.search(clean_query, top_k=request.limit) + formatted = [format_result_by_collection(r) for r in results] return SearchResponse( - data=results, + data=formatted, meta={"total": len(results), "query": request.query, "routed": request.use_router}, ) except Exception as e: diff --git a/src/api/v1/documents/files.py b/src/api/v1/documents/files.py index 4f53a14b..c448c85c 100644 --- a/src/api/v1/documents/files.py +++ b/src/api/v1/documents/files.py @@ -8,6 +8,8 @@ from pathlib import Path from typing import Any +from src.shared.constants import NULL_UUID + from fastapi import APIRouter, Depends, UploadFile from fastapi import File as FastAPIFile from pydantic import BaseModel @@ -144,7 +146,7 @@ async def add_local_file( "file_size": file_size, "original_url": f"file://{dest_path.as_posix()}", "normalized_url": f"file://{dest_path.as_posix()}", - "rule_id": "00000000-0000-0000-0000-000000000000", + "rule_id": NULL_UUID, "title": title, "timestamp": iso_now(), "last_seen": iso_now(), diff --git a/src/api/v1/documents/spec.py b/src/api/v1/documents/spec.py index e886ed12..e2f5bd48 100644 --- a/src/api/v1/documents/spec.py +++ b/src/api/v1/documents/spec.py @@ -6,6 +6,7 @@ from src.api.v1.base.repo_router import create_repo_router from src.config.settings import get_config from src.infrastructure.database import DatabaseService +from src.workers.enums import WorkerName logger = logging.getLogger(__name__) @@ -28,4 +29,7 @@ def _cleanup_sigma_spec(org: str, name: str) -> None: include_outdated_check=False, on_delete_cleanup=_cleanup_sigma_spec, use_get_for_sync=True, + select_dirs_worker=WorkerName.SPEC_DISCOVERY, + scan_workers=[WorkerName.SPEC_DISCOVERY], + source_type="spec", ) diff --git a/src/api/v1/infrastructure/github.py b/src/api/v1/infrastructure/github.py index 14f6b4c5..9d5e2a25 100644 --- a/src/api/v1/infrastructure/github.py +++ b/src/api/v1/infrastructure/github.py @@ -12,4 +12,5 @@ include_detail_endpoint=True, include_status_endpoint=True, include_outdated_check=True, + source_type="github", ) diff --git a/src/api/v1/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index af95d936..637e0ed5 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -64,8 +64,18 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N collection_name=collection_name, vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE), sparse_vectors_config={ - "text-sparse": models.SparseVectorParams(index=models.SparseIndexParams()) + "text-sparse": models.SparseVectorParams( + index=models.SparseIndexParams(), + modifier=models.Modifier.IDF, + ) }, + quantization_config=models.ScalarQuantization( + scalar=models.ScalarQuantizationConfig( + type=models.ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) + ), ) diff --git a/src/api/v1/models/models_embedding.py b/src/api/v1/models/models_embedding.py index de74fced..740b4116 100644 --- a/src/api/v1/models/models_embedding.py +++ b/src/api/v1/models/models_embedding.py @@ -113,9 +113,13 @@ async def search_embedding_models( manager: EmbeddingManager = Depends(get_embedding_manager), ) -> JSONResponse: """Search for embedding models on HuggingFace.""" + from src.application.models.exceptions import DownloadError + try: results = await manager.search_models(query, limit=limit) return JSONResponse(content={"models": [{"repo_id": r.full_id} for r in results]}) + except DownloadError as e: + return JSONResponse(status_code=400, content={"error": str(e)}) except Exception as e: logger.error(f"Search failed: {e}") return JSONResponse(status_code=500, content={"error": "An internal error occurred"}) diff --git a/src/api/v1/models/models_llm.py b/src/api/v1/models/models_llm.py index 868a45b0..9c8d2d12 100644 --- a/src/api/v1/models/models_llm.py +++ b/src/api/v1/models/models_llm.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os from fastapi import APIRouter from fastapi.responses import JSONResponse @@ -74,10 +75,14 @@ async def search_llm_models(query: str, limit: int = 20) -> JSONResponse: No pipeline tag filter is applied because many GGUF repos (e.g. MaziyarPanahi/*, bartowski/*) do not set a pipeline_tag. """ + from src.application.models.exceptions import DownloadError + try: service = HFDownloadService() results = await service.list_models(query) return JSONResponse(content={"models": [{"repo_id": r.full_id} for r in results[:limit]]}) + except DownloadError as e: + return JSONResponse(status_code=400, content={"error": str(e)}) except Exception as e: logger.error(f"LLM search failed: {e}") return JSONResponse(status_code=500, content={"error": "An internal error occurred"}) @@ -91,7 +96,7 @@ async def list_llm_model_files(repo_id: str | None = None) -> JSONResponse: try: mm = get_embedding_manager() - files = mm.download_service.list_gguf_files(HFRepo.from_string(repo_id)) + files = await mm.download_service.list_gguf_files(HFRepo.from_string(repo_id)) return JSONResponse(content={"files": files}) except Exception as e: logger.error(f"Failed to list files for {repo_id}: {e}") @@ -171,7 +176,7 @@ async def download_in_background(): resolved_filename = filename if not resolved_filename: mm = get_embedding_manager() - files = mm.download_service.list_gguf_files(HFRepo.from_string(repo_id)) + files = await mm.download_service.list_gguf_files(HFRepo.from_string(repo_id)) if files: resolved_filename = files[0]["filename"] if not resolved_filename: @@ -183,11 +188,20 @@ async def download_in_background(): dest_dir = LLM_DIR / repo.owner / repo.name dest_dir.mkdir(parents=True, exist_ok=True) - hf_hub_download( - repo_id=repo_id, - filename=resolved_filename, - local_dir=dest_dir, - ) + import huggingface_hub.constants as hc + + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False + try: + _raw_token = os.environ.get("HF_TOKEN") + hf_hub_download( + repo_id=repo_id, + filename=resolved_filename, + local_dir=dest_dir, + token=_raw_token if _raw_token else None, + ) + finally: + hc.HF_HUB_OFFLINE = was_offline db = get_database_service() reg = get_unified_registry() @@ -202,8 +216,14 @@ async def download_in_background(): from src.infrastructure.llm.llamacpp.auto_start import start_llamacpp as _start_llm from src.infrastructure.vectorstore.auto_start import start_qdrant as _start_qd - await _start_qd() - await _start_llm() + try: + await _start_qd() + except Exception: + logger.exception("Failed to auto-start Qdrant after download") + try: + await _start_llm() + except Exception: + logger.exception("Failed to auto-start llama.cpp after download") asyncio.create_task(download_in_background()) diff --git a/src/api/v1/models/models_sparse.py b/src/api/v1/models/models_sparse.py index 2263da32..02b4df96 100644 --- a/src/api/v1/models/models_sparse.py +++ b/src/api/v1/models/models_sparse.py @@ -4,7 +4,6 @@ import asyncio import logging -import os from typing import Any from fastapi import APIRouter @@ -42,7 +41,10 @@ async def download_in_background() -> None: SPARSE_MODEL_DIR.mkdir(parents=True, exist_ok=True) - was_offline = os.environ.pop("HF_HUB_OFFLINE", None) + import huggingface_hub.constants as hc + + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False try: from transformers import AutoModelForMaskedLM, AutoTokenizer @@ -56,8 +58,7 @@ async def download_in_background() -> None: tokenizer.save_pretrained(str(SPARSE_MODEL_DIR)) model.save_pretrained(str(SPARSE_MODEL_DIR)) finally: - if was_offline is not None: - os.environ["HF_HUB_OFFLINE"] = was_offline + hc.HF_HUB_OFFLINE = was_offline _download_progress["sparse"] = {"progress": 100, "status": "completed"} except Exception as e: diff --git a/src/api/v1/sigma/translate.py b/src/api/v1/sigma/translate.py index 8698aa32..38aeacaa 100644 --- a/src/api/v1/sigma/translate.py +++ b/src/api/v1/sigma/translate.py @@ -17,6 +17,17 @@ DEFAULT_PROMPT_ID = "vulgarisation-english" +# Module-level singleton to share SearchEngine/LlamaClient/Cache with ChatService +_rag_pipeline: RAGPipeline | None = None + + +def get_rag_pipeline() -> RAGPipeline: + """Return the shared RAGPipeline singleton.""" + global _rag_pipeline + if _rag_pipeline is None: + _rag_pipeline = RAGPipeline() + return _rag_pipeline + class TranslateDetectionRequest(BaseModel): """Request body for translating a Sigma detection block.""" @@ -57,7 +68,7 @@ async def translate_detection_endpoint(req: TranslateDetectionRequest) -> JSONRe content={"error": "yaml is required and cannot be empty"}, ) - rag = RAGPipeline() + rag = get_rag_pipeline() try: translation = await translate_detection( diff --git a/src/application/chat/rag.py b/src/application/chat/rag.py index 99804fbe..c93d6d3c 100644 --- a/src/application/chat/rag.py +++ b/src/application/chat/rag.py @@ -11,6 +11,7 @@ from jinja2 import Template from src.core.search.engine import SearchEngine +from src.core.sigma.models import SigmaRule from src.infrastructure.llm.llamacpp import LlamaClient from src.application.system.cache import ResponseCache @@ -96,25 +97,25 @@ def _fallback_prompt(self) -> str: async def explain_rule( self, - rule_data: dict[str, Any], + rule: SigmaRule, related_results: list[dict[str, Any]] | None = None, system_prompt_id: str = "", ) -> str: """Generate explanation for an uploaded Sigma rule.""" related_text = self._format_search_results(related_results or []) cache_key = self.cache.generate_key( - query=rule_data.get("name", ""), + query=rule.name, context=related_text, ) cached = self.cache.get(cache_key) if cached: - logger.info(f"Cache hit for rule explanation: {rule_data.get('name')}") + logger.info(f"Cache hit for rule explanation: {rule.name}") return cached try: prompt_content = self._resolve_prompt(system_prompt_id, mode="explain") - rule_yaml = await self._format_rule_yaml(rule_data) + rule_yaml = await self._format_rule_yaml(rule) prompt = Template(prompt_content).render( uploaded_rule=rule_yaml, related_rules=related_text, @@ -128,31 +129,31 @@ async def explain_rule( return response except Exception as e: logger.error(f"LLM generation failed: {e}") - return self._fallback_explanation(rule_data) + return self._fallback_explanation(rule) async def explain_rule_stream( self, - rule_data: dict[str, Any], + rule: SigmaRule, related_results: list[dict[str, Any]] | None = None, system_prompt_id: str = "", ) -> AsyncGenerator[str, None]: """Stream explanation for an uploaded Sigma rule.""" related_text = self._format_search_results(related_results or []) cache_key = self.cache.generate_key( - query=rule_data.get("name", ""), + query=rule.name, context=related_text, ) cached = self.cache.get(cache_key) if cached: - logger.info(f"Cache hit for rule explanation: {rule_data.get('name')}") + logger.info(f"Cache hit for rule explanation: {rule.name}") for token in cached: yield token return try: prompt_content = self._resolve_prompt(system_prompt_id, mode="explain") - rule_yaml = await self._format_rule_yaml(rule_data) + rule_yaml = await self._format_rule_yaml(rule) prompt = Template(prompt_content).render( uploaded_rule=rule_yaml, related_rules=related_text, @@ -166,7 +167,7 @@ async def explain_rule_stream( yield token except Exception as e: logger.error(f"LLM generation failed: {e}") - fallback = self._fallback_explanation(rule_data) + fallback = self._fallback_explanation(rule) for token in fallback: yield token @@ -268,25 +269,25 @@ async def answer_search_query( async def analyze_coverage( self, - rule_data: dict[str, Any], + rule: SigmaRule, related_results: list[dict[str, Any]], system_prompt_id: str = "", ) -> str: """Analyze detection coverage gaps.""" related_text = self._format_search_results(related_results) cache_key = self.cache.generate_key( - query=rule_data.get("name", ""), + query=rule.name, context=related_text, ) cached = self.cache.get(cache_key) if cached: - logger.info(f"Cache hit for coverage analysis: {rule_data.get('name')}") + logger.info(f"Cache hit for coverage analysis: {rule.name}") return cached try: prompt_content = self._resolve_prompt(system_prompt_id, mode="coverage") - rule_yaml = await self._format_rule_yaml(rule_data) + rule_yaml = await self._format_rule_yaml(rule) prompt = Template(prompt_content).render( uploaded_rule=rule_yaml, related_rules=related_text, @@ -304,27 +305,27 @@ async def analyze_coverage( async def analyze_coverage_stream( self, - rule_data: dict[str, Any], + rule: SigmaRule, related_results: list[dict[str, Any]], system_prompt_id: str = "", ) -> AsyncGenerator[str, None]: """Stream coverage analysis for an uploaded Sigma rule.""" related_text = self._format_search_results(related_results) cache_key = self.cache.generate_key( - query=rule_data.get("name", ""), + query=rule.name, context=related_text, ) cached = self.cache.get(cache_key) if cached: - logger.info(f"Cache hit for coverage analysis: {rule_data.get('name')}") + logger.info(f"Cache hit for coverage analysis: {rule.name}") for token in cached: yield token return try: prompt_content = self._resolve_prompt(system_prompt_id, mode="coverage") - rule_yaml = await self._format_rule_yaml(rule_data) + rule_yaml = await self._format_rule_yaml(rule) prompt = Template(prompt_content).render( uploaded_rule=rule_yaml, related_rules=related_text, @@ -377,21 +378,22 @@ def _format_search_results(self, results: list[dict[str, Any]]) -> str: return "\n\n".join(lines) - async def _format_rule_yaml(self, rule: dict[str, Any]) -> str: + async def _format_rule_yaml(self, rule: SigmaRule) -> str: """Format rule data as readable YAML-like text.""" loop = asyncio.get_event_loop() return await loop.run_in_executor( - None, lambda: yaml.dump(rule, default_flow_style=False, allow_unicode=True) + None, + lambda: yaml.dump(rule.to_dict(), default_flow_style=False, allow_unicode=True), ) - def _fallback_explanation(self, rule_data: dict[str, Any]) -> str: + def _fallback_explanation(self, rule: SigmaRule) -> str: """Fallback explanation without LLM.""" parts = [ - f"**Rule:** {rule_data.get('name', 'Unknown')}", - f"**ID:** {rule_data.get('id', 'N/A')}", + f"**Rule:** {rule.name}", + f"**ID:** {rule.id}", ] - if desc := rule_data.get("description"): - parts.append(f"**Description:** {desc}") + if rule.description: + parts.append(f"**Description:** {rule.description}") return "\n".join(parts) def _fallback_search_results(self, results: list[dict[str, Any]]) -> str: diff --git a/src/application/chat/service.py b/src/application/chat/service.py index 88cf4d69..5c653084 100644 --- a/src/application/chat/service.py +++ b/src/application/chat/service.py @@ -16,26 +16,36 @@ translate_detection, ) from src.application.tools import ToolContext, ToolDispatcher, get_tools -from src.core.search.engine import SearchEngine +from src.core.search.engine import SearchEngine, format_result_by_collection +from src.core.sigma.models import SigmaRule from src.api.v1.chat.schemas import ChatMode +from src.shared.session import SessionStore logger = logging.getLogger(__name__) MAX_HISTORY = 50 MAX_TOOL_CALLS = 5 +_session_store = SessionStore() + + +def get_session_store() -> SessionStore: + return _session_store + class ChatService: - """Service layer for chat operations.""" + """Service layer for chat operations. + + Session-scoped state (history, uploaded rule, citations) is stored in + a shared :class:`SessionStore` keyed by *session_id*. Callers should + pass a unique *session_id* (e.g. from ``X-Session-ID`` header) to + avoid state leaking between users. + """ def __init__(self, use_router: bool = True) -> None: self.search_engine = SearchEngine(use_router=use_router) self.rag_pipeline = RAGPipeline() self.validator = SigmaValidator() - self._history: list[dict[str, str]] = [] - self._uploaded_rule: dict[str, Any] | None = None - self._last_citations: list[str] = [] - self._current_prompt_id: str = "" # Tool-calling setup self._tool_executor = ToolDispatcher(get_tools()) @@ -45,6 +55,31 @@ def __init__(self, use_router: bool = True) -> None: rag_pipeline=self.rag_pipeline, ) + # ------------------------------------------------------------------ + # Session helpers + # ------------------------------------------------------------------ + + def _sid(self, session_id: str | None) -> str: + return session_id or "_default" + + def _get_rule(self, session_id: str | None) -> SigmaRule | None: + return get_session_store().get(self._sid(session_id), "_uploaded_rule") + + def _set_rule(self, session_id: str | None, rule: SigmaRule | None) -> None: + get_session_store().set(self._sid(session_id), "_uploaded_rule", rule) + + def _get_history(self, session_id: str | None) -> list[dict[str, str]]: + return get_session_store().get(self._sid(session_id), "_history", []) + + def _set_history(self, session_id: str | None, history: list[dict[str, str]]) -> None: + get_session_store().set(self._sid(session_id), "_history", history) + + def _get_citations(self, session_id: str | None) -> list[str]: + return get_session_store().get(self._sid(session_id), "_last_citations", []) + + def _set_citations(self, session_id: str | None, citations: list[str]) -> None: + get_session_store().set(self._sid(session_id), "_last_citations", citations) + def _get_tool_schemas(self) -> list[dict[str, Any]]: """Return OpenAI-compatible tool schemas for the LLM.""" return self._tool_executor.list_tools() @@ -55,6 +90,7 @@ async def process_message( mode: str = ChatMode.SEARCH.value, model: str = "", prompt_id: str = "", + session_id: str | None = None, ) -> str: """Process a chat message based on the current mode. @@ -63,6 +99,7 @@ async def process_message( mode: Chat mode (search, coverage, explain) model: Selected LLM model path prompt_id: Selected system prompt ID + session_id: Session identifier for state isolation. Returns: AI response text @@ -70,21 +107,21 @@ async def process_message( if not message.strip(): return "" - self._add_to_history("user", message) - self._current_prompt_id = prompt_id + sid = self._sid(session_id) + self._add_to_history("user", message, sid) try: if mode == ChatMode.EXPLAIN.value: - response = await self._handle_explain(message) + response = await self._handle_explain(message, prompt_id, sid) elif mode == ChatMode.COVERAGE.value: - response = await self._handle_coverage(message) + response = await self._handle_coverage(message, prompt_id, sid) else: - response = await self._handle_search(message) + response = await self._handle_search(message, prompt_id) except Exception as e: logger.error(f"Chat processing error: {e}") response = f"Error processing message: {str(e)}" - self._add_to_history("assistant", response) + self._add_to_history("assistant", response, sid) return response async def process_message_stream( @@ -93,13 +130,14 @@ async def process_message_stream( mode: str = ChatMode.SEARCH.value, model: str = "", prompt_id: str = "", + session_id: str | None = None, ) -> AsyncGenerator[str, None]: """Stream a chat message response based on the current mode.""" if not message.strip(): return - self._add_to_history("user", message) - self._current_prompt_id = prompt_id + sid = self._sid(session_id) + self._add_to_history("user", message, sid) if model: logger.info("Selected model: %s", model) if prompt_id: @@ -108,15 +146,15 @@ async def process_message_stream( accumulated: list[str] = [] try: if mode == ChatMode.EXPLAIN.value: - async for token in self._handle_explain_stream(message): + async for token in self._handle_explain_stream(message, prompt_id, sid): accumulated.append(token) yield token elif mode == ChatMode.COVERAGE.value: - async for token in self._handle_coverage_stream(message): + async for token in self._handle_coverage_stream(message, prompt_id, sid): accumulated.append(token) yield token else: - async for token in self._handle_search_stream(message): + async for token in self._handle_search_stream(message, prompt_id, sid): accumulated.append(token) yield token except Exception as e: @@ -125,7 +163,7 @@ async def process_message_stream( accumulated.append(error_text) yield error_text - self._add_to_history("assistant", "".join(accumulated)) + self._add_to_history("assistant", "".join(accumulated), sid) async def _execute_tool_calls(self, messages: list[dict[str, Any]]) -> str: """Execute tool calls returned by the LLM in a multi-turn loop. @@ -208,13 +246,9 @@ async def _send_chat_with_tools( logger.error("Tool-calling chat failed: %s", e) raise - async def _handle_search(self, message: str) -> str: - """Handle search mode: multi-turn tool-calling loop. - - Uses search_sigma, filter_metadata, explain_detection, explain_rule, - and summarize tools to answer questions about Sigma detection rules. - """ - prompt_content = self.rag_pipeline._resolve_prompt(self._current_prompt_id, mode="search") + async def _handle_search(self, message: str, prompt_id: str) -> str: + """Handle search mode: multi-turn tool-calling loop.""" + prompt_content = self.rag_pipeline._resolve_prompt(prompt_id, mode="search") assistant_message = { "role": "assistant", "content": message, @@ -226,34 +260,37 @@ async def _handle_search(self, message: str) -> str: messages = [system_msg, assistant_message] return await self._execute_tool_calls(messages) - async def _handle_explain(self, message: str) -> str: + async def _handle_explain(self, message: str, prompt_id: str, sid: str) -> str: """Handle explain mode: analyze uploaded Sigma rule.""" - if not self._uploaded_rule: + rule = self._get_rule(sid) + if not rule: return "No Sigma rule uploaded. Please upload a .yaml file first." - related = await self.search_engine.search(self._uploaded_rule.get("name", "")) - return await self.rag_pipeline.explain_rule( - self._uploaded_rule, related, system_prompt_id=self._current_prompt_id - ) + related = await self.search_engine.search(rule.name) + return await self.rag_pipeline.explain_rule(rule, related, system_prompt_id=prompt_id) async def _handle_explain_stream( self, message: str, + prompt_id: str, + sid: str, ) -> AsyncGenerator[str, None]: """Handle explain mode with streaming.""" - if not self._uploaded_rule: + rule = self._get_rule(sid) + if not rule: yield "No Sigma rule uploaded. Please upload a .yaml file first." return - related = await self.search_engine.search(self._uploaded_rule.get("name", "")) + related = await self.search_engine.search(rule.name) async for token in self.rag_pipeline.explain_rule_stream( - self._uploaded_rule, related, system_prompt_id=self._current_prompt_id + rule, related, system_prompt_id=prompt_id ): yield token - async def _handle_coverage(self, message: str) -> str: + async def _handle_coverage(self, message: str, prompt_id: str, sid: str) -> str: """Handle coverage mode.""" - if not self._uploaded_rule: + rule = self._get_rule(sid) + if not rule: return "No Sigma rule uploaded. Upload a .yaml file to check coverage." results = await self.search_engine.search(message) @@ -262,7 +299,7 @@ async def _handle_coverage(self, message: str) -> str: try: return await self.rag_pipeline.analyze_coverage( - self._uploaded_rule, results, system_prompt_id=self._current_prompt_id + rule, results, system_prompt_id=prompt_id ) except Exception as e: logger.error(f"RAG pipeline failed: {e}") @@ -271,6 +308,8 @@ async def _handle_coverage(self, message: str) -> str: async def _handle_search_stream( self, message: str, + prompt_id: str, + sid: str, ) -> AsyncGenerator[str, None]: """Handle search mode with streaming: semantic search + LLM token stream. @@ -295,14 +334,15 @@ async def _handle_search_stream( yield "No matching Sigma rules found. Try a different query." return - self._last_citations = [ + citations = [ self.search_engine.get_citation(r) for r in results[:5] if self.search_engine.get_citation(r) ] + self._set_citations(sid, citations) - if self._last_citations: - yield f"__CITATIONS__:{json.dumps(self._last_citations)}" + if citations: + yield f"__CITATIONS__:{json.dumps(citations)}" # If we have a translation, prepend it as context augmented_message = message @@ -311,10 +351,12 @@ async def _handle_search_stream( f"[Auto-translated detection]\n{translation}\n\n[Original YAML]\n{message}" ) + formatted_results = [format_result_by_collection(r) for r in results] + try: found = False async for token in self.rag_pipeline.answer_search_query_stream( - augmented_message, results, system_prompt_id=self._current_prompt_id + augmented_message, formatted_results, system_prompt_id=prompt_id ): yield token found = True @@ -326,9 +368,12 @@ async def _handle_search_stream( async def _handle_coverage_stream( self, message: str, + prompt_id: str, + sid: str, ) -> AsyncGenerator[str, None]: """Handle coverage mode with streaming.""" - if not self._uploaded_rule: + rule = self._get_rule(sid) + if not rule: yield "No Sigma rule uploaded. Upload a .yaml file to check coverage." return @@ -337,53 +382,61 @@ async def _handle_coverage_stream( yield "No related rules found for coverage analysis." return + formatted_results = [format_result_by_collection(r) for r in results] + try: async for token in self.rag_pipeline.analyze_coverage_stream( - self._uploaded_rule, results, system_prompt_id=self._current_prompt_id + rule, formatted_results, system_prompt_id=prompt_id ): yield token except Exception as e: logger.error(f"RAG pipeline failed: {e}") - async def validate_and_store_yaml(self, content: bytes) -> dict[str, Any]: + async def validate_and_store_yaml( + self, content: bytes, session_id: str | None = None + ) -> SigmaRule: """Validate YAML content and store rule data in session. Args: content: Raw YAML file content + session_id: Session identifier for state isolation. Returns: - Parsed and validated rule dictionary + Parsed and validated SigmaRule """ - rule_data = self.validator.validate(content) - self._uploaded_rule = rule_data + rule = self.validator.validate(content) + self._set_rule(session_id, rule) self.rag_pipeline.cache.invalidate() - return rule_data + return rule - def get_last_citations(self) -> list[str]: + def get_last_citations(self, session_id: str | None = None) -> list[str]: """Get citations from the last search response.""" - return self._last_citations.copy() + return self._get_citations(session_id).copy() - def _add_to_history(self, role: str, content: str) -> None: + def _add_to_history(self, role: str, content: str, sid: str) -> None: """Add a message to chat history (max 50 messages).""" - self._history.append( + history = self._get_history(sid) + history.append( { "role": role, "content": content, "timestamp": datetime.now(UTC).isoformat(), } ) - if len(self._history) > MAX_HISTORY: - self._history = self._history[-MAX_HISTORY:] + if len(history) > MAX_HISTORY: + history = history[-MAX_HISTORY:] + self._set_history(sid, history) - def get_history(self) -> list[dict[str, str]]: + def get_history(self, session_id: str | None = None) -> list[dict[str, str]]: """Get chat history.""" - return self._history.copy() + return self._get_history(session_id).copy() - async def clear_history(self) -> None: + async def clear_history(self, session_id: str | None = None) -> None: """Clear chat history, uploaded rule, and llama.cpp KV cache.""" - self._history.clear() - self._uploaded_rule = None - self._last_citations.clear() + sid = self._sid(session_id) + self._set_history(sid, []) + self._set_rule(sid, None) + self._set_citations(sid, []) try: await self.rag_pipeline.llm_client.erase_slot_cache() except Exception: diff --git a/src/application/documents/indexing.py b/src/application/documents/indexing.py index 61b40c77..f8bf5832 100644 --- a/src/application/documents/indexing.py +++ b/src/application/documents/indexing.py @@ -1,4 +1,4 @@ -"""Qdrant indexing for Sigma rules.""" +"""Qdrant indexing for Sigma rules and reference documents.""" from __future__ import annotations @@ -30,6 +30,7 @@ async def index_sigma_rules( rules: list[SigmaRule], collection_name: str | None = None, mode: str = "flat", + index_after_upload: bool = False, ) -> dict[str, Any]: """Index Sigma rules in Qdrant. @@ -37,6 +38,7 @@ async def index_sigma_rules( rules: List of SigmaRule to index collection_name: Optional collection name override mode: 'flat' (one chunk per rule) or 'rich' (multiple chunks per rule) + index_after_upload: If True, index immediately after upload (incremental) Returns: Dict with indexing results @@ -47,16 +49,17 @@ async def index_sigma_rules( config = get_config() collection = collection_name or config.qdrant_collection_name - service_client = get_qdrant_client(host=config.qdrant_host, port=config.qdrant_port) - try: - # Delete + full reindex strategy: clear old vectors before pipeline insert - service_client.delete_collection(collection) - except Exception as e: - logger.warning( - "Collection %s not found or cannot be deleted (first run?): %s", collection, e - ) - finally: - service_client.close() + if not index_after_upload: + service_client = get_qdrant_client(host=config.qdrant_host, port=config.qdrant_port) + try: + # Delete + full reindex strategy: clear old vectors before pipeline insert + service_client.delete_collection(collection) + except Exception as e: + logger.warning( + "Collection %s not found or cannot be deleted (first run?): %s", collection, e + ) + finally: + service_client.close() nodes = [] for rule in rules: @@ -123,6 +126,9 @@ def _sigma_rule_to_text(rule: SigmaRule) -> str: if rule.level: parts.append(f"Level: {rule.level}") + if rule.references: + parts.append(f"References: {', '.join(rule.references)}") + return "\n".join(parts) @@ -137,6 +143,7 @@ def _sigma_rule_to_metadata(rule: SigmaRule) -> dict[str, Any]: "status": rule.status, "tags": rule.tags, "logsource": rule.logsource, + "references": rule.references, "chunk_type": "full_rule", } diff --git a/src/application/documents/models.py b/src/application/documents/models.py index a3abc097..190c88d3 100644 --- a/src/application/documents/models.py +++ b/src/application/documents/models.py @@ -4,24 +4,6 @@ from pydantic import BaseModel, Field -from src.core.sigma.models import SigmaRule - - -class ValidationError(BaseModel): - """Validation error for a single field.""" - - field: str - message: str - - -class ValidationResult(BaseModel): - """Result of validating a Sigma rule.""" - - valid: bool - rule: SigmaRule | None = None - errors: list[ValidationError] = Field(default_factory=list) - file_path: str | None = None - class IngestRequest(BaseModel): """Request to ingest Sigma rules.""" diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index f832ee3e..2a3720b5 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -2,362 +2,125 @@ from __future__ import annotations -import ipaddress import logging -import re import threading -import time -import urllib.parse from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor, as_completed from pathlib import Path from typing import Any, cast -import httpx import yaml +from src.shared.constants import NULL_UUID +from src.shared.http import download_file as http_download_file +from src.shared.http import head_url as http_head_url +from src.shared.utils.registry_utils import build_registry_entry +from src.shared.utils.crypto_utils import ( + compute_sha256_bytes, + compute_sha256_file, + compute_sha256_str, +) from src.shared.utils.identify_file_type import ( SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, + filetype_ext, ) +from src.shared.utils.url_utils import is_private_url, normalize_url from src.infrastructure.database import DatabaseService from src.core.sigma.models import is_sigma_rule_dict -from src.shared.utils import iso_now +from src.shared.utils.sigma_utils import extract_sigma_references +from src.config.settings import get_config + +# Local helper modules +from .sigma_ref_paths import ( + resolve_rule_path as _resolve_rule_path, + sigmaref_resolve_path as _sigmaref_resolve_path, + sigmaref_write_path as _sigmaref_write_path, +) +from .sigma_ref_url import detect_url_type as _detect_url_type +from .sigma_ref_registry import ( + load_registry as _load_registry, + load_error_registry as _load_error_registry, + maybe_record_error as _maybe_record_error, + save_registry as _save_registry, +) logger = logging.getLogger(__name__) - - -def _is_private_url(url: str) -> bool: - """Check if a URL points to a private/reserved IP to prevent SSRF.""" - parsed = urllib.parse.urlparse(url) - host = parsed.hostname or "" - if host in ("localhost", "127.0.0.1", "::1"): - return True - try: - ip = ipaddress.ip_address(host) - return ip.is_private or ip.is_loopback or ip.is_link_local - except ValueError: - return False - - -GITHUB_BLOB_PATTERN = re.compile( - r"^https?://(?:www\.)?github\.com/([^/]+/[^/]+)/blob/([^#?]+)", - re.IGNORECASE, -) -MAX_RETRIES = 3 -BACKOFF_DELAYS = [1, 4, 9] -RETRY_STATUSES = {429, 500, 502, 503, 504} DEFAULT_REQUEST_DELAY = 0.5 DEFAULT_MAX_WORKERS = 5 SUPPORTED_EXTENSIONS: dict[str, str] = { ext: ft.value for ext, ft in SUPPORTED_DOC_EXTENSION_MAP.items() } - _registry_lock = threading.Lock() -def normalize_url(url: str) -> str: - """Normalize a reference URL for deduplication.""" - match = GITHUB_BLOB_PATTERN.match(url) - if match: - repo = match.group(1) - path_part = match.group(2) - path_part = re.sub(r"^refs/heads/", "", path_part) - path_part = re.sub(r"^refs/tags/", "", path_part) - result = f"https://raw.githubusercontent.com/{repo}/{path_part}" - parsed = urllib.parse.urlparse(result) - if parsed.fragment: - result = urllib.parse.urlunparse(parsed._replace(fragment="")) - return result - - parsed = urllib.parse.urlparse(url) - clean = parsed._replace(fragment="") - return urllib.parse.urlunparse(clean) - - -def _detect_url_type(url: str, content_type: str | None = None) -> str | None: - """Detect the document type of a reference URL. +def download_sigma_references( + db: DatabaseService, + output_dir: str, + mode: str = "scan", + rules_dir: str | None = None, + supported_types: set[str] | None = None, + request_delay: float = DEFAULT_REQUEST_DELAY, + progress_callback: Callable[[int, int, str], None] | None = None, + max_workers: int = DEFAULT_MAX_WORKERS, + selected_dirs: list[str] | None = None, +) -> dict[str, Any]: + """Download Sigma rule references using the specified mode. - Checks URL extension first, then falls back to HTTP Content-Type. + Two modes are supported: - Args: - url: The reference URL. - content_type: Optional HTTP Content-Type header value. + - **scan** (default): Scans a local directory of Sigma rule YAML files, + extracts reference URLs, and downloads matching documents. + Requires ``rules_dir``. - Returns: - The FileType value string (e.g. "markdown") or None if unsupported. - """ - parsed = urllib.parse.urlparse(url) - path = parsed.path.rstrip("/") - ext = Path(path).suffix.lower() - - if ext in SUPPORTED_EXTENSIONS: - return SUPPORTED_EXTENSIONS[ext] - - if content_type: - ct = content_type.lower() - if ct.startswith("text/markdown"): - return "markdown" - if ct.startswith("text/html"): - return "html" - if ct.startswith("text/plain"): - if ext in {".md", ".markdown"}: - return "markdown" - return "plain_text" - if ct.startswith("application/pdf"): - return "pdf" - if ct.startswith("application/vnd.openxmlformats-officedocument"): - return "office_document" - if ct.startswith("application/vnd.oasis.opendocument"): - return "office_document" - if ct.startswith("application/msword"): - return "office_document" - if ct.startswith("application/rtf"): - return "office_document" - - return None - - -def _head_content_type(url: str, timeout: int = 10) -> str | None: - """Do a HEAD request to discover the Content-Type of a URL. + - **registry**: Reads pending Sigma rule entries from the doc_registry, + resolves their local file paths, and downloads referenced documents. + Does **not** require ``rules_dir``. Args: - url: The URL to check. - timeout: HTTP request timeout in seconds. - - Returns: - The Content-Type header value, or None if the request failed. - """ - try: - with httpx.Client(timeout=httpx.Timeout(timeout), follow_redirects=True) as client: - response = client.head(url) - response.raise_for_status() - ct = response.headers.get("content-type") - return str(ct) if ct else None - except Exception: - return None - - -def _download_file( - url: str, - output_path: Path, - timeout: int = 30, - max_retries: int = MAX_RETRIES, -) -> tuple[bool, int | None]: - """Download a single file with retry and exponential backoff. - - Args: - url: The URL to download. - output_path: Local filesystem path to save the file. - timeout: HTTP request timeout in seconds. - max_retries: Maximum number of retry attempts. 0 means no retries. + db: Database service instance. + output_dir: Path to the output directory for downloaded files. + mode: ``"scan"`` or ``"registry"``. + rules_dir: Path to the directory containing Sigma rule YAML files + (only used in ``"scan"`` mode). + supported_types: Set of FileType values to accept (e.g. {"markdown"}). + Defaults to all reference doc types. + request_delay: Seconds to wait between download requests. + progress_callback: Optional callback(current, total, phase). + max_workers: Max concurrent HTTP download threads. + selected_dirs: Optional list of relative directory paths to scan + (only used in ``"scan"`` mode). Returns: - True if download succeeded, False otherwise. - """ - for attempt in range(1, max_retries + 1): - try: - with httpx.Client(timeout=httpx.Timeout(timeout), follow_redirects=True) as client: - response = client.get(url) - response.raise_for_status() - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_bytes(response.content) - return True, None - - except OSError as exc: - logger.warning("Filesystem error for %s: %s — skipping", url, exc) - if output_path.exists(): - try: - output_path.unlink() - except OSError: - pass - return False, None - - except httpx.HTTPStatusError as exc: - status = exc.response.status_code - if status in RETRY_STATUSES and attempt < max_retries: - retry_after = _get_retry_after(exc.response) - if retry_after is not None: - wait: float = min(retry_after, 120) - else: - wait = _backoff_delay(attempt) - logger.warning( - "HTTP %d on attempt %d/%d for %s — waiting %ds", - status, - attempt, - max_retries, - url, - wait, - ) - time.sleep(wait) - continue - logger.warning("HTTP %d for %s — giving up", status, url) - return False, status - - except ( - httpx.TimeoutException, - httpx.NetworkError, - httpx.ConnectError, - httpx.RemoteProtocolError, - ) as exc: - if attempt < max_retries: - wait: float = _backoff_delay(attempt) # type: ignore[no-redef] - logger.warning( - "Network error on attempt %d/%d for %s: %s — waiting %ds", - attempt, - max_retries, - url, - exc, - wait, - ) - time.sleep(wait) - continue - logger.warning("Network error for %s after %d attempts: %s", url, max_retries, exc) - return False, None - - return False, None - - -def _backoff_delay(attempt: int) -> float: - """Return the backoff delay for the given attempt number (1-indexed). - - Falls back to the last configured delay value if attempt exceeds the list. - """ - idx = attempt - 1 - if idx < len(BACKOFF_DELAYS): - return BACKOFF_DELAYS[idx] - return BACKOFF_DELAYS[-1] - - -def _get_retry_after(response: httpx.Response) -> int | None: - """Extract Retry-After header value as seconds.""" - raw = response.headers.get("Retry-After") - if raw is None: - return None - try: - return int(raw) - except ValueError: - return None - - -def _load_registry(path: Path, db: DatabaseService) -> dict[str, Any]: - """Load the registry from doc_registry for sigmaref org. - - Returns an empty dict if DB not available. + Dict with summary stats: total_rules, total_refs, downloaded, skipped, failed. """ - entries = db.get_entries_by_org("sigmaref", limit=0) - registry = {} - for entry in entries: - url_hash = entry["url_hash"] - registry[url_hash] = { - "original_url": entry.get("original_url", ""), - "normalized_url": entry.get("normalized_url"), - "content_type": entry.get("content_type"), - "rule_id": entry.get("rule_id"), - "title": entry.get("title"), - "timestamp": entry.get("timestamp"), - "content_sha256": entry.get("content_sha256"), - "embed_status": entry.get("embed_status"), - "last_seen": entry.get("last_seen"), - "file_name": entry.get("file_name", ""), - } - return registry - - -def _save_registry(registry: dict[str, Any], path: Path, db: DatabaseService) -> None: - """Save the registry to doc_registry atomically in a single batch.""" - rows = [] - now = iso_now() - for url_hash, entry in registry.items(): - if isinstance(entry, dict): - rows.append( - { - "url_hash": url_hash, - "original_url": entry.get("original_url", ""), - "normalized_url": entry.get("normalized_url"), - "content_type": entry.get("content_type"), - "rule_id": entry.get("rule_id"), - "title": entry.get("title"), - "timestamp": entry.get("timestamp"), - "content_sha256": entry.get("content_sha256"), - "org": entry.get("org", "sigmaref"), - "repo": entry.get("repo", "references"), - "file_name": entry.get("file_name", ""), - "file_size": entry.get("file_size"), - "embed_status": entry.get("embed_status", "discovery"), - "last_seen": entry.get("last_seen", now), - } - ) - db.batch_upsert_doc_registry(rows) - + if supported_types is None: + supported_types = SUPPORTED_REFERENCE_DOC_TYPES -def _load_error_registry(db: DatabaseService) -> set[str]: - """Load the set of url_hash values that have previously failed (30x/40x).""" - try: - entries = db.get_doc_errors() - return {e["url_hash"] for e in entries} - except Exception: - logger.warning("Failed to load error registry from DuckDB — proceeding without it") - return set() - - -def _maybe_record_error( - db: DatabaseService, - url_hash: str, - original_url: str, - normalized_url: str, - status_code: int | None, - rule_id: str, - rule_title: str, -) -> None: - """Record a 30x/40x download error in doc_error so it is skipped on retry.""" - if status_code is None: - return - if 300 <= status_code < 500 or (status_code >= 500 and status_code not in RETRY_STATUSES): - try: - db.upsert_doc_error( - { - "url_hash": url_hash, - "original_url": original_url, - "normalized_url": normalized_url, - "error_code": status_code, - "error_message": f"HTTP {status_code}", - "org": "sigmaref", - "repo": "references", - "timestamp": iso_now(), - } - ) - except Exception: - logger.warning("Failed to record error for %s", normalized_url) - - -def _make_entry( - url_hash: str, - original_url: str, - normalized_url: str, - content_type: str, - rule_id: str, - title: str, - timestamp: str, - content_sha256: str, - file_name: str = "", - file_size: int | None = None, -) -> dict[str, Any]: - """Build a registry entry dict with all fields expected by _save_registry.""" - return { - "original_url": original_url, - "normalized_url": normalized_url, - "content_type": content_type, - "rule_id": rule_id, - "title": title, - "timestamp": timestamp, - "content_sha256": content_sha256, - "org": "sigmaref", - "repo": "references", - "file_name": file_name, - "file_size": file_size, - "last_seen": timestamp, - } + if mode == "scan": + if not rules_dir: + raise ValueError("rules_dir is required in scan mode") + return _download_scan_mode( + rules_dir=rules_dir, + output_dir=output_dir, + db=db, + supported_types=supported_types, + request_delay=request_delay, + progress_callback=progress_callback, + max_workers=max_workers, + selected_dirs=selected_dirs, + ) + if mode == "registry": + return _download_registry_mode( + output_dir=output_dir, + db=db, + supported_types=supported_types, + request_delay=request_delay, + progress_callback=progress_callback, + max_workers=max_workers, + ) + msg = f"Unknown mode: {mode!r} (expected 'scan' or 'registry')" + raise ValueError(msg) def download_references( @@ -375,23 +138,32 @@ def download_references( Scans all Sigma rules in the given directory, extracts reference URLs, filters by supported document types, and downloads matching files. - Args: - rules_dir: Path to the directory containing Sigma rule YAML files. - output_dir: Path to the output directory for downloaded files. - db: Database service instance. - supported_types: Set of FileType values to accept (e.g. {"markdown"}). - Defaults to {"markdown"}. - request_delay: Seconds to wait between download requests (sequential - phase only; parallel phase uses max_workers instead). - progress_callback: Optional callback(current, total) called after each file. - max_workers: Max concurrent HTTP download threads. - selected_dirs: Optional list of relative directory paths to scan. - If provided, only files within these directories are processed. - Directories not in this list are excluded from scanning. - - Returns: - Dict with summary stats: total_rules, total_refs, downloaded, skipped, failed. + Delegates to :func:`download_sigma_references` with ``mode="scan"``. """ + return download_sigma_references( + db=db, + output_dir=output_dir, + mode="scan", + rules_dir=rules_dir, + supported_types=supported_types, + request_delay=request_delay, + progress_callback=progress_callback, + max_workers=max_workers, + selected_dirs=selected_dirs, + ) + + +def _download_scan_mode( + rules_dir: str, + output_dir: str, + db: DatabaseService, + supported_types: set[str] | None = None, + request_delay: float = DEFAULT_REQUEST_DELAY, + progress_callback: Callable[[int, int, str], None] | None = None, + max_workers: int = DEFAULT_MAX_WORKERS, + selected_dirs: list[str] | None = None, +) -> dict[str, Any]: + """Scan-mode implementation — see :func:`download_sigma_references`.""" if supported_types is None: supported_types = SUPPORTED_REFERENCE_DOC_TYPES @@ -449,6 +221,7 @@ def _collect_yaml_files() -> list[Path]: # Phase 1: scan YAML files, classify refs download_queue: list[dict[str, Any]] = [] head_pending: list[dict[str, Any]] = [] + rule_refs: list[dict[str, str]] = [] total_files = len(yml_files) for file_idx, yml_file in enumerate(yml_files): @@ -481,7 +254,7 @@ def _collect_yaml_files() -> list[Path]: logger.debug("Non-HTTP ref skipped: %s", ref) continue - if _is_private_url(ref): + if is_private_url(ref): logger.warning("Skipping private URL ref: %s", ref) skipped += 1 continue @@ -489,7 +262,9 @@ def _collect_yaml_files() -> list[Path]: total_refs += 1 normalized = normalize_url(ref) - url_hash = _sha256(normalized) + url_hash = compute_sha256_str(normalized) + + rule_refs.append({"rule_id": rule_id, "url_hash": url_hash, "ref_url": ref}) if url_hash in error_registry: logger.debug("Skipping previously failed URL: %s", normalized) @@ -500,16 +275,24 @@ def _collect_yaml_files() -> list[Path]: if url_hash in registry: fname = registry[url_hash].get("file_name", "") if fname: - output_file = output_path / fname + content_type_for_subdir = registry[url_hash].get("content_type", "") + output_file = _sigmaref_resolve_path( + output_path, content_type_for_subdir, fname + ) if output_file.exists(): existing_sha = registry[url_hash].get("content_sha256") - if existing_sha is not None and _sha256_file(output_file) != existing_sha: + if ( + existing_sha is not None + and compute_sha256_file(output_file) != existing_sha + ): download_queue.append( { "url_hash": url_hash, "original_url": ref, "normalized_url": normalized, - "output_file": output_file, + "output_file": _sigmaref_write_path( + output_path, content_type_for_subdir, fname + ), "content_type": registry[url_hash].get( "content_type", "markdown" ), @@ -524,7 +307,7 @@ def _collect_yaml_files() -> list[Path]: # File missing or no file_name — fall through to re-download # Determine extension and content type - ext = _url_ext(normalized) + ext = filetype_ext(normalized) ftype = _detect_url_type(normalized) if ftype is None and url_hash in registry: ct = registry[url_hash].get("content_type") @@ -545,22 +328,16 @@ def _collect_yaml_files() -> list[Path]: continue if not ext and ftype is not None: - _TYPE_TO_EXT: dict[str, str] = { - "html": ".html", - "markdown": ".md", - "plain_text": ".txt", - "pdf": ".pdf", - "office_document": ".docx", - } - ext = _TYPE_TO_EXT.get(ftype, ".md") + ext = filetype_ext(ftype) if not ext: ext = ".md" - output_file = output_path / f"{url_hash}{ext}" + fname = f"{url_hash}{ext}" + output_file = _sigmaref_resolve_path(output_path, ftype, fname) if output_file.exists(): - content_hash = _sha256_file(output_file) + content_hash = compute_sha256_file(output_file) if url_hash in registry: existing_sha = registry[url_hash].get("content_sha256") if existing_sha is not None and content_hash != existing_sha: @@ -569,7 +346,7 @@ def _collect_yaml_files() -> list[Path]: "url_hash": url_hash, "original_url": ref, "normalized_url": normalized, - "output_file": output_file, + "output_file": _sigmaref_write_path(output_path, ftype, fname), "content_type": ftype or "markdown", "rule_id": rule_id, "rule_title": rule_title, @@ -577,14 +354,13 @@ def _collect_yaml_files() -> list[Path]: } ) continue - registry[url_hash] = _make_entry( + registry[url_hash] = build_registry_entry( url_hash=url_hash, original_url=ref, normalized_url=normalized, content_type=ftype or "markdown", rule_id=rule_id, title=rule_title, - timestamp=iso_now(), content_sha256=content_hash, file_name=output_file.name, file_size=output_file.stat().st_size, @@ -601,7 +377,7 @@ def _collect_yaml_files() -> list[Path]: "url_hash": url_hash, "original_url": ref, "normalized_url": normalized, - "output_file": output_file, + "output_file": _sigmaref_write_path(output_path, ftype, fname), "content_type": ftype, "rule_id": rule_id, "rule_title": rule_title, @@ -615,7 +391,7 @@ def _collect_yaml_files() -> list[Path]: with ThreadPoolExecutor(max_workers=max_workers) as executor: future_map: dict[Future, dict[str, Any]] = {} for item in head_pending: - future = executor.submit(_head_content_type, item["normalized"]) + future = executor.submit(http_head_url, item["normalized"], 10, check_ssrf=False) future_map[future] = item resolved = 0 for future in as_completed(future_map): @@ -623,21 +399,15 @@ def _collect_yaml_files() -> list[Path]: item = future_map[future] if progress_callback: progress_callback(resolved, len(head_pending), "downloading") - head_ct = future.result() + head_ct, _, _ = future.result() ftype = _detect_url_type(item["normalized"], content_type=head_ct) ext = item["ext"] if not ext and ftype is not None: - _TYPE_TO_EXT = { - "html": ".html", - "markdown": ".md", - "plain_text": ".txt", - "pdf": ".pdf", - "office_document": ".docx", - } - ext = _TYPE_TO_EXT.get(ftype, ".md") + ext = filetype_ext(ftype) if not ext: ext = ".md" - output_file = output_path / f"{item['url_hash']}{ext}" + fname = f"{item['url_hash']}{ext}" + output_file = _sigmaref_resolve_path(output_path, ftype, fname) if ftype is None or ftype not in supported_types: skipped += 1 continue @@ -646,7 +416,7 @@ def _collect_yaml_files() -> list[Path]: "url_hash": item["url_hash"], "original_url": item["original_url"], "normalized_url": item["normalized"], - "output_file": output_file, + "output_file": _sigmaref_write_path(output_path, ftype, fname), "content_type": ftype, "rule_id": item["rule_id"], "rule_title": item["rule_title"], @@ -662,9 +432,10 @@ def _collect_yaml_files() -> list[Path]: future_map = {} for item in download_queue: future = executor.submit( - _download_file, # type: ignore[arg-type] + http_download_file, # type: ignore[arg-type] item["normalized_url"], item["output_file"], + check_ssrf=False, ) future_map[future] = item @@ -675,21 +446,26 @@ def _collect_yaml_files() -> list[Path]: success, status_code = cast(tuple[bool, int | None], future.result()) if success: output_file = item["output_file"] - content_hash = _sha256_file(output_file) if output_file.exists() else "" - registry[item["url_hash"]] = _make_entry( + content_hash = compute_sha256_file(output_file) if output_file.exists() else "" + registry[item["url_hash"]] = build_registry_entry( url_hash=item["url_hash"], original_url=item["original_url"], normalized_url=item["normalized_url"], content_type=item["content_type"], rule_id=item["rule_id"], title=item["rule_title"], - timestamp=iso_now(), content_sha256=content_hash, file_name=output_file.name, file_size=output_file.stat().st_size if output_file.exists() else None, ) downloaded += 1 else: + logger.debug( + "Reference download failed | url=%s rule_id=%s status=%s", + item["original_url"], + item["rule_id"], + status_code, + ) _maybe_record_error( db, item["url_hash"], @@ -704,6 +480,9 @@ def _collect_yaml_files() -> list[Path]: with _registry_lock: _save_registry(registry, output_path, db) + if rule_refs: + db.batch_upsert_rule_references(rule_refs) + summary: dict[str, Any] = { "total_rules": total_rules, "total_refs": total_refs, @@ -711,35 +490,24 @@ def _collect_yaml_files() -> list[Path]: "skipped": skipped, "failed": failed, } - logger.info("Download complete: %s", summary) + _log_download_summary(summary) return summary -def _sha256(text: str) -> str: - """Compute SHA256 hex digest of a string.""" - import hashlib - - return hashlib.sha256(text.encode("utf-8")).hexdigest() - - -def _sha256_file(path: Path) -> str: - """Compute SHA256 hex digest of a file's contents.""" - import hashlib - - h = hashlib.sha256() - try: - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(65536), b""): - h.update(chunk) - except OSError: - return "" - return h.hexdigest() - - -def _url_ext(url: str) -> str: - """Extract the file extension from a URL path.""" - parsed = urllib.parse.urlparse(url) - return Path(parsed.path).suffix.lower() +def _log_download_summary(summary: dict[str, Any]) -> None: + """Log download summary and warn if failure rate exceeds threshold.""" + logger.info("Download complete: %s", summary) + total_refs = summary.get("total_refs", 0) + failed = summary.get("failed", 0) + if total_refs > 0: + fail_rate = failed / total_refs + if fail_rate > 0.05: + logger.warning( + "High reference failure rate: %.1f%% (%d/%d) — check network or URL validity", + fail_rate * 100, + failed, + total_refs, + ) def _empty_summary() -> dict[str, Any]: @@ -751,3 +519,257 @@ def _empty_summary() -> dict[str, Any]: "skipped": 0, "failed": 0, } + + +# ------------------------------------------------------------------ +# Registry mode — reads pending Sigma rules from doc_registry and +# downloads their referenced documents. +# ------------------------------------------------------------------ + + +def _download_registry_mode( + output_dir: str, + db: DatabaseService, + supported_types: set[str] | None = None, + request_delay: float = DEFAULT_REQUEST_DELAY, + progress_callback: Callable[[int, int, str], None] | None = None, + max_workers: int = DEFAULT_MAX_WORKERS, +) -> dict[str, Any]: + """Registry-mode implementation — see :func:`download_sigma_references`.""" + if supported_types is None: + supported_types = SUPPORTED_REFERENCE_DOC_TYPES + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + entries = db.get_pending_registry_all() + sigma_entries = [ + e + for e in entries + if e.get("content_type") == "sigma_rule" and e.get("embed_status") == "discovery" + ] + + total_rules = len(sigma_entries) + total_refs = 0 + downloaded = 0 + skipped = 0 + failed = 0 + + cfg = get_config() + + with _registry_lock: + registry = _load_registry(output_path, db) + + head_queue: list[dict[str, Any]] = [] + download_ready: list[dict[str, Any]] = [] + seen_urls: set[str] = set() + rule_refs: list[dict[str, str]] = [] + + for rule_entry in sigma_entries: + rule_id: str = rule_entry.get("rule_id", NULL_UUID) + original_url = rule_entry.get("original_url", "") + file_name = rule_entry.get("file_name", "") + + file_path = _resolve_rule_path(rule_entry, cfg) + if not file_path or not file_path.exists(): + logger.warning("Rule file not found, skipping: %s (url=%s)", file_name, original_url) + skipped += 1 + continue + + refs = extract_sigma_references(file_path) + rule_title = rule_entry.get("title", file_name) + logger.info("Rule %s processed: %d reference(s) found", rule_id, len(refs)) + + if not refs: + continue + + total_refs += len(refs) + + for ref_url in refs: + ref_url_clean = ref_url.strip() + if not ref_url_clean: + continue + + norm_url = normalize_url(ref_url_clean) + url_hash = compute_sha256_str(norm_url) + + if url_hash in seen_urls: + logger.debug("Reference already queued this run: %s", ref_url_clean) + skipped += 1 + continue + seen_urls.add(url_hash) + + rule_refs.append({"rule_id": rule_id, "url_hash": url_hash, "ref_url": ref_url_clean}) + + existing = registry.get(url_hash) + + if existing and existing.get("content_sha256"): + skipped += 1 + continue + + if existing: + download_ready.append( + { + "url": ref_url_clean, + "rule_id": rule_id, + "rule_title": rule_title, + "url_hash": url_hash, + "final_url": existing.get("normalized_url", ref_url_clean), + "content_type": existing.get("content_type", ""), + } + ) + else: + head_queue.append( + { + "url": ref_url_clean, + "rule_id": rule_id, + "rule_title": rule_title, + } + ) + + total_head = len(head_queue) + total_download_ready = len(download_ready) + + # Phase 1: parallel HEAD requests + head_pending: list[dict[str, Any]] = [] + head_completed = 0 + if head_queue: + with ThreadPoolExecutor(max_workers=max_workers) as head_executor: + head_futures: dict[Future, dict[str, Any]] = {} + for item in head_queue: + future = head_executor.submit(http_head_url, item["url"], 15.0) + head_futures[future] = item + + for future in as_completed(head_futures): + head_completed += 1 + if progress_callback: + progress_callback( + head_completed, + total_head + total_download_ready, + "resolving URLs", + ) + item = head_futures[future] + url = item["url"] + try: + content_type, size, final_url = future.result() + norm_url = normalize_url(final_url or url) + url_hash = compute_sha256_str(norm_url) + + if content_type not in supported_types: + db.batch_upsert_doc_registry( + [ + build_registry_entry( + normalized_url=norm_url, + content_type=content_type or "unknown", + rule_id=item["rule_id"], + title=item["rule_title"], + embed_status="head_verified", + ) + ] + ) + skipped += 1 + continue + + db.batch_upsert_doc_registry( + [ + build_registry_entry( + normalized_url=norm_url, + content_type=content_type, + rule_id=item["rule_id"], + title=item["rule_title"], + file_size=size, + embed_status="head_verified", + ) + ] + ) + head_pending.append( + { + **item, + "content_type": content_type, + "size": size, + "final_url": final_url or url, + "url_hash": url_hash, + } + ) + except Exception as e: + logger.warning("HEAD failed for %s: %s", url, e) + failed += 1 + + # Phase 2: merge & download + all_to_download = head_pending + download_ready + total_to_download = len(all_to_download) + + def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: + url = item["final_url"] + content_type = item.get("content_type", "") + ext = filetype_ext(content_type) + url_hash = item.get("url_hash") or compute_sha256_str(normalize_url(url)) + fname = f"{url_hash}{ext}" + existing_path = _sigmaref_resolve_path(output_path, content_type, fname) + if existing_path.exists(): + existing_entry = registry.get(url_hash) + if existing_entry and existing_entry.get("content_sha256"): + return None + + file_path = _sigmaref_write_path(output_path, content_type, fname) + + ok, _ = http_download_file(url, file_path, check_ssrf=False) + if ok: + content = file_path.read_bytes() + content_hash = compute_sha256_bytes(content) + return ("ok", url_hash, content_hash, len(content)) + logger.error("Reference download failed: %s", url) + return ("fail", "", "", 0) + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures: dict[Future, dict[str, Any]] = { + executor.submit(_download_one, item): item for item in all_to_download + } + + for future in as_completed(futures): + item = futures[future] + try: + result = future.result() + if result is None: + skipped += 1 + elif result is not None and result[0] == "ok": + _, url_hash, content_hash, size = cast(tuple[str, str, str, int], result) + entry = build_registry_entry( + url_hash=url_hash, + normalized_url=item.get("final_url", item["url"]), + content_type=item["content_type"], + rule_id=item["rule_id"], + title=item["rule_title"], + content_sha256=content_hash, + file_name=f"{url_hash}{filetype_ext(item['content_type'])}", + file_size=size, + embed_status="discovery", + ) + db.batch_upsert_doc_registry([entry]) + downloaded += 1 + else: + failed += 1 + except Exception as e: + logger.error("Download task failed: %s", e) + failed += 1 + + if progress_callback: + completed = downloaded + failed + progress_callback( + total_head + completed, + total_head + total_to_download, + "downloading", + ) + + if rule_refs: + db.batch_upsert_rule_references(rule_refs) + + summary = { + "total_rules": total_rules, + "total_refs": total_refs, + "downloaded": downloaded, + "skipped": skipped, + "failed": failed, + } + _log_download_summary(summary) + return summary diff --git a/src/application/documents/sigma_ref_paths.py b/src/application/documents/sigma_ref_paths.py new file mode 100644 index 00000000..f2d8b917 --- /dev/null +++ b/src/application/documents/sigma_ref_paths.py @@ -0,0 +1,86 @@ +"""Path resolution helpers for Sigma reference documents. + +Handles mapping between logical paths (content_type, file_name) and +physical filesystem paths for the sigmaref storage layout. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from src.shared.utils.identify_file_type import filetype_subdir + + +def subdir_for(content_type: str | None) -> str: + """Return the subdirectory name for a given content type. + + Args: + content_type: Content type string (e.g. "markdown", "pdf"). + + Returns: + Subdirectory name (e.g. "markdown"). + """ + return filetype_subdir(content_type or "") + + +def sigmaref_write_path(output_path: Path, content_type: str | None, file_name: str) -> Path: + """Compute the write path for a sigmaref document. + + Args: + output_path: Base output directory. + content_type: Content type of the document. + file_name: Name of the file. + + Returns: + Full path where the file should be written. + """ + return output_path / subdir_for(content_type) / file_name + + +def sigmaref_resolve_path(output_path: Path, content_type: str | None, file_name: str) -> Path: + """Resolve the actual path for a sigmaref document (handles existing files). + + Args: + output_path: Base output directory. + content_type: Content type of the document. + file_name: Name of the file. + + Returns: + Path that exists or should be created. + """ + candidate = sigmaref_write_path(output_path, content_type, file_name) + if candidate.exists(): + return candidate + return sigmaref_write_path(output_path, content_type, file_name) + + +def resolve_rule_path(entry: dict, cfg: Any) -> Path | None: + """Resolve the on-disk path for a rule entry. + + Args: + entry: Registry entry with org, repo, file_name fields. + cfg: Configuration object. + + Returns: + Path to the rule file, or None if not found. + """ + org: str = entry.get("org", "") or "" + repo: str = entry.get("repo", "") or "" + file_name: str = entry.get("file_name", "") or "" + + if org == "local": + return Path(cfg.local_documents_path).resolve() / file_name + + if org == "sigmaref": + base = Path(cfg.sigmaref_documents_path).resolve() + subdir = filetype_subdir(entry.get("content_type", "")) + candidate = base / subdir / file_name + if candidate.exists(): + return candidate + return candidate + + if org and repo: + return Path(cfg.paths_github_dir).resolve() / org / repo / file_name + + return None diff --git a/src/application/documents/sigma_ref_processor.py b/src/application/documents/sigma_ref_processor.py index 83c66f5d..30c0ec71 100644 --- a/src/application/documents/sigma_ref_processor.py +++ b/src/application/documents/sigma_ref_processor.py @@ -1,413 +1,66 @@ -"""Process Sigma rules already in doc_registry and download their reference documents.""" - -from __future__ import annotations - -import logging -from collections.abc import Callable -from concurrent.futures import Future, ThreadPoolExecutor, as_completed -from pathlib import Path -from typing import Any, cast - -import httpx - -from src.shared.utils.identify_file_type import SUPPORTED_REFERENCE_DOC_TYPES -from src.shared.utils.sigma_utils import extract_sigma_references -from src.config.settings import get_config -from src.shared.utils import iso_now - -logger = logging.getLogger(__name__) - -DEFAULT_REQUEST_DELAY = 0.5 -DEFAULT_MAX_WORKERS = 5 -GITHUB_BLOB_PATTERN: Any = None # imported lazily to avoid circular imports - - -def _build_head_entry( - normalized_url: str, - content_type: str, - rule_id: str, - title: str, - file_size: int | None = None, -) -> dict[str, Any]: - now = iso_now() - return { - "url_hash": _sha256_bytes(normalized_url.encode()), - "org": "sigmaref", - "repo": "references", - "content_type": content_type, - "file_name": Path(normalized_url).name, - "content_sha256": "", - "file_size": file_size or 0, - "original_url": normalized_url, - "normalized_url": normalized_url, - "rule_id": rule_id, - "title": title, - "timestamp": now, - "last_seen": now, - "embed_status": "head_verified", - } - - -def _build_download_entry( - normalized_url: str, - content_type: str, - rule_id: str, - title: str, - content_sha256: str, - file_name: str = "", - file_size: int | None = None, -) -> dict[str, Any]: - """Build a doc_registry entry for a downloaded reference document.""" - now = iso_now() - return { - "url_hash": _sha256_bytes(normalized_url.encode()), - "org": "sigmaref", - "repo": "references", - "content_type": content_type, - "file_name": file_name or Path(normalized_url).name, - "content_sha256": content_sha256, - "file_size": file_size or 0, - "original_url": normalized_url, - "normalized_url": normalized_url, - "rule_id": rule_id, - "title": title, - "timestamp": now, - "last_seen": now, - "embed_status": "discovery", - } - - -def process_sigma_refs( - db: Any, - output_dir: str, - supported_types: set[str] | None = None, - request_delay: float = DEFAULT_REQUEST_DELAY, - progress_callback: Callable[[int, int, str], None] | None = None, - max_workers: int = DEFAULT_MAX_WORKERS, -) -> dict[str, Any]: - """Process Sigma rules already in doc_registry and download their references. - - Parameters - ---------- - db : DatabaseService - Database service for writing registry entries. - output_dir : str - Directory where downloaded reference files are stored. - supported_types : set[str] | None - Allowed content types. Defaults to all reference doc types. - request_delay : float - Delay between sequential HEAD requests. - progress_callback : callable | None - ``progress_callback(current, total, phase)``. - max_workers : int - Max concurrent download threads. - - Returns - ------- - dict - Summary with keys: ``total_rules``, ``total_refs``, ``downloaded``, - ``skipped``, ``failed``. - """ - if supported_types is None: - supported_types = SUPPORTED_REFERENCE_DOC_TYPES - - output_path = Path(output_dir) - output_path.mkdir(parents=True, exist_ok=True) - - entries = db.get_pending_registry_all() - sigma_entries = [ - e - for e in entries - if e.get("content_type") == "sigma_rule" and e.get("embed_status") == "discovery" - ] - - total_rules = len(sigma_entries) - total_refs = 0 - downloaded = 0 - skipped = 0 - failed = 0 - - head_queue: list[dict[str, Any]] = [] # needs HEAD + download - download_ready: list[dict[str, Any]] = [] # has cached HEAD, needs download - seen_urls: set[str] = set() # deduplicate across rules within one run - - cfg = get_config() - for rule_entry in sigma_entries: - rule_id = rule_entry.get("rule_id", "00000000-0000-0000-0000-000000000000") - original_url = rule_entry.get("original_url", "") - file_name = rule_entry.get("file_name", "") - - # Resolve the actual file path from org/repo/file_name - file_path = _resolve_rule_path(rule_entry, cfg) - if not file_path or not file_path.exists(): - logger.warning("Rule file not found, skipping: %s (url=%s)", file_name, original_url) - skipped += 1 - continue - - # Extract reference URLs from Sigma rule - refs = extract_sigma_references(file_path) - rule_title = rule_entry.get("title", file_name) - logger.info("Rule %s processed: %d reference(s) found", rule_id, len(refs)) - - if not refs: - continue - - total_refs += len(refs) - - for ref_url in refs: - ref_url_clean = ref_url.strip() - if not ref_url_clean: - continue - - norm_url = _normalize_url(ref_url_clean) - url_hash = _sha256_bytes(norm_url.encode()) - - # Deduplicate within a single run (same URL across multiple rules) - if url_hash in seen_urls: - logger.debug("Reference already queued this run: %s", ref_url_clean) - skipped += 1 - continue - seen_urls.add(url_hash) - - existing = db.get_entry(url_hash) - - if existing and existing.get("content_sha256"): - # Already fully downloaded - logger.debug("Reference already registered: %s", ref_url_clean) - skipped += 1 - continue - - base = {"url": ref_url_clean, "rule_id": rule_id, "rule_title": rule_title} - if existing: - # Cached HEAD result — skip HEAD, go straight to download - download_ready.append( - { - **base, - "url_hash": url_hash, - "final_url": existing.get("normalized_url", ref_url_clean), - "content_type": existing.get("content_type", ""), - } - ) - else: - head_queue.append(base) - - total_head = len(head_queue) - total_download_ready = len(download_ready) - - # Phase 1: Parallel HEAD requests to resolve content types - head_pending: list[dict[str, Any]] = [] - head_completed = 0 - head_futures: dict[Future, dict[str, Any]] = {} - - total_head = len(head_queue) - with ThreadPoolExecutor(max_workers=max_workers) as head_executor: - for item in head_queue: - future = head_executor.submit(_head_request, item["url"]) - head_futures[future] = item - - for future in as_completed(head_futures): - head_completed += 1 - if progress_callback: - progress_callback( - head_completed, - total_head + total_download_ready, - "resolving URLs", - ) - - item = head_futures[future] - url = item["url"] - try: - content_type, size, final_url = future.result() - norm_url = _normalize_url(final_url or url) - url_hash = _sha256_bytes(norm_url.encode()) - - if content_type not in supported_types: - logger.info("Reference skipped (unsupported type): %s (%s)", url, content_type) - db.batch_upsert_doc_registry( - [ - _build_head_entry( - normalized_url=norm_url, - content_type=content_type or "unknown", - rule_id=item["rule_id"], - title=item["rule_title"], - ) - ] - ) - skipped += 1 - continue - - # Cache HEAD result immediately - db.batch_upsert_doc_registry( - [ - _build_head_entry( - normalized_url=norm_url, - content_type=content_type, - rule_id=item["rule_id"], - title=item["rule_title"], - file_size=size, - ) - ] - ) - head_pending.append( - { - **item, - "content_type": content_type, - "size": size, - "final_url": final_url or url, - "url_hash": url_hash, - } - ) - except Exception as e: - logger.warning("HEAD failed for %s: %s", url, e) - failed += 1 - - total_to_download = len(head_pending) + total_download_ready - - # Merge head_pending + download_ready (refs with cached HEAD) - all_to_download = head_pending + download_ready - - # Phase 2: Parallel downloads - def _download_one(item: dict[str, Any]) -> tuple[str, str, int] | None: - url = item["final_url"] - output_path = Path(output_dir) - file_path = output_path / _sanitize_filename(url) - - if file_path.exists(): - url_hash = _sha256_bytes(url.encode()) - existing = db.get_entry(url_hash) - if existing and existing.get("content_sha256"): - logger.info("Reference already present: %s", url) - return None - - try: - with httpx.Client( - timeout=httpx.Timeout(30.0), headers={"User-Agent": "SigmaRAG/1.0"} - ) as client: - resp = client.get(url) - resp.raise_for_status() - content = resp.content - file_path.write_bytes(content) - logger.info("Reference downloaded: %s", url) - return ("ok", _sha256_bytes(content), len(content)) - except Exception as e: - logger.error("Reference download failed: %s - %s", url, e) - return ("fail", "", 0) - - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures: dict[Future, dict[str, Any]] = { - executor.submit(_download_one, item): item for item in all_to_download - } - - for future in as_completed(futures): - item = futures[future] - try: - result = future.result() - if result is None: - skipped += 1 - elif result[0] == "ok": - final_url = item["final_url"] - head_result = cast("tuple[str, str, int]", result) - entry = _build_download_entry( - normalized_url=final_url, - content_type=item["content_type"], - rule_id=item["rule_id"], - title=item["rule_title"], - content_sha256=head_result[1], - file_size=head_result[2], - ) - db.batch_upsert_doc_registry([entry]) - downloaded += 1 - else: - failed += 1 - except Exception as e: - logger.error("Download task failed: %s", e) - failed += 1 - - if progress_callback: - completed = downloaded + failed - progress_callback( - total_head + completed, - total_head + total_to_download, - "downloading", - ) - - return { - "total_rules": total_rules, - "total_refs": total_refs, - "downloaded": downloaded, - "skipped": skipped, - "failed": failed, - } - - -# ------------------------------------------------------------------ -# Helpers -# ------------------------------------------------------------------ - - -def _resolve_rule_path(entry: dict, cfg: Any) -> Path | None: - """Resolve the local file path for a sigma rule entry.""" - org = entry.get("org", "") - repo = entry.get("repo", "") - file_name = entry.get("file_name", "") - - if not file_name: - return None - - if org == "local": - base = Path(str(cfg.local_documents_path)) - return Path(base, file_name) - - if org == "sigmaref": - base = Path(str(cfg.sigmaref_documents_path)) - return Path(base, file_name) - - if org and repo: - base = Path(str(cfg.paths_github_dir)) - return Path(base, org, repo, file_name) - - return None - - -def _normalize_url(url: str) -> str: - """Simple URL normalizer - keeps the URL as-is for now.""" - return url.strip().rstrip("/") - - -def _head_request(url: str, delay: float = 0.0) -> tuple[str | None, int | None, str | None]: - """HEAD request to resolve content type and size.""" - try: - with httpx.Client( - timeout=httpx.Timeout(15.0), headers={"User-Agent": "SigmaRAG/1.0"} - ) as client: - resp = client.head(url, follow_redirects=True) - resp.raise_for_status() - ctype = resp.headers.get("content-type", "").split(";")[0].strip() - size = int(resp.headers.get("content-length", 0)) - return ctype, size, str(resp.url) - except Exception: - return None, None, None - - -def _sha256_file(path: Path) -> str: - import hashlib - - h = hashlib.sha256() - with open(path, "rb") as f: - for chunk in iter(lambda: f.read(65536), b""): - h.update(chunk) - return h.hexdigest() - - -def _sha256_bytes(data: bytes) -> str: - import hashlib - - return hashlib.sha256(data).hexdigest() - - -def _sanitize_filename(url: str) -> str: - import re - - name = Path(url).name - name = re.sub(r"[^a-zA-Z0-9._-]", "_", name) - return name if name else "downloaded_file" +"""Process Sigma rules already in doc_registry and download their reference documents. + +This module delegates to :func:`download_sigma_references` with ``mode="registry"`` +for the actual download logic. The legacy ``process_sigma_refs`` entry point is +kept for backward compatibility. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from src.application.documents.sigma_ref_downloader import download_sigma_references +from src.shared.utils.identify_file_type import SUPPORTED_REFERENCE_DOC_TYPES + +logger = logging.getLogger(__name__) + +DEFAULT_REQUEST_DELAY = 0.5 +DEFAULT_MAX_WORKERS = 5 + + +def process_sigma_refs( + db: Any, + output_dir: str, + supported_types: set[str] | None = None, + request_delay: float = DEFAULT_REQUEST_DELAY, + progress_callback: Callable[[int, int, str], None] | None = None, + max_workers: int = DEFAULT_MAX_WORKERS, +) -> dict[str, Any]: + """Process Sigma rules already in doc_registry and download their references. + + Delegates to :func:`download_sigma_references` with ``mode="registry"``. + + Parameters + ---------- + db : DatabaseService + Database service for writing registry entries. + output_dir : str + Directory where downloaded reference files are stored. + supported_types : set[str] | None + Allowed content types. Defaults to all reference doc types. + request_delay : float + Delay between sequential HEAD requests. + progress_callback : callable | None + ``progress_callback(current, total, phase)``. + max_workers : int + Max concurrent download threads. + + Returns + ------- + dict + Summary with keys: ``total_rules``, ``total_refs``, ``downloaded``, + ``skipped``, ``failed``. + """ + if supported_types is None: + supported_types = SUPPORTED_REFERENCE_DOC_TYPES + return download_sigma_references( + db=db, + output_dir=output_dir, + mode="registry", + supported_types=supported_types, + request_delay=request_delay, + progress_callback=progress_callback, + max_workers=max_workers, + ) diff --git a/src/application/documents/sigma_ref_registry.py b/src/application/documents/sigma_ref_registry.py new file mode 100644 index 00000000..23b0138a --- /dev/null +++ b/src/application/documents/sigma_ref_registry.py @@ -0,0 +1,135 @@ +"""Registry management for Sigma reference documents. + +Handles loading/saving of the registry from the database. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from src.infrastructure.database import DatabaseService +from src.shared.http import RETRY_STATUSES +from src.shared.utils import iso_now + +logger = logging.getLogger(__name__) + + +def load_registry(path: Any, db: DatabaseService) -> dict[str, Any]: + """Load the registry from doc_registry for sigmaref org. + + Args: + path: Output path (unused, kept for API compatibility). + db: Database service instance. + + Returns: + Registry dict mapping url_hash → entry. + """ + entries = db.get_entries_by_org("sigmaref", limit=0) + registry: dict[str, Any] = {} + for entry in entries: + url_hash = entry["url_hash"] + registry[url_hash] = { + "original_url": entry.get("original_url", ""), + "normalized_url": entry.get("normalized_url"), + "content_type": entry.get("content_type"), + "rule_id": entry.get("rule_id"), + "title": entry.get("title"), + "timestamp": entry.get("timestamp"), + "content_sha256": entry.get("content_sha256"), + "embed_status": entry.get("embed_status"), + "last_seen": entry.get("last_seen"), + "file_name": entry.get("file_name", ""), + } + return registry + + +def save_registry(registry: dict[str, Any], path: Any, db: DatabaseService) -> None: + """Save the registry to doc_registry atomically in a single batch. + + Args: + registry: Registry dict to save. + path: Output path (unused, kept for API compatibility). + db: Database service instance. + """ + rows = [] + now = iso_now() + for url_hash, entry in registry.items(): + if isinstance(entry, dict): + rows.append( + { + "url_hash": url_hash, + "original_url": entry.get("original_url", ""), + "normalized_url": entry.get("normalized_url"), + "content_type": entry.get("content_type"), + "rule_id": entry.get("rule_id"), + "title": entry.get("title"), + "timestamp": entry.get("timestamp"), + "content_sha256": entry.get("content_sha256"), + "org": entry.get("org", "sigmaref"), + "repo": entry.get("repo", "references"), + "file_name": entry.get("file_name", ""), + "file_size": entry.get("file_size"), + "embed_status": entry.get("embed_status", "discovery"), + "last_seen": entry.get("last_seen", now), + } + ) + if rows: + db.batch_upsert_doc_registry(rows) + + +def load_error_registry(db: DatabaseService) -> set[str]: + """Load the set of url_hash values that have previously failed (30x/40x). + + Args: + db: Database service instance. + + Returns: + Set of url_hash values that have errors. + """ + try: + entries = db.get_doc_errors() + return {e["url_hash"] for e in entries} + except Exception: + logger.warning("Failed to load error registry from DuckDB — proceeding without it") + return set() + + +def maybe_record_error( + db: DatabaseService, + url_hash: str, + original_url: str, + normalized_url: str, + status_code: int | None, + rule_id: str, + rule_title: str, +) -> None: + """Record a 30x/40x download error in doc_error so it is skipped on retry. + + Args: + db: Database service instance. + url_hash: Hash of the normalized URL. + original_url: Original URL from the rule. + normalized_url: Normalized URL. + status_code: HTTP status code (None for network errors). + rule_id: ID of the Sigma rule that referenced this URL. + rule_title: Title of the Sigma rule. + """ + if status_code is None: + return + if 300 <= status_code < 500 or (status_code >= 500 and status_code not in RETRY_STATUSES): + try: + db.upsert_doc_error( + { + "url_hash": url_hash, + "original_url": original_url, + "normalized_url": normalized_url, + "error_code": status_code, + "error_message": f"HTTP {status_code}", + "org": "sigmaref", + "repo": "references", + "timestamp": iso_now(), + } + ) + except Exception: + logger.warning("Failed to record error for %s", normalized_url) diff --git a/src/application/documents/sigma_ref_url.py b/src/application/documents/sigma_ref_url.py new file mode 100644 index 00000000..1c43609a --- /dev/null +++ b/src/application/documents/sigma_ref_url.py @@ -0,0 +1,99 @@ +"""URL type detection for Sigma reference documents. + +Detects the type of a URL based on its content type header or URL pattern. +""" + +from __future__ import annotations + +import logging +import re + +from src.shared.utils.identify_file_type import ( + SUPPORTED_REFERENCE_DOC_TYPES, + filetype_ext, +) +from src.shared.utils.url_utils import url_ext + +logger = logging.getLogger(__name__) + + +def detect_url_type(url: str, content_type: str | None = None) -> str | None: + """Detect the type of a URL. + + Args: + url: The URL to detect. + content_type: Optional content type from HEAD request. + + Returns: + Detected file type string, or None if not supported. + """ + # Try content type first + if content_type: + ctype = content_type.split(";")[0].strip().lower() + if "markdown" in ctype: + return "markdown" + + # Fall back to URL extension + ext = url_ext(url) + if ext: + ext = ext.lower().lstrip(".") + # Direct match against file type names (e.g. "markdown" -> "markdown") + if ext in SUPPORTED_REFERENCE_DOC_TYPES: + return ext + # Map short extension to file type (e.g. "md" -> "markdown") + for ft in SUPPORTED_REFERENCE_DOC_TYPES: + if filetype_ext(ft).lstrip(".") == ext: + return ft + + return None + + +def resolve_ext(url: str, ftype: str | None) -> str: + """Resolve the file extension for a URL and content type. + + Args: + url: The URL. + ftype: Detected content type. + + Returns: + File extension (e.g. ".md", ".pdf"). + """ + if ftype: + ext = filetype_ext(ftype) + if ext: + return ext + + # Fall back to URL extension + ext = url_ext(url) + if ext: + return ext + + return ".md" + + +# Pattern to match common reference URL patterns +REFERENCE_PATTERNS: list[tuple[str, re.Pattern[str]]] = [ + ("github_raw", re.compile(r"github\.com.*\/raw\/", re.IGNORECASE)), + ("github_blob", re.compile(r"github\.com.*\/blob\/", re.IGNORECASE)), + ("gitlab_raw", re.compile(r"gitlab\.com.*\/raw\/", re.IGNORECASE)), + ("bitbucket", re.compile(r"bitbucket\.org.*\/raw\/", re.IGNORECASE)), + ("rawcdn", re.compile(r"rawcdn\.com", re.IGNORECASE)), + ("pastebin", re.compile(r"pastebin\.com", re.IGNORECASE)), + ("hastebin", re.compile(r"hastebin\.com", re.IGNORECASE)), + ("dpaste", re.compile(r"dpaste\.org", re.IGNORECASE)), +] + + +def is_reference_url(url: str) -> bool: + """Check if a URL looks like a reference document URL. + + Args: + url: URL to check. + + Returns: + True if the URL matches reference patterns. + """ + for _, pattern in REFERENCE_PATTERNS: + if pattern.search(url): + return True + return False diff --git a/src/application/documents/validator.py b/src/application/documents/validator.py index 457d8bdd..d8fe4f57 100644 --- a/src/application/documents/validator.py +++ b/src/application/documents/validator.py @@ -1,80 +1,3 @@ -"""Sigma rule validator.""" - -from __future__ import annotations - -import logging - -from src.application.documents.models import ValidationError, ValidationResult -from src.core.sigma.models import SigmaRule - -logger = logging.getLogger(__name__) - -REQUIRED_FIELDS = {"title", "detection", "condition"} -OPTIONAL_FIELDS = { - "description", - "author", - "date", - "modified", - "references", - "tags", - "level", - "falsepositives", - "logsource", - "status", - "license", -} - - -def validate_sigma_rule(rule: SigmaRule) -> ValidationResult: - """Validate a Sigma rule against specification. - - Args: - rule: SigmaRule to validate - - Returns: - ValidationResult with validation status and errors - """ - errors: list[ValidationError] = [] - - if not rule.title or not rule.title.strip(): - errors.append(ValidationError(field="title", message="Title cannot be empty")) - - if not rule.condition or not str(rule.condition).strip(): - errors.append(ValidationError(field="condition", message="Condition cannot be empty")) - - if not rule.detection or not isinstance(rule.detection, dict): - errors.append( - ValidationError(field="detection", message="Detection must be a non-empty dict") - ) - - if rule.level is not None: - valid_levels = [ - "informational", - "low", - "medium", - "high", - "critical", - ] - if rule.level.lower() not in valid_levels: - errors.append( - ValidationError( - field="level", - message=f"Invalid level '{rule.level}'. Must be one of: {valid_levels}", - ) - ) - - if rule.status is not None: - valid_statuses = ["experimental", "stable", "testing", "deprecated", "test", "unsupported"] - if rule.status.lower() not in valid_statuses: - errors.append( - ValidationError( - field="status", - message=f"Invalid status '{rule.status}'. Must be one of: {valid_statuses}", - ) - ) - - return ValidationResult( - valid=len(errors) == 0, - rule=rule, - errors=errors, - ) +"""Sigma rule validator (deprecated — validation moved to SigmaValidator).""" + +from __future__ import annotations diff --git a/src/application/models/download.py b/src/application/models/download.py index 776b108a..aefda9d8 100644 --- a/src/application/models/download.py +++ b/src/application/models/download.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import hashlib import os from pathlib import Path @@ -22,7 +23,8 @@ def __init__( token: str | None = None, ) -> None: self.temp_dir = temp_dir or TEMP_DIR - self.token = token or os.environ.get("HF_TOKEN") + raw_token = token or os.environ.get("HF_TOKEN") + self.token = raw_token if raw_token else None self.temp_dir.mkdir(parents=True, exist_ok=True) self._metadata: dict[str, dict] = {} self._load_metadata() @@ -36,16 +38,21 @@ def _load_metadata(self) -> None: except RuntimeError: pass - def list_gguf_files( + async def list_gguf_files( self, repo: HFRepo, ) -> list[dict[str, Any]]: """List all .gguf files in a repository with metadata.""" + import huggingface_hub.constants as hc from huggingface_hub import HfApi - api = HfApi(token=self.token) + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False try: - info = api.model_info(repo_id=repo.full_id, files_metadata=True) + api = HfApi(token=self.token) + info = await asyncio.to_thread( + api.model_info, repo_id=repo.full_id, files_metadata=True + ) siblings = info.siblings or [] results = [] @@ -60,24 +67,38 @@ def list_gguf_files( return results except Exception as e: raise DownloadError(f"Failed to list GGUF files for {repo.full_id}: {e}") from e + finally: + hc.HF_HUB_OFFLINE = was_offline - def get_model_info(self, repo: HFRepo): + async def get_model_info(self, repo: HFRepo): """Get model info from HuggingFace.""" + import huggingface_hub.constants as hc from huggingface_hub import HfApi - api = HfApi(token=self.token) - return api.model_info(repo_id=repo.full_id) + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False + try: + api = HfApi(token=self.token) + return await asyncio.to_thread(api.model_info, repo_id=repo.full_id) + finally: + hc.HF_HUB_OFFLINE = was_offline def download_repo(self, repo: HFRepo, target_dir: Path) -> Path: """Download an entire repository.""" + import huggingface_hub.constants as hc from huggingface_hub import snapshot_download - path = snapshot_download( - repo_id=repo.full_id, - local_dir=target_dir, - token=self.token, - ) - return Path(path) + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False + try: + path = snapshot_download( + repo_id=repo.full_id, + local_dir=target_dir, + token=self.token, + ) + return Path(path) + finally: + hc.HF_HUB_OFFLINE = was_offline def verify_checksum(self, file_path: Path, expected_sha256: str) -> bool: """Verify file checksum.""" @@ -102,22 +123,36 @@ async def list_models(self, query: str, task: str | None = None) -> list[HFRepo] query: Search query string. task: Optional pipeline tag filter (e.g. ``"feature-extraction"`` for embedding models, ``"text-generation"`` for LLMs). + When ``None``, no pipeline tag filter is applied. Returns: List of matching :class:`HFRepo` instances. """ + import huggingface_hub.constants as hc from huggingface_hub import HfApi - api = HfApi(token=self.token) - kwargs: dict[str, Any] = {"search": query, "sort": "downloads"} - pipeline_tags = [task] if task else ["feature-extraction"] - if task == "feature-extraction": - pipeline_tags = ["feature-extraction", "sentence-similarity"] + was_offline = hc.HF_HUB_OFFLINE + hc.HF_HUB_OFFLINE = False try: + api = HfApi(token=self.token) + kwargs: dict[str, Any] = {"search": query, "sort": "downloads"} + if task is None: + pipeline_tags: list[str] = [] + elif task == "feature-extraction": + pipeline_tags = ["feature-extraction", "sentence-similarity"] + else: + pipeline_tags = [task] seen: set[str] = set() results: list[HFRepo] = [] - for tag in pipeline_tags: - kwargs["pipeline_tag"] = tag + if pipeline_tags: + for tag in pipeline_tags: + kwargs["pipeline_tag"] = tag + for r in api.list_models(**kwargs): + if r.id not in seen: + seen.add(r.id) + results.append(HFRepo.from_string(r.id)) + else: + kwargs.pop("pipeline_tag", None) for r in api.list_models(**kwargs): if r.id not in seen: seen.add(r.id) @@ -125,3 +160,5 @@ async def list_models(self, query: str, task: str | None = None) -> list[HFRepo] return results except Exception as e: raise DownloadError(f"Failed to search models: {e}") from e + finally: + hc.HF_HUB_OFFLINE = was_offline diff --git a/src/application/models/embedding.py b/src/application/models/embedding.py index 645912d1..315ade43 100644 --- a/src/application/models/embedding.py +++ b/src/application/models/embedding.py @@ -76,7 +76,7 @@ async def get_model_info(self, repo_id: str) -> Any: service = HFDownloadService() repo = HFRepo.from_string(repo_id) - return service.get_model_info(repo) + return await service.get_model_info(repo) async def search_models( self, query: str = "sentence-transformers", limit: int = 10 @@ -212,7 +212,7 @@ async def get_repo_files(self, repo_id: str) -> list[str]: repo = HFRepo.from_string(repo_id) except ValueError as e: raise DownloadError(f"Invalid repo_id '{repo_id}': {e}") from e - api = self.download_service.get_model_info(repo) + api = await self.download_service.get_model_info(repo) if api.siblings: return [f.rfilename for f in api.siblings] return [] diff --git a/src/application/sigma/validator.py b/src/application/sigma/validator.py index dcd427f9..cd74b3f3 100644 --- a/src/application/sigma/validator.py +++ b/src/application/sigma/validator.py @@ -7,26 +7,32 @@ import yaml +from src.core.sigma.models import SigmaRule from src.shared.exceptions import ValidationError logger = logging.getLogger(__name__) -REQUIRED_FIELDS = ["id", "name", "description", "detection"] +REQUIRED_FIELDS = ["id", "description", "detection"] DEPRECATED_FIELDS = ["level", "falsepositives"] # Sigma v2 deprecated MAX_FILE_SIZE = 1024 * 1024 +_VALID_LEVELS = frozenset({"informational", "low", "medium", "high", "critical"}) +_VALID_STATUSES = frozenset( + {"experimental", "stable", "testing", "deprecated", "test", "unsupported"} +) + class SigmaValidator: """Validates Sigma rule YAML files.""" - def validate(self, content: bytes) -> dict[str, Any]: + def validate(self, content: bytes) -> SigmaRule: """Validate Sigma rule YAML content. Args: content: Raw YAML file content as bytes Returns: - Parsed and validated rule dictionary + Parsed and validated SigmaRule Raises: ValidationError: If validation fails @@ -41,29 +47,49 @@ def validate(self, content: bytes) -> dict[str, Any]: ) try: - rule_data = yaml.safe_load(content) + data = yaml.safe_load(content) except yaml.YAMLError as e: raise ValidationError( field="yaml_syntax", message=f"Invalid YAML syntax: {str(e)}", ) from None - if not isinstance(rule_data, dict): + if not isinstance(data, dict): raise ValidationError( field="yaml_structure", message="YAML content must be a mapping (dictionary)", ) - self._validate_required_fields(rule_data) - self._validate_field_types(rule_data) - self._validate_detection_section(rule_data) - self._check_deprecated_fields(rule_data) - self._validate_condition_syntax(rule_data) - return rule_data + # Normalize name/title duality — Sigma YAML uses "name" or "title" + self._normalize_name_title(data) + + self._validate_required_fields(data) + self._validate_field_types(data) + self._validate_detection_section(data) + self._check_deprecated_fields(data) + self._validate_condition_syntax(data) + self._validate_level(data) + self._validate_status(data) + + return SigmaRule.from_dict(data) + + def _normalize_name_title(self, data: dict[str, Any]) -> None: + """Normalize 'name' / 'title' duality. + + Sigma YAML can use either key. The model uses 'title'. + """ + if "title" in data and "name" not in data: + data["name"] = data["title"] + elif "name" in data and "title" not in data: + data["title"] = data["name"] def _validate_required_fields(self, data: dict[str, Any]) -> None: """Validate that all required Sigma fields are present.""" + # At least one of 'name' or 'title' is required + has_name_field = "name" in data or "title" in data missing = [field for field in REQUIRED_FIELDS if field not in data] + if not has_name_field: + missing.append("name") if missing: raise ValidationError( field="required_fields", @@ -75,7 +101,8 @@ def _validate_field_types(self, data: dict[str, Any]) -> None: if not isinstance(data.get("id"), str) or not data["id"].strip(): raise ValidationError(field="id", message="Rule ID must be a non-empty string") - if not isinstance(data.get("name"), str) or not data["name"].strip(): + name_val = data.get("name") or data.get("title", "") + if not isinstance(name_val, str) or not name_val.strip(): raise ValidationError(field="name", message="Rule name must be a non-empty string") if not isinstance(data.get("description"), str) or not data["description"].strip(): @@ -108,10 +135,8 @@ def _validate_condition_syntax(self, data: dict[str, Any]) -> None: """Validate condition syntax if present.""" condition = data.get("condition") if condition and isinstance(condition, str): - # Basic validation: condition should reference detection keys detection_keys = set(data.get("detection", {}).keys()) condition_words = set(condition.replace("(", " ").replace(")", " ").split()) - # Check if condition references non-existent detection keys invalid_refs = ( condition_words - detection_keys - {"and", "or", "not", "1", "of", "them"} ) @@ -119,3 +144,23 @@ def _validate_condition_syntax(self, data: dict[str, Any]) -> None: logger.warning( f"Condition may reference non-existent detection keys: {invalid_refs}" ) + + def _validate_level(self, data: dict[str, Any]) -> None: + """Validate level field if present.""" + level = data.get("level") + if level is not None and level not in _VALID_LEVELS: + raise ValidationError( + field="level", + message=f"Invalid level '{level}'. " + f"Must be one of: {', '.join(sorted(_VALID_LEVELS))}", + ) + + def _validate_status(self, data: dict[str, Any]) -> None: + """Validate status field if present.""" + status = data.get("status") + if status is not None and status not in _VALID_STATUSES: + raise ValidationError( + field="status", + message=f"Invalid status '{status}'. " + f"Must be one of: {', '.join(sorted(_VALID_STATUSES))}", + ) diff --git a/src/application/system/duckdb.py b/src/application/system/duckdb.py index eca4a05a..812a7350 100644 --- a/src/application/system/duckdb.py +++ b/src/application/system/duckdb.py @@ -38,6 +38,7 @@ class DuckDbStatus(NamedTuple): "git_metadata", "git_selected_dirs", "sigma_spec", + "rule_references", "doc_error", "worker_state", "release_cache", diff --git a/src/core/pipeline/__init__.py b/src/core/pipeline/__init__.py index 6b825187..9d546185 100644 --- a/src/core/pipeline/__init__.py +++ b/src/core/pipeline/__init__.py @@ -5,7 +5,6 @@ build_embed_model, get_embedding_dimension, ) -from .orchestrator import RAGPipeline __all__ = [ "UnifiedIndexer", @@ -13,5 +12,4 @@ "build_embed_model", "DEFAULT_MODEL", "get_embedding_dimension", - "RAGPipeline", ] diff --git a/src/core/pipeline/indexer.py b/src/core/pipeline/indexer.py index 7340d808..906ea971 100644 --- a/src/core/pipeline/indexer.py +++ b/src/core/pipeline/indexer.py @@ -16,6 +16,7 @@ from src.core.document.parser.generic_parser import GenericTransform from src.core.pipeline.ingestion import IngestionPipelineBuilder from src.infrastructure.database import DatabaseService +from src.shared.utils.identify_file_type import filetype_subdir logger = logging.getLogger(__name__) @@ -106,6 +107,26 @@ async def index(self, route: IndexRoute) -> IndexResult: if not docs: continue + # Inject doc_registry metadata (rule_id, original_url, etc.) + # into each document so the Qdrant point carries the provenance + # of the reference (which rule referenced it, where from). + if route.table_name == "doc_registry": + source_meta = { + k: row.get(k) + for k in ( + "rule_id", + "original_url", + "normalized_url", + "url_hash", + "content_type", + "title", + ) + if row.get(k) + } + if source_meta: + for doc in docs: + doc.metadata.update(source_meta) + # Some transforms (notably SigmaParser → SigmaChunker) may emit # chunks with empty text (missing fields). The pipeline filters # those out internally, but we skip empties early to avoid the @@ -138,7 +159,12 @@ async def index(self, route: IndexRoute) -> IndexResult: return result async def index_all(self, group: str | None = None) -> list[IndexResult]: - """Execute configured routes, optionally filtered by group ("spec" or "docs").""" + """Execute configured routes, optionally filtered by group ("spec" or "docs"). + + Only processes entries with ``embed_status = 'discovery'`` (new or + changed files). Entries that are already ``'embedded'`` and whose + content has not changed are skipped — see Q3.4. + """ if group == "spec": routes = [r for r in ROUTES if r.table_name == "sigma_spec"] elif group == "docs": @@ -151,6 +177,16 @@ async def index_all(self, group: str | None = None) -> list[IndexResult]: results.append(r) return results + async def index_incremental(self, group: str | None = None) -> list[IndexResult]: + """Alias for :meth:`index_all` — only pending/changed entries are indexed. + + This is the preferred method for routine sync operations. Full rebuild + (all entries) requires an explicit ``reset_embed_status_for_collection`` + call via the API before calling this method. + """ + logger.info("Starting incremental index (group=%s)", group) + return await self.index_all(group=group) + def _get_pending(self, route: IndexRoute) -> list[dict]: """Fetch pending entries for a given route.""" if route.table_name == "sigma_spec": @@ -175,7 +211,12 @@ def _resolve_path(self, table_name: str, row: dict) -> Path | None: return Path(cfg.local_documents_path).resolve() / file_name if org == "sigmaref": - return Path(cfg.sigmaref_documents_path).resolve() / file_name + base = Path(cfg.sigmaref_documents_path).resolve() + subdir = filetype_subdir(row.get("content_type", "")) + candidate = base / subdir / file_name + if candidate.exists(): + return candidate + return base / file_name repo = row.get("repo", "") or "" if org and repo: diff --git a/src/core/pipeline/ingestion.py b/src/core/pipeline/ingestion.py index 37766885..4dd9ba0a 100644 --- a/src/core/pipeline/ingestion.py +++ b/src/core/pipeline/ingestion.py @@ -32,8 +32,8 @@ DEFAULT_MODEL = "intfloat/multilingual-e5-small" DEFAULT_CHUNK_SIZE = 1024 DEFAULT_CHUNK_OVERLAP = 100 -DEFAULT_EMBED_BATCH_SIZE = 8 -DEFAULT_NUM_WORKERS = 0 +DEFAULT_EMBED_BATCH_SIZE = 64 +DEFAULT_NUM_WORKERS = 4 DEFAULT_SIMILARITY_TOP_K = 5 # Collection names that use the transform system instead of SentenceSplitter. @@ -115,23 +115,18 @@ def build_embed_model(model_name: str) -> BaseEmbedding: model_path = str(local_path) if local_path.exists() else model_name global _embed_dim - try: - logger.info("Loading embedding model from %s", model_path) - # Air-gap mode: use only local files, no network calls to HF Hub - import os as _os - - _was_offline = _os.environ.get("HF_HUB_OFFLINE") == "1" + # Air-gap mode is already enforced by main.py at startup — + # huggingface_hub.constants.HF_HUB_OFFLINE is True. + # No os.environ manipulation needed here. - if not Path(model_path).exists(): - # Force offline mode when model is not found locally (air-gap) - _os.environ["HF_HUB_OFFLINE"] = "1" + import sys as _sys + from io import StringIO as _StringIO - # Suppress tqdm/progress bar output during model loading - import sys as _sys - from io import StringIO as _StringIO + _old_stderr = _sys.stderr + _sys.stderr = _StringIO() - _old_stderr = _sys.stderr - _sys.stderr = _StringIO() + try: + logger.info("Loading embedding model from %s", model_path) model = HuggingFaceEmbedding( model_name=model_path, @@ -141,23 +136,10 @@ def build_embed_model(model_name: str) -> BaseEmbedding: text_instruction="passage: ", ) - # Restore stderr and previous offline state after loading - _sys.stderr = _old_stderr - if not _was_offline and "HF_HUB_OFFLINE" in _os.environ: - del _os.environ["HF_HUB_OFFLINE"] - _embed_dim = _detect_embed_dim(model) logger.info("Detected embedding dimension: %d", _embed_dim) return model except Exception as e: - # Restore stderr on error too - if " _sys" in dir() and hasattr(_sys, "stderr"): - try: - _sys.stderr = _old_stderr - except Exception: - pass - if not _was_offline and "HF_HUB_OFFLINE" in _os.environ: - del _os.environ["HF_HUB_OFFLINE"] logger.error( "Embedding model %s failed to load (path: %s): %s", model_name, @@ -165,6 +147,8 @@ def build_embed_model(model_name: str) -> BaseEmbedding: e, ) raise + finally: + _sys.stderr = _old_stderr class IngestionPipelineBuilder: diff --git a/src/core/pipeline/orchestrator.py b/src/core/pipeline/orchestrator.py index b4d77a1a..158c98a9 100644 --- a/src/core/pipeline/orchestrator.py +++ b/src/core/pipeline/orchestrator.py @@ -1,22 +1,3 @@ -"""RAG pipeline implementation — delegates to IngestionPipelineBuilder.""" +"""RAG pipeline — (module kept for backwards compatibility, core logic moved to SearchEngine).""" from __future__ import annotations - -import logging -from typing import Any - -logger = logging.getLogger(__name__) - - -class RAGPipeline: - """RAG pipeline for semantic search backed by IngestionPipelineBuilder.""" - - def __init__(self) -> None: - """Initialize the pipeline.""" - from src.core.search.engine import SearchEngine - - self.search_engine = SearchEngine() - - async def search(self, query: str, limit: int = 10) -> list[dict[str, Any]]: - """Search for similar documents.""" - return await self.search_engine.search(query, top_k=limit) diff --git a/src/core/registry.py b/src/core/registry.py index 88175b80..f793984e 100644 --- a/src/core/registry.py +++ b/src/core/registry.py @@ -4,7 +4,8 @@ import logging from pathlib import Path -from typing import TYPE_CHECKING, Iterator, Type +from collections.abc import Iterator +from typing import TYPE_CHECKING, Type from .base import DocumentTransform diff --git a/src/core/search/engine.py b/src/core/search/engine.py index 37227fc6..4dbb47d8 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -37,6 +37,7 @@ "chunk_type", "collection", "tags", + "references", } ) @@ -95,6 +96,16 @@ def _get_search_embed_model() -> Any: DEFAULT_COLLECTIONS = ["sigma_rules", "sigma_docs", "sigma_spec"] +# Per-collection hybrid alpha values (1.0 = pure dense, 0.0 = pure sparse) +# Tuned for each collection's usage pattern — see Q2.1 +ALPHA_BY_COLLECTION: dict[str, float] = { + "sigma_rules": 0.5, # balanced — rules mix technical keywords & semantics + "sigma_docs": 0.7, # keyword-leaning — reference docs use precise terms + "sigma_spec": 0.3, # semantic-leaning — specs describe abstract concepts +} + +RRF_K_DEFAULT = 60 + def _node_to_result(node: Any) -> dict[str, Any]: """Convert a LlamaIndex NodeWithScore to the legacy dict format. @@ -244,6 +255,9 @@ def format_result_by_collection(result: dict[str, Any]) -> dict[str, Any]: "score": result.get("score", 0.0), "collection": collection, "source_file": meta.get("source_file", ""), + "file_path": meta.get("file_path", ""), + "line_start": meta.get("line_start", ""), + "metadata": meta, } if collection == "sigma_rules": @@ -264,6 +278,8 @@ def format_result_by_collection(result: dict[str, Any]) -> dict[str, Any]: "doc_type": meta.get("doc_type", ""), "heading_text": meta.get("heading_text", ""), "heading_level": meta.get("heading_level", 0), + "original_url": meta.get("original_url", ""), + "source_rule_id": meta.get("rule_id", ""), } # sigma_spec or unknown @@ -298,7 +314,10 @@ def __init__( similarity_threshold: float = SIMILARITY_THRESHOLD, use_router: bool = False, llm_client: LlamaClient | None = None, - alpha: float = 0.3, + alpha: float | None = None, + alpha_by_collection: dict[str, float] | None = None, + rrf_k: int = 60, + rrf_weights: dict[str, float] | None = None, ) -> None: """Initialize search engine. @@ -311,14 +330,27 @@ def __init__( relevant collections instead of all three. llm_client: Optional LlamaClient for the router. When not provided, creates a new LlamaClient. - alpha: Hybrid search weight (1.0=pure dense, 0.0=pure sparse, 0.3=keyword-leaning). + alpha: Fallback hybrid weight when a collection has no per-collection + override. Defaults to 0.3. + alpha_by_collection: Per-collection hybrid weights. Takes precedence + over the global ``alpha``. See ``ALPHA_BY_COLLECTION`` for defaults. + rrf_k: RRF constant (default 60). Lower values give top ranks more + weight; higher values flatten the score distribution. + rrf_weights: Per-collection RRF boost factors (default 1.0 for all). + Weighted RRF formula: ``score = weight / (rrf_k + rank)``. """ self.collection_names = collection_names or list(DEFAULT_COLLECTIONS) self.top_k = top_k self.similarity_threshold = similarity_threshold self.use_router = use_router self._llm_client = llm_client - self.alpha = alpha + self._fallback_alpha = alpha if alpha is not None else 0.3 + merged = dict(ALPHA_BY_COLLECTION) + if alpha_by_collection: + merged.update(alpha_by_collection) + self._alpha_by_collection = merged + self._rrf_k = rrf_k + self._rrf_weights = rrf_weights or {} async def search( self, @@ -339,6 +371,15 @@ async def search( limit = top_k if top_k is not None else self.top_k per_collection_k = max(limit * 2, 10) + # Parse inline key:value filters (e.g. references:url) from query + inline_filters, clean_query = parse_query_filters(query) + if inline_filters: + metadata_filter = {**(metadata_filter or {}), **inline_filters} + query = clean_query if clean_query else query + + # When filtering by references, always include sigma_docs + has_ref_filter = metadata_filter and "references" in metadata_filter + # Determine which collections to search cols_to_search: list[str] = [] if self.use_router: @@ -353,6 +394,11 @@ async def search( else: cols_to_search = self.collection_names + # When filtering by references, always include sigma_docs even if + # the router chose otherwise — reference docs live there. + if has_ref_filter and "sigma_docs" not in cols_to_search: + cols_to_search.append("sigma_docs") + # Build Qdrant filter for metadata filtering qdrant_filter = build_qdrant_filter(metadata_filter) if metadata_filter else None @@ -364,7 +410,7 @@ async def search( collection_name=col, top_k=per_collection_k, metadata_filter=qdrant_filter, - alpha=self.alpha, + alpha=self._alpha_by_collection.get(col, self._fallback_alpha), ) retrievers.append(retriever) except Exception as e: @@ -428,16 +474,21 @@ async def _retrieve_from(retriever: Any) -> list[dict[str, Any]]: ] all_results = await asyncio.gather(*tasks, return_exceptions=True) + rrf_k = self._rrf_k + rrf_weights = self._rrf_weights + rrf_scores: dict[tuple[str, str], dict[str, Any]] = {} - for col_results in all_results: + for col_idx, col_results in enumerate(all_results): if isinstance(col_results, Exception): logger.warning( "Collection search failed during RRF fusion, skipping: %s", col_results ) continue results_list: list[dict[str, Any]] = col_results # type: ignore[assignment] + col_name = collections[col_idx] if col_idx < len(collections) else "" + weight = rrf_weights.get(col_name, 1.0) for rank, result in enumerate(results_list, start=1): - rrf_score = 1.0 / (60 + rank) + rrf_score = weight / (rrf_k + rank) text = result.get("text", "") meta = result.get("metadata", {}) file_path = meta.get("file_path", "") if isinstance(meta, dict) else "" diff --git a/src/core/search/eval/__init__.py b/src/core/search/eval/__init__.py new file mode 100644 index 00000000..c349f82c --- /dev/null +++ b/src/core/search/eval/__init__.py @@ -0,0 +1,28 @@ +"""Search quality evaluation infrastructure. + +Provides golden set management, recall/precision metrics, evaluation runner, +and exact-search baseline comparison for Qdrant hybrid search tuning. +""" + +from src.core.search.eval.golden_set import GoldenQuery, GoldenSet, load_golden_set, save_golden_set +from src.core.search.eval.metrics import ( + context_precision, + context_recall, + mean_reciprocal_rank, + recall_at_k, +) +from src.core.search.eval.runner import SearchEvaluator +from src.core.search.eval.exact_search import compare_exact_vs_approximate + +__all__ = [ + "GoldenQuery", + "GoldenSet", + "load_golden_set", + "save_golden_set", + "recall_at_k", + "mean_reciprocal_rank", + "context_precision", + "context_recall", + "SearchEvaluator", + "compare_exact_vs_approximate", +] diff --git a/src/core/search/eval/exact_search.py b/src/core/search/eval/exact_search.py new file mode 100644 index 00000000..7338c8a4 --- /dev/null +++ b/src/core/search/eval/exact_search.py @@ -0,0 +1,229 @@ +"""Exact vs approximate search comparison for Qdrant collections. + +Provides utilities to: +- Run dense exact search (flat scan) on a Qdrant collection +- Compare exact results against the default HNSW approximate search +- Compute recall@k of approximate search relative to exact search + +This is used as a baseline (Q0.2) to determine whether HNSW parameters +need tuning — if approximate recall drops >5% below exact, ``ef`` and ``m`` +should be adjusted. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any + +from qdrant_client import QdrantClient + +from src.core.search.eval.metrics import recall_at_k + +logger = logging.getLogger(__name__) + + +@dataclass +class SearchComparisonResult: + """Result of comparing exact vs approximate search.""" + + collection: str + query: str + exact_ids: list[str] + approx_ids: list[str] + recall_at_10: float + recall_at_20: float + recall_at_50: float + exact_count: int + approx_count: int + + def summary(self) -> dict[str, Any]: + return { + "collection": self.collection, + "query": self.query, + "exact_count": self.exact_count, + "approx_count": self.approx_count, + "recall@10": self.recall_at_10, + "recall@20": self.recall_at_20, + "recall@50": self.recall_at_50, + "approx_below_exact_threshold": self.recall_at_10 < 0.95, + } + + +def _extract_doc_ids(points: list[Any]) -> list[str]: + """Extract document IDs from Qdrant scored points.""" + ids: list[str] = [] + for point in points: + payload = point.payload or {} + doc_id = payload.get("doc_id") or payload.get("file_path") or payload.get("id") + if doc_id: + ids.append(str(doc_id)) + return ids + + +def run_exact_search( + client: QdrantClient, + collection: str, + query_embedding: list[float], + limit: int = 100, +) -> list[str]: + """Run dense exact search (flat scan, no HNSW approximation). + + Uses ``search()`` with a large limit to get a near-complete result set, + which serves as the ground truth for recall computation. + + Args: + client: Active Qdrant client. + collection: Collection name. + query_embedding: Query vector. + limit: Maximum number of results to return. + + Returns: + List of document IDs in relevance order. + """ + points = client.search( # type: ignore[attr-defined] + collection_name=collection, + query_vector=query_embedding, + limit=limit, + ) + return _extract_doc_ids(points) + + +def run_approximate_search( + client: QdrantClient, + collection: str, + query_embedding: list[float], + limit: int = 50, +) -> list[str]: + """Run dense approximate search (HNSW, default parameters). + + Args: + client: Active Qdrant client. + collection: Collection name. + query_embedding: Query vector. + limit: Maximum number of results to return. + + Returns: + List of document IDs in relevance order. + """ + points = client.query_points( + collection_name=collection, + query=query_embedding, + limit=limit, + ) + return _extract_doc_ids(points.points) + + +def compare_exact_vs_approximate( + client: QdrantClient, + collection: str, + query_embedding: list[float], + exact_limit: int = 100, + approx_limit: int = 50, + recall_threshold: float = 0.95, +) -> SearchComparisonResult: + """Compare exact search against approximate (HNSW) search. + + Runs both searches and computes recall@k of approximate relative to exact. + + Args: + client: Active Qdrant client. + collection: Collection name. + query_embedding: Query vector. + exact_limit: Number of results for exact search. + approx_limit: Number of results for approximate search. + recall_threshold: If recall@10 falls below this, the result flags + ``approx_below_exact_threshold`` as ``True``. + + Returns: + ``SearchComparisonResult`` with recall scores at k=10, 20, 50. + """ + exact_ids = run_exact_search(client, collection, query_embedding, limit=exact_limit) + approx_ids = run_approximate_search(client, collection, query_embedding, limit=approx_limit) + + r10 = recall_at_k(approx_ids, exact_ids, k=10) + r20 = recall_at_k(approx_ids, exact_ids, k=20) + r50 = recall_at_k(approx_ids, exact_ids, k=min(50, len(exact_ids))) + + result = SearchComparisonResult( + collection=collection, + query="", + exact_ids=exact_ids, + approx_ids=approx_ids, + recall_at_10=r10, + recall_at_20=r20, + recall_at_50=r50, + exact_count=len(exact_ids), + approx_count=len(approx_ids), + ) + + if r10 < recall_threshold: + logger.warning( + "Approximate search recall@10=%.3f below threshold %.2f on '%s'", + r10, + recall_threshold, + collection, + ) + else: + logger.info( + "Approximate search recall@10=%.3f on '%s' (OK, threshold=%.2f)", + r10, + collection, + recall_threshold, + ) + + return result + + +def compare_collection_baseline( + client: QdrantClient, + collection: str, + query_embeddings: list[list[float]], + recall_threshold: float = 0.95, +) -> dict[str, float]: + """Run exact vs approximate comparison across multiple queries. + + Aggregates recall@10 across all queries to produce a single baseline + metric for the collection. + + Args: + client: Active Qdrant client. + collection: Collection name. + query_embeddings: List of query vectors to test. + recall_threshold: Threshold for flagging poor recall. + + Returns: + Dict with aggregated metrics: + ``mean_recall_at_10``, ``min_recall_at_10``, ``max_recall_at_10``, + ``queries_below_threshold``. + """ + recalls_10: list[float] = [] + below_count = 0 + + for i, embedding in enumerate(query_embeddings): + try: + result = compare_exact_vs_approximate( + client, collection, embedding, recall_threshold=recall_threshold + ) + recalls_10.append(result.recall_at_10) + if result.recall_at_10 < recall_threshold: + below_count += 1 + except Exception as e: + logger.warning("Comparison failed for query %d on '%s': %s", i, collection, e) + + if not recalls_10: + return { + "mean_recall_at_10": 0.0, + "min_recall_at_10": 0.0, + "max_recall_at_10": 0.0, + "queries_below_threshold": 0, + "total_queries": 0, + } + + return { + "mean_recall_at_10": sum(recalls_10) / len(recalls_10), + "min_recall_at_10": min(recalls_10), + "max_recall_at_10": max(recalls_10), + "queries_below_threshold": below_count, + "total_queries": len(query_embeddings), + } diff --git a/src/core/search/eval/golden_set.py b/src/core/search/eval/golden_set.py new file mode 100644 index 00000000..2b9caf0a --- /dev/null +++ b/src/core/search/eval/golden_set.py @@ -0,0 +1,182 @@ +"""Golden set data model and I/O for search quality evaluation. + +A golden set is a curated collection of ``(query, relevant_doc_ids)`` pairs +used to measure retrieval quality (recall@k, precision@k, MRR). + +File format (JSON):: + + { + "metadata": { + "version": 1, + "description": "Sigma rules retrieval — 50 queries", + "created_at": "2025-01-15T10:30:00Z" + }, + "queries": [ + { + "id": "q001", + "query": "detect powershell execution", + "collection": "sigma_rules", + "relevant_doc_ids": ["rule-abc-123", "rule-def-456"], + "k": 10 + } + ] + } +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +@dataclass +class GoldenQuery: + """A single golden-set query with its ground-truth relevant documents.""" + + id: str + query: str + collection: str + relevant_doc_ids: list[str] + k: int = 10 + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "GoldenQuery": + return cls( + id=data["id"], + query=data["query"], + collection=data["collection"], + relevant_doc_ids=data["relevant_doc_ids"], + k=data.get("k", 10), + ) + + +@dataclass +class GoldenSet: + """A complete golden set with metadata and queries.""" + + description: str = "" + queries: list[GoldenQuery] = field(default_factory=list) + + def add( + self, + query: str, + collection: str, + relevant_doc_ids: list[str], + k: int = 10, + query_id: str | None = None, + ) -> GoldenQuery: + q = GoldenQuery( + id=query_id or _make_id(len(self.queries)), + query=query, + collection=collection, + relevant_doc_ids=relevant_doc_ids, + k=k, + ) + self.queries.append(q) + return q + + def __len__(self) -> int: + return len(self.queries) + + def to_dict(self) -> dict[str, Any]: + return { + "metadata": { + "version": 1, + "description": self.description, + "created_at": datetime.now(timezone.utc).isoformat(), + "num_queries": len(self.queries), + }, + "queries": [q.to_dict() for q in self.queries], + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "GoldenSet": + queries = [GoldenQuery.from_dict(q) for q in data.get("queries", [])] + return cls( + description=data.get("metadata", {}).get("description", ""), + queries=queries, + ) + + +def _make_id(index: int) -> str: + return f"q{index:04d}" + + +def load_golden_set(path: str | Path) -> GoldenSet: + """Load a golden set from a JSON file. + + Args: + path: Path to the golden set JSON file. + + Returns: + Parsed GoldenSet. + + Raises: + FileNotFoundError: If the file does not exist. + ValueError: If the JSON structure is invalid. + """ + path = Path(path) + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + + if "queries" not in data: + raise ValueError(f"Golden set file missing 'queries' key: {path}") + return GoldenSet.from_dict(data) + + +def save_golden_set(golden_set: GoldenSet, path: str | Path) -> Path: + """Save a golden set to a JSON file. + + Args: + golden_set: The golden set to save. + path: Output file path. + + Returns: + The resolved output path. + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(golden_set.to_dict(), f, ensure_ascii=False, indent=2) + return path + + +def create_sample_golden_set() -> GoldenSet: + """Create a small sample golden set for development / testing. + + These queries use doc_ids that correspond to typical Sigma rule IDs + found in the sigma_rules collection. Replace with real ground truth + data for production evaluation. + """ + gs = GoldenSet(description="Sample golden set for development") + gs.add( + query="detect powershell execution", + collection="sigma_rules", + relevant_doc_ids=["sample-rule-001", "sample-rule-002"], + k=10, + ) + gs.add( + query="CMD.exe suspicious command line", + collection="sigma_rules", + relevant_doc_ids=["sample-rule-003"], + k=10, + ) + gs.add( + query="network connection to C2 server", + collection="sigma_docs", + relevant_doc_ids=["sample-doc-001"], + k=10, + ) + gs.add( + query="Sigma rule format specification", + collection="sigma_spec", + relevant_doc_ids=["sample-spec-001"], + k=10, + ) + return gs diff --git a/src/core/search/eval/metrics.py b/src/core/search/eval/metrics.py new file mode 100644 index 00000000..bfbc18f2 --- /dev/null +++ b/src/core/search/eval/metrics.py @@ -0,0 +1,178 @@ +"""Retrieval quality metrics for search evaluation. + +Implements standard information retrieval metrics: +- ``recall_at_k``: fraction of relevant docs found in top-k +- ``precision_at_k``: fraction of top-k results that are relevant +- ``mean_reciprocal_rank``: average of 1/rank of first relevant result +- ``context_precision``: weighted precision by position (AP-style) +- ``context_recall``: recall estimated via retrieved contexts vs ground truth +""" + +from __future__ import annotations + + +def recall_at_k(retrieved_ids: list[str], relevant_ids: list[str], k: int | None = None) -> float: + """Fraction of relevant documents found in the top-k retrieved results. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + k: Number of results to consider. Defaults to ``len(retrieved_ids)``. + + Returns: + Recall score in ``[0.0, 1.0]``. Returns ``0.0`` if no relevant docs + are known (avoids division by zero). + """ + if not relevant_ids: + return 0.0 + cutoff = retrieved_ids[:k] if k is not None else retrieved_ids + hits = sum(1 for doc_id in cutoff if doc_id in relevant_ids) + return hits / len(relevant_ids) + + +def precision_at_k( + retrieved_ids: list[str], relevant_ids: list[str], k: int | None = None +) -> float: + """Fraction of top-k retrieved results that are relevant. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + k: Number of results to consider. Defaults to ``len(retrieved_ids)``. + + Returns: + Precision score in ``[0.0, 1.0]``. Returns ``0.0`` if k=0. + """ + cutoff = retrieved_ids[:k] if k is not None else retrieved_ids + if not cutoff: + return 0.0 + hits = sum(1 for doc_id in cutoff if doc_id in relevant_ids) + return hits / len(cutoff) + + +def mean_reciprocal_rank(retrieved_ids: list[str], relevant_ids: list[str]) -> float: + """Mean Reciprocal Rank — average of ``1/rank`` for the first relevant doc. + + If no relevant document appears in the retrieved list, returns ``0.0``. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + + Returns: + MRR score in ``[0.0, 1.0]``. + """ + relevant_set = set(relevant_ids) + for rank, doc_id in enumerate(retrieved_ids, start=1): + if doc_id in relevant_set: + return 1.0 / rank + return 0.0 + + +def average_precision(retrieved_ids: list[str], relevant_ids: list[str]) -> float: + """Average Precision — precision averaged at each relevant document position. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + + Returns: + AP score in ``[0.0, 1.0]``. Returns ``0.0`` if no relevant docs. + """ + if not relevant_ids: + return 0.0 + relevant_set = set(relevant_ids) + hits = 0 + precision_sum = 0.0 + for rank, doc_id in enumerate(retrieved_ids, start=1): + if doc_id in relevant_set: + hits += 1 + precision_sum += hits / rank + return precision_sum / len(relevant_ids) + + +def context_precision(retrieved_ids: list[str], relevant_ids: list[str]) -> float: + """Position-weighted precision (similar to AP but normalised by min(k, |relevant|)). + + Unlike AP which divides by ``|relevant_ids|``, this divides by the number of + relevant docs that actually appear in the retrieved list (or all relevant + docs if they all appear), giving a score in ``[0, 1]`` even when only a + subset of relevant docs is retrieved. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + + Returns: + Context precision score in ``[0.0, 1.0]``. + """ + if not retrieved_ids or not relevant_ids: + return 0.0 + relevant_set = set(relevant_ids) + weighted_sum = 0.0 + for rank, doc_id in enumerate(retrieved_ids, start=1): + if doc_id in relevant_set: + weighted_sum += 1.0 / rank + divisor = min(len(retrieved_ids), len(relevant_ids)) + return weighted_sum / divisor + + +def context_recall(retrieved_ids: list[str], relevant_ids: list[str]) -> float: + """Estimates recall from the retriever's perspective. + + Measures what fraction of known relevant documents the retriever actually + returned. This is the standard recall@all (k = full retrieved list). + + Args: + retrieved_ids: Document IDs returned by the retriever. + relevant_ids: Ground-truth relevant document IDs. + + Returns: + Context recall score in ``[0.0, 1.0]``. + """ + if not relevant_ids: + return 0.0 + retrieved_set = set(retrieved_ids) + hits = sum(1 for doc_id in relevant_ids if doc_id in retrieved_set) + return hits / len(relevant_ids) + + +def evaluate_query( + retrieved_ids: list[str], + relevant_ids: list[str], + k: int | None = None, +) -> dict[str, float]: + """Compute all metrics for a single query. + + Args: + retrieved_ids: Document IDs returned by the retriever, in rank order. + relevant_ids: Ground-truth relevant document IDs. + k: Optional cutoff for recall@k and precision@k. + + Returns: + Dict with keys: ``recall_at_k``, ``precision_at_k``, ``mrr``, + ``average_precision``, ``context_precision``, ``context_recall``. + """ + return { + "recall_at_k": recall_at_k(retrieved_ids, relevant_ids, k), + "precision_at_k": precision_at_k(retrieved_ids, relevant_ids, k), + "mrr": mean_reciprocal_rank(retrieved_ids, relevant_ids), + "average_precision": average_precision(retrieved_ids, relevant_ids), + "context_precision": context_precision(retrieved_ids, relevant_ids), + "context_recall": context_recall(retrieved_ids, relevant_ids), + } + + +def aggregate_metrics(metrics_list: list[dict[str, float]]) -> dict[str, float]: + """Average a list of per-query metric dicts. + + Args: + metrics_list: Per-query results from :func:`evaluate_query`. + + Returns: + Dict of mean metric values. + """ + if not metrics_list: + return {} + keys = metrics_list[0].keys() + return {key: sum(m[key] for m in metrics_list) / len(metrics_list) for key in keys} diff --git a/src/core/search/eval/runner.py b/src/core/search/eval/runner.py new file mode 100644 index 00000000..5918336f --- /dev/null +++ b/src/core/search/eval/runner.py @@ -0,0 +1,179 @@ +"""Search evaluation runner. + +Orchestrates golden set evaluation against a ``SearchEngine`` (or any +callable that returns ranked document IDs for a query). + +Usage:: + + evaluator = SearchEvaluator(search_engine) + results = evaluator.run(golden_set) + print(results.summary()) +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Callable + +from src.core.search.eval.golden_set import GoldenSet +from src.core.search.eval.metrics import aggregate_metrics, evaluate_query + +logger = logging.getLogger(__name__) + +# Type alias for a search function: (query, collection) -> list[doc_id] +SearchFn = Callable[[str, str], list[str]] + + +@dataclass +class QueryResult: + """Evaluation result for a single query.""" + + query_id: str + query: str + collection: str + retrieved_ids: list[str] + relevant_ids: list[str] + metrics: dict[str, float] + elapsed_ms: float + + +@dataclass +class EvaluationResults: + """Aggregated evaluation results across all queries.""" + + queries: list[QueryResult] = field(default_factory=list) + description: str = "" + + def add(self, result: QueryResult) -> None: + self.queries.append(result) + + @property + def num_queries(self) -> int: + return len(self.queries) + + def summary(self) -> dict[str, Any]: + """Return aggregated metrics and per-collection breakdown.""" + if not self.queries: + return {"num_queries": 0} + + per_collection: dict[str, list[dict[str, float]]] = {} + for qr in self.queries: + per_collection.setdefault(qr.collection, []).append(qr.metrics) + + overall = aggregate_metrics([qr.metrics for qr in self.queries]) + per_collection_agg = { + col: aggregate_metrics(metrics_list) for col, metrics_list in per_collection.items() + } + + return { + "description": self.description, + "num_queries": self.num_queries, + "overall": overall, + "per_collection": per_collection_agg, + } + + def report(self) -> str: + """Return a human-readable evaluation report.""" + summary = self.summary() + lines = [ + "=== Search Evaluation Report ===", + f"Queries: {summary['num_queries']}", + f"Description: {summary.get('description', '')}", + "", + "--- Overall Metrics ---", + ] + for key, value in summary.get("overall", {}).items(): + if isinstance(value, float): + lines.append(f" {key}: {value:.4f}") + else: + lines.append(f" {key}: {value}") + + for col, metrics in summary.get("per_collection", {}).items(): + lines.append("") + lines.append(f"--- {col} ---") + for key, value in metrics.items(): + if isinstance(value, float): + lines.append(f" {key}: {value:.4f}") + else: + lines.append(f" {key}: {value}") + + return "\n".join(lines) + + +class SearchEvaluator: + """Evaluates a search function against a golden set. + + Args: + search_fn: Callable that takes ``(query, collection)`` and returns + a list of document IDs in rank order. Typically wraps + ``SearchEngine.search()`` or ``SearchEngine.search_collection()``. + """ + + def __init__(self, search_fn: SearchFn) -> None: + self._search_fn = search_fn + + def _search(self, query: str, collection: str) -> list[str]: + """Execute search and extract document IDs.""" + try: + result = self._search_fn(query, collection) + if isinstance(result, list): + return result + return [] + except Exception as e: + logger.warning("Search failed for '%s' on '%s': %s", query, collection, e) + return [] + + def evaluate_query(self, golden_query: Any) -> QueryResult: + """Evaluate a single golden query. + + Args: + golden_query: A ``GoldenQuery`` dataclass instance. + + Returns: + ``QueryResult`` with metrics. + """ + start = time.perf_counter() + retrieved_ids = self._search(golden_query.query, golden_query.collection) + elapsed_ms = (time.perf_counter() - start) * 1000 + + metrics = evaluate_query(retrieved_ids, golden_query.relevant_doc_ids, k=golden_query.k) + + return QueryResult( + query_id=golden_query.id, + query=golden_query.query, + collection=golden_query.collection, + retrieved_ids=retrieved_ids, + relevant_ids=golden_query.relevant_doc_ids, + metrics=metrics, + elapsed_ms=elapsed_ms, + ) + + def run( + self, + golden_set: GoldenSet, + description: str = "", + ) -> EvaluationResults: + """Run evaluation across all queries in the golden set. + + Args: + golden_set: The golden set to evaluate against. + description: Optional description for the report. + + Returns: + ``EvaluationResults`` with per-query and aggregated metrics. + """ + results = EvaluationResults(description=description) + for q in golden_set.queries: + qr = self.evaluate_query(q) + results.add(qr) + logger.info( + " [%s] recall@%d=%.3f MRR=%.3f (%.1fms)", + q.id, + q.k, + qr.metrics["recall_at_k"], + qr.metrics["mrr"], + qr.elapsed_ms, + ) + return results diff --git a/src/core/search/router.py b/src/core/search/router.py index 5d2aab0d..f05a5afe 100644 --- a/src/core/search/router.py +++ b/src/core/search/router.py @@ -21,9 +21,13 @@ Collections: - sigma_rules: Detection rules (YARA, Sigma, Splunk queries, MITRE ATT&CK techniques, threat hunting, IOCs, log sources) -- sigma_docs: Documentation (architecture, setup, configuration, how-to guides, explanations) +- sigma_docs: Reference documents downloaded from external URLs linked to Sigma rules (blog posts, Microsoft docs, security research, LOLBas pages, PDFs, troubleshooting guides, external articles, references cited by rules) - sigma_spec: Specification reference (YAML format, field definitions, schema, syntax, encoding rules, modifiers, tags, logsource taxonomy, correlation rules, filters, FAQ level/status definitions) +Search sigma_docs when the user asks about external references, citations, blog posts, articles, +or documentation linked from Sigma rules. Queries mentioning specific URLs or reference +documents should include sigma_docs. + Examples of sigma_spec queries: "What severity levels exist", "How do I write a correlation rule", "What does the contains modifier do", "How to tag a MITRE technique", "What logsource for Windows Security", "How are maps evaluated", "What filename conventions for Sigma rules", "How does group-by work", "What is the difference between temporal and temporal_ordered" Return ONLY a JSON object with a single key "collections" containing a list of collection names. diff --git a/src/core/search/sparse_encoder.py b/src/core/search/sparse_encoder.py index 27431ec8..ce8166ed 100644 --- a/src/core/search/sparse_encoder.py +++ b/src/core/search/sparse_encoder.py @@ -110,6 +110,10 @@ _MIN_WORD_LENGTH = 3 _MAX_HASH_ID = 2**24 +# BM25 default parameters +_K1 = 1.2 +_B = 0.75 + def _tokenize(text: str) -> list[str]: """Lowercase tokenisation — alphanumeric words of at least *min* chars.""" @@ -121,8 +125,23 @@ def _token_id(token: str) -> int: return int(hashlib.md5(token.encode()).hexdigest()[:8], 16) % _MAX_HASH_ID -def _encode_single(text: str) -> tuple[list[int], list[float]]: - """Encode *text* into a sparse vector (indices, values).""" +def _tf_weight(freq: int) -> float: + """Sublinear TF saturation with BM25's k1.""" + return (_K1 + 1.0) * freq / (_K1 + freq) + + +def _encode_single( + text: str, + idf_map: dict[str, float] | None = None, + avg_doc_len: float | None = None, +) -> tuple[list[int], list[float]]: + """Encode *text* into a sparse vector (indices, values). + + When ``idf_map`` is provided, weights are BM25-style: + ``score = IDF * (k1 + 1) * tf / (k1 * (1 - b + b * Ld / Lavg) + tf)`` + + Otherwise falls back to sublinear TF (``1 + log(tf)``). + """ tokens = [t for t in _tokenize(text) if t not in STOP_WORDS] if not tokens: return [], [] @@ -131,19 +150,65 @@ def _encode_single(text: str) -> tuple[list[int], list[float]]: for t in tokens: tf[t] = tf.get(t, 0) + 1 + doc_len = len(tokens) + length_norm = 1.0 + if avg_doc_len is not None and avg_doc_len > 0: + length_norm = 1.0 - _B + _B * (doc_len / avg_doc_len) * _K1 + indices: list[int] = [] values: list[float] = [] for term, freq in tf.items(): + if idf_map is not None: + idf = idf_map.get(term, 1.5) + bm25_score = idf * _tf_weight(freq) / (length_norm + freq) + values.append(bm25_score) + else: + values.append(1.0 + math.log(freq)) indices.append(_token_id(term)) - values.append(1.0 + math.log(freq)) return indices, values +class IDFCalculator: + """Accumulates corpus-level term frequencies to compute IDF. + + Usage: + calc = IDFCalculator() + for text in corpus: + calc.add_document(text) + idf_map = calc.idf() + """ + + def __init__(self) -> None: + self._df: dict[str, int] = {} + self._num_docs: int = 0 + + def add_document(self, text: str) -> int: + """Add a single document to the corpus. Returns token count.""" + tokens = set(_tokenize(text)) + tokens.discard("") + for t in tokens: + self._df[t] = self._df.get(t, 0) + 1 + self._num_docs += 1 + return len(tokens) + + def idf(self, smooth: bool = True) -> dict[str, float]: + """Return ``{term: idf}`` map using BM25's IDF formula. + + With smoothing: ``idf = log(1 + (N - df + 0.5) / (df + 0.5))`` + """ + n = self._num_docs + if n == 0: + return {} + if smooth: + return {t: math.log(1.0 + (n - df + 0.5) / (df + 0.5)) for t, df in self._df.items()} + return {t: math.log(n / df) for t, df in self._df.items()} + + def bm25_sparse_encoder( texts: list[str], ) -> tuple[list[list[int]], list[list[float]]]: - """BM25-style sparse encoder for a batch of texts. + """TF-only sparse encoder for a batch of texts. Each unique term gets weight = 1 + log(term_frequency). Token IDs are deterministic hashes, making the encoder stateless. @@ -166,8 +231,36 @@ def bm25_sparse_encoder( return all_indices, all_values +def bm25_idf_sparse_encoder( + texts: list[str], + *, + idf_map: dict[str, float], + avg_doc_len: float, +) -> tuple[list[list[int]], list[list[float]]]: + """Full BM25 sparse encoder with IDF + document length normalisation. + + Args: + texts: Batch of input strings. + idf_map: Pre-computed ``{term: idf}`` map. + avg_doc_len: Average document length for length normalisation. + + Returns: + Tuple of (all_indices, all_values) where each element corresponds + to one input text. + """ + all_indices: list[list[int]] = [] + all_values: list[list[float]] = [] + + for text in texts: + indices, values = _encode_single(text, idf_map=idf_map, avg_doc_len=avg_doc_len) + all_indices.append(indices) + all_values.append(values) + + return all_indices, all_values + + def create_sparse_encoder(): - """Return a BM25-based sparse encoder compatible with LlamaIndex/Qdrant. + """Return a TF-based sparse encoder compatible with LlamaIndex/Qdrant. The returned callable accepts ``List[str]`` (batched) and returns ``Tuple[List[List[int]], List[List[float]]]``. diff --git a/src/core/sigma/chunker.py b/src/core/sigma/chunker.py index 8a3657ec..f035ba94 100644 --- a/src/core/sigma/chunker.py +++ b/src/core/sigma/chunker.py @@ -86,167 +86,134 @@ def post_process(self, documents: list[Document]) -> list[Document]: return documents - def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> list[dict]: - """Chunk a single Sigma rule dict into enriched chunk dicts. - - This is the refactored version of the legacy chunk_sigma_rules_rich function. - - Args: - rule: Raw Sigma rule dict. - llm_client: Optional LLM client for keyword extraction. - """ - title = rule.get("title", "Untitled Sigma rule") - rule_id = rule.get("id") - description = rule.get("description", "") - level = rule.get("level", "unknown") - status = rule.get("status", "unknown") - tags = rule.get("tags", []) - logsource = rule.get("logsource", {}) - detection = rule.get("detection", {}) - condition = detection.get("condition", "") - falsepositives = rule.get("falsepositives", []) - references = rule.get("references", []) - author = rule.get("author", "") - date = rule.get("date", "") - modified = rule.get("modified", "") - - product = logsource.get("product", "unknown") - category = logsource.get("category", "unknown") - service = logsource.get("service", "unknown") - - chunks: list[dict] = [] + def _extract_fields(self, rule: dict) -> dict: + return { + "title": rule.get("title", "Untitled Sigma rule"), + "rule_id": rule.get("id"), + "description": rule.get("description", ""), + "level": rule.get("level", "unknown"), + "status": rule.get("status", "unknown"), + "tags": rule.get("tags", []), + "detection": rule.get("detection", {}), + "condition": rule.get("detection", {}).get("condition", ""), + "falsepositives": rule.get("falsepositives", []), + "references": rule.get("references", []), + "author": rule.get("author", ""), + "date": rule.get("date", ""), + "modified": rule.get("modified", ""), + "product": rule.get("logsource", {}).get("product", "unknown"), + "category": rule.get("logsource", {}).get("category", "unknown"), + "service": rule.get("logsource", {}).get("service", "unknown"), + } - # Executive summary - chunks.append( - make_chunk( - rule, - "executive_summary", - ( - f"Sigma rule: {title}\n" - f"Rule ID: {rule_id}\n" - f"Purpose: {description}\n" - f"This rule is designed for {product} logs with " - f"category={category} and service={service}.\n" - f"Severity: {level}. Status: {status}.\n" - f"Main detection logic: {condition}" - ), - eval_questions=self._summarize_questions(title), - ) + def _build_executive_summary(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "executive_summary", + ( + f"Sigma rule: {f['title']}\n" + f"Rule ID: {f['rule_id']}\n" + f"Purpose: {f['description']}\n" + f"This rule is designed for {f['product']} logs with " + f"category={f['category']} and service={f['service']}.\n" + f"Severity: {f['level']}. Status: {f['status']}.\n" + f"Main detection logic: {f['condition']}" + ), + eval_questions=self._summarize_questions(f["title"]), ) - # Rule metadata and lifecycle - chunks.append( - make_chunk( - rule, - "rule_metadata_lifecycle", - ( - f"Rule metadata for {title}.\n" - f"Rule ID: {rule_id}\n" - f"Author: {author}\n" - f"Created date: {date}\n" - f"Modified date: {modified}\n" - f"Status: {status}\n" - f"Level: {level}" - ), - extra_meta={"references": references}, - eval_questions=self._lifecycle_questions(title), - ) + def _build_metadata_lifecycle(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "rule_metadata_lifecycle", + ( + f"Rule metadata for {f['title']}.\n" + f"Rule ID: {f['rule_id']}\n" + f"Author: {f['author']}\n" + f"Created date: {f['date']}\n" + f"Modified date: {f['modified']}\n" + f"Status: {f['status']}\n" + f"Level: {f['level']}" + ), + extra_meta={"references": f["references"]}, + eval_questions=self._lifecycle_questions(f["title"]), ) - # Logsource context - chunks.append( - make_chunk( - rule, - "logsource_context", - ( - f"Logsource context for Sigma rule {title}.\n" - f"Product: {product}\n" - f"Category: {category}\n" - f"Service: {service}\n" - f"The rule expects telemetry from product={product}, " - f"category={category}, service={service}." - ), - eval_questions=self._logsource_questions(title, product, category, service), - ) + def _build_logsource_context(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "logsource_context", + ( + f"Logsource context for Sigma rule {f['title']}.\n" + f"Product: {f['product']}\n" + f"Category: {f['category']}\n" + f"Service: {f['service']}\n" + f"The rule expects telemetry from product={f['product']}, " + f"category={f['category']}, service={f['service']}." + ), + eval_questions=self._logsource_questions( + f["title"], f["product"], f["category"], f["service"] + ), ) - # MITRE ATT&CK mapping - attack_tags = [tag for tag in tags if str(tag).startswith("attack.")] - if attack_tags: - chunks.append( - make_chunk( - rule, - "mitre_attack_mapping", - (f"MITRE ATT&CK mapping for {title}.\nTags:\n{format_value(attack_tags)}"), - extra_meta={"attack_tags": attack_tags}, - eval_questions=self._attck_questions(title, attack_tags), - ) - ) + def _build_mitre_attack_mapping(self, rule: dict, f: dict) -> dict | None: + attack_tags = [tag for tag in f["tags"] if str(tag).startswith("attack.")] + if not attack_tags: + return None + return make_chunk( + rule, + "mitre_attack_mapping", + f"MITRE ATT&CK mapping for {f['title']}.\nTags:\n{format_value(attack_tags)}", + extra_meta={"attack_tags": attack_tags}, + eval_questions=self._attck_questions(f["title"], attack_tags), + ) - # Detection condition - chunks.append( - make_chunk( - rule, - "detection_condition", - ( - f"Detection condition for {title}.\n" - f"Condition: {condition}\n" - f"Interpretation: this condition defines how selection and " - f"filter blocks are combined to trigger the rule." - ), - extra_meta={"condition": condition}, - eval_questions=[ - f"What is the Sigma condition of {title}?", - f"How are selections and filters combined in {title}?", - f"Que signifie la condition de détection de {title} ?", - ], - ) + def _build_detection_condition(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "detection_condition", + ( + f"Detection condition for {f['title']}.\n" + f"Condition: {f['condition']}\n" + f"Interpretation: this condition defines how selection and " + f"filter blocks are combined to trigger the rule." + ), + extra_meta={"condition": f["condition"]}, + eval_questions=[ + f"What is the Sigma condition of {f['title']}?", + f"How are selections and filters combined in {f['title']}?", + f"Que signifie la condition de détection de {f['title']} ?", + ], ) - # Iteration over detection blocks + def _build_detection_block_chunks(self, rule: dict, f: dict) -> tuple[list[dict], list[dict]]: + chunks: list[dict] = [] all_atomic_facts: list[dict] = [] - for detection_name, detection_value in detection.items(): + for detection_name, detection_value in f["detection"].items(): if detection_name == "condition": continue - is_filter = detection_name.startswith("filter") chunk_type = "detection_filter_block" if is_filter else "detection_selection_block" - facts = flatten_detection_values(detection_value) all_atomic_facts.extend( - [ - { - "detection_name": detection_name, - **fact, - "is_filter": is_filter, - } - for fact in facts - ] + {"detection_name": detection_name, **fact, "is_filter": is_filter} for fact in facts ) - - # Detection block chunk chunks.append( make_chunk( rule, chunk_type, ( - f"Detection block {detection_name} in Sigma rule {title}.\n" + f"Detection block {detection_name} in Sigma rule {f['title']}.\n" f"Block role: {'exclusion / false positive reduction' if is_filter else 'positive detection indicator'}\n" f"Raw block:\n{format_value(detection_value)}" ), - extra_meta={ - "detection_name": detection_name, - "is_filter": is_filter, - }, - eval_questions=self._block_questions(title, detection_name, is_filter), + extra_meta={"detection_name": detection_name, "is_filter": is_filter}, + eval_questions=self._block_questions(f["title"], detection_name, is_filter), ) ) - - # Field/operator groups by_field_operator: dict[str, list] = {} for fact in facts: by_field_operator.setdefault(fact["field_operator"], []).append(fact["value"]) - for field_operator, values in by_field_operator.items(): field, operator = split_field_operator(field_operator) chunks.append( @@ -254,7 +221,7 @@ def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> li rule, "field_operator_group", ( - f"Field/operator group in {title}.\n" + f"Field/operator group in {f['title']}.\n" f"Detection block: {detection_name}\n" f"Field: {field}\n" f"Operator: {operator}\n" @@ -267,11 +234,9 @@ def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> li "operator": operator, "is_filter": is_filter, }, - eval_questions=self._field_questions(title, field, operator), + eval_questions=self._field_questions(f["title"], field, operator), ) ) - - # Atomic indicators for fact in facts: field, operator = split_field_operator(fact["field_operator"]) value = fact["value"] @@ -280,13 +245,13 @@ def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> li rule, "atomic_indicator", ( - f"Atomic Sigma indicator for {title}.\n" + f"Atomic Sigma indicator for {f['title']}.\n" f"Detection block: {detection_name}\n" f"Field: {field}\n" f"Operator: {operator}\n" f"Value: {value}\n" f"Role: {'legitimate exclusion / filter' if is_filter else 'suspicious or monitored indicator'}\n" - f"A match on this value contributes to the rule condition: {condition}" + f"A match on this value contributes to the rule condition: {f['condition']}" ), extra_meta={ "detection_name": detection_name, @@ -295,141 +260,178 @@ def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> li "value": value, "is_filter": is_filter, }, - eval_questions=self._indicator_questions(title, value), + eval_questions=self._indicator_questions(f["title"], value), ) ) - - # Indicator inventory - suspicious_values = [f for f in all_atomic_facts if not f["is_filter"]] - filter_values = [f for f in all_atomic_facts if f["is_filter"]] - chunks.append( - make_chunk( - rule, - "indicator_inventory", - ( - f"Indicator inventory for {title}.\n" - f"Suspicious or monitored values:\n" - f"{format_value([f['value'] for f in suspicious_values])}\n\n" - f"Filtered legitimate values:\n" - f"{format_value([f['value'] for f in filter_values])}" - ), - eval_questions=[ - f"List all monitored values in {title}.", - f"What indicators are monitored by {title}?", - f"Which values are excluded by {title}?", - ], - ) + return chunks, all_atomic_facts + + def _build_indicator_inventory(self, rule: dict, f: dict, all_facts: list[dict]) -> dict: + suspicious = [x for x in all_facts if not x["is_filter"]] + filters = [x for x in all_facts if x["is_filter"]] + return make_chunk( + rule, + "indicator_inventory", + ( + f"Indicator inventory for {f['title']}.\n" + f"Suspicious or monitored values:\n" + f"{format_value([x['value'] for x in suspicious])}\n\n" + f"Filtered legitimate values:\n" + f"{format_value([x['value'] for x in filters])}" + ), + eval_questions=[ + f"List all monitored values in {f['title']}.", + f"What indicators are monitored by {f['title']}?", + f"Which values are excluded by {f['title']}?", + ], ) - # Investigation guidance - chunks.append( - make_chunk( - rule, - "investigation_guidance", - ( - f"Investigation guidance for alerts from {title}.\n" - f"Investigate the entity, user, process, host, timestamp, " - f"and event context that matched the Sigma rule.\n" - f"Review whether the matched values are expected in the environment." - ), - eval_questions=[ - f"How should an analyst investigate alerts from {title}?", - f"What context matters for {title} alerts?", - f"Que faut-il vérifier lors d'une alerte {title} ?", - ], - ) + def _build_investigation_guidance(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "investigation_guidance", + ( + f"Investigation guidance for alerts from {f['title']}.\n" + f"Investigate the entity, user, process, host, timestamp, " + f"and event context that matched the Sigma rule.\n" + f"Review whether the matched values are expected in the environment." + ), + eval_questions=[ + f"How should an analyst investigate alerts from {f['title']}?", + f"What context matters for {f['title']} alerts?", + f"Que faut-il vérifier lors d'une alerte {f['title']} ?", + ], ) - # False positive context - chunks.append( - make_chunk( - rule, - "false_positive_context", - ( - f"False positive context for {title}.\n" - f"False positives:\n{format_value(falsepositives)}\n" - f"Common benign causes may include administrative activity, " - f"automation, package managers, security tools." - ), - eval_questions=[ - f"What false positives can occur for {title}?", - f"How can false positives be reduced for {title}?", - f"Quels faux positifs sont attendus pour {title} ?", - ], - ) + def _build_false_positive_context(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "false_positive_context", + ( + f"False positive context for {f['title']}.\n" + f"False positives:\n{format_value(f['falsepositives'])}\n" + f"Common benign causes may include administrative activity, " + f"automation, package managers, security tools." + ), + eval_questions=[ + f"What false positives can occur for {f['title']}?", + f"How can false positives be reduced for {f['title']}?", + f"Quels faux positifs sont attendus pour {f['title']} ?", + ], ) - # Natural language queries - chunks.append( - make_chunk( - rule, - "natural_language_queries", - ( - f"Natural language retrieval hints for {title}.\n" - f"This rule is relevant for questions such as:\n" - f"- What does {title} detect?\n" - f"- Which detection fields and values are used?\n" - f"- What logsource is required?" - ), - eval_questions=[ - f"Which rule detects behavior described by {title}?", - f"What fields are used by {title}?", - f"What ATT&CK mapping is associated with {title}?", - ], - ) + def _build_natural_language_queries(self, rule: dict, f: dict) -> dict: + return make_chunk( + rule, + "natural_language_queries", + ( + f"Natural language retrieval hints for {f['title']}.\n" + f"This rule is relevant for questions such as:\n" + f"- What does {f['title']} detect?\n" + f"- Which detection fields and values are used?\n" + f"- What logsource is required?" + ), + eval_questions=[ + f"Which rule detects behavior described by {f['title']}?", + f"What fields are used by {f['title']}?", + f"What ATT&CK mapping is associated with {f['title']}?", + ], ) - # Backend mapping hints - detection_summary_lines: list[str] = [] - for det_name, det_val in detection.items(): + def _build_backend_mapping_hints(self, rule: dict, f: dict) -> dict: + lines: list[str] = [] + for det_name, det_val in f["detection"].items(): if det_name == "condition": continue - is_filt = det_name.startswith("filter") - role = "exclusion" if is_filt else "selection" - detection_summary_lines.append(f" - {det_name} ({role}): {format_value(det_val)}") - detection_summary = "\n".join(detection_summary_lines) if detection_summary_lines else "N/A" - - chunks.append( - make_chunk( - rule, - "backend_mapping_hints", - ( - f"Backend mapping hints for {title}.\n" - f"The detection fields should be mapped to the corresponding " - f"SIEM, EDR, or log backend schema.\n" - f"Operators from Sigma such as contains, startswith, endswith, " - f"all, and equals should be preserved during translation.\n" - f"Condition: {condition}\n\n" - f"Detection blocks:\n{detection_summary}" - ), - eval_questions=[ - f"What fields should be mapped for {title}?", - f"How should {title} be translated to a SIEM query?", - f"What operators are used in {title}?", - ], - ) + role = "exclusion" if det_name.startswith("filter") else "selection" + lines.append(f" - {det_name} ({role}): {format_value(det_val)}") + summary = "\n".join(lines) if lines else "N/A" + return make_chunk( + rule, + "backend_mapping_hints", + ( + f"Backend mapping hints for {f['title']}.\n" + f"The detection fields should be mapped to the corresponding " + f"SIEM, EDR, or log backend schema.\n" + f"Operators from Sigma such as contains, startswith, endswith, " + f"all, and equals should be preserved during translation.\n" + f"Condition: {f['condition']}\n\n" + f"Detection blocks:\n{summary}" + ), + eval_questions=[ + f"What fields should be mapped for {f['title']}?", + f"How should {f['title']} be translated to a SIEM query?", + f"What operators are used in {f['title']}?", + ], ) - # Enrich all chunks with LLM-generated summaries and keywords + def _enrich_chunks(self, chunks: list[dict], llm_client: LLMClientLike) -> list[dict]: + enriched: list[dict] = [] + for chunk in chunks: + try: + result = enrich_by_llm(chunk["text"], llm_client) + except Exception: + logger.debug("LLM enrichment failed for chunk %s", chunk.get("chunk_type", "?")) + result = {"summary": None, "keywords": None, "error": "enrichment_failed"} + summary = result.get("summary") or "" + keywords = result.get("keywords") or "" + if summary or keywords: + enrichment = "\n\n---\n" + if summary: + enrichment += f"Summary: {summary}\n\n" + if keywords: + enrichment += f"Keywords: {keywords}\n" + chunk["text"] = chunk["text"] + enrichment + enriched.append(chunk) + return enriched + + def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> list[dict]: + f = self._extract_fields(rule) + chunks = self._assemble_chunks(rule, f) + if llm_client is not None: - enriched_chunks: list[dict] = [] + chunks = self._enrich_chunks(chunks, llm_client) + + return chunks + + def _assemble_chunks(self, rule: dict, f: dict) -> list[dict]: + """Assemble all chunks for a Sigma rule. + + Args: + rule: The raw Sigma rule dict. + f: Extracted fields dict from ``_extract_fields()``. + + Returns: + List of chunk dicts ready for document conversion. + """ + chunks: list[dict] = [ + self._build_executive_summary(rule, f), + self._build_metadata_lifecycle(rule, f), + self._build_logsource_context(rule, f), + ] + + if (attack := self._build_mitre_attack_mapping(rule, f)) is not None: + chunks.append(attack) + + chunks.append(self._build_detection_condition(rule, f)) + + det_chunks, all_facts = self._build_detection_block_chunks(rule, f) + chunks.extend(det_chunks) + + chunks.extend( + [ + self._build_indicator_inventory(rule, f, all_facts), + self._build_investigation_guidance(rule, f), + self._build_false_positive_context(rule, f), + self._build_natural_language_queries(rule, f), + self._build_backend_mapping_hints(rule, f), + ] + ) + + refs = f.get("references", []) + if refs: + ref_text = "\n\nReferences:\n" + "\n".join(f" - {r}" for r in refs) for chunk in chunks: - try: - result = enrich_by_llm(chunk["text"], llm_client) - except Exception: - logger.debug("LLM enrichment failed for chunk %s", chunk.get("chunk_type", "?")) - result = {"summary": None, "keywords": None, "error": "enrichment_failed"} - summary = result.get("summary") or "" - keywords = result.get("keywords") or "" - if summary or keywords: - enrichment = "\n\n---\n" - if summary: - enrichment += f"Summary: {summary}\n\n" - if keywords: - enrichment += f"Keywords: {keywords}\n" - chunk["text"] = chunk["text"] + enrichment - enriched_chunks.append(chunk) - chunks = enriched_chunks + chunk["text"] += ref_text return chunks @@ -500,20 +502,6 @@ def _indicator_questions(self, title: str, value: str) -> list[str]: f"What field contains {value} in {title}?", ] - def _generate_eval_questions(self, rule: dict) -> dict[str, list[str]]: - """Generate eval question groups for a rule (for flat-mode post_process).""" - title = rule.get("title", "Untitled Sigma rule") - return { - "summary": self._summarize_questions(title), - "lifecycle": self._lifecycle_questions(title), - "logsource": self._logsource_questions( - title, - rule.get("logsource", {}).get("product", "unknown"), - rule.get("logsource", {}).get("category", "unknown"), - rule.get("logsource", {}).get("service", "unknown"), - ), - } - # Register SigmaChunker for rich mode. TransformRegistry.register(SigmaChunker) diff --git a/src/core/sigma/models.py b/src/core/sigma/models.py index e19fb5c2..a059f608 100644 --- a/src/core/sigma/models.py +++ b/src/core/sigma/models.py @@ -57,12 +57,20 @@ def from_dict( rule_data["file_path"] = str(file_path) if line_number is not None: rule_data["line_number"] = line_number + # Normalize non-string condition values + if "condition" in rule_data and not isinstance(rule_data["condition"], str): + rule_data["condition"] = str(rule_data["condition"]) return cls(**rule_data) def to_dict(self) -> dict[str, Any]: """Convert to dictionary.""" return self.model_dump(exclude_none=True) + @property + def name(self) -> str: + """Alias for title — Sigma YAML uses 'name', model uses 'title'.""" + return self.title + @property def path(self) -> Path | None: """Get file path as Path object.""" diff --git a/src/infrastructure/database/core.py b/src/infrastructure/database/core.py index 324e61d6..c20cff0e 100644 --- a/src/infrastructure/database/core.py +++ b/src/infrastructure/database/core.py @@ -27,6 +27,7 @@ def _default_db_path() -> str: "doc_registry", "sigma_spec", "doc_error", + "rule_references", "git_metadata", "git_selected_dirs", "worker_state", @@ -71,12 +72,27 @@ def initialize(self) -> None: initdb_path = Path(__file__).parent / "initdb.sql" self._writer_conn.execute(initdb_path.read_text(encoding="utf-8")) self._writer_conn.commit() + self._apply_migrations() if self.db_path.exists(): self._load_from_file() self._initialized = True self._conn = self._writer_conn logger.info("Schema initialized successfully") + def _apply_migrations(self) -> None: + """Apply schema migrations idempotently.""" + migrations = [ + "ALTER TABLE git_selected_dirs ADD COLUMN IF NOT EXISTS source_type TEXT", + "UPDATE git_selected_dirs SET source_type = '' WHERE source_type IS NULL", + "UPDATE config SET value = '20260704' WHERE key = 'schema_version'", + ] + for sql in migrations: + try: + self._writer_conn.execute(sql) + self._writer_conn.commit() + except Exception as e: + logger.warning("Migration failed (may already be applied): %s", e) + def _load_from_file(self) -> None: with self._lock: logger.info("Loading database from %s", self.db_path) @@ -88,6 +104,15 @@ def _load_from_file(self) -> None: f"INSERT OR REPLACE INTO {table} SELECT * FROM file_db.{table}" ) except Exception: + if table == "git_selected_dirs": + try: + self._writer_conn.execute( + "INSERT OR REPLACE INTO git_selected_dirs (repo_key, dir_path, updated) " + "SELECT repo_key, dir_path, updated FROM file_db.git_selected_dirs" + ) + continue + except Exception: + pass logger.warning("Table %s not found in existing database — skipping", table) self._writer_conn.execute("DETACH file_db") logger.info("Loaded tables from disk") diff --git a/src/infrastructure/database/doc_ops.py b/src/infrastructure/database/doc_ops.py index ebd787b8..7b1b1368 100644 --- a/src/infrastructure/database/doc_ops.py +++ b/src/infrastructure/database/doc_ops.py @@ -5,7 +5,10 @@ import hashlib import logging import os +from datetime import datetime, timezone, timedelta from pathlib import Path + +from src.shared.constants import NULL_UUID from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -99,7 +102,12 @@ def upsert_doc_registry(self, data: dict) -> None: title = EXCLUDED.title, timestamp = EXCLUDED.timestamp, last_seen = EXCLUDED.last_seen, - embed_status = EXCLUDED.embed_status""", + embed_status = CASE + WHEN excluded.content_sha256 IS NOT NULL + AND excluded.content_sha256 != doc_registry.content_sha256 + THEN 'discovery' + ELSE doc_registry.embed_status + END""", ( data.get("url_hash"), data.get("org"), @@ -110,7 +118,7 @@ def upsert_doc_registry(self, data: dict) -> None: data.get("file_size"), data.get("original_url"), data.get("normalized_url"), - data.get("rule_id", "00000000-0000-0000-0000-000000000000"), + data.get("rule_id", NULL_UUID), data.get("title"), data.get("timestamp"), data.get("last_seen"), @@ -141,7 +149,12 @@ def batch_upsert_doc_registry(self, rows: list[dict]) -> None: title = EXCLUDED.title, timestamp = EXCLUDED.timestamp, last_seen = EXCLUDED.last_seen, - embed_status = EXCLUDED.embed_status""", + embed_status = CASE + WHEN excluded.content_sha256 IS NOT NULL + AND excluded.content_sha256 != doc_registry.content_sha256 + THEN 'discovery' + ELSE doc_registry.embed_status + END""", [ ( r.get("url_hash"), @@ -153,7 +166,7 @@ def batch_upsert_doc_registry(self, rows: list[dict]) -> None: r.get("file_size"), r.get("original_url"), r.get("normalized_url"), - r.get("rule_id", "00000000-0000-0000-0000-000000000000"), + r.get("rule_id", NULL_UUID), r.get("title"), r.get("timestamp"), r.get("last_seen"), @@ -252,7 +265,12 @@ def batch_upsert_sigma_spec(self, rows: list[dict]) -> None: title = EXCLUDED.title, timestamp = EXCLUDED.timestamp, last_seen = EXCLUDED.last_seen, - embed_status = EXCLUDED.embed_status""", + embed_status = CASE + WHEN excluded.content_sha256 IS NOT NULL + AND excluded.content_sha256 != sigma_spec.content_sha256 + THEN 'discovery' + ELSE sigma_spec.embed_status + END""", [ ( r.get("url_hash"), @@ -431,6 +449,25 @@ def delete_doc_registry_by_url(self, original_url: str) -> None: ) self._writer_conn.commit() + def delete_doc_registry_by_url_hashes(self, url_hashes: list[str]) -> None: + if not url_hashes: + return + with self._lock: + placeholders = ",".join("?" for _ in url_hashes) + self._writer_conn.execute( + f"DELETE FROM doc_registry WHERE url_hash IN ({placeholders})", + url_hashes, + ) + self._writer_conn.commit() + + def get_doc_registry_url_hashes_by_repo(self, org: str, repo: str) -> list[str]: + with self._lock: + results = self._writer_conn.execute( + "SELECT url_hash FROM doc_registry WHERE org = ? AND repo = ?", + (org, repo), + ).fetchall() + return [row[0] for row in results] + def get_local_files(self, limit: int = 1000, offset: int = 0) -> list[dict]: with self._lock: results = self._writer_conn.execute( @@ -564,3 +601,177 @@ def _hash_file(path: Path) -> str | None: logger.error(f"[resync_local_file_sizes] Batch update failed, rolling back: {e}") return {"updated": updated, "skipped": skipped, "error": errors, "incomplete": incomplete} + + # ------------------------------------------------------------------ + # RULE_REFERENCES table (junction rule_id ↔ url_hash) + # ------------------------------------------------------------------ + + def upsert_rule_reference(self, rule_id: str, url_hash: str, ref_url: str) -> None: + """Record that *rule_id* references *ref_url* (identified by *url_hash*).""" + with self._lock: + self._writer_conn.execute( + "INSERT INTO rule_references (rule_id, url_hash, ref_url) " + "VALUES (?, ?, ?) " + "ON CONFLICT (rule_id, url_hash) DO NOTHING", + (rule_id, url_hash, ref_url), + ) + self._writer_conn.commit() + + def batch_upsert_rule_references(self, rows: list[dict[str, str]]) -> None: + """Batch upsert rule_references rows. + + Each row must have keys: rule_id, url_hash, ref_url. + """ + with self._lock: + self._writer_conn.executemany( + "INSERT INTO rule_references (rule_id, url_hash, ref_url) " + "VALUES (?, ?, ?) " + "ON CONFLICT (rule_id, url_hash) DO NOTHING", + [(r["rule_id"], r["url_hash"], r["ref_url"]) for r in rows], + ) + self._writer_conn.commit() + + def get_referencing_rules(self, url_hash: str) -> list[dict]: + """Return all rules that reference the document identified by *url_hash*.""" + with self._lock: + results = self._writer_conn.execute( + "SELECT rule_id, ref_url, created FROM rule_references WHERE url_hash = ? ORDER BY rule_id", + (url_hash,), + ).fetchall() + col_names = [desc[0] for desc in self._writer_conn.description] + return [dict(zip(col_names, row)) for row in results] + + def get_rule_references(self, rule_id: str) -> list[dict]: + """Return all reference documents for a given *rule_id*.""" + with self._lock: + results = self._writer_conn.execute( + "SELECT r.url_hash, r.ref_url, r.created, " + "d.content_type, d.file_name, d.embed_status, d.content_sha256 " + "FROM rule_references r " + "LEFT JOIN doc_registry d ON r.url_hash = d.url_hash " + "WHERE r.rule_id = ? ORDER BY r.ref_url", + (rule_id,), + ).fetchall() + col_names = [desc[0] for desc in self._writer_conn.description] + return [dict(zip(col_names, row)) for row in results] + + def delete_rule_references_by_url_hash(self, url_hash: str) -> None: + """Delete all rule_references entries for a given *url_hash*.""" + with self._lock: + self._writer_conn.execute("DELETE FROM rule_references WHERE url_hash = ?", (url_hash,)) + self._writer_conn.commit() + + def get_rule_id_by_url_hash(self, url_hash: str) -> str | None: + """Return the rule_id associated with a doc_registry entry.""" + with self._lock: + result = self._writer_conn.execute( + "SELECT rule_id FROM doc_registry WHERE url_hash = ?", (url_hash,) + ).fetchone() + return result[0] if result else None + + def get_rule_reference_paths(self, rule_id: str) -> list[Path]: + """Return the on-disk paths for all reference documents of a rule.""" + with self._lock: + results = self._writer_conn.execute( + "SELECT r.ref_url, d.file_name, d.content_type " + "FROM rule_references r " + "JOIN doc_registry d ON r.url_hash = d.url_hash " + "WHERE r.rule_id = ?", + (rule_id,), + ).fetchall() + paths: list[Path] = [] + for ref_url, file_name, content_type in results: + try: + from src.application.documents.sigma_ref_paths import sigmaref_resolve_path + from src.config.settings import get_config + + cfg = get_config() + path = sigmaref_resolve_path( + Path(cfg.sigmaref_documents_path).resolve(), + content_type, + file_name or "", + ) + if path.exists(): + paths.append(path) + except Exception: + continue + return paths + + # ------------------------------------------------------------------ + # R1.4 — cleanup orphaned head_verified entries (no content_sha256) + # ------------------------------------------------------------------ + + def delete_head_verified_orphans(self, grace_days: int = 7) -> int: + """Delete doc_registry entries stuck in 'head_verified' without content. + + These are entries created by a HEAD request whose content type was not + in the supported set — they will never transition to 'embedded'. + Also removes corresponding rule_references rows. + + Args: + grace_days: Delete entries older than N days (default 7). + + Returns: + Number of deleted entries. + """ + cutoff = (datetime.now(timezone.utc) - timedelta(days=grace_days)).strftime( + "%Y-%m-%dT%H:%M:%SZ" + ) + with self._lock: + # Find orphan url_hashes + orphans = self._writer_conn.execute( + "SELECT url_hash FROM doc_registry " + "WHERE embed_status = 'head_verified' " + "AND (content_sha256 IS NULL OR content_sha256 = '') " + "AND (last_seen IS NULL " + " OR last_seen < ?)", + [cutoff], + ).fetchall() + orphan_hashes = [row[0] for row in orphans] + + if not orphan_hashes: + return 0 + + # Delete from rule_references first (FK-like cleanup) + placeholders = ",".join("?" for _ in orphan_hashes) + self._writer_conn.execute( + f"DELETE FROM rule_references WHERE url_hash IN ({placeholders})", + orphan_hashes, + ) + # Delete from doc_registry + self._writer_conn.execute( + f"DELETE FROM doc_registry WHERE url_hash IN ({placeholders})", + orphan_hashes, + ) + self._writer_conn.commit() + + logger.info("Deleted %d orphan head_verified entries", len(orphan_hashes)) + return len(orphan_hashes) + + def delete_unreferenced_entries(self) -> int: + """Delete sigmaref entries whose url_hash is no longer in rule_references. + + These are documents downloaded for rules that no longer exist or whose + references have been removed. Only affects entries with ``org='sigmaref'`` + to avoid touching local/GitHub entries. + """ + with self._lock: + orphans = self._writer_conn.execute( + "SELECT url_hash FROM doc_registry " + "WHERE org = 'sigmaref' " + "AND url_hash NOT IN (SELECT DISTINCT url_hash FROM rule_references)", + ).fetchall() + orphan_hashes = [row[0] for row in orphans] + + if not orphan_hashes: + return 0 + + placeholders = ",".join("?" for _ in orphan_hashes) + self._writer_conn.execute( + f"DELETE FROM doc_registry WHERE url_hash IN ({placeholders})", + orphan_hashes, + ) + self._writer_conn.commit() + + logger.info("Deleted %d unreferenced sigmaref entries", len(orphan_hashes)) + return len(orphan_hashes) diff --git a/src/infrastructure/database/initdb.sql b/src/infrastructure/database/initdb.sql index acda7b8f..dbf4a033 100644 --- a/src/infrastructure/database/initdb.sql +++ b/src/infrastructure/database/initdb.sql @@ -65,6 +65,7 @@ CREATE TABLE IF NOT EXISTS git_metadata ( CREATE TABLE IF NOT EXISTS git_selected_dirs ( repo_key TEXT NOT NULL, dir_path TEXT NOT NULL, + source_type TEXT NOT NULL DEFAULT '', updated TEXT, PRIMARY KEY (repo_key, dir_path) ); @@ -89,6 +90,17 @@ CREATE TABLE IF NOT EXISTS sigma_spec ( CREATE INDEX IF NOT EXISTS idx_sigma_spec_embed ON sigma_spec(embed_status); CREATE INDEX IF NOT EXISTS idx_sigma_spec_org_repo ON sigma_spec(org, repo); +-- rule_references (junction table: rule_id ↔ reference url_hash — M:N) +CREATE TABLE IF NOT EXISTS rule_references ( + rule_id TEXT NOT NULL, + url_hash TEXT NOT NULL, + ref_url TEXT NOT NULL, + created TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', current_timestamp)), + PRIMARY KEY (rule_id, url_hash) +); + +CREATE INDEX IF NOT EXISTS idx_rule_ref_url_hash ON rule_references(url_hash); + -- doc_error (failed URLs — 30x/40x errors to skip on retry) CREATE TABLE IF NOT EXISTS doc_error ( url_hash TEXT PRIMARY KEY, @@ -139,4 +151,4 @@ CREATE INDEX IF NOT EXISTS idx_doc_registry_org_repo ON doc_registry(org, repo); INSERT OR IGNORE INTO config (key, value) VALUES ('app_version', '"0.1.0"'), ('theme', '"dark"'), - ('schema_version', '20260619'); + ('schema_version', '20260704'); diff --git a/src/infrastructure/database/table_ops.py b/src/infrastructure/database/table_ops.py index c0470eb2..10dc1a4c 100644 --- a/src/infrastructure/database/table_ops.py +++ b/src/infrastructure/database/table_ops.py @@ -191,22 +191,26 @@ def get_selected_dirs(self, repo_key: str) -> list[str]: ).fetchall() return [row[0] for row in results] - def get_repos_with_selected_dirs(self) -> list[str]: + def get_repos_with_selected_dirs(self, source_type: str | None = None) -> list[str]: + query = "SELECT DISTINCT repo_key FROM git_selected_dirs" + params: list[Any] = [] + if source_type: + query += " WHERE source_type = ?" + params.append(source_type) + query += " ORDER BY repo_key" with self._lock: - results = self._writer_conn.execute( - "SELECT DISTINCT repo_key FROM git_selected_dirs ORDER BY repo_key" - ).fetchall() + results = self._writer_conn.execute(query, params).fetchall() return [row[0] for row in results] - def set_selected_dirs(self, repo_key: str, dirs: list[str]) -> None: + def set_selected_dirs(self, repo_key: str, dirs: list[str], source_type: str = "") -> None: with self._lock: self._writer_conn.execute( "DELETE FROM git_selected_dirs WHERE repo_key = ?", (repo_key,) ) for d in dirs: self._writer_conn.execute( - "INSERT INTO git_selected_dirs (repo_key, dir_path, updated) VALUES (?, ?, ?)", - (repo_key, d, iso_now()), + "INSERT INTO git_selected_dirs (repo_key, dir_path, source_type, updated) VALUES (?, ?, ?, ?)", + (repo_key, d, source_type, iso_now()), ) self._writer_conn.commit() diff --git a/src/infrastructure/github/git.py b/src/infrastructure/github/git.py index 1ec4f33c..524cd005 100644 --- a/src/infrastructure/github/git.py +++ b/src/infrastructure/github/git.py @@ -384,7 +384,11 @@ def _walk_dir(path: Path, current_depth: int) -> list[dict[str, Any]]: def save_selected_dirs( - org: str, name: str, selected: list[str], repos_dir: Path | None = None + org: str, + name: str, + selected: list[str], + repos_dir: Path | None = None, + source_type: str = "", ) -> dict[str, Any]: """Save selected directories for a repository to DuckDB. @@ -393,6 +397,7 @@ def save_selected_dirs( name: Repository name selected: List of relative folder paths repos_dir: Base directory for cloned repos + source_type: Source type identifier (e.g. 'github', 'spec') Returns: Result dict with success status @@ -400,7 +405,7 @@ def save_selected_dirs( db = DatabaseService.get_instance() repo_key = _get_repo_key(org, name) try: - db.set_selected_dirs(repo_key, selected) + db.set_selected_dirs(repo_key, selected, source_type=source_type) db.persist() logger.info(f"Saved selection for {org}/{name}: {selected}") return {"success": True} diff --git a/src/infrastructure/llm/llamacpp/auto_start.py b/src/infrastructure/llm/llamacpp/auto_start.py index 08ccbae2..a9c4754e 100644 --- a/src/infrastructure/llm/llamacpp/auto_start.py +++ b/src/infrastructure/llm/llamacpp/auto_start.py @@ -121,7 +121,6 @@ async def start_llamacpp() -> None: pass await asyncio.sleep(1) - _llamacpp_started_by_us = True raise ServiceStartError("llama.cpp process started but health check timed out after 10s") diff --git a/src/infrastructure/vectorstore/__init__.py b/src/infrastructure/vectorstore/__init__.py index 526f83ae..3a220faa 100644 --- a/src/infrastructure/vectorstore/__init__.py +++ b/src/infrastructure/vectorstore/__init__.py @@ -226,7 +226,9 @@ async def health_check(self) -> bool: finally: client.close() - async def create_collection(self, enable_hybrid: bool = True) -> None: + async def create_collection( + self, enable_hybrid: bool = True, enable_quantization: bool = True + ) -> None: """Create the collection if it doesn't exist.""" if self._vector_store is None: await self.initialize() @@ -239,8 +241,12 @@ async def create_collection(self, enable_hybrid: bool = True) -> None: self.collection_name, self.vector_size, enable_hybrid=enable_hybrid, + enable_quantization=enable_quantization, + ) + logger.info( + f"Collection {self.collection_name} ready " + f"(hybrid={enable_hybrid}, quantization={enable_quantization})" ) - logger.info(f"Collection {self.collection_name} ready (hybrid={enable_hybrid})") def __repr__(self) -> str: return f"QdrantService(collection={self.collection_name}, host={self.host}:{self.port})" diff --git a/src/infrastructure/vectorstore/auto_start.py b/src/infrastructure/vectorstore/auto_start.py index edbbcc2f..026a2178 100644 --- a/src/infrastructure/vectorstore/auto_start.py +++ b/src/infrastructure/vectorstore/auto_start.py @@ -102,7 +102,6 @@ async def start_qdrant( pass await asyncio.sleep(1) - _qdrant_started_by_us = True raise ServiceStartError("Qdrant process started but health check timed out after 10s") diff --git a/src/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index f798deab..6ad10846 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -15,6 +15,32 @@ logger = logging.getLogger(__name__) +def collection_hnsw_config( + collection_name: str, +) -> qdrant_client.models.HnswConfigDiff | None: + """Return per-collection HNSW config. + + ``sigma_rules`` is kept in-RAM with higher ``ef_construct`` for + better recall. ``sigma_docs`` and ``sigma_spec`` (cold collections) + use on-disk storage to save RAM. + """ + if collection_name == "sigma_rules": + return qdrant_client.models.HnswConfigDiff( + m=16, + ef_construct=200, + full_scan_threshold=10000, + on_disk=False, + ) + if collection_name in ("sigma_docs", "sigma_spec"): + return qdrant_client.models.HnswConfigDiff( + m=16, + ef_construct=100, + full_scan_threshold=10000, + on_disk=True, + ) + return None + + def _get_collections_sync(client) -> Any: """Synchronous wrapper for client.get_collections().""" return client.get_collections() @@ -32,15 +58,24 @@ def _count_sync(client, collection_name: str) -> int: def _create_collection_sync( - client, collection_name: str, vectors_config, sparse_vectors_config=None + client, + collection_name: str, + vectors_config, + sparse_vectors_config=None, + quantization_config=None, + hnsw_config=None, ): """Synchronous wrapper for client.create_collection().""" - kwargs = { + kwargs: dict[str, Any] = { "collection_name": collection_name, "vectors_config": vectors_config, } if sparse_vectors_config is not None: kwargs["sparse_vectors_config"] = sparse_vectors_config + if quantization_config is not None: + kwargs["quantization_config"] = quantization_config + if hnsw_config is not None: + kwargs["hnsw_config"] = hnsw_config client.create_collection(**kwargs) @@ -58,6 +93,7 @@ def _create_collection_sync( ("chunk_type", "keyword"), ("collection", "keyword"), ("tags", "keyword"), + ("references", "keyword"), ] @@ -104,9 +140,24 @@ async def list_collections(host: str, port: int) -> list[dict[str, Any]]: async def create_collection( - host: str, port: int, collection_name: str, vector_size: int = 384, enable_hybrid: bool = True + host: str, + port: int, + collection_name: str, + vector_size: int = 384, + enable_hybrid: bool = True, + enable_quantization: bool = True, ) -> bool: - """Create a new collection (with optional sparse vector support for hybrid search).""" + """Create a new collection (with optional sparse vector support for hybrid search). + + When *enable_quantization* is True, ScalarQuantization INT8 is applied + to reduce vector memory footprint by ~4x with minimal recall loss when + combined with rescore. + + HNSW config is selected per-collection via :func:`collection_hnsw_config`: + ``sigma_rules`` uses in-RAM HNSW with higher ``ef_construct`` for + better recall, while ``sigma_docs`` and ``sigma_spec`` use on-disk + storage to save RAM. + """ client = get_qdrant_client(host=host, port=port) try: vectors_config = qdrant_client.models.VectorParams( @@ -117,11 +168,28 @@ async def create_collection( if enable_hybrid: sparse_vectors_config = { "text-sparse": qdrant_client.models.SparseVectorParams( - index=qdrant_client.models.SparseIndexParams() + index=qdrant_client.models.SparseIndexParams(), + modifier=qdrant_client.models.Modifier.IDF, ) } + quantization_config = None + if enable_quantization: + quantization_config = qdrant_client.models.ScalarQuantization( + scalar=qdrant_client.models.ScalarQuantizationConfig( + type=qdrant_client.models.ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) + ) + hnsw_config = collection_hnsw_config(collection_name) await asyncio.to_thread( - _create_collection_sync, client, collection_name, vectors_config, sparse_vectors_config + _create_collection_sync, + client, + collection_name, + vectors_config, + sparse_vectors_config, + quantization_config, + hnsw_config, ) await asyncio.to_thread(_create_payload_indexes_sync, client, collection_name) logger.info( diff --git a/src/infrastructure/vectorstore/downloader.py b/src/infrastructure/vectorstore/downloader.py index 85715b20..cc938fb9 100644 --- a/src/infrastructure/vectorstore/downloader.py +++ b/src/infrastructure/vectorstore/downloader.py @@ -195,16 +195,6 @@ async def download_binary(self, progress_callback: ProgressCallback = None) -> d archive_path = self.bin_dir / f"qdrant{ext}" binary_path = self.get_binary_path() - # Clean stale files before extraction to avoid version mix - for f in list(self.bin_dir.iterdir()): - if f.name == archive_path.name or f.is_dir(): - continue - if f.suffix == ".exe" or f.name == "qdrant": - try: - f.unlink() - except OSError: - pass - try: if progress_callback: progress_callback(5, "Downloading binary...") @@ -221,6 +211,16 @@ async def download_binary(self, progress_callback: ProgressCallback = None) -> d else: self._safe_extract_tar_gz(archive_path, self.bin_dir) + # Clean stale files only after successful download + extraction + for f in list(self.bin_dir.iterdir()): + if f.name == archive_path.name or f.is_dir(): + continue + if f.suffix == ".exe" or f.name == "qdrant": + try: + f.unlink() + except OSError: + pass + if progress_callback: progress_callback(80, "Cleaning up...") diff --git a/src/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index f3866a6f..6c8079e1 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -7,7 +7,7 @@ from typing import Any import qdrant_client -from qdrant_client.models import FieldCondition, Filter, MatchValue +from qdrant_client.models import FieldCondition, Filter, MatchValue, Modifier from src.infrastructure.vectorstore.client import get_qdrant_client @@ -97,15 +97,17 @@ async def store_embeddings( ) ) + client = get_qdrant_client() try: - client = get_qdrant_client() - # Auto-create collection if missing (include sparse vectors so hybrid # search works even when the collection is created via this code path). existing_collections = [c.name for c in client.get_collections().collections] if collection_name not in existing_collections: from qdrant_client.models import ( Distance, + ScalarQuantization, + ScalarQuantizationConfig, + ScalarType, SparseIndexParams, SparseVectorParams, VectorParams, @@ -115,8 +117,18 @@ async def store_embeddings( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), sparse_vectors_config={ - "text-sparse": SparseVectorParams(index=SparseIndexParams()), + "text-sparse": SparseVectorParams( + index=SparseIndexParams(), + modifier=Modifier.IDF, + ), }, + quantization_config=ScalarQuantization( + scalar=ScalarQuantizationConfig( + type=ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) + ), ) # Delete old points for each source before upserting @@ -137,6 +149,8 @@ async def store_embeddings( except Exception as e: logger.error("Failed to store embeddings: %s", e) return False + finally: + client.close() async def search( diff --git a/src/presentation/static/js/internal/chat.js b/src/presentation/static/js/internal/chat.js index 4b7b23e9..145ef324 100644 --- a/src/presentation/static/js/internal/chat.js +++ b/src/presentation/static/js/internal/chat.js @@ -403,6 +403,38 @@ } if (data.indexOf("__CITATIONS__:") === 0) { + const citations = (() => { + try { + return JSON.parse(data.slice(15)); + } catch { + return []; + } + })(); + if (citations.length > 0) { + if (!bubbleInfo) { + hideTyping(); + bubbleInfo = makeMessageEl("assistant", "", true); + } + const existing = + bubbleInfo.el.querySelector(".citations-block"); + if (existing) existing.remove(); + const block = document.createElement("div"); + block.className = "citations-block"; + const title = document.createElement("div"); + title.className = "citations-title"; + title.textContent = "Sources"; + block.appendChild(title); + const list = document.createElement("div"); + list.className = "citations-list"; + citations.forEach((c) => { + const item = document.createElement("div"); + item.className = "citation-item"; + item.textContent = c; + list.appendChild(item); + }); + block.appendChild(list); + bubbleInfo.el.appendChild(block); + } continue; } diff --git a/src/presentation/static/js/internal/config.js b/src/presentation/static/js/internal/config.js index 2f262105..5a99aad7 100644 --- a/src/presentation/static/js/internal/config.js +++ b/src/presentation/static/js/internal/config.js @@ -1630,9 +1630,13 @@ function searchEmbModel() { encodeURIComponent(query) + "&limit=10", ) - .then((r) => - r.ok ? r.json() : Promise.reject(new Error(`HTTP ${r.status}`)), - ) + .then(async (r) => { + if (!r.ok) { + const body = await r.json().catch(() => ({})); + throw new Error(body?.error || `HTTP ${r.status}`); + } + return r.json(); + }) .then((data) => { const resultsEl = document.getElementById("emb-search-results"); if (!data.models || data.models.length === 0) { @@ -1656,7 +1660,7 @@ function searchEmbModel() { .catch((e) => { console.error(e); document.getElementById("emb-search-results").innerHTML = - '

Search error.

'; + `

Search error: ${escHtml(e.message)}.

`; }); } @@ -1665,9 +1669,13 @@ function searchLlmModel() { if (!query) return; fetch(`${CONFIG.llm.search}?query=${encodeURIComponent(query)}&limit=10`) - .then((r) => - r.ok ? r.json() : Promise.reject(new Error(`HTTP ${r.status}`)), - ) + .then(async (r) => { + if (!r.ok) { + const body = await r.json().catch(() => ({})); + throw new Error(body?.error || `HTTP ${r.status}`); + } + return r.json(); + }) .then((data) => { const resultsEl = document.getElementById("llm-search-results"); if (!data.models || data.models.length === 0) { @@ -1693,7 +1701,7 @@ function searchLlmModel() { .catch((e) => { console.error(e); document.getElementById("llm-search-results").innerHTML = - '

Search error.

'; + `

Search error: ${escHtml(e.message)}.

`; }); } diff --git a/src/presentation/static/js/internal/repo-browser.js b/src/presentation/static/js/internal/repo-browser.js index 7d2bf6ba..5e86ccce 100644 --- a/src/presentation/static/js/internal/repo-browser.js +++ b/src/presentation/static/js/internal/repo-browser.js @@ -288,6 +288,7 @@ class RepoBrowser { if (btn) btn.disabled = false; }, 2000); this.loadRepos(); + this.closeRepoDetail(); return; } diff --git a/src/presentation/templates/data/base.html.j2 b/src/presentation/templates/data/base.html.j2 index 9d821045..ed9615c6 100644 --- a/src/presentation/templates/data/base.html.j2 +++ b/src/presentation/templates/data/base.html.j2 @@ -16,8 +16,8 @@ {% block sidebar_subtitle %}Manage your data{% endblock %} {% block sidebar_nav %} -GitHub Sigma Spec +GitHub Local Files Vector DB Prompts diff --git a/src/shared/constants.py b/src/shared/constants.py new file mode 100644 index 00000000..80feb9d3 --- /dev/null +++ b/src/shared/constants.py @@ -0,0 +1,3 @@ +"""Shared project-wide constants.""" + +NULL_UUID = "00000000-0000-0000-0000-000000000000" diff --git a/src/shared/http.py b/src/shared/http.py new file mode 100644 index 00000000..a136ff20 --- /dev/null +++ b/src/shared/http.py @@ -0,0 +1,286 @@ +"""Shared HTTP utilities — HEAD, download with retry/backoff, and client factory. + +Consolidates the HTTP logic duplicated across: +- sigma_ref_downloader.py (``_head_content_type``, ``_download_file``) +- sigma_ref_processor.py (``_head_request``, ``_download_one``) +""" + +from __future__ import annotations + +import logging +import threading +import time +from pathlib import Path +from typing import Any + +import httpx + +from src.shared.utils.url_utils import is_private_url + +logger = logging.getLogger(__name__) + +DEFAULT_TIMEOUT = 30.0 +DEFAULT_USER_AGENT = "SigmaRAG/1.0" +MAX_RETRIES = 3 +BACKOFF_DELAYS = [1.0, 4.0, 9.0] +RETRY_STATUSES = {429, 500, 502, 503, 504} + + +# ------------------------------------------------------------------ +# Connection pool +# ------------------------------------------------------------------ + +_pool_lock = threading.Lock() +_pool: dict[str, httpx.Client] = {} + + +def _pool_key(timeout: float, follow_redirects: bool) -> str: + return f"{timeout:.1f}:{follow_redirects}" + + +def get_pooled_client( + timeout: float = DEFAULT_TIMEOUT, + headers: dict[str, str] | None = None, + follow_redirects: bool = True, +) -> httpx.Client: + """Return a pooled ``httpx.Client`` keyed by (timeout, follow_redirects). + + Connections are reused across requests to the same host, reducing TCP + handshake overhead for repeated downloads (e.g. reference documents). + """ + key = _pool_key(timeout, follow_redirects) + with _pool_lock: + client = _pool.get(key) + if client is not None: + return client + + merged: dict[str, str] = {"User-Agent": DEFAULT_USER_AGENT} + if headers: + merged.update(headers) + + transport = httpx.HTTPTransport( + retries=2, + limits=httpx.Limits( + max_connections=100, + max_keepalive_connections=20, + keepalive_expiry=30.0, + ), + ) + + new_client = httpx.Client( + timeout=httpx.Timeout(timeout), + headers=merged, + follow_redirects=follow_redirects, + transport=transport, + ) + + with _pool_lock: + _pool[key] = new_client + return new_client + + +def close_all_pooled_clients() -> None: + """Close all pooled HTTP clients. Call at shutdown.""" + with _pool_lock: + for client in _pool.values(): + try: + client.close() + except Exception: + pass + _pool.clear() + + +def create_client( + timeout: float = DEFAULT_TIMEOUT, + headers: dict[str, str] | None = None, + follow_redirects: bool = True, +) -> httpx.Client: + """Create a properly configured ``httpx.Client``. + + Parameters + ---------- + timeout : + Request timeout in seconds. + headers : + Extra HTTP headers. ``User-Agent`` is set automatically. + follow_redirects : + Whether to follow redirects automatically. + + Returns + ------- + httpx.Client + """ + merged: dict[str, Any] = {"User-Agent": DEFAULT_USER_AGENT} + if headers: + merged.update(headers) + return httpx.Client( + timeout=httpx.Timeout(timeout), + headers=merged, + follow_redirects=follow_redirects, + ) + + +def head_url( + url: str, + timeout: float = 10.0, + *, + check_ssrf: bool = True, +) -> tuple[str | None, int | None, str | None]: + """HEAD request to discover content type, size, and final URL. + + Parameters + ---------- + url : + The URL to check. + timeout : + Request timeout in seconds. + check_ssrf : + If ``True`` (default), reject private/reserved IPs before connecting. + + Returns + ------- + tuple of (content_type, content_length, final_url) + ``(None, None, None)`` on failure or SSRF rejection. + """ + if check_ssrf and is_private_url(url): + logger.warning("Skipping private URL: %s", url) + return None, None, None + + try: + client = get_pooled_client(timeout=timeout) + resp = client.head(url) + resp.raise_for_status() + ctype = resp.headers.get("content-type") + if ctype: + ctype = ctype.split(";")[0].strip() + size_str = resp.headers.get("content-length", "0") + size = int(size_str) if size_str else 0 + return ctype, size, str(resp.url) + except Exception: + return None, None, None + + +def download_file( + url: str, + output_path: str | Path, + timeout: float = DEFAULT_TIMEOUT, + max_retries: int = MAX_RETRIES, + *, + check_ssrf: bool = True, +) -> tuple[bool, int | None]: + """Download a file with retry and backoff. + + Parameters + ---------- + url : + The URL to download. + output_path : + Local filesystem path to save the file. + timeout : + HTTP request timeout in seconds. + max_retries : + Maximum number of retry attempts. 0 means no retries. + check_ssrf : + If ``True`` (default), reject private/reserved IPs before connecting. + + Returns + ------- + tuple of (success, http_status_or_None) + ``http_status`` is ``None`` for non-HTTP errors. + """ + if check_ssrf and is_private_url(url): + logger.warning("Skipping private URL (SSRF): %s", url) + return False, None + + path = Path(output_path) + client = get_pooled_client(timeout=timeout) + + for attempt in range(1, max_retries + 1): + try: + resp = client.get(url) + resp.raise_for_status() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(resp.content) + return True, None + + except OSError as exc: + logger.warning("Filesystem error for %s: %s — skipping", url, exc) + if path.exists(): + try: + path.unlink() + except OSError: + pass + return False, None + + except httpx.HTTPStatusError as exc: + status = exc.response.status_code + if status in RETRY_STATUSES and attempt < max_retries: + retry_after = _get_retry_after(exc.response) + if retry_after is not None: + wait = min(float(retry_after), 120.0) + else: + wait = _backoff_delay(attempt) + logger.warning( + "HTTP %d for %s — retrying in %.1fs (attempt %d/%d)", + status, + url, + wait, + attempt, + max_retries, + ) + time.sleep(wait) + continue + logger.warning("HTTP %d for %s — giving up", status, url) + return False, status + + except ( + httpx.TimeoutException, + httpx.NetworkError, + httpx.ConnectError, + httpx.RemoteProtocolError, + ) as exc: + if attempt < max_retries: + wait = _backoff_delay(attempt) + logger.warning( + "Network error for %s — retrying in %.1fs (attempt %d/%d): %s", + url, + wait, + attempt, + max_retries, + exc, + ) + time.sleep(wait) + continue + logger.warning( + "Network error for %s after %d attempts: %s", + url, + max_retries, + exc, + ) + return False, None + + return False, None + + +# ------------------------------------------------------------------ +# Internal helpers +# ------------------------------------------------------------------ + + +def _backoff_delay(attempt: int) -> float: + """Return the backoff delay for the given attempt number (1-indexed).""" + idx = attempt - 1 + if idx < len(BACKOFF_DELAYS): + return BACKOFF_DELAYS[idx] + return BACKOFF_DELAYS[-1] + + +def _get_retry_after(response: httpx.Response) -> int | None: + """Extract Retry-After header value as seconds.""" + raw = response.headers.get("Retry-After") + if raw is None: + return None + try: + return int(raw) + except ValueError: + return None diff --git a/src/shared/session.py b/src/shared/session.py new file mode 100644 index 00000000..7e66d5e0 --- /dev/null +++ b/src/shared/session.py @@ -0,0 +1,45 @@ +"""In-memory session store with LRU eviction.""" + +from __future__ import annotations + +import time +from typing import Any + + +class SessionStore: + """Simple in-memory session store. + + Each session holds a dict of key-value pairs. Stale sessions are + evicted LRU-style when the store exceeds *max_sessions*. + """ + + def __init__(self, max_sessions: int = 100) -> None: + self._max = max_sessions + self._data: dict[str, dict[str, Any]] = {} + self._accessed: dict[str, float] = {} + + def get(self, session_id: str, key: str, default: Any = None) -> Any: + self._touch(session_id) + return self._data.get(session_id, {}).get(key, default) + + def set(self, session_id: str, key: str, value: Any) -> None: + if session_id not in self._data and len(self._data) >= self._max: + self._evict() + self._data.setdefault(session_id, {})[key] = value + self._touch(session_id) + + def delete(self, session_id: str) -> None: + self._data.pop(session_id, None) + self._accessed.pop(session_id, None) + + def clear(self, session_id: str) -> None: + if session_id in self._data: + self._data[session_id].clear() + self._touch(session_id) + + def _touch(self, session_id: str) -> None: + self._accessed[session_id] = time.monotonic() + + def _evict(self) -> None: + oldest = min(self._accessed, key=self._accessed.get) + self.delete(oldest) diff --git a/src/shared/utils/crypto_utils.py b/src/shared/utils/crypto_utils.py new file mode 100644 index 00000000..3048bfb2 --- /dev/null +++ b/src/shared/utils/crypto_utils.py @@ -0,0 +1,28 @@ +"""Shared cryptographic utility functions.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + + +def compute_sha256_str(data: str) -> str: + """Compute SHA256 hex digest of a string.""" + return hashlib.sha256(data.encode("utf-8")).hexdigest() + + +def compute_sha256_file(path: Path) -> str: + """Compute SHA256 hex digest of a file's contents.""" + h = hashlib.sha256() + try: + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(65536), b""): + h.update(chunk) + except OSError: + return "" + return h.hexdigest() + + +def compute_sha256_bytes(data: bytes) -> str: + """Compute SHA256 hex digest of bytes.""" + return hashlib.sha256(data).hexdigest() diff --git a/src/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index c1dadf39..59ed2ae3 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -56,12 +56,26 @@ class FileType(Enum): ".odp": FileType.OFFICE_DOCUMENT, } +FILETYPE_INFO: dict[str, dict[str, str]] = { + FileType.MARKDOWN.value: {"ext": ".md", "subdir": "markdown"}, + FileType.SIGMA_RULE.value: {"ext": ".yml", "subdir": "sigma"}, + FileType.OFFICE_DOCUMENT.value: {"ext": ".docx", "subdir": "office"}, +} + + +def filetype_ext(content_type: str) -> str: + """Return the file extension for a given content type.""" + return FILETYPE_INFO.get(content_type, {}).get("ext", ".md") + + +def filetype_subdir(content_type: str) -> str: + """Return the subdirectory name for a given content type.""" + return FILETYPE_INFO.get(content_type, {}).get("subdir", "misc") + + SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, - FileType.PDF.value, - FileType.PLAIN_TEXT.value, - FileType.HTML.value, - FileType.OFFICE_DOCUMENT.value, + FileType.SIGMA_RULE.value, } PUREMAGIC_TYPE_MAP: dict[str, FileType] = { diff --git a/src/shared/utils/registry_utils.py b/src/shared/utils/registry_utils.py new file mode 100644 index 00000000..068145bc --- /dev/null +++ b/src/shared/utils/registry_utils.py @@ -0,0 +1,81 @@ +"""Shared registry entry builder for Sigma reference document tracking. + +Consolidates the entry-building logic duplicated across: +- sigma_ref_downloader.py (``_make_entry``) +- sigma_ref_processor.py (``_build_head_entry``, ``_build_download_entry``) +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from src.shared.utils import iso_now +from src.shared.utils.crypto_utils import compute_sha256_bytes + + +def build_registry_entry( + normalized_url: str, + content_type: str, + rule_id: str, + title: str, + content_sha256: str = "", + file_name: str = "", + file_size: int | None = None, + embed_status: str = "discovery", + url_hash: str = "", + original_url: str = "", +) -> dict[str, Any]: + """Build a doc_registry entry for a Sigma reference document. + + Parameters + ---------- + normalized_url : + Normalized URL of the reference document. + content_type : + MIME content type (e.g. ``"markdown"``, ``"html"``). + rule_id : + ID of the Sigma rule that references this document. + title : + Title of the referencing Sigma rule. + content_sha256 : + SHA-256 hex digest of the downloaded content. Empty string if not + yet downloaded. + file_name : + Local filename. Defaults to the last path component of the URL. + file_size : + File size in bytes. + embed_status : + Embedding status (``"discovery"``, ``"head_verified"``, ``"embedded"``). + url_hash : + Pre-computed hash of ``normalized_url``. Computed automatically if + empty. + original_url : + Original (non-normalized) URL. Falls back to ``normalized_url``. + + Returns + ------- + dict + A dictionary compatible with ``_save_registry`` and + ``batch_upsert_doc_registry``. + """ + now = iso_now() + hash_value = url_hash or compute_sha256_bytes(normalized_url.encode()) + name = file_name or Path(normalized_url).name + + return { + "url_hash": hash_value, + "org": "sigmaref", + "repo": "references", + "content_type": content_type, + "file_name": name, + "content_sha256": content_sha256, + "file_size": file_size or 0, + "original_url": original_url or normalized_url, + "normalized_url": normalized_url, + "rule_id": rule_id, + "title": title, + "timestamp": now, + "last_seen": now, + "embed_status": embed_status, + } diff --git a/src/shared/utils/sigma_utils.py b/src/shared/utils/sigma_utils.py index 991a2b12..81db258c 100644 --- a/src/shared/utils/sigma_utils.py +++ b/src/shared/utils/sigma_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations from pathlib import Path -from typing import Optional + import yaml @@ -33,7 +33,7 @@ def extract_sigma_references(file_path: Path) -> list[str]: return [] -def get_sigma_rule_id(file_path: Path) -> Optional[str]: +def get_sigma_rule_id(file_path: Path) -> str | None: """Extract the Sigma rule UUID from the ``id`` field if present.""" try: data = yaml.safe_load(file_path.read_text(encoding="utf-8")) diff --git a/src/shared/utils/url_utils.py b/src/shared/utils/url_utils.py new file mode 100644 index 00000000..11ac0e3b --- /dev/null +++ b/src/shared/utils/url_utils.py @@ -0,0 +1,55 @@ +"""Shared URL normalization utilities.""" + +from __future__ import annotations + +import ipaddress +import re +import urllib.parse +from pathlib import Path + +_GITHUB_BLOB_PATTERN = re.compile( + r"^https?://(?:www\.)?github\.com/([^/]+/[^/]+)/blob/([^#?]+)", + re.IGNORECASE, +) + + +def normalize_url(url: str) -> str: + """Normalize a reference URL for deduplication. + + Converts GitHub blob URLs to raw URLs, strips fragments, + and removes refs/heads/ or refs/tags/ prefixes. + """ + match = _GITHUB_BLOB_PATTERN.match(url) + if match: + repo = match.group(1) + path_part = match.group(2) + path_part = re.sub(r"^refs/heads/", "", path_part) + path_part = re.sub(r"^refs/tags/", "", path_part) + result = f"https://raw.githubusercontent.com/{repo}/{path_part}" + parsed = urllib.parse.urlparse(result) + if parsed.fragment: + result = urllib.parse.urlunparse(parsed._replace(fragment="")) + return result + + parsed = urllib.parse.urlparse(url) + clean = parsed._replace(fragment="") + return urllib.parse.urlunparse(clean) + + +def is_private_url(url: str) -> bool: + """Check if a URL points to a private/reserved IP to prevent SSRF.""" + parsed = urllib.parse.urlparse(url) + host = parsed.hostname or "" + if host in ("localhost", "127.0.0.1", "::1"): + return True + try: + ip = ipaddress.ip_address(host) + return ip.is_private or ip.is_loopback or ip.is_link_local + except ValueError: + return False + + +def url_ext(url: str) -> str: + """Extract the file extension from a URL path.""" + parsed = urllib.parse.urlparse(url) + return Path(parsed.path).suffix.lower() diff --git a/src/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 9253eb63..712425f3 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -4,12 +4,22 @@ from pathlib import Path from src.config.settings import get_config +from src.shared.utils.identify_file_type import filetype_subdir from src.workers.base import BaseWorker from src.workers.enums import WorkerName, WorkerStatus logger = logging.getLogger(__name__) +def _is_orphan_candidate(stem: str, known_hashes: set[str]) -> bool: + """Return True if *stem* looks like a SHA256 hex hash not in *known_hashes*.""" + if len(stem) != 64: + return False + if not all(c in "0123456789abcdef" for c in stem): + return False + return stem not in known_hashes + + class DocGCWorker(BaseWorker): """Garbage-collects document entries whose files are permanently deleted. @@ -44,17 +54,25 @@ def process(self, task: dict) -> None: ) deleted = self._gc_entries( - table="doc_registry", local_base=Path(cfg.local_documents_path), github_base=Path(cfg.paths_github_dir), grace_days=grace_days, ) skipped_found = self._cleanup_orphaned_error_entries() + + removed_head = self.db.delete_head_verified_orphans(grace_days=grace_days) + removed_unref = self.db.delete_unreferenced_entries() + removed_orphan_files = self._gc_orphaned_sigmaref_files() + removed_trash = self._cleanup_trash(Path(cfg.sigmaref_documents_path) / ".trash") + scanned = deleted logger.info( - f"[DocGCWorker] Complete: scanned={scanned}, removed={deleted}, reappears={skipped_found}" + f"[DocGCWorker] Complete: scanned={scanned}, removed={deleted}, " + f"reappears={skipped_found}, head_verified={removed_head}, " + f"unreferenced={removed_unref}, orphaned_files={removed_orphan_files}, " + f"trash_cleaned={removed_trash}" ) except Exception as e: @@ -70,7 +88,6 @@ def process(self, task: dict) -> None: def _gc_entries( self, - table: str, local_base: Path, github_base: Path, grace_days: int, @@ -81,11 +98,12 @@ def _gc_entries( try: with self.db._lock: rows = self.db._writer_conn.execute( - f"""SELECT url_hash, org, repo, file_name, embed_status, last_seen, content_type - FROM {table} - WHERE embed_status IN ('error', 'skipped') - AND last_seen < (CURRENT_TIMESTAMP - INTERVAL {grace_days} DAY) - ORDER BY url_hash""" + "SELECT url_hash, org, repo, file_name, embed_status, last_seen, content_type " + "FROM doc_registry " + "WHERE embed_status IN ('error', 'skipped') " + "AND last_seen < (CURRENT_TIMESTAMP - INTERVAL ? DAY) " + "ORDER BY url_hash", + (grace_days,), ).fetchall() for row in rows: @@ -94,6 +112,7 @@ def _gc_entries( org or "", repo or "", file_name or "", + content_type or "", local_base, github_base, ) @@ -101,15 +120,14 @@ def _gc_entries( if scanned: continue - with self.db._lock: - self.db._writer_conn.execute( - f"DELETE FROM {table} WHERE url_hash = ?", - (url_hash,), - ) + self.db._writer_conn.execute( + "DELETE FROM doc_registry WHERE url_hash = ?", + (url_hash,), + ) deleted += 1 except Exception as e: - logger.error(f"[DocGCWorker] Error processing table {table}: {e}", exc_info=True) + logger.error(f"[DocGCWorker] Error processing table doc_registry: {e}", exc_info=True) return deleted @@ -118,6 +136,7 @@ def _file_exists_locally( org: str, repo: str, file_name: str, + content_type: str, local_base: Path, github_base: Path, ) -> bool: @@ -127,19 +146,21 @@ def _file_exists_locally( candidates: list[Path] = [] - # Case 1: local files → {local_base}/{file_name} + # Case 1: local files → {local_base}/{file_name} if org == "local": candidates.append(local_base / file_name) - # Case 2: GitHub files → {github_base}/{org}/{repo}/{file_name} + # Case 2: GitHub files → {github_base}/{org}/{repo}/{file_name} if org and org not in ("local", "sigmaref") and repo: candidates.append(github_base / org / repo / file_name) - # Case 3: sigmaref files ��' local base with file_name or hash prefix + # Case 3: sigmaref files → subdir/{file_name} with flat fallback if org == "sigmaref": from src.config.settings import get_config sigmaref_base = get_config().sigmaref_documents_path + subdir = filetype_subdir(content_type or "") + candidates.append(Path(sigmaref_base) / subdir / file_name) candidates.append(Path(sigmaref_base) / file_name) if "." in file_name and not file_name.startswith("."): base, ext = file_name.rsplit(".", 1) @@ -156,6 +177,98 @@ def _file_exists_locally( return False + def _gc_orphaned_sigmaref_files(self) -> int: + """Move files in ``sigmaref/`` whose ``url_hash`` no longer exists in + ``doc_registry`` into ``.trash/``. + + Returns the number of files removed. + """ + cfg = get_config() + base = Path(cfg.sigmaref_documents_path) + if not base.exists(): + return 0 + + # Collect all known sigmaref url_hashes in a single query + known: set[str] = set() + try: + with self.db._lock: + rows = self.db._writer_conn.execute( + "SELECT url_hash FROM doc_registry WHERE org = 'sigmaref'" + ).fetchall() + known = {r[0] for r in rows} + except Exception: + logger.warning("Failed to query doc_registry for orphaned file scan") + return 0 + + trash_dir = base / ".trash" + removed = 0 + + for entry in sorted(base.iterdir()): + if entry.name.startswith("."): + continue + + if entry.is_dir(): + for f in sorted(entry.iterdir()): + if f.is_file() and not f.name.startswith("."): + if _is_orphan_candidate(f.stem, known): + self._trash_file(f, trash_dir) + removed += 1 + if entry.is_dir() and not any(entry.iterdir()): + try: + entry.rmdir() + except OSError: + pass + elif entry.is_file(): + if _is_orphan_candidate(entry.stem, known): + self._trash_file(entry, trash_dir) + removed += 1 + + return removed + + def _cleanup_trash(self, trash_dir: Path, max_age_days: int = 7) -> int: + """Permanently delete files in *trash_dir* older than *max_age_days*. + + Returns the number of files deleted. + """ + import time + + deleted = 0 + if not trash_dir.exists(): + return 0 + + cutoff = time.time() - (max_age_days * 86400) + + for entry in sorted(trash_dir.iterdir()): + if entry.is_file() and not entry.name.startswith("."): + try: + stat = entry.stat() + if stat.st_mtime < cutoff: + entry.unlink() + deleted += 1 + logger.info("Permanently deleted trashed file: %s", entry) + except OSError: + logger.warning("Failed to delete trashed file: %s", entry) + + # Remove empty subdirectories + for entry in sorted(trash_dir.iterdir()): + if entry.is_dir() and not any(entry.iterdir()): + try: + entry.rmdir() + except OSError: + pass + + return deleted + + def _trash_file(self, file_path: Path, trash_dir: Path) -> None: + """Move *file_path* to *trash_dir*, skipping if already gone.""" + try: + trash_dir.mkdir(parents=True, exist_ok=True) + dest = trash_dir / file_path.name + file_path.rename(dest) + logger.info("Moved orphaned sigmaref file to trash: %s", dest) + except OSError: + logger.warning("Failed to move orphaned file: %s", file_path) + def _cleanup_orphaned_error_entries(self) -> int: """Remove error entries whose URLs are no longer in doc_registry.""" deleted = 0 diff --git a/src/workers/processor.py b/src/workers/processor.py index fbaf0797..6f5d7668 100644 --- a/src/workers/processor.py +++ b/src/workers/processor.py @@ -2,7 +2,7 @@ import threading import uuid from concurrent.futures import Future, ThreadPoolExecutor -from typing import Dict, Type +from typing import Type from src.infrastructure.database.service import DatabaseService from src.workers.base import BaseWorker @@ -22,7 +22,7 @@ class TaskDispatcher: The dispatcher alone controls when a WAITING worker transitions to RUNNING. """ - _WORKER_TYPES: Dict[WorkerName, Type[BaseWorker]] = { + _WORKER_TYPES: dict[WorkerName, Type[BaseWorker]] = { WorkerName.SIGMAREF_DISCOVERY: SigmaRefProcessor, WorkerName.GITHUB_DISCOVERY: GenericDiscoveryWorker, WorkerName.LOCAL_DISCOVERY: GenericDiscoveryWorker, @@ -37,12 +37,12 @@ def __init__(self, poll_interval: float = 1.0, max_workers: int = 1): self._running = False self._stop_event = threading.Event() self._lock = threading.Lock() - self._pending_tasks: Dict[WorkerName, dict] = {} + self._pending_tasks: dict[WorkerName, dict] = {} self._executor: ThreadPoolExecutor | None = None self._thread: threading.Thread | None = None self._db: DatabaseService | None = None - self._workers: Dict[WorkerName, BaseWorker] = {} - self._worker_states: Dict[WorkerName, dict] = {} + self._workers: dict[WorkerName, BaseWorker] = {} + self._worker_states: dict[WorkerName, dict] = {} # ------------------------------------------------------------------ # Public API diff --git a/src/workers/sigma/discovery_base.py b/src/workers/sigma/discovery_base.py index 2de885e1..766ebca5 100644 --- a/src/workers/sigma/discovery_base.py +++ b/src/workers/sigma/discovery_base.py @@ -1,7 +1,8 @@ -import hashlib import logging from pathlib import Path +from src.shared.constants import NULL_UUID +from src.shared.utils.crypto_utils import compute_sha256_file, compute_sha256_str from src.shared.utils.identify_file_type import identify from src.shared.utils import iso_now from src.workers.base import BaseWorker @@ -14,16 +15,14 @@ class DiscoveryWorker(BaseWorker): def _compute_sha256(self, file_path: Path) -> tuple[str, int]: """Returns (content_hash, file_size) using streaming hash.""" - h = hashlib.sha256() - size = 0 - try: - with open(file_path, "rb") as f: - for chunk in iter(lambda: f.read(65536), b""): - h.update(chunk) - size += len(chunk) - return h.hexdigest(), size - except Exception: + content_hash = compute_sha256_file(file_path) + if not content_hash: return "", 0 + try: + size = file_path.stat().st_size + except OSError: + size = 0 + return content_hash, size def _identify_content_type(self, file_path: Path) -> str: try: @@ -42,9 +41,9 @@ def _make_doc_registry_entry( original_url: str, normalized_url: str, title: str, - rule_id: str = "00000000-0000-0000-0000-000000000000", + rule_id: str = NULL_UUID, ) -> dict: - url_hash = hashlib.sha256(normalized_url.encode()).hexdigest() + url_hash = compute_sha256_str(normalized_url) now = iso_now() return { "url_hash": url_hash, @@ -75,7 +74,7 @@ def _make_sigma_spec_entry( normalized_url: str, title: str, ) -> dict: - url_hash = hashlib.sha256(normalized_url.encode()).hexdigest() + url_hash = compute_sha256_str(normalized_url) now = iso_now() return { "url_hash": url_hash, diff --git a/src/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index b6b8323d..74e7548f 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -5,9 +5,12 @@ import logging from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Optional +from collections.abc import Callable +from typing import TYPE_CHECKING +from src.shared.constants import NULL_UUID from src.shared.utils.identify_file_type import SIGMA_RULE_EXTENSIONS, SUPPORTED_DOC_EXTENSION_MAP +from src.shared.utils.sigma_utils import get_sigma_rule_id from src.infrastructure.database.service import DatabaseService from src.workers.enums import WorkerName from src.workers.sigma.discovery_base import DiscoveryWorker @@ -51,14 +54,14 @@ class GenericDiscoveryWorker(DiscoveryWorker): def __init__( self, - db: Optional["DatabaseService"] = None, - dispatcher: Optional["TaskDispatcher"] = None, + db: DatabaseService | None = None, + dispatcher: TaskDispatcher | None = None, *, source_type: SourceType = SourceType.LOCAL, - base_dir: Optional[Path] = None, - github_base_dir: Optional[Path] = None, - spec_repos_dir: Optional[Path] = None, - selected_dirs: Optional[list[str]] = None, + base_dir: Path | None = None, + github_base_dir: Path | None = None, + spec_repos_dir: Path | None = None, + selected_dirs: list[str] | None = None, ) -> None: super().__init__(db or DatabaseService.get_instance(), dispatcher) self.source_type = source_type @@ -113,6 +116,32 @@ def _process_local(self, task: dict, worker_name: WorkerName) -> None: self._write_entries( entries, worker_name, len(files_to_process), processed_count, skipped_count ) + self._garbage_collect_local(collection_name, entries) + + def _garbage_collect_local(self, collection_name: str, current_entries: list[dict]) -> None: + current_hashes = {e["url_hash"] for e in current_entries} + if not current_hashes: + return + try: + existing = self.db.get_doc_registry_url_hashes_by_repo("local", collection_name) + except Exception: + return + stale_hashes = [h for h in existing if h not in current_hashes] + if stale_hashes: + self._cleanup_rule_references(stale_hashes) + try: + self.db.delete_doc_registry_by_url_hashes(stale_hashes) + logger.info( + "[GenericDiscoveryWorker] Garbage collected %d stale local entries for %s", + len(stale_hashes), + collection_name, + ) + except Exception as e: + logger.error( + "[GenericDiscoveryWorker] Failed to garbage collect local %s: %s", + collection_name, + e, + ) # ------------------------------------------------------------------ # GitHub source @@ -122,64 +151,131 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: gh_base = self.github_base_dir or Path(task.get("github_base_dir", "data/github")) gh_base = gh_base.resolve() - try: - repo_keys = self.db.get_repos_with_selected_dirs() - except Exception as e: - logger.error(f"[GenericDiscoveryWorker] Failed to query repo keys: {e}") - return - - if not repo_keys: - self._update_progress(worker_name, 100, "") - logger.info("[GenericDiscoveryWorker] No repos with selected dirs") - return + st = "github" if worker_name == WorkerName.GITHUB_DISCOVERY else "" - all_files: list[tuple[Path, Path, str, str]] = [] - for repo_key in repo_keys: + repo_key = task.get("repo_key") + if repo_key: parts = repo_key.split("/") - if len(parts) != 2: + if len(parts) == 2: + repo_items = [(parts[0], parts[1])] + else: logger.warning(f"[GenericDiscoveryWorker] Invalid repo key: {repo_key}") - continue - - org, repo = parts - repo_path = gh_base / org / repo - - if not repo_path.exists(): - logger.warning(f"[GenericDiscoveryWorker] Repo not found: {repo_path}") - continue - - selected = self.selected_dirs or self._get_selected_dirs(repo_key) - - for found_file in repo_path.rglob("*"): - if ( - not found_file.is_file() - or found_file.suffix.lower() not in SUPPORTED_EXTENSIONS - ): + return + else: + try: + repo_keys = self.db.get_repos_with_selected_dirs(source_type=st) + except Exception as e: + logger.error(f"[GenericDiscoveryWorker] Failed to query repo keys: {e}") + return + + if not repo_keys: + self._update_progress(worker_name, 100, "") + logger.info("[GenericDiscoveryWorker] No repos with selected dirs") + return + + repo_items = [] + for rk in repo_keys: + parts = rk.split("/") + if len(parts) != 2: + logger.warning(f"[GenericDiscoveryWorker] Invalid repo key: {rk}") continue + repo_items.append((parts[0], parts[1])) - if selected: - rel_to_repo = found_file.relative_to(repo_path).as_posix() - if not any( - rel_to_repo == sd.lstrip("./") - or rel_to_repo.startswith(sd.lstrip("./") + "/") - for sd in selected - if sd - ): - continue + # Always garbage collect from all repos with active selections + # plus the repo being scanned (selections may have just been cleared) + gc_repos: list[tuple[str, str]] = [] + gc_seen: set[str] = set() + try: + all_selected = self.db.get_repos_with_selected_dirs(source_type=st) + for rk in all_selected: + parts = rk.split("/") + if len(parts) == 2: + key = f"{parts[0]}/{parts[1]}" + if key not in gc_seen: + gc_repos.append((parts[0], parts[1])) + gc_seen.add(key) + except Exception: + pass + if repo_key and repo_key not in gc_seen: + parts = repo_key.split("/") + if len(parts) == 2: + gc_repos.append((parts[0], parts[1])) - all_files.append((found_file, repo_path, org, repo)) + all_files = self._collect_repo_files(repo_items, gh_base) if self.dispatcher: self._update_progress(worker_name, 1, f"{len(all_files)} files found") logger.info( - f"[GenericDiscoveryWorker] Found {len(all_files)} files across {len(repo_keys)} repos" + f"[GenericDiscoveryWorker] Found {len(all_files)} files across {len(repo_items)} repos" ) if all_files: - entries, processed_count, skipped_count = self._scan_all_github(all_files, worker_name) + entries, processed_count, skipped_count = self._scan_all( + all_files, + worker_name, + lambda fp, bp, o, r: self._prepare_entry(fp, bp, o, r, is_github=True), + ) self._write_entries( entries, worker_name, len(all_files), processed_count, skipped_count ) + self._garbage_collect_github(gc_repos, entries) + + def _garbage_collect_github( + self, + repo_items: list[tuple[str, str]], + current_entries: list[dict], + ) -> None: + current_hashes = {e["url_hash"] for e in current_entries} + if not current_hashes: + return + for org, repo in repo_items: + try: + existing = self.db.get_doc_registry_url_hashes_by_repo(org, repo) + except Exception: + continue + stale_hashes = [h for h in existing if h not in current_hashes] + if stale_hashes: + self._cleanup_rule_references(stale_hashes) + try: + self.db.delete_doc_registry_by_url_hashes(stale_hashes) + logger.info( + "[GenericDiscoveryWorker] Garbage collected %d stale entries for %s/%s", + len(stale_hashes), + org, + repo, + ) + except Exception as e: + logger.error( + "[GenericDiscoveryWorker] Failed to garbage collect for %s/%s: %s", + org, + repo, + e, + ) + + def _cleanup_rule_references(self, url_hashes: list[str]) -> None: + """Delete on-disk reference files for stale doc_registry entries.""" + for url_hash in url_hashes: + try: + rule_id = self.db.get_rule_id_by_url_hash(url_hash) + if not rule_id or rule_id == "00000000-0000-0000-0000-000000000000": + continue + paths = self.db.get_rule_reference_paths(rule_id) + for path in paths: + try: + path.unlink(missing_ok=True) + except OSError as e: + logger.warning( + "[GenericDiscoveryWorker] Failed to delete ref %s: %s", + path, + e, + ) + except Exception as e: + logger.warning( + "[GenericDiscoveryWorker] Failed to cleanup refs for %s: %s", + url_hash, + e, + ) # ------------------------------------------------------------------ # Specification repository source @@ -205,47 +301,20 @@ def _process_spec(self, task: dict, worker_name: WorkerName) -> None: logger.info("[GenericDiscoveryWorker] No spec repos found") return - all_files: list[tuple[Path, Path, str, str]] = [] + repo_items: list[tuple[str, str]] = [] for repo_info in repos: org = repo_info.get("org", "") repo_name = repo_info.get("name", "") - if not org or not repo_name: - continue + if org and repo_name: + repo_items.append((org, repo_name)) - repo_path = spec_base / org / repo_name - if not repo_path.exists(): - logger.warning(f"[GenericDiscoveryWorker] Repo path not found: {repo_path}") - continue - - repo_key = f"{org}/{repo_name}" - selected = self.selected_dirs or self._get_selected_dirs(repo_key) - - try: - for found_file in repo_path.rglob("*"): - if ( - not found_file.is_file() - or found_file.suffix.lower() not in SUPPORTED_EXTENSIONS - ): - continue - - if selected: - rel_to_repo = found_file.relative_to(repo_path).as_posix() - if not any( - rel_to_repo == sd.lstrip("./") - or rel_to_repo.startswith(sd.lstrip("./") + "/") - for sd in selected - if sd - ): - continue - - all_files.append((found_file, repo_path, org, repo_name)) - except PermissionError as e: - logger.warning( - "[GenericDiscoveryWorker] Permission denied scanning %s: %s", - repo_path, - e, - ) - continue + try: + all_files = self._collect_repo_files(repo_items, spec_base) + except PermissionError as e: + logger.warning( + "[GenericDiscoveryWorker] Permission denied scanning %s: %s", spec_base, e + ) + all_files = [] if self.dispatcher: self._update_progress(worker_name, 1, f"{len(all_files)} files found") @@ -255,15 +324,60 @@ def _process_spec(self, task: dict, worker_name: WorkerName) -> None: ) if all_files: - entries, processed_count, skipped_count = self._scan_all_spec(all_files, worker_name) - self._write_spec_entries( - entries, worker_name, len(all_files), processed_count, skipped_count + entries, processed_count, skipped_count = self._scan_all( + all_files, worker_name, self._prepare_spec_entry ) + self._write_entries( + entries, + worker_name, + len(all_files), + processed_count, + skipped_count, + batch_upsert_fn=self.db.batch_upsert_sigma_spec, + ) + self._garbage_collect_spec(repos, entries) - def _scan_all_spec( + def _garbage_collect_spec( + self, + repos: list[dict], + current_entries: list[dict], + ) -> None: + current_hashes = {e["url_hash"] for e in current_entries} + if not current_hashes: + return + for repo_info in repos: + org = repo_info.get("org", "") + repo_name = repo_info.get("name", "") + if not org or not repo_name: + continue + try: + existing = self.db.get_doc_registry_url_hashes_by_repo(org, repo_name) + except Exception: + continue + stale_hashes = [h for h in existing if h not in current_hashes] + if stale_hashes: + self._cleanup_rule_references(stale_hashes) + try: + self.db.delete_doc_registry_by_url_hashes(stale_hashes) + logger.info( + "[GenericDiscoveryWorker] Garbage collected %d stale spec entries for %s/%s", + len(stale_hashes), + org, + repo_name, + ) + except Exception as e: + logger.error( + "[GenericDiscoveryWorker] Failed to garbage collect spec %s/%s: %s", + org, + repo_name, + e, + ) + + def _scan_all( self, files: list[tuple[Path, Path, str, str]], worker_name: WorkerName, + prepare_fn: Callable, ) -> tuple[list[dict], int, int]: entries: list[dict] = [] processed_count = 0 @@ -271,7 +385,7 @@ def _scan_all_spec( for idx, (file_path, base_path, org, repo) in enumerate(files): try: - entry = self._prepare_spec_entry(file_path, base_path, org, repo) + entry = prepare_fn(file_path, base_path, org, repo) if entry is not None: entries.append(entry) processed_count += 1 @@ -327,26 +441,6 @@ def _prepare_spec_entry( logger.error(f"[GenericDiscoveryWorker] Cannot prepare spec entry for {file_path}: {e}") return None - def _write_spec_entries( - self, entries: list[dict], worker_name: WorkerName, total: int, processed: int, skipped: int - ) -> None: - if not entries: - self._update_progress(worker_name, 100, "") - return - - try: - self.db.batch_upsert_sigma_spec(entries) - except Exception as e: - logger.error( - f"[GenericDiscoveryWorker] Batch upsert sigma_spec failed: {e}", exc_info=True - ) - - if total > 0 and self.dispatcher: - pct = int((processed + skipped) / total * 100) - self._update_progress(worker_name, pct, "") - - self._update_progress(worker_name, 100, "") - # ------------------------------------------------------------------ # Shared scanning logic # ------------------------------------------------------------------ @@ -376,33 +470,6 @@ def _scan( return entries, processed_count, skipped_count - def _scan_all_github( - self, - files: list[tuple[Path, Path, str, str]], - worker_name: WorkerName, - ) -> tuple[list[dict], int, int]: - entries: list[dict] = [] - processed_count = 0 - skipped_count = 0 - - for idx, (file_path, base_path, org, repo) in enumerate(files): - try: - entry = self._prepare_entry(file_path, base_path, org, repo, is_github=True) - if entry is not None: - entries.append(entry) - processed_count += 1 - else: - skipped_count += 1 - except Exception as e: - logger.error(f"[GenericDiscoveryWorker] Unexpected error on {file_path}: {e}") - skipped_count += 1 - - if idx % 50 == 0 and self.dispatcher: - pct = int((idx + 1) / len(files) * 100) if files else 0 - self._update_progress(worker_name, pct, str(file_path)) - - return entries, processed_count, skipped_count - # ------------------------------------------------------------------ # Entry preparation # ------------------------------------------------------------------ @@ -438,10 +505,8 @@ def _prepare_entry( content_type = self._identify_content_type(file_path) title = file_path.stem - rule_id = "00000000-0000-0000-0000-000000000000" + rule_id = NULL_UUID if content_type == "sigma_rule": - from src.shared.utils.sigma_utils import get_sigma_rule_id - rid = get_sigma_rule_id(file_path) if rid: rule_id = rid @@ -462,6 +527,55 @@ def _prepare_entry( logger.error(f"[GenericDiscoveryWorker] Cannot prepare entry for {file_path}: {e}") return None + # ------------------------------------------------------------------ + # Shared repo scanning + # ------------------------------------------------------------------ + + def _collect_repo_files( + self, + repo_items: list[tuple[str, str]], + repo_base: Path, + ) -> list[tuple[Path, Path, str, str]]: + """Iterate over repos and collect supported files with selected_dirs filtering. + + Args: + repo_items: List of ``(org, repo_name)`` tuples. + repo_base: Base path under which ``{org}/{repo}`` directories live. + + Returns: + List of ``(file_path, repo_path, org, repo_name)`` tuples. + """ + all_files: list[tuple[Path, Path, str, str]] = [] + for org, repo in repo_items: + repo_path = repo_base / org / repo + if not repo_path.exists(): + logger.warning(f"[GenericDiscoveryWorker] Repo not found: {repo_path}") + continue + + repo_key = f"{org}/{repo}" + selected = self.selected_dirs or self._get_selected_dirs(repo_key) + + for found_file in repo_path.rglob("*"): + if ( + not found_file.is_file() + or found_file.suffix.lower() not in SUPPORTED_EXTENSIONS + ): + continue + + if selected: + rel_to_repo = found_file.relative_to(repo_path).as_posix() + if not any( + rel_to_repo == sd.lstrip("./") + or rel_to_repo.startswith(sd.lstrip("./") + "/") + for sd in selected + if sd + ): + continue + + all_files.append((found_file, repo_path, org, repo)) + + return all_files + # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ @@ -476,14 +590,21 @@ def _get_selected_dirs(self, repo_key: str) -> list[str]: return [] def _write_entries( - self, entries: list[dict], worker_name: WorkerName, total: int, processed: int, skipped: int + self, + entries: list[dict], + worker_name: WorkerName, + total: int, + processed: int, + skipped: int, + batch_upsert_fn: Callable | None = None, ) -> None: if not entries: self._update_progress(worker_name, 100, "") return + upsert_fn = batch_upsert_fn or self.db.batch_upsert_doc_registry try: - self.db.batch_upsert_doc_registry(entries) + upsert_fn(entries) except Exception as e: logger.error(f"[GenericDiscoveryWorker] Batch upsert failed: {e}", exc_info=True) diff --git a/tests/integration/test_chat_flow.py b/tests/integration/test_chat_flow.py index 15123e96..1032736a 100644 --- a/tests/integration/test_chat_flow.py +++ b/tests/integration/test_chat_flow.py @@ -1,144 +1,136 @@ -"""Integration tests for complete chat flows.""" - -from __future__ import annotations - -from unittest.mock import AsyncMock, MagicMock - -import pytest -from src.application.chat.service import ChatService - - -@pytest.fixture -def chat_service() -> ChatService: - """Create ChatService with mocked dependencies.""" - # Create service first - service = ChatService() - - # Mock search_engine - service.search_engine = MagicMock() - service.search_engine.search = AsyncMock(return_value=[]) - - # Mock the rag_pipeline - service.rag_pipeline = MagicMock() - service.rag_pipeline.explain_rule = AsyncMock( - return_value="This rule detects failed login attempts..." - ) - service.rag_pipeline.answer_search_query = AsyncMock(return_value="I found 2 relevant rules...") - service.rag_pipeline.analyze_coverage = AsyncMock(return_value="Coverage analysis: ...") - service.rag_pipeline.cache = MagicMock() - service.rag_pipeline.cache.invalidate = MagicMock() - - # Mock the validator - service.validator = MagicMock() - service.validator.validate.return_value = { - "id": "test_001", - "name": "Test Rule", - "description": "A test rule", - "detection": {"selection": {"EventID": 4625}}, - } - - # Mock the tool execution to avoid HTTP calls to LLM - service._execute_tool_calls = AsyncMock(return_value="I found 2 relevant rules...") - - return service - - -@pytest.mark.asyncio -async def test_upload_then_explain_flow(chat_service: ChatService) -> None: - """Test complete flow: upload YAML → explain rule.""" - # Upload a valid Sigma rule - yaml_content = b""" -id: test_001 -name: Test Rule -description: A test rule -detection: - selection: - EventID: 4625 -""" - rule_data = await chat_service.validate_and_store_yaml(yaml_content) - assert rule_data["id"] == "test_001" - - # Explain the rule - response = await chat_service._handle_explain("explain this rule") - assert "detects failed login" in response.lower() - chat_service.rag_pipeline.explain_rule.assert_called_once() - - -@pytest.mark.asyncio -async def test_search_flow(chat_service: ChatService) -> None: - """Test search mode flow with LLM response.""" - # Mock search results - chat_service.search_engine.search = AsyncMock( - return_value=[ - {"text": "Rule 1 content", "citation": "sigma:rule_1", "score": 0.95}, - {"text": "Rule 2 content", "citation": "sigma:rule_2", "score": 0.85}, - ] - ) - - # Mock RAG pipeline response - chat_service.rag_pipeline.answer_search_query = AsyncMock( - return_value="I found 2 relevant rules..." - ) - - response = await chat_service._handle_search("failed logon events") - assert "found" in response.lower() - chat_service._execute_tool_calls.assert_called_once() - - -@pytest.mark.asyncio -async def test_coverage_flow(chat_service: ChatService) -> None: - """Test coverage analysis flow.""" - # Setup uploaded rule - chat_service._uploaded_rule = { - "id": "test_001", - "name": "Test Rule", - "detection": {"selection": {"EventID": 4625}}, - } - - # Mock search results - chat_service.search_engine.search = AsyncMock( - return_value=[ - {"text": "Related rule 1", "citation": "sigma:rule_2"}, - ] - ) - - # Mock RAG pipeline response - chat_service.rag_pipeline.analyze_coverage = AsyncMock(return_value="Coverage analysis: ...") - - response = await chat_service._handle_coverage("check coverage") - assert "coverage" in response.lower() - chat_service.rag_pipeline.analyze_coverage.assert_called_once() - - -@pytest.mark.asyncio -async def test_cache_invalidation_on_upload(chat_service: ChatService) -> None: - """Test that cache is invalidated when new rule is uploaded.""" - yaml_content = b""" -id: test_002 -name: New Rule -description: Another test -detection: - selection: - EventID: 4648 -""" - await chat_service.validate_and_store_yaml(yaml_content) - chat_service.rag_pipeline.cache.invalidate.assert_called_once() - - -@pytest.mark.asyncio -async def test_fallback_when_llm_unavailable(chat_service: ChatService) -> None: - """Test fallback responses when LLM fails.""" - # Mock RAG pipeline to raise exception - chat_service.rag_pipeline.answer_search_query = AsyncMock(side_effect=Exception("LLM down")) - - # Mock search results - chat_service.search_engine.search = AsyncMock( - return_value=[ - {"text": "Some rule content", "score": 0.9}, - ] - ) - - response = await chat_service._handle_search("test query") - # Should return fallback (not exception) - assert response != "" - assert "Some rule content" in response or "rule" in response.lower() +"""Integration tests for complete chat flows.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from src.application.chat.service import ChatService +from src.core.sigma.models import SigmaRule + + +def _make_test_rule(**overrides: str) -> SigmaRule: + data = { + "id": "test_001", + "title": "Test Rule", + "description": "A test rule", + "detection": {"selection": {"EventID": 4625}}, + "logsource": {"category": "process_creation", "product": "windows"}, + **overrides, + } + return SigmaRule(**data) + + +@pytest.fixture +def chat_service() -> ChatService: + """Create ChatService with mocked dependencies.""" + service = ChatService() + + service.search_engine = MagicMock() + service.search_engine.search = AsyncMock(return_value=[]) + + service.rag_pipeline = MagicMock() + service.rag_pipeline.explain_rule = AsyncMock( + return_value="This rule detects failed login attempts..." + ) + service.rag_pipeline.answer_search_query = AsyncMock(return_value="I found 2 relevant rules...") + service.rag_pipeline.analyze_coverage = AsyncMock(return_value="Coverage analysis: ...") + service.rag_pipeline.cache = MagicMock() + service.rag_pipeline.cache.invalidate = MagicMock() + + service.validator = MagicMock() + service.validator.validate.return_value = _make_test_rule() + + service._execute_tool_calls = AsyncMock(return_value="I found 2 relevant rules...") + + return service + + +_TEST_SID = "test-session" + + +@pytest.mark.asyncio +async def test_upload_then_explain_flow(chat_service: ChatService) -> None: + """Test complete flow: upload YAML to session → explain rule.""" + yaml_content = b""" +id: test_001 +name: Test Rule +description: A test rule +detection: + selection: + EventID: 4625 +""" + rule_data = await chat_service.validate_and_store_yaml(yaml_content, _TEST_SID) + assert rule_data.id == "test_001" + + response = await chat_service._handle_explain("explain this rule", "", _TEST_SID) + assert "detects failed login" in response.lower() + chat_service.rag_pipeline.explain_rule.assert_called_once() + + +@pytest.mark.asyncio +async def test_search_flow(chat_service: ChatService) -> None: + """Test search mode flow with LLM response.""" + chat_service.search_engine.search = AsyncMock( + return_value=[ + {"text": "Rule 1 content", "citation": "sigma:rule_1", "score": 0.95}, + {"text": "Rule 2 content", "citation": "sigma:rule_2", "score": 0.85}, + ] + ) + + chat_service.rag_pipeline.answer_search_query = AsyncMock( + return_value="I found 2 relevant rules..." + ) + + response = await chat_service._handle_search("failed logon events", "") + assert "found" in response.lower() + chat_service._execute_tool_calls.assert_called_once() + + +@pytest.mark.asyncio +async def test_coverage_flow(chat_service: ChatService) -> None: + """Test coverage analysis flow.""" + chat_service._set_rule(_TEST_SID, _make_test_rule()) + + chat_service.search_engine.search = AsyncMock( + return_value=[ + {"text": "Related rule 1", "citation": "sigma:rule_2"}, + ] + ) + + chat_service.rag_pipeline.analyze_coverage = AsyncMock(return_value="Coverage analysis: ...") + + response = await chat_service._handle_coverage("check coverage", "", _TEST_SID) + assert "coverage" in response.lower() + chat_service.rag_pipeline.analyze_coverage.assert_called_once() + + +@pytest.mark.asyncio +async def test_cache_invalidation_on_upload(chat_service: ChatService) -> None: + """Test that cache is invalidated when new rule is uploaded.""" + yaml_content = b""" +id: test_002 +name: New Rule +description: Another test +detection: + selection: + EventID: 4648 +""" + await chat_service.validate_and_store_yaml(yaml_content, _TEST_SID) + chat_service.rag_pipeline.cache.invalidate.assert_called_once() + + +@pytest.mark.asyncio +async def test_fallback_when_llm_unavailable(chat_service: ChatService) -> None: + """Test fallback responses when LLM fails.""" + chat_service.rag_pipeline.answer_search_query = AsyncMock(side_effect=Exception("LLM down")) + + chat_service.search_engine.search = AsyncMock( + return_value=[ + {"text": "Some rule content", "score": 0.9}, + ] + ) + + response = await chat_service._handle_search("test query", "") + assert response != "" + assert "Some rule content" in response or "rule" in response.lower() diff --git a/tests/unit/api/v1/documents/test_spec.py b/tests/unit/api/v1/documents/test_spec.py index a993931e..16c2d849 100644 --- a/tests/unit/api/v1/documents/test_spec.py +++ b/tests/unit/api/v1/documents/test_spec.py @@ -13,11 +13,14 @@ @pytest.fixture def client(): """Create a test client with the spec router.""" + from unittest.mock import MagicMock + from fastapi import FastAPI from fastapi.testclient import TestClient app = FastAPI() app.include_router(router) + app.state.dispatcher = MagicMock() return TestClient(app, raise_server_exceptions=False) @@ -187,7 +190,7 @@ def test_select_dirs_success(self, client, tmp_path): repo_dir = tmp_path / "SigmaHQ" / "sigma-specification" repo_dir.mkdir(parents=True) - def mock_save(org, name, selected, repos_dir=None): + def mock_save(org, name, selected, repos_dir=None, source_type=""): return {"success": True} repo_path = tmp_path / "SigmaHQ" / "sigma-specification" diff --git a/tests/unit/api/v1/test_admin_download.py b/tests/unit/api/v1/test_admin_download.py index bb072db6..a22e165b 100644 --- a/tests/unit/api/v1/test_admin_download.py +++ b/tests/unit/api/v1/test_admin_download.py @@ -56,8 +56,13 @@ def test_download_with_idempotency_key_returns_same_result( assert response1.status_code == 202 assert response2.status_code == 202 - def test_download_returns_202_with_service_info(self, client: TestClient) -> None: + @patch("src.api.v1.system.orchestration.start_download", new_callable=AsyncMock) + def test_download_returns_202_with_service_info( + self, mock_start: AsyncMock, client: TestClient + ) -> None: """Given download endpoint called, when service name provided, then returns 202 with job info.""" + mock_start.return_value = {"job_id": "job-456", "status": "started"} + response = client.post( "/api/v1/orchestration/download", json={"service": "qdrant"}, diff --git a/tests/unit/api/v1/test_github.py b/tests/unit/api/v1/test_github.py index 090b3a56..6b044a95 100644 --- a/tests/unit/api/v1/test_github.py +++ b/tests/unit/api/v1/test_github.py @@ -15,8 +15,11 @@ @pytest.fixture def client(): """Create a test client with mocked dependencies.""" + from unittest.mock import MagicMock + app = FastAPI() app.include_router(router) + app.state.dispatcher = MagicMock() return TestClient(app, raise_server_exceptions=False) diff --git a/tests/unit/api/v1/test_orchestration_download.py b/tests/unit/api/v1/test_orchestration_download.py index 4b60eb2c..4293ebc0 100644 --- a/tests/unit/api/v1/test_orchestration_download.py +++ b/tests/unit/api/v1/test_orchestration_download.py @@ -81,8 +81,13 @@ def test_download_with_idempotency_key_returns_same_result( assert response1.status_code == 202 assert response2.status_code == 202 - def test_download_returns_202_with_service_info(self, client: TestClient) -> None: + @patch("src.api.v1.system.orchestration.start_download", new_callable=AsyncMock) + def test_download_returns_202_with_service_info( + self, mock_start: AsyncMock, client: TestClient + ) -> None: """Given download endpoint called, when service name provided, then returns 202 with job info.""" + mock_start.return_value = {"job_id": "job-456", "status": "started"} + response = client.post( "/api/v1/orchestration/download", json={"service": "qdrant"}, diff --git a/tests/unit/api/v1/test_search.py b/tests/unit/api/v1/test_search.py index 84fac99d..e1e9dcab 100644 --- a/tests/unit/api/v1/test_search.py +++ b/tests/unit/api/v1/test_search.py @@ -38,13 +38,37 @@ def test_search_returns_empty_data(self, mock_search: AsyncMock, client: TestCli @patch("src.core.search.engine.SearchEngine.search", new_callable=AsyncMock) def test_search_returns_results(self, mock_search: AsyncMock, client: TestClient) -> None: """Test search returns results.""" - mock_search.return_value = [{"id": "rule-001"}, {"id": "rule-002"}] + mock_search.return_value = [ + { + "text": "rule text", + "score": 0.9, + "metadata": { + "collection": "sigma_rules", + "source_file": "/rules/test.yml", + "rule_id": "abc-123", + "title": "Test Rule", + }, + }, + { + "text": "doc text", + "score": 0.8, + "metadata": { + "collection": "sigma_docs", + "source_file": "/docs/test.md", + "original_url": "https://example.com", + }, + }, + ] response = client.post("/api/v1/search", json={"query": "test", "limit": 10}) assert response.status_code == 200 data = response.json() - assert data["data"] == [{"id": "rule-001"}, {"id": "rule-002"}] + assert len(data["data"]) == 2 + assert data["data"][0]["collection"] == "sigma_rules" + assert data["data"][0]["rule_id"] == "abc-123" + assert data["data"][1]["collection"] == "sigma_docs" + assert data["data"][1]["source_file"] == "/docs/test.md" assert data["meta"]["total"] == 2 def test_search_empty_query(self, client: TestClient) -> None: diff --git a/tests/unit/application/documents/test_documents.py b/tests/unit/application/documents/test_documents.py index b6e4ab76..f8cb8ba5 100644 --- a/tests/unit/application/documents/test_documents.py +++ b/tests/unit/application/documents/test_documents.py @@ -4,9 +4,7 @@ from pathlib import Path -from src.core.sigma.models import SigmaRule from src.application.documents.parser import parse_sigma_rule, scan_directory -from src.application.documents.validator import validate_sigma_rule FIXTURES_DIR = (Path(__file__).parent / ".." / ".." / ".." / "fixtures").resolve() @@ -51,82 +49,3 @@ def test_scan_empty_directory(self) -> None: files = scan_directory("/nonexistent/directory") assert files == [] - - -class TestSigmaRuleValidator: - """Test Sigma rule validator.""" - - def test_validate_valid_rule(self) -> None: - """Test validating a valid rule.""" - rule = SigmaRule( - id="test-001", - title="Test Rule", - detection={"selection": {"EventID": 4688}}, - condition="selection", - level="high", - ) - - result = validate_sigma_rule(rule) - - assert result.valid is True - assert result.rule is not None - assert len(result.errors) == 0 - - def test_validate_missing_title(self) -> None: - """Test validating rule with missing title.""" - rule = SigmaRule( - id="test-001", - title="", - detection={"selection": {"EventID": 4688}}, - condition="selection", - ) - - result = validate_sigma_rule(rule) - - assert result.valid is False - assert any(e.field == "title" for e in result.errors) - - def test_validate_missing_condition(self) -> None: - """Test validating rule with missing condition.""" - rule = SigmaRule( - id="test-001", - title="Test", - detection={"selection": {"EventID": 4688}}, - condition="", - ) - - result = validate_sigma_rule(rule) - - assert result.valid is False - assert any(e.field == "condition" for e in result.errors) - - def test_validate_invalid_level(self) -> None: - """Test validating rule with invalid level.""" - rule = SigmaRule( - id="test-001", - title="Test", - detection={"selection": {"EventID": 4688}}, - condition="selection", - level="invalid_level", - ) - - result = validate_sigma_rule(rule) - - assert result.valid is False - assert any(e.field == "level" for e in result.errors) - - def test_validate_valid_levels(self) -> None: - """Test validating rule with valid levels.""" - valid_levels = ["informational", "low", "medium", "high", "critical"] - - for level in valid_levels: - rule = SigmaRule( - id="test-001", - title="Test", - detection={"selection": {"EventID": 4688}}, - condition="selection", - level=level, - ) - - result = validate_sigma_rule(rule) - assert result.valid is True, f"Level {level} should be valid" diff --git a/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 93db3e66..f355f59a 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -3,25 +3,20 @@ from __future__ import annotations from pathlib import Path -from unittest.mock import ANY, MagicMock, patch - -import httpx +from unittest.mock import MagicMock, patch from src.shared.utils import iso_now +from src.shared.utils.crypto_utils import compute_sha256_file as _sha256_file +from src.shared.utils.crypto_utils import compute_sha256_str as _sha256 +from src.shared.utils.identify_file_type import filetype_subdir +from src.shared.utils.url_utils import is_private_url as _is_private_url, normalize_url from src.application.documents.sigma_ref_downloader import ( - _backoff_delay, _detect_url_type, - _download_file, - _get_retry_after, - _is_private_url, _load_registry, _registry_lock, _save_registry, - _sha256, - _sha256_file, download_references, - normalize_url, ) @@ -152,108 +147,6 @@ def test_text_plain_with_md_extension(self) -> None: ) -class TestDownloadFile: - def test_successful_download(self, tmp_path: Path) -> None: - url = "https://raw.githubusercontent.com/user/repo/main/test.md" - output = tmp_path / "test.md" - - with patch("httpx.Client") as mock_client: - mock_response = mock_client.return_value.__enter__.return_value.get.return_value - mock_response.raise_for_status.return_value = None - mock_response.content = b"# Hello" - mock_response.status_code = 200 - - success, status_code = _download_file(url, output, timeout=30) - assert success is True - assert status_code is None - assert output.read_text() == "# Hello" - - def test_retry_then_success(self, tmp_path: Path) -> None: - url = "https://example.com/doc.md" - output = tmp_path / "doc.md" - attempts: list[int] = [] - - class FakeResponse: - status_code = 200 - headers = {} - content = b"success" - - def raise_for_status(self) -> None: - pass - - class FakeErrorResponse: - status_code = 500 - headers = {} - content = b"" - - def mock_get(client_self, url: str, **kwargs: object) -> FakeResponse: - attempts.append(len(attempts) + 1) - if len(attempts) < 3: - resp = FakeErrorResponse() - raise httpx.HTTPStatusError("Server error", request=ANY, response=resp) # type: ignore[arg-type] - return FakeResponse() - - with patch.object(httpx.Client, "get", mock_get), patch("time.sleep"): - success, status_code = _download_file(url, output, timeout=30) - assert success is True - assert status_code is None - assert output.read_text() == "success" - assert len(attempts) == 3 - - def test_all_retries_fail(self, tmp_path: Path) -> None: - url = "https://example.com/fail.md" - output = tmp_path / "fail.md" - - class ErrorResp: - status_code = 500 - headers = {} - - def failing_get(self, url, **kwargs): - raise httpx.HTTPStatusError("500", request=ANY, response=ErrorResp()) - - with patch.object(httpx.Client, "get", failing_get): - with patch("time.sleep"): - success, status_code = _download_file(url, output, timeout=30) - assert success is False - assert status_code == 500 - assert not output.exists() - - def test_zero_retries(self, tmp_path: Path) -> None: - url = "https://example.com/zero.md" - output = tmp_path / "zero.md" - - call_count = 0 - - def failing_get(self, url, **kwargs): - nonlocal call_count - call_count += 1 - raise httpx.ConnectError("connection failed") - - with patch.object(httpx.Client, "get", failing_get): - success, status_code = _download_file(url, output, max_retries=0) - assert success is False - assert status_code is None - assert call_count == 0 - - def test_network_error_retries(self, tmp_path: Path) -> None: - url = "https://example.com/net.md" - output = tmp_path / "net.md" - - call_count = 0 - - def failing_get(self, url, **kwargs): - nonlocal call_count - call_count += 1 - raise httpx.ConnectError("connection refused") - - with patch.object(httpx.Client, "get", failing_get): - with patch("time.sleep"): - success, status_code = _download_file(url, output, timeout=30) - assert success is False - assert status_code is None - assert call_count == 3 - - class TestRegistry: def test_load_empty(self, tmp_path: Path) -> None: db = _make_db() @@ -361,7 +254,7 @@ def test_skip_non_http_refs(self, tmp_path: Path) -> None: db = _make_db() with patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", return_value=(True, None), ): result = download_references( @@ -425,7 +318,7 @@ def write_file(url: str, output_path: Path, **kwargs: object) -> tuple[bool, int db = _make_db() with patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", write_file, ): first = download_references(str(rules_dir), str(output_dir), db, selected_dirs=[""]) @@ -464,7 +357,7 @@ def check_normalized_url( db = _make_db() with patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", check_normalized_url, ): result = download_references(str(rules_dir), str(output_dir), db, selected_dirs=[""]) @@ -506,7 +399,7 @@ def test_output_dir_created(self, tmp_path: Path) -> None: db = _make_db() with patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", return_value=(True, None), ): result = download_references(str(rules_dir), str(output_dir), db, selected_dirs=[""]) @@ -534,7 +427,7 @@ def test_partial_download_cleanup(self, tmp_path: Path) -> None: db = _make_db() with patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", return_value=(False, None), ): result = download_references(str(rules_dir), str(output_dir), db, selected_dirs=[""]) @@ -572,7 +465,9 @@ def capture_url(url: str, output_path: Path, **kwargs: object) -> tuple[bool, in return True, None db = _make_db() - with patch("src.application.documents.sigma_ref_downloader._download_file", capture_url): + with patch( + "src.application.documents.sigma_ref_downloader.http_download_file", capture_url + ): result = download_references(str(rules_dir), str(output_dir), db, selected_dirs=[""]) assert result["downloaded"] == 2 assert len(urls_downloaded) == 2 @@ -614,7 +509,7 @@ def write_on_download( db = _make_db() with ( patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", write_on_download, ), patch("time.sleep"), @@ -670,7 +565,7 @@ def test_content_changed_re_downloads(self, tmp_path: Path) -> None: download_calls: list[str] = [] def tracking_download( - url: str, output_path: Path, timeout: int = 30 + url: str, output_path: Path, **kwargs: object ) -> tuple[bool, int | None]: download_calls.append(url) output_path.parent.mkdir(parents=True, exist_ok=True) @@ -679,7 +574,7 @@ def tracking_download( with ( patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", tracking_download, ), patch("time.sleep"), @@ -740,7 +635,7 @@ def tracking_download( with ( patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", tracking_download, ), ): @@ -798,7 +693,7 @@ def tracking_download( with ( patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", tracking_download, ), ): @@ -836,7 +731,7 @@ def test_download_uppercase_scheme(self, tmp_path: Path) -> None: db = _make_db() with ( patch( - "src.application.documents.sigma_ref_downloader._download_file", + "src.application.documents.sigma_ref_downloader.http_download_file", return_value=(True, None), ), patch("time.sleep"), @@ -866,38 +761,6 @@ def test_github_url_with_query_and_fragment(self) -> None: assert "#" not in result -class TestBackoffDelay: - def test_first_attempt(self) -> None: - assert _backoff_delay(1) == 1 - - def test_second_attempt(self) -> None: - assert _backoff_delay(2) == 4 - - def test_third_attempt(self) -> None: - assert _backoff_delay(3) == 9 - - def test_beyond_list_length(self) -> None: - assert _backoff_delay(4) == 9 - assert _backoff_delay(10) == 9 - - -class TestGetRetryAfter: - def test_no_header(self) -> None: - response = MagicMock() - response.headers = {} - assert _get_retry_after(response) is None - - def test_valid_int(self) -> None: - response = MagicMock() - response.headers = {"Retry-After": "30"} - assert _get_retry_after(response) == 30 - - def test_invalid_value(self) -> None: - response = MagicMock() - response.headers = {"Retry-After": "not-a-number"} - assert _get_retry_after(response) is None - - class TestSha256File: def test_file_sha(self, tmp_path: Path) -> None: f = tmp_path / "test.txt" @@ -912,3 +775,272 @@ class TestConcurrencyLock: def test_lock_has_acquire_release(self) -> None: assert hasattr(_registry_lock, "acquire") assert hasattr(_registry_lock, "release") + + +class TestDownloadSigmaReferencesContract: + """Contract tests: ``download_sigma_references`` with mode="scan" and + mode="registry" must use the same ``{url_hash}{ext}`` naming convention + for downloaded files. + """ + + def test_scan_mode_naming_convention(self, tmp_path: Path) -> None: + """Scan mode downloads to ``{url_hash}.md``.""" + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + output_dir = tmp_path / "output" + ref_url = "https://example.com/test-doc.md" + + rule = rules_dir / "test_rule.yml" + rule.write_text(f""" +title: Test Rule +id: contract-001 +logsource: + category: process_creation + product: windows +detection: + selection: + EventID: 4688 + condition: selection +references: + - {ref_url} +""") + + expected_hash = _sha256(normalize_url(ref_url)) + expected_filename = f"{expected_hash}.md" + + def _fake_download( + url: str, output_path: Path, **kwargs: object + ) -> tuple[bool, int | None]: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("# test") + return True, None + + db = _make_db() + with patch( + "src.application.documents.sigma_ref_downloader.http_download_file", + _fake_download, + ): + from src.application.documents.sigma_ref_downloader import ( + download_sigma_references, + ) + + result = download_sigma_references( + db=db, + output_dir=str(output_dir), + mode="scan", + rules_dir=str(rules_dir), + ) + + assert result["downloaded"] == 1 + assert (output_dir / filetype_subdir("markdown") / expected_filename).exists() + + def test_registry_mode_runs_without_error(self, tmp_path: Path) -> None: + """Registry mode runs without error (naming convention is inherited + from the shared download helpers).""" + output_dir = tmp_path / "output" + output_dir.mkdir(parents=True) + + ref_url = "https://example.com/registry-doc.md" + norm_url = normalize_url(ref_url) + expected_hash = _sha256(norm_url) + + rule_file = tmp_path / "rule.yml" + rule_file.write_text(f"""title: Registry Rule +id: contract-002 +logsource: + category: process_creation + product: windows +detection: + selection: + EventID: 4688 + condition: selection +references: + - {ref_url} +""") + + db = MagicMock() + db.get_pending_registry_all.return_value = [ + { + "org": "local", + "repo": "references", + "file_name": "rule.yml", + "rule_id": "contract-002", + "original_url": "", + "title": "Registry Rule", + "content_type": "sigma_rule", + "embed_status": "discovery", + "url_hash": "dummy", + "normalized_url": "", + } + ] + db.get_entry.return_value = None + db.batch_upsert_doc_registry = MagicMock() + + def _fake_download( + url: str, output_path: Path, **kwargs: object + ) -> tuple[bool, int | None]: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("# downloaded") + return True, None + + with ( + patch( + "src.application.documents.sigma_ref_downloader.get_config", + ) as mock_cfg, + patch( + "src.application.documents.sigma_ref_downloader.http_head_url", + return_value=("markdown", 1024, ref_url), + ), + patch( + "src.application.documents.sigma_ref_downloader.http_download_file", + _fake_download, + ), + ): + from src.application.documents.sigma_ref_downloader import ( + download_sigma_references, + ) + + cfg = MagicMock() + cfg.local_documents_path = str(tmp_path) + cfg.sigmaref_documents_path = str(tmp_path) + cfg.paths_github_dir = str(tmp_path) + mock_cfg.return_value = cfg + + result = download_sigma_references( + db=db, + output_dir=str(output_dir), + mode="registry", + ) + + assert result["total_rules"] == 1 + assert result["total_refs"] == 1 + # File should be written with {url_hash}{ext} naming + expected_filename = f"{expected_hash}.md" + assert (output_dir / filetype_subdir("markdown") / expected_filename).exists() + + +class TestRuleReferences: + """Tests for rule↔reference tracking.""" + + def test_scan_mode_populates_rule_references(self, tmp_path: Path) -> None: + """Scan mode inserts rule_references rows via batch_upsert_rule_references.""" + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + output_dir = tmp_path / "output" + ref_url = "https://example.com/tracked-doc.md" + + rule = rules_dir / "test_rule.yml" + rule.write_text(f""" +title: Tracked Rule +id: tracked-001 +logsource: + category: process_creation + product: windows +detection: + selection: + EventID: 4688 + condition: selection +references: + - {ref_url} +""") + + db = _make_db() + db.batch_upsert_rule_references = MagicMock() + + def _fake_download( + url: str, output_path: Path, **kwargs: object + ) -> tuple[bool, int | None]: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("# tracked") + return True, None + + with patch( + "src.application.documents.sigma_ref_downloader.http_download_file", + _fake_download, + ): + from src.application.documents.sigma_ref_downloader import ( + download_sigma_references, + ) + + result = download_sigma_references( + db=db, + output_dir=str(output_dir), + mode="scan", + rules_dir=str(rules_dir), + ) + + assert result["downloaded"] == 1 + db.batch_upsert_rule_references.assert_called_once() + args = db.batch_upsert_rule_references.call_args[0][0] + assert len(args) == 1 + assert args[0]["rule_id"] == "tracked-001" + assert args[0]["ref_url"] == ref_url + + def test_scan_mode_dedup_urls_in_rule_refs(self, tmp_path: Path) -> None: + """Two rules referencing the same URL get two rule_references rows.""" + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + output_dir = tmp_path / "output" + ref_url = "https://example.com/shared-doc.md" + + rule_a = rules_dir / "rule_a.yml" + rule_a.write_text(f""" +title: Rule A +id: rule-a +logsource: + category: process_creation + product: windows +detection: + selection: + EventID: 4625 + condition: selection +references: + - {ref_url} +""") + + rule_b = rules_dir / "rule_b.yml" + rule_b.write_text(f""" +title: Rule B +id: rule-b +logsource: + category: process_creation + product: windows +detection: + selection: + EventID: 4688 + condition: selection +references: + - {ref_url} +""") + + db = _make_db() + db.batch_upsert_rule_references = MagicMock() + + def _fake_download( + url: str, output_path: Path, **kwargs: object + ) -> tuple[bool, int | None]: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("# shared") + return True, None + + with patch( + "src.application.documents.sigma_ref_downloader.http_download_file", + _fake_download, + ): + from src.application.documents.sigma_ref_downloader import ( + download_sigma_references, + ) + + download_sigma_references( + db=db, + output_dir=str(output_dir), + mode="scan", + rules_dir=str(rules_dir), + ) + + db.batch_upsert_rule_references.assert_called_once() + rows = db.batch_upsert_rule_references.call_args[0][0] + assert len(rows) == 2 + rule_ids = {r["rule_id"] for r in rows} + assert rule_ids == {"rule-a", "rule-b"} + assert rows[0]["url_hash"] == rows[1]["url_hash"] diff --git a/tests/unit/application/documents/test_sigma_ref_paths.py b/tests/unit/application/documents/test_sigma_ref_paths.py new file mode 100644 index 00000000..1401747b --- /dev/null +++ b/tests/unit/application/documents/test_sigma_ref_paths.py @@ -0,0 +1,110 @@ +"""Tests for sigma_ref_paths module.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +from src.application.documents.sigma_ref_paths import ( + resolve_rule_path, + sigmaref_resolve_path, + sigmaref_write_path, + subdir_for, +) + + +class TestSubdirFor: + def test_markdown(self) -> None: + assert subdir_for("markdown") == "markdown" + + def test_pdf(self) -> None: + assert subdir_for("pdf") == "pdf" + + def test_html(self) -> None: + assert subdir_for("html") == "html" + + def test_none(self) -> None: + assert subdir_for(None) == "misc" + + def test_empty_string(self) -> None: + assert subdir_for("") == "misc" + + def test_plain_text(self) -> None: + assert subdir_for("plain_text") == "plain_text" + + def test_office_document(self) -> None: + assert subdir_for("office_document") == "office" + + +class TestSigmarefWritePath: + def test_basic_path(self, tmp_path: Path) -> None: + result = sigmaref_write_path(tmp_path, "markdown", "doc.md") + assert result == tmp_path / "markdown" / "doc.md" + + def test_pdf_path(self, tmp_path: Path) -> None: + result = sigmaref_write_path(tmp_path, "pdf", "doc.pdf") + assert result == tmp_path / "pdf" / "doc.pdf" + + def test_none_content_type(self, tmp_path: Path) -> None: + result = sigmaref_write_path(tmp_path, None, "doc.txt") + assert result == tmp_path / "misc" / "doc.txt" + + def test_html_path(self, tmp_path: Path) -> None: + result = sigmaref_write_path(tmp_path, "html", "doc.html") + assert result == tmp_path / "html" / "doc.html" + + +class TestSigmarefResolvePath: + def test_existing_candidate_returns_candidate(self, tmp_path: Path) -> None: + existing = tmp_path / "markdown" / "doc.md" + existing.parent.mkdir(parents=True, exist_ok=True) + existing.touch() + + result = sigmaref_resolve_path(tmp_path, "markdown", "doc.md") + assert result == existing + + def test_non_existing_returns_default(self, tmp_path: Path) -> None: + result = sigmaref_resolve_path(tmp_path, "markdown", "doc.md") + assert result == tmp_path / "doc.md" + + def test_non_existing_pdf(self, tmp_path: Path) -> None: + result = sigmaref_resolve_path(tmp_path, "pdf", "doc.pdf") + assert result == tmp_path / "doc.pdf" + + +class TestResolveRulePath: + def test_local_org(self) -> None: + mock_cfg = MagicMock() + mock_cfg.local_documents_path = "/tmp/local" + result = resolve_rule_path({"org": "local", "file_name": "rule.yml"}, mock_cfg) + assert result == Path("/tmp/local/rule.yml").resolve() + + def test_sigmaref_org_existing(self, tmp_path: Path) -> None: + candidate = tmp_path / "markdown" / "rule.yml" + candidate.parent.mkdir(parents=True, exist_ok=True) + candidate.touch() + + mock_cfg = MagicMock() + mock_cfg.sigmaref_documents_path = str(tmp_path) + entry = {"org": "sigmaref", "file_name": "rule.yml", "content_type": "markdown"} + result = resolve_rule_path(entry, mock_cfg) + assert result == candidate + + def test_sigmaref_org_not_existing(self, tmp_path: Path) -> None: + mock_cfg = MagicMock() + mock_cfg.sigmaref_documents_path = str(tmp_path) + entry = {"org": "sigmaref", "file_name": "rule.yml", "content_type": "markdown"} + result = resolve_rule_path(entry, mock_cfg) + assert result == tmp_path / "rule.yml" + + def test_github_org_repo(self) -> None: + mock_cfg = MagicMock() + mock_cfg.sigmaref_documents_path = "/tmp/sigmaref" + entry = {"org": "sigma-project", "repo": "rules", "file_name": "rule.yml"} + result = resolve_rule_path(entry, mock_cfg) + assert result == Path("/tmp/sigmaref/sigma-project/rules/rule.yml") + + def test_empty_org_returns_none(self) -> None: + mock_cfg = MagicMock() + result = resolve_rule_path({"org": "", "file_name": "rule.yml"}, mock_cfg) + assert result is None diff --git a/tests/unit/application/documents/test_sigma_ref_registry.py b/tests/unit/application/documents/test_sigma_ref_registry.py new file mode 100644 index 00000000..01f5a956 --- /dev/null +++ b/tests/unit/application/documents/test_sigma_ref_registry.py @@ -0,0 +1,256 @@ +"""Tests for sigma_ref_registry module.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from src.application.documents.sigma_ref_registry import ( + load_error_registry, + load_registry, + maybe_record_error, + save_registry, +) + + +class TestLoadRegistry: + def test_loads_from_db(self) -> None: + mock_db = MagicMock() + mock_db.get_entries_by_org.return_value = [ + { + "url_hash": "abc123", + "original_url": "https://example.com/doc.md", + "normalized_url": "https://example.com/doc.md", + "content_type": "markdown", + "rule_id": "test-rule", + "title": "Test Rule", + "timestamp": "2024-01-01T00:00:00Z", + "content_sha256": "sha256hash", + "embed_status": "discovery", + "last_seen": "2024-01-01T00:00:00Z", + "file_name": "doc.md", + } + ] + + result = load_registry("/tmp/output", mock_db) + + assert "abc123" in result + assert result["abc123"]["content_type"] == "markdown" + assert result["abc123"]["embed_status"] == "discovery" + mock_db.get_entries_by_org.assert_called_once_with("sigmaref", limit=0) + + def test_empty_registry(self) -> None: + mock_db = MagicMock() + mock_db.get_entries_by_org.return_value = [] + + result = load_registry("/tmp/output", mock_db) + + assert result == {} + + +class TestSaveRegistry: + def test_saves_to_db(self) -> None: + mock_db = MagicMock() + registry = { + "abc123": { + "original_url": "https://example.com/doc.md", + "normalized_url": "https://example.com/doc.md", + "content_type": "markdown", + "rule_id": "test-rule", + "title": "Test Rule", + "timestamp": "2024-01-01T00:00:00Z", + "content_sha256": "sha256hash", + "embed_status": "discovery", + "last_seen": "2024-01-01T00:00:00Z", + "file_name": "doc.md", + "file_size": 1024, + "org": "sigmaref", + "repo": "references", + } + } + + save_registry(registry, "/tmp/output", mock_db) + + mock_db.batch_upsert_doc_registry.assert_called_once() + rows = mock_db.batch_upsert_doc_registry.call_args[0][0] + assert len(rows) == 1 + assert rows[0]["url_hash"] == "abc123" + assert rows[0]["content_type"] == "markdown" + assert rows[0]["embed_status"] == "discovery" + assert rows[0]["file_size"] == 1024 + + def test_empty_registry_noop(self) -> None: + mock_db = MagicMock() + save_registry({}, "/tmp/output", mock_db) + mock_db.batch_upsert_doc_registry.assert_not_called() + + def test_non_dict_entry_skipped(self) -> None: + mock_db = MagicMock() + registry = {"abc123": "not a dict"} + save_registry(registry, "/tmp/output", mock_db) + mock_db.batch_upsert_doc_registry.assert_not_called() + + +class TestLoadErrorRegistry: + def test_loads_errors(self) -> None: + mock_db = MagicMock() + mock_db.get_doc_errors.return_value = [ + {"url_hash": "abc123", "error_code": 404}, + {"url_hash": "def456", "error_code": 403}, + ] + + result = load_error_registry(mock_db) + + assert result == {"abc123", "def456"} + + def test_empty_errors(self) -> None: + mock_db = MagicMock() + mock_db.get_doc_errors.return_value = [] + + result = load_error_registry(mock_db) + + assert result == set() + + def test_db_error_returns_empty_set(self) -> None: + mock_db = MagicMock() + mock_db.get_doc_errors.side_effect = Exception("DB error") + + result = load_error_registry(mock_db) + + assert result == set() + + +class TestMaybeRecordError: + def test_records_404(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=404, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_called_once() + args = mock_db.upsert_doc_error.call_args[0][0] + assert args["error_code"] == 404 + assert args["error_message"] == "HTTP 404" + + def test_records_403(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=403, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_called_once() + + def test_records_301_redirect(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/old", + normalized_url="https://example.com/new", + status_code=301, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_called_once() + + def test_no_record_for_200(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=200, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_not_called() + + def test_no_record_for_500_retryable(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=500, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_not_called() + + def test_no_record_for_none_status(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=None, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_not_called() + + def test_no_record_for_502_retryable(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=502, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_not_called() + + def test_records_503_non_retryable(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=503, + rule_id="test-rule", + rule_title="Test Rule", + ) + # 503 is in RETRY_STATUSES, so should NOT be recorded + mock_db.upsert_doc_error.assert_not_called() + + def test_records_500_non_retryable(self) -> None: + mock_db = MagicMock() + # 500 is in RETRY_STATUSES, so should NOT be recorded + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=500, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_not_called() + + def test_records_501_not_in_retry_statuses(self) -> None: + mock_db = MagicMock() + maybe_record_error( + mock_db, + url_hash="abc123", + original_url="https://example.com/doc.md", + normalized_url="https://example.com/doc.md", + status_code=501, + rule_id="test-rule", + rule_title="Test Rule", + ) + mock_db.upsert_doc_error.assert_called_once() diff --git a/tests/unit/application/documents/test_sigma_ref_url.py b/tests/unit/application/documents/test_sigma_ref_url.py new file mode 100644 index 00000000..f577344f --- /dev/null +++ b/tests/unit/application/documents/test_sigma_ref_url.py @@ -0,0 +1,109 @@ +"""Tests for sigma_ref_url module.""" + +from __future__ import annotations + +from src.application.documents.sigma_ref_url import ( + detect_url_type, + is_reference_url, + resolve_ext, +) + + +class TestDetectUrlType: + def test_markdown_extension(self) -> None: + assert detect_url_type("https://example.com/doc.md") == "markdown" + + def test_markdown_extension_alt(self) -> None: + assert detect_url_type("https://example.com/doc.markdown") == "markdown" + + def test_pdf_extension(self) -> None: + assert detect_url_type("https://example.com/doc.pdf") == "pdf" + + def test_html_extension(self) -> None: + assert detect_url_type("https://example.com/doc.html") == "html" + + def test_unsupported_extension(self) -> None: + assert detect_url_type("https://example.com/doc.exe") is None + + def test_no_extension_with_markdown_content_type(self) -> None: + assert ( + detect_url_type("https://example.com/doc", content_type="text/markdown") == "markdown" + ) + + def test_no_extension_with_pdf_content_type(self) -> None: + assert detect_url_type("https://example.com/doc", content_type="application/pdf") == "pdf" + + def test_no_extension_with_html_content_type(self) -> None: + assert detect_url_type("https://example.com/doc", content_type="text/html") == "html" + + def test_no_extension_with_text_content_type(self) -> None: + assert detect_url_type("https://example.com/doc", content_type="text/plain") == "markdown" + + def test_no_extension_no_content_type(self) -> None: + assert detect_url_type("https://example.com/doc") is None + + def test_content_type_takes_precedence(self) -> None: + assert ( + detect_url_type("https://example.com/doc.md", content_type="application/pdf") == "pdf" + ) + + def test_empty_url(self) -> None: + assert detect_url_type("") is None + + def test_plain_text_file_type(self) -> None: + assert detect_url_type("https://example.com/doc.txt") == "plain_text" + + def test_office_document(self) -> None: + assert detect_url_type("https://example.com/doc.docx") == "office_document" + + +class TestResolveExt: + def test_with_content_type(self) -> None: + assert resolve_ext("https://example.com/doc", "markdown") == ".md" + + def test_with_pdf_content_type(self) -> None: + assert resolve_ext("https://example.com/doc", "pdf") == ".pdf" + + def test_fallback_to_url_extension(self) -> None: + assert resolve_ext("https://example.com/doc.md", None) == ".md" + + def test_fallback_to_md(self) -> None: + assert resolve_ext("https://example.com/doc", None) == ".md" + + def test_content_type_overrides_url(self) -> None: + assert resolve_ext("https://example.com/doc.md", "pdf") == ".pdf" + + +class TestIsReferenceUrl: + def test_github_raw(self) -> None: + assert is_reference_url("https://github.com/user/repo/raw/main/doc.md") is True + + def test_github_blob(self) -> None: + assert is_reference_url("https://github.com/user/repo/blob/main/doc.md") is True + + def test_gitlab_raw(self) -> None: + assert is_reference_url("https://gitlab.com/user/repo/raw/main/doc.md") is True + + def test_bitbucket_raw(self) -> None: + assert is_reference_url("https://bitbucket.org/user/repo/raw/main/doc.md") is True + + def test_rawcdn(self) -> None: + assert is_reference_url("https://rawcdn.com/user/doc.md") is True + + def test_pastebin(self) -> None: + assert is_reference_url("https://pastebin.com/abc123") is True + + def test_hastebin(self) -> None: + assert is_reference_url("https://hastebin.com/abc123") is True + + def test_dpaste(self) -> None: + assert is_reference_url("https://dpaste.org/abc123") is True + + def test_normal_url(self) -> None: + assert is_reference_url("https://example.com/doc.md") is False + + def test_empty_url(self) -> None: + assert is_reference_url("") is False + + def test_local_file(self) -> None: + assert is_reference_url("file:///tmp/doc.md") is False diff --git a/tests/unit/application/feedback/test_feedback.py b/tests/unit/application/feedback/test_feedback.py index c62c36f9..72dfdfd8 100644 --- a/tests/unit/application/feedback/test_feedback.py +++ b/tests/unit/application/feedback/test_feedback.py @@ -105,8 +105,8 @@ def test_feedback_stats_calculation(self) -> None: class TestResolveDbPath: def test_absolute_path(self) -> None: - repo = FeedbackRepository(db_path="C:\\absolute\\path.db") - assert str(repo._db_path) == "C:\\absolute\\path.db" + repo = FeedbackRepository(db_path="/tmp/absolute/path.db") + assert str(repo._db_path) == "/tmp/absolute/path.db" def test_relative_path(self) -> None: repo = FeedbackRepository(db_path="relative.db") diff --git a/tests/unit/application/services/test_chat_service_cache.py b/tests/unit/application/services/test_chat_service_cache.py index 1df81ae6..372971de 100644 --- a/tests/unit/application/services/test_chat_service_cache.py +++ b/tests/unit/application/services/test_chat_service_cache.py @@ -26,7 +26,7 @@ async def test_erase_slot_cache_called_after_translate(self): with patch.object( svc, "_execute_tool_calls", new_callable=AsyncMock, return_value="Answer" ): - result = await svc._handle_search(message) + result = await svc._handle_search(message, "") svc.rag_pipeline.llm_client.erase_slot_cache.assert_not_awaited() assert result == "Answer" @@ -37,7 +37,7 @@ async def test_erase_slot_cache_not_called_without_yaml(self): with patch.object( svc, "_execute_tool_calls", new_callable=AsyncMock, return_value="Answer" ): - await svc._handle_search("simple question") + await svc._handle_search("simple question", "") svc.rag_pipeline.llm_client.erase_slot_cache.assert_not_awaited() @@ -50,7 +50,7 @@ async def test_erase_slot_cache_exception_does_not_break_flow(self): with patch.object( svc, "_execute_tool_calls", new_callable=AsyncMock, return_value="Answer" ): - result = await svc._handle_search("detection:\n condition: selection") + result = await svc._handle_search("detection:\n condition: selection", "") assert result == "Answer" @@ -83,7 +83,7 @@ async def fake_stream(*args: object, **kwargs: object): ), ): tokens = [] - async for t in svc._handle_search_stream(message): + async for t in svc._handle_search_stream(message, "", ""): tokens.append(t) svc.rag_pipeline.llm_client.erase_slot_cache.assert_awaited_once() @@ -96,11 +96,11 @@ async def test_erase_slot_cache_called_in_endpoint(self): patch( "src.api.v1.sigma.translate.translate_detection", new_callable=AsyncMock ) as mock_translate, - patch("src.api.v1.sigma.translate.RAGPipeline") as mock_rag_cls, + patch("src.api.v1.sigma.translate.get_rag_pipeline") as mock_get_rag, ): mock_rag = MagicMock() mock_rag.llm_client.erase_slot_cache = AsyncMock() - mock_rag_cls.return_value = mock_rag + mock_get_rag.return_value = mock_rag mock_translate.return_value = "Translated text" from fastapi.testclient import TestClient diff --git a/tests/unit/application/services/test_rag_pipeline.py b/tests/unit/application/services/test_rag_pipeline.py index 84616243..c4eb92b7 100644 --- a/tests/unit/application/services/test_rag_pipeline.py +++ b/tests/unit/application/services/test_rag_pipeline.py @@ -1,84 +1,92 @@ -"""Tests for RAG pipeline.""" - -from __future__ import annotations - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest -from src.application.chat.rag import RAGPipeline - - -@pytest.fixture -def rag_pipeline() -> RAGPipeline: - """Create RAG pipeline with fully mocked dependencies.""" - with ( - patch("src.application.chat.rag.SearchEngine") as mock_search, - patch("src.application.chat.rag.LlamaClient") as mock_llm, - ): - pipeline = RAGPipeline() - pipeline.search_engine = mock_search.return_value - pipeline.llm_client = mock_llm.return_value - # Mock _resolve_prompt to avoid DB access - pipeline._resolve_prompt = MagicMock(return_value="You are a helpful assistant.") - # Mock the jinja2 environment to avoid template loading - pipeline.env = MagicMock() - yield pipeline - - -@pytest.mark.asyncio -async def test_explain_rule(rag_pipeline: RAGPipeline) -> None: - """Test rule explanation with LLM.""" - rule_data = { - "id": "test_001", - "name": "Test Rule", - "description": "A test rule", - "detection": {"selection": {"EventID": 4625}}, - } - - rag_pipeline.llm_client.generate = AsyncMock(return_value="Rule explanation here") - - result = await rag_pipeline.explain_rule(rule_data) - assert result == "Rule explanation here" - rag_pipeline.llm_client.generate.assert_called_once() - - -@pytest.mark.asyncio -async def test_explain_rule_fallback(rag_pipeline: RAGPipeline) -> None: - """Test fallback when LLM fails.""" - rule_data = { - "id": "test_001", - "name": "Test Rule", - "description": "A test rule", - } - - rag_pipeline.llm_client.generate = AsyncMock(side_effect=Exception("LLM down")) - - result = await rag_pipeline.explain_rule(rule_data) - assert "**Rule:** Test Rule" in result - - -@pytest.mark.asyncio -async def test_answer_search_query(rag_pipeline: RAGPipeline) -> None: - """Test search query answering.""" - results = [ - {"text": "Rule 1 content", "citation": "sigma:rule1"}, - {"text": "Rule 2 content", "citation": "sigma:rule2"}, - ] - - rag_pipeline.llm_client.generate = AsyncMock(return_value="Search answer here") - - result = await rag_pipeline.answer_search_query("What about EventID?", results) - assert result == "Search answer here" - rag_pipeline.llm_client.generate.assert_called_once() - - -@pytest.mark.asyncio -async def test_analyze_coverage(rag_pipeline: RAGPipeline) -> None: - """Test coverage analysis.""" - rule_data = {"id": "test_001", "name": "Test Rule"} - results = [{"text": "Related rule", "citation": "sigma:related"}] - - rag_pipeline.llm_client.generate = AsyncMock(return_value="Coverage analysis here") - - result = await rag_pipeline.analyze_coverage(rule_data, results) - assert result == "Coverage analysis here" +"""Tests for RAG pipeline.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from src.application.chat.rag import RAGPipeline +from src.core.sigma.models import SigmaRule + + +@pytest.fixture +def rag_pipeline() -> RAGPipeline: + """Create RAG pipeline with fully mocked dependencies.""" + with ( + patch("src.application.chat.rag.SearchEngine") as mock_search, + patch("src.application.chat.rag.LlamaClient") as mock_llm, + ): + pipeline = RAGPipeline() + pipeline.search_engine = mock_search.return_value + pipeline.llm_client = mock_llm.return_value + pipeline._resolve_prompt = MagicMock(return_value="You are a helpful assistant.") + pipeline.env = MagicMock() + yield pipeline + + +@pytest.fixture +def sigma_rule() -> SigmaRule: + """Create a test SigmaRule.""" + return SigmaRule( + id="test_001", + title="Test Rule", + description="A test rule", + detection={"selection": {"EventID": 4625}}, + ) + + +@pytest.mark.asyncio +async def test_explain_rule(rag_pipeline: RAGPipeline, sigma_rule: SigmaRule) -> None: + """Test rule explanation with LLM.""" + rag_pipeline.llm_client.generate = AsyncMock(return_value="Rule explanation here") + + result = await rag_pipeline.explain_rule(sigma_rule) + assert result == "Rule explanation here" + rag_pipeline.llm_client.generate.assert_called_once() + + +@pytest.mark.asyncio +async def test_explain_rule_fallback(rag_pipeline: RAGPipeline) -> None: + """Test fallback when LLM fails.""" + rule = SigmaRule( + id="test_001", + title="Test Rule", + description="A test rule", + detection={"selection": {"EventID": 4625}}, + ) + + rag_pipeline.llm_client.generate = AsyncMock(side_effect=Exception("LLM down")) + + result = await rag_pipeline.explain_rule(rule) + assert "**Rule:** Test Rule" in result + + +@pytest.mark.asyncio +async def test_answer_search_query(rag_pipeline: RAGPipeline) -> None: + """Test search query answering.""" + results = [ + {"text": "Rule 1 content", "citation": "sigma:rule1"}, + {"text": "Rule 2 content", "citation": "sigma:rule2"}, + ] + + rag_pipeline.llm_client.generate = AsyncMock(return_value="Search answer here") + + result = await rag_pipeline.answer_search_query("What about EventID?", results) + assert result == "Search answer here" + rag_pipeline.llm_client.generate.assert_called_once() + + +@pytest.mark.asyncio +async def test_analyze_coverage(rag_pipeline: RAGPipeline) -> None: + """Test coverage analysis.""" + rule = SigmaRule( + id="test_001", + title="Test Rule", + detection={"selection": {"EventID": 4625}}, + ) + results = [{"text": "Related rule", "citation": "sigma:related"}] + + rag_pipeline.llm_client.generate = AsyncMock(return_value="Coverage analysis here") + + result = await rag_pipeline.analyze_coverage(rule, results) + assert result == "Coverage analysis here" diff --git a/tests/unit/application/services/test_sigma_validator.py b/tests/unit/application/services/test_sigma_validator.py index c4f99f6c..85dca067 100644 --- a/tests/unit/application/services/test_sigma_validator.py +++ b/tests/unit/application/services/test_sigma_validator.py @@ -19,9 +19,10 @@ def test_validate_valid_yaml() -> None: EventID: 4625 """ result = validator.validate(content) - assert result["id"] == "test_rule_001" - assert result["name"] == "Test Rule" - assert "detection" in result + assert result.id == "test_rule_001" + assert result.name == "Test Rule" + assert result.title == "Test Rule" + assert "selection" in result.detection def test_validate_missing_fields() -> None: @@ -211,11 +212,11 @@ def test_validate_condition_ok() -> None: condition: selection """ result = validator.validate(content) - assert result["condition"] == "selection" + assert result.condition == "selection" def test_validate_condition_non_string() -> None: - """Test non-string condition is silently accepted.""" + """Test non-string condition is cast to string.""" validator = SigmaValidator() content = b""" id: rule_001 @@ -227,4 +228,4 @@ def test_validate_condition_non_string() -> None: condition: true """ result = validator.validate(content) - assert result["condition"] is True + assert result.condition == "True" diff --git a/tests/unit/application/services/test_sigma_validator_advanced.py b/tests/unit/application/services/test_sigma_validator_advanced.py index 7dd3cde1..2602d0c5 100644 --- a/tests/unit/application/services/test_sigma_validator_advanced.py +++ b/tests/unit/application/services/test_sigma_validator_advanced.py @@ -55,7 +55,7 @@ def test_validate_deprecated_fields() -> None: """ # Should not raise - deprecated fields are warnings only result = validator.validate(content) - assert result["level"] == "high" + assert result.level == "high" def test_validate_condition_syntax() -> None: @@ -72,7 +72,7 @@ def test_validate_condition_syntax() -> None: """ # Should not raise - valid condition result = validator.validate(content) - assert result.get("condition") == "selection" + assert result.condition == "selection" def test_validate_condition_invalid_ref() -> None: @@ -89,7 +89,7 @@ def test_validate_condition_invalid_ref() -> None: """ # Should log warning but not raise result = validator.validate(content) - assert result.get("condition") == "nonexistent" + assert result.condition == "nonexistent" def test_validate_empty_description() -> None: diff --git a/tests/unit/application/services/test_sigma_validator_regression.py b/tests/unit/application/services/test_sigma_validator_regression.py new file mode 100644 index 00000000..35c1d912 --- /dev/null +++ b/tests/unit/application/services/test_sigma_validator_regression.py @@ -0,0 +1,226 @@ +"""Regression tests for SigmaValidator contract. + +These tests capture the contract of SigmaValidator.validate() +after the refactoring from dict[str, Any] to SigmaRule. +""" + +from __future__ import annotations + +import pytest +import yaml + +from src.application.sigma.validator import SigmaValidator +from src.core.sigma.models import SigmaRule +from src.shared.exceptions import ValidationError + +_VALID_YAML = b""" +id: regr_test_001 +title: Regression Test Rule +description: A rule for regression testing +detection: + selection: + EventID: 4625 +condition: selection +""" + +_VALID_YAML_WITH_NAME = b""" +id: regr_test_002 +name: Name-Based Rule +description: A rule using name instead of title +detection: + selection: + EventID: 4625 +condition: selection +""" + + +class TestSigmaValidatorRegression: + """Regression tests: updated contract of SigmaValidator.""" + + def setup_method(self) -> None: + self.validator = SigmaValidator() + + # --- Return type contract --- + + def test_validate_returns_sigma_rule(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert isinstance(result, SigmaRule), "Must return a SigmaRule" + + def test_validate_rule_has_expected_attributes(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert result.id == "regr_test_001" + assert result.title == "Regression Test Rule" + assert result.name == "Regression Test Rule" + assert result.description == "A rule for regression testing" + assert "selection" in result.detection + + def test_validate_rule_supports_name_property(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert result.name == result.title + + def test_validate_rule_yaml_dumpable(self) -> None: + result = self.validator.validate(_VALID_YAML) + dumped = yaml.dump(result.to_dict(), default_flow_style=False) + assert isinstance(dumped, str) + assert "regr_test_001" in dumped + + def test_validate_accepts_name_instead_of_title(self) -> None: + result = self.validator.validate(_VALID_YAML_WITH_NAME) + assert isinstance(result, SigmaRule) + assert result.title == "Name-Based Rule" + assert result.name == "Name-Based Rule" + + # --- Error contract --- + + def test_validate_raises_shared_validation_error(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b"") + assert exc_info.type is ValidationError + assert exc_info.value.details["field"] == "file" + + def test_validate_error_has_field_in_details(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b"") + assert "field" in exc_info.value.details + + def test_validate_error_has_message(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b"") + assert exc_info.value.message + + # --- Deprecated fields (warnings, not errors) --- + + def test_deprecated_fields_still_in_rule(self) -> None: + yaml_with_deprecated = b""" +id: regr_test_003 +title: Rule with deprecated +description: Contains level and falsepositives +level: high +falsepositives: + - FP1 +detection: + selection: + EventID: 4625 +""" + result = self.validator.validate(yaml_with_deprecated) + assert result.level == "high" + assert result.falsepositives == ["FP1"] + + # --- Level validation --- + + def test_invalid_level_raises_error(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b""" +id: test-001 +title: Test +description: Invalid level +level: invalid_level +detection: + selection: + EventID: 1 +""") + assert exc_info.value.details["field"] == "level" + + def test_valid_levels_accepted(self) -> None: + for level in ("informational", "low", "medium", "high", "critical"): + result = self.validator.validate( + f"""id: test-001 +title: Test {level} +description: Testing level {level} +level: {level} +detection: + selection: + EventID: 1 +condition: selection +""".encode() + ) + assert result.level == level + + # --- Status validation --- + + def test_invalid_status_raises_error(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b""" +id: test-001 +title: Test +description: Invalid status +status: unknown_status +detection: + selection: + EventID: 1 +""") + assert exc_info.value.details["field"] == "status" + + # --- Edge cases --- + + def test_validate_large_file_raises_error(self) -> None: + large_content = b"id: test\nname: test\ndescription: x\n" + b"detection:\n x: 1\n" + large_content = large_content + b" " * (1024 * 1024) + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(large_content) + assert exc_info.value.details["field"] == "file" + + def test_validate_non_dict_yaml(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b"[1, 2, 3]") + assert exc_info.value.details["field"] == "yaml_structure" + + def test_validate_invalid_yaml_syntax(self) -> None: + with pytest.raises(ValidationError) as exc_info: + self.validator.validate(b"{invalid: yaml: too many: colons: }") + assert exc_info.value.details["field"] == "yaml_syntax" + + +class TestChatServiceConsumerRegression: + """Regression: ChatService usage patterns on SigmaRule. + + After the refactoring, ChatService._uploaded_rule is SigmaRule | None + and is accessed via: + - .name (was .get("name", "")) + - .id (was .get("id", "N/A")) + - .description (was .get("description")) + """ + + def _make_rule(self) -> SigmaRule: + return SigmaRule( + id="regr_test_001", + title="Regression Test Rule", + description="A rule for regression testing", + detection={"selection": {"EventID": 4625}}, + condition="selection", + ) + + def test_uploaded_rule_name(self) -> None: + rule = self._make_rule() + assert rule.name == "Regression Test Rule" + + def test_uploaded_rule_id(self) -> None: + rule = self._make_rule() + assert rule.id == "regr_test_001" + + def test_uploaded_rule_description(self) -> None: + rule = self._make_rule() + assert rule.description == "A rule for regression testing" + + def test_rag_pipeline_format_rule_yaml(self) -> None: + rule = self._make_rule() + dumped = yaml.dump(rule.to_dict(), default_flow_style=False, allow_unicode=True) + assert isinstance(dumped, str) + assert "regr_test_001" in dumped + + def test_rag_pipeline_fallback_explanation(self) -> None: + rule = self._make_rule() + parts = [ + f"**Rule:** {rule.name}", + f"**ID:** {rule.id}", + ] + if rule.description: + parts.append(f"**Description:** {rule.description}") + text = "\n".join(parts) + assert "Regression Test Rule" in text + assert "regr_test_001" in text + + def test_search_engine_search_uses_name(self) -> None: + rule = self._make_rule() + query = rule.name + assert query == "Regression Test Rule" diff --git a/tests/unit/application/system/test_datadir.py b/tests/unit/application/system/test_datadir.py index 360aa354..21b32e32 100644 --- a/tests/unit/application/system/test_datadir.py +++ b/tests/unit/application/system/test_datadir.py @@ -8,7 +8,7 @@ class TestOfficialDirs: def test_returns_expected_count(self, tmp_path: Path) -> None: dirs = src.application.system.datadir.DataDirManager.official_dirs(str(tmp_path)) - assert len(dirs) == 14 # 12 leaves + documents + models parents + assert len(dirs) == 15 # 13 leaves + documents + models parents def test_base_is_resolved(self, tmp_path: Path) -> None: dirs = src.application.system.datadir.DataDirManager.official_dirs(str(tmp_path)) diff --git a/tests/unit/application/system/test_duckdb.py b/tests/unit/application/system/test_duckdb.py index ab166107..afcd62a9 100644 --- a/tests/unit/application/system/test_duckdb.py +++ b/tests/unit/application/system/test_duckdb.py @@ -72,7 +72,7 @@ def test_dirty_tables_state(self, manager: DuckDbManager, tmp_path: Path) -> Non assert status["state"] == "dirty_tables" assert status["needs_fix"] is True assert status["needs_clean"] is False - assert len(status["tables_missing"]) == 10 + assert len(status["tables_missing"]) == 11 assert status["tables_excess"] == [] def test_excess_tables_state(self, manager: DuckDbManager, tmp_path: Path) -> None: diff --git a/tests/unit/back/database/test_service.py b/tests/unit/back/database/test_service.py index 035f807f..a6748483 100644 --- a/tests/unit/back/database/test_service.py +++ b/tests/unit/back/database/test_service.py @@ -1131,6 +1131,21 @@ def test_with_dirs(self, db: DatabaseService) -> None: assert "a/r1" in repos assert "b/r2" in repos + def test_filter_by_source_type(self, db: DatabaseService) -> None: + db.set_selected_dirs("a/r1", ["rules"], source_type="github") + db.set_selected_dirs("b/r2", ["docs"], source_type="spec") + github_repos = db.get_repos_with_selected_dirs(source_type="github") + spec_repos = db.get_repos_with_selected_dirs(source_type="spec") + assert github_repos == ["a/r1"] + assert spec_repos == ["b/r2"] + + def test_no_filter_returns_all(self, db: DatabaseService) -> None: + db.set_selected_dirs("a/r1", ["rules"], source_type="github") + db.set_selected_dirs("b/r2", ["docs"], source_type="spec") + repos = db.get_repos_with_selected_dirs() + assert "a/r1" in repos + assert "b/r2" in repos + class TestGetInstance: def test_get_instance_without_init_raises(self) -> None: diff --git a/tests/unit/back/github/test_git_advanced.py b/tests/unit/back/github/test_git_advanced.py index 0c2360fa..4dcf64a1 100644 --- a/tests/unit/back/github/test_git_advanced.py +++ b/tests/unit/back/github/test_git_advanced.py @@ -116,7 +116,9 @@ def test_saves_successfully(self) -> None: ): result = save_selected_dirs("org", "repo", ["src/", "docs/"]) assert result["success"] is True - mock_db.set_selected_dirs.assert_called_once_with("org/repo", ["src/", "docs/"]) + mock_db.set_selected_dirs.assert_called_once_with( + "org/repo", ["src/", "docs/"], source_type="" + ) def test_handles_exception(self) -> None: mock_db = MagicMock() diff --git a/tests/unit/core/search/eval/test_exact_search.py b/tests/unit/core/search/eval/test_exact_search.py new file mode 100644 index 00000000..5e62c439 --- /dev/null +++ b/tests/unit/core/search/eval/test_exact_search.py @@ -0,0 +1,131 @@ +"""Tests for exact vs approximate search comparison.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from src.core.search.eval.exact_search import ( + SearchComparisonResult, + _extract_doc_ids, + compare_collection_baseline, + compare_exact_vs_approximate, + run_approximate_search, + run_exact_search, +) + + +class TestExtractDocIds: + def test_flat_payload(self) -> None: + p1 = MagicMock() + p1.payload = {"doc_id": "id1"} + p2 = MagicMock() + p2.payload = {"doc_id": "id2"} + assert _extract_doc_ids([p1, p2]) == ["id1", "id2"] + + def test_file_path_fallback(self) -> None: + p = MagicMock() + p.payload = {"file_path": "/rules/test.yml"} + assert _extract_doc_ids([p]) == ["/rules/test.yml"] + + def test_empty_payload(self) -> None: + p = MagicMock() + p.payload = {} + assert _extract_doc_ids([p]) == [] + + def test_none_payload(self) -> None: + p = MagicMock() + p.payload = None + assert _extract_doc_ids([p]) == [] + + +class TestRunExactSearch: + def test_returns_doc_ids(self) -> None: + mock_point = MagicMock() + mock_point.payload = {"doc_id": "exact-1"} + mock_client = MagicMock() + mock_client.search.return_value = [mock_point] + result = run_exact_search(mock_client, "sigma_rules", [0.1] * 384, limit=10) + assert result == ["exact-1"] + mock_client.search.assert_called_once() + + def test_empty_results(self) -> None: + mock_client = MagicMock() + mock_client.search.return_value = [] + result = run_exact_search(mock_client, "sigma_rules", [0.1] * 384) + assert result == [] + + +class TestRunApproximateSearch: + def test_returns_doc_ids(self) -> None: + mock_point = MagicMock() + mock_point.payload = {"doc_id": "approx-1"} + mock_points = MagicMock() + mock_points.points = [mock_point] + mock_client = MagicMock() + mock_client.query_points.return_value = mock_points + result = run_approximate_search(mock_client, "sigma_rules", [0.1] * 384, limit=10) + assert result == ["approx-1"] + + +class TestCompareExactVsApproximate: + def test_perfect_recall(self) -> None: + mock_point = MagicMock() + mock_point.payload = {"doc_id": "d1"} + mock_client = MagicMock() + mock_client.search.return_value = [mock_point] + mock_points = MagicMock() + mock_points.points = [mock_point] + mock_client.query_points.return_value = mock_points + result = compare_exact_vs_approximate(mock_client, "sigma_rules", [0.1] * 384) + assert isinstance(result, SearchComparisonResult) + assert result.recall_at_10 == 1.0 + + def test_zero_recall(self) -> None: + exact_point = MagicMock() + exact_point.payload = {"doc_id": "d1"} + approx_point = MagicMock() + approx_point.payload = {"doc_id": "d99"} + mock_client = MagicMock() + mock_client.search.return_value = [exact_point] + mock_points = MagicMock() + mock_points.points = [approx_point] + mock_client.query_points.return_value = mock_points + result = compare_exact_vs_approximate(mock_client, "sigma_rules", [0.1] * 384) + assert result.recall_at_10 == 0.0 + + +class TestCompareCollectionBaseline: + def test_single_query(self) -> None: + mock_point = MagicMock() + mock_point.payload = {"doc_id": "d1"} + mock_client = MagicMock() + mock_client.search.return_value = [mock_point] + mock_points = MagicMock() + mock_points.points = [mock_point] + mock_client.query_points.return_value = mock_points + result = compare_collection_baseline(mock_client, "sigma_rules", [[0.1] * 384]) + assert result["mean_recall_at_10"] == 1.0 + assert result["total_queries"] == 1 + + def test_multiple_queries(self) -> None: + mock_point = MagicMock() + mock_point.payload = {"doc_id": "d1"} + mock_client = MagicMock() + mock_client.search.return_value = [mock_point] + mock_points = MagicMock() + mock_points.points = [mock_point] + mock_client.query_points.return_value = mock_points + result = compare_collection_baseline(mock_client, "sigma_rules", [[0.1] * 384, [0.2] * 384]) + assert result["total_queries"] == 2 + + def test_empty_embeddings(self) -> None: + mock_client = MagicMock() + result = compare_collection_baseline(mock_client, "sigma_rules", []) + assert result["total_queries"] == 0 + assert result["mean_recall_at_10"] == 0.0 + + def test_search_failure_handled(self) -> None: + mock_client = MagicMock() + mock_client.search.side_effect = RuntimeError("connection lost") + result = compare_collection_baseline(mock_client, "sigma_rules", [[0.1] * 384]) + assert result["total_queries"] == 0 diff --git a/tests/unit/core/search/eval/test_golden_set.py b/tests/unit/core/search/eval/test_golden_set.py new file mode 100644 index 00000000..7cf4dac6 --- /dev/null +++ b/tests/unit/core/search/eval/test_golden_set.py @@ -0,0 +1,124 @@ +"""Tests for the golden set module.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from src.core.search.eval.golden_set import ( + GoldenQuery, + GoldenSet, + create_sample_golden_set, + load_golden_set, + save_golden_set, +) + + +class TestGoldenQuery: + def test_to_dict_roundtrip(self) -> None: + q = GoldenQuery( + id="q001", + query="test query", + collection="sigma_rules", + relevant_doc_ids=["doc1", "doc2"], + k=10, + ) + d = q.to_dict() + q2 = GoldenQuery.from_dict(d) + assert q2.id == q.id + assert q2.query == q.query + assert q2.collection == q.collection + assert q2.relevant_doc_ids == q.relevant_doc_ids + assert q2.k == q.k + + def test_default_k(self) -> None: + q = GoldenQuery( + id="q001", + query="test", + collection="sigma_rules", + relevant_doc_ids=["doc1"], + ) + assert q.k == 10 + + def test_from_dict_missing_k_defaults(self) -> None: + d = { + "id": "q001", + "query": "test", + "collection": "sigma_rules", + "relevant_doc_ids": ["doc1"], + } + q = GoldenQuery.from_dict(d) + assert q.k == 10 + + +class TestGoldenSet: + def test_add_query(self) -> None: + gs = GoldenSet(description="test") + q = gs.add("query", "sigma_rules", ["doc1", "doc2"], k=5) + assert len(gs) == 1 + assert q.id == "q0000" + assert q.query == "query" + + def test_add_query_custom_id(self) -> None: + gs = GoldenSet() + q = gs.add("query", "sigma_rules", ["doc1"], query_id="custom-id") + assert q.id == "custom-id" + + def test_to_dict_roundtrip(self) -> None: + gs = GoldenSet(description="test set") + gs.add("q1", "sigma_rules", ["d1"]) + gs.add("q2", "sigma_docs", ["d2", "d3"], k=20) + d = gs.to_dict() + gs2 = GoldenSet.from_dict(d) + assert gs2.description == "test set" + assert len(gs2) == 2 + assert gs2.queries[0].query == "q1" + assert gs2.queries[1].k == 20 + + def test_empty_set(self) -> None: + gs = GoldenSet() + assert len(gs) == 0 + d = gs.to_dict() + assert d["metadata"]["num_queries"] == 0 + assert d["queries"] == [] + + +class TestSaveLoadGoldenSet: + def test_save_and_load(self, tmp_path: Path) -> None: + gs = GoldenSet(description="roundtrip test") + gs.add("test query", "sigma_rules", ["doc1", "doc2"]) + path = save_golden_set(gs, tmp_path / "golden.json") + assert path.exists() + + loaded = load_golden_set(path) + assert loaded.description == "roundtrip test" + assert len(loaded) == 1 + assert loaded.queries[0].query == "test query" + + def test_load_missing_file_raises(self) -> None: + with pytest.raises(FileNotFoundError): + load_golden_set("/nonexistent/path/golden.json") + + def test_load_invalid_structure_raises(self, tmp_path: Path) -> None: + bad_file = tmp_path / "bad.json" + bad_file.write_text(json.dumps({"not_queries": []})) + with pytest.raises(ValueError, match="missing 'queries' key"): + load_golden_set(bad_file) + + def test_creates_parent_dirs(self, tmp_path: Path) -> None: + gs = GoldenSet(description="nested") + gs.add("q", "sigma_rules", ["d1"]) + path = save_golden_set(gs, tmp_path / "a" / "b" / "golden.json") + assert path.exists() + + +class TestCreateSampleGoldenSet: + def test_sample_has_queries(self) -> None: + gs = create_sample_golden_set() + assert len(gs) == 4 + assert gs.queries[0].collection == "sigma_rules" + assert gs.queries[1].collection == "sigma_rules" + assert gs.queries[2].collection == "sigma_docs" + assert gs.queries[3].collection == "sigma_spec" diff --git a/tests/unit/core/search/eval/test_metrics.py b/tests/unit/core/search/eval/test_metrics.py new file mode 100644 index 00000000..cc72d4d8 --- /dev/null +++ b/tests/unit/core/search/eval/test_metrics.py @@ -0,0 +1,141 @@ +"""Tests for retrieval quality metrics.""" + +from __future__ import annotations + +import pytest + +from src.core.search.eval.metrics import ( + aggregate_metrics, + average_precision, + context_precision, + context_recall, + evaluate_query, + mean_reciprocal_rank, + precision_at_k, + recall_at_k, +) + + +class TestRecallAtK: + def test_perfect_recall(self) -> None: + assert recall_at_k(["d1", "d2", "d3"], ["d1", "d2"]) == 1.0 + + def test_partial_recall(self) -> None: + assert recall_at_k(["d1", "d4"], ["d1", "d2"]) == 0.5 + + def test_no_hits(self) -> None: + assert recall_at_k(["d4", "d5"], ["d1", "d2"]) == 0.0 + + def test_k_cutoff(self) -> None: + # d1, d2 in top-2, 3 relevant docs total -> 2/3 + assert recall_at_k(["d1", "d2", "d3", "d4"], ["d1", "d2", "d3"], k=2) == pytest.approx( + 2.0 / 3 + ) + + def test_empty_relevant(self) -> None: + assert recall_at_k(["d1"], []) == 0.0 + + def test_empty_retrieved(self) -> None: + assert recall_at_k([], ["d1"]) == 0.0 + + +class TestPrecisionAtK: + def test_perfect_precision(self) -> None: + assert precision_at_k(["d1", "d2"], ["d1", "d2"]) == 1.0 + + def test_partial_precision(self) -> None: + assert precision_at_k(["d1", "d3"], ["d1", "d2"]) == 0.5 + + def test_no_hits(self) -> None: + assert precision_at_k(["d3", "d4"], ["d1", "d2"]) == 0.0 + + def test_k_cutoff(self) -> None: + assert precision_at_k(["d1", "d2", "d3"], ["d1"], k=2) == 0.5 + + def test_empty(self) -> None: + assert precision_at_k([], ["d1"]) == 0.0 + + +class TestMeanReciprocalRank: + def test_first_rank(self) -> None: + assert mean_reciprocal_rank(["d1", "d2"], ["d1"]) == 1.0 + + def test_second_rank(self) -> None: + assert mean_reciprocal_rank(["d2", "d1"], ["d1"]) == 0.5 + + def test_no_hit(self) -> None: + assert mean_reciprocal_rank(["d3", "d4"], ["d1"]) == 0.0 + + def test_multiple_relevant(self) -> None: + assert mean_reciprocal_rank(["d1", "d2", "d3"], ["d1", "d2"]) == 1.0 + + +class TestAveragePrecision: + def test_perfect(self) -> None: + assert average_precision(["d1", "d2"], ["d1", "d2"]) == 1.0 + + def test_partial(self) -> None: + # d1 at rank 1, d2 at rank 3 -> AP = (1/1 + 2/3) / 2 = 5/6 + assert average_precision(["d1", "d3", "d2"], ["d1", "d2"]) == pytest.approx(5 / 6) + + def test_no_hits(self) -> None: + assert average_precision(["d3", "d4"], ["d1", "d2"]) == 0.0 + + def test_empty_relevant(self) -> None: + assert average_precision(["d1"], []) == 0.0 + + +class TestContextPrecision: + def test_perfect(self) -> None: + assert context_precision(["d1", "d2"], ["d1", "d2"]) > 0 + + def test_no_hits(self) -> None: + assert context_precision(["d3", "d4"], ["d1"]) == 0.0 + + def test_empty(self) -> None: + assert context_precision([], ["d1"]) == 0.0 + assert context_precision(["d1"], []) == 0.0 + + +class TestContextRecall: + def test_perfect_recall(self) -> None: + assert context_recall(["d1", "d2", "d3"], ["d1", "d2"]) == 1.0 + + def test_partial_recall(self) -> None: + assert context_recall(["d1"], ["d1", "d2"]) == 0.5 + + def test_no_recall(self) -> None: + assert context_recall(["d3"], ["d1", "d2"]) == 0.0 + + def test_empty_relevant(self) -> None: + assert context_recall(["d1"], []) == 0.0 + + +class TestEvaluateQuery: + def test_returns_all_metrics(self) -> None: + metrics = evaluate_query(["d1", "d2"], ["d1", "d2"]) + assert "recall_at_k" in metrics + assert "precision_at_k" in metrics + assert "mrr" in metrics + assert "average_precision" in metrics + assert "context_precision" in metrics + assert "context_recall" in metrics + + def test_with_k_cutoff(self) -> None: + metrics = evaluate_query(["d1", "d2", "d3", "d4"], ["d1", "d2", "d3"], k=2) + assert metrics["recall_at_k"] == pytest.approx(2.0 / 3) + assert metrics["precision_at_k"] == 1.0 + + +class TestAggregateMetrics: + def test_aggregate(self) -> None: + lists = [ + {"recall_at_k": 1.0, "mrr": 1.0}, + {"recall_at_k": 0.5, "mrr": 0.5}, + ] + agg = aggregate_metrics(lists) + assert agg["recall_at_k"] == pytest.approx(0.75) + assert agg["mrr"] == pytest.approx(0.75) + + def test_empty(self) -> None: + assert aggregate_metrics([]) == {} diff --git a/tests/unit/core/search/eval/test_runner.py b/tests/unit/core/search/eval/test_runner.py new file mode 100644 index 00000000..db791ed5 --- /dev/null +++ b/tests/unit/core/search/eval/test_runner.py @@ -0,0 +1,121 @@ +"""Tests for the search evaluation runner.""" + +from __future__ import annotations + +from src.core.search.eval.golden_set import GoldenSet +from src.core.search.eval.runner import EvaluationResults, QueryResult, SearchEvaluator + + +def _mock_search(query: str, collection: str) -> list[str]: + """Deterministic mock search returning fixed IDs based on query.""" + return [f"result-for-{query}"] + + +class TestQueryResult: + def test_creation(self) -> None: + qr = QueryResult( + query_id="q001", + query="test", + collection="sigma_rules", + retrieved_ids=["r1"], + relevant_ids=["r1"], + metrics={"recall_at_k": 1.0}, + elapsed_ms=42.0, + ) + assert qr.query_id == "q001" + assert qr.elapsed_ms == 42.0 + + +class TestEvaluationResults: + def test_empty_summary(self) -> None: + results = EvaluationResults(description="empty") + summary = results.summary() + assert summary["num_queries"] == 0 + + def test_add_and_summary(self) -> None: + results = EvaluationResults(description="test") + results.add( + QueryResult( + query_id="q001", + query="q", + collection="sigma_rules", + retrieved_ids=["r1"], + relevant_ids=["r1"], + metrics={"recall_at_k": 1.0, "mrr": 1.0}, + elapsed_ms=10.0, + ) + ) + summary = results.summary() + assert summary["num_queries"] == 1 + assert summary["overall"]["recall_at_k"] == 1.0 + assert "sigma_rules" in summary["per_collection"] + + def test_report_format(self) -> None: + results = EvaluationResults(description="report test") + results.add( + QueryResult( + query_id="q001", + query="q", + collection="sigma_rules", + retrieved_ids=["r1"], + relevant_ids=["r1"], + metrics={"recall_at_k": 1.0, "mrr": 1.0}, + elapsed_ms=10.0, + ) + ) + report = results.report() + assert "Search Evaluation Report" in report + assert "recall_at_k" in report + + +class TestSearchEvaluator: + def test_evaluate_query(self) -> None: + evaluator = SearchEvaluator(_mock_search) + gs = GoldenSet() + q = gs.add("test", "sigma_rules", ["result-for-test"]) + qr = evaluator.evaluate_query(q) + assert qr.query_id == "q0000" + assert qr.metrics["recall_at_k"] == 1.0 + assert qr.elapsed_ms >= 0 + + def test_evaluate_query_no_hit(self) -> None: + evaluator = SearchEvaluator(_mock_search) + gs = GoldenSet() + q = gs.add("test", "sigma_rules", ["nonexistent"]) + qr = evaluator.evaluate_query(q) + assert qr.metrics["recall_at_k"] == 0.0 + assert qr.metrics["mrr"] == 0.0 + + def test_run_full_set(self) -> None: + evaluator = SearchEvaluator(_mock_search) + gs = GoldenSet(description="full test") + gs.add("test", "sigma_rules", ["result-for-test"]) + gs.add("test2", "sigma_docs", ["result-for-test2"]) + results = evaluator.run(gs) + assert results.num_queries == 2 + summary = results.summary() + assert summary["num_queries"] == 2 + + def test_search_failure_returns_empty(self) -> None: + def failing_search(query: str, collection: str) -> list[str]: + raise RuntimeError("search failed") + + evaluator = SearchEvaluator(failing_search) + gs = GoldenSet() + q = gs.add("test", "sigma_rules", ["doc1"]) + qr = evaluator.evaluate_query(q) + assert qr.retrieved_ids == [] + assert qr.metrics["recall_at_k"] == 0.0 + + def test_per_collection_breakdown(self) -> None: + def multi_collection_search(query: str, collection: str) -> list[str]: + return [f"hit-{collection}"] + + evaluator = SearchEvaluator(multi_collection_search) + gs = GoldenSet() + gs.add("q1", "sigma_rules", ["hit-sigma_rules"]) + gs.add("q2", "sigma_docs", ["hit-sigma_docs"]) + results = evaluator.run(gs) + summary = results.summary() + assert "sigma_rules" in summary["per_collection"] + assert "sigma_docs" in summary["per_collection"] diff --git a/tests/unit/core/search/test_sparse_encoder.py b/tests/unit/core/search/test_sparse_encoder.py new file mode 100644 index 00000000..f234708c --- /dev/null +++ b/tests/unit/core/search/test_sparse_encoder.py @@ -0,0 +1,164 @@ +"""Tests for the sparse encoder module (Q3.2 — IDF-aware BM25).""" + +from __future__ import annotations + +import math + +import pytest + +from src.core.search.sparse_encoder import ( + IDFCalculator, + _encode_single, + _token_id, + _tokenize, + bm25_idf_sparse_encoder, + bm25_sparse_encoder, +) + + +class TestTokenize: + def test_basic(self) -> None: + assert _tokenize("Process Creation with Image") == [ + "process", + "creation", + "with", + "image", + ] + + def test_short_words_filtered(self) -> None: + assert _tokenize("a an the of to be") == ["the"] + + def test_alphanumeric(self) -> None: + assert _tokenize("T1059.001") == ["t1059"] + + def test_empty(self) -> None: + assert _tokenize("") == [] + + +class TestTokenId: + def test_deterministic(self) -> None: + assert _token_id("process") == _token_id("process") + assert _token_id("process") != _token_id("Process") + + def test_within_range(self) -> None: + tid = _token_id("process_creation") + assert 0 <= tid < 2**24 + + +class TestEncodeSingle: + def test_basic_tf(self) -> None: + indices, values = _encode_single("the process process creation") + token_ids = {_token_id("process"): 0, _token_id("creation"): 0} + for idx, val in zip(indices, values): + token_ids[idx] = val + assert _token_id("process") in token_ids + assert _token_id("creation") in token_ids + assert values == [1.0 + math.log(2), 1.0 + math.log(1)] + + def test_empty_returns_empty(self) -> None: + assert _encode_single("") == ([], []) + assert _encode_single("a an") == ([], []) + + def test_with_idf(self) -> None: + idf_map = {"process": 2.5, "creation": 1.5} + indices, values = _encode_single("the process process creation", idf_map=idf_map) + id_by_token = {_token_id("process"): None, _token_id("creation"): None} + for idx, val in zip(indices, values): + id_by_token[idx] = val + + proc_val = id_by_token[_token_id("process")] + creat_val = id_by_token[_token_id("creation")] + assert proc_val is not None + assert creat_val is not None + assert proc_val != creat_val + + +class TestIDFCalculator: + def test_empty_corpus(self) -> None: + calc = IDFCalculator() + assert calc.idf() == {} + + def test_single_doc(self) -> None: + calc = IDFCalculator() + calc.add_document("the process creation") + idf_map = calc.idf() + assert "process" in idf_map + assert "creation" in idf_map + + def test_idf_higher_for_rare_terms(self) -> None: + calc = IDFCalculator() + for _ in range(100): + calc.add_document("process creation") + for _ in range(5): + calc.add_document("rare term specific") + idf_map = calc.idf() + assert idf_map["rare"] > idf_map["process"] + + def test_smooth_idf(self) -> None: + calc = IDFCalculator() + calc.add_document("term") + assert calc.idf(smooth=True)["term"] == pytest.approx( + math.log(1.0 + (1 - 1 + 0.5) / (1 + 0.5)) + ) + + +class TestBm25SparseEncoder: + def test_basic(self) -> None: + indices, values = bm25_sparse_encoder(["process creation"]) + assert len(indices) == 1 + assert len(values) == 1 + assert len(indices[0]) == 2 + + def test_batch(self) -> None: + texts = ["process creation", "network connection", ""] + indices, values = bm25_sparse_encoder(texts) + assert len(indices) == 3 + assert len(values) == 3 + assert indices[2] == [] # empty text + assert values[2] == [] + + +class TestBm25IdfSparseEncoder: + def test_basic(self) -> None: + idf_map = {"process": 2.0, "creation": 1.5} + indices, values = bm25_idf_sparse_encoder( + ["process creation"], + idf_map=idf_map, + avg_doc_len=10.0, + ) + assert len(indices) == 1 + assert len(indices[0]) == 2 + + def test_unknown_term_gets_default_idf(self) -> None: + indices, values = bm25_idf_sparse_encoder( + ["process creation"], + idf_map={"foo": 1.0}, + avg_doc_len=10.0, + ) + assert len(indices[0]) == 2 + assert all(v > 0 for v in values[0]) + + def test_empty_text(self) -> None: + indices, values = bm25_idf_sparse_encoder( + [""], + idf_map={"term": 1.0}, + avg_doc_len=10.0, + ) + assert indices == [[]] + assert values == [[]] + + def test_length_normalization(self) -> None: + """Short doc should get higher BM25 weight per term.""" + idf_map = {"process": 2.0} + _, values_short = bm25_idf_sparse_encoder( + ["process"], + idf_map=idf_map, + avg_doc_len=100.0, + ) + _, values_long = bm25_idf_sparse_encoder( + ["process"], + idf_map=idf_map, + avg_doc_len=1.0, + ) + # Shorter-than-avg doc gets slightly different weight + assert values_short[0] != values_long[0] diff --git a/tests/unit/shared/test_http.py b/tests/unit/shared/test_http.py new file mode 100644 index 00000000..860ed6f3 --- /dev/null +++ b/tests/unit/shared/test_http.py @@ -0,0 +1,313 @@ +"""Tests for shared HTTP utilities.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import httpx + +from src.shared.http import ( + close_all_pooled_clients, + create_client, + download_file, + get_pooled_client, + head_url, + _backoff_delay, + _get_retry_after, +) + + +class TestCreateClient: + def test_returns_httpx_client(self) -> None: + client = create_client() + assert isinstance(client, httpx.Client) + client.close() + + def test_sets_user_agent(self) -> None: + client = create_client() + assert client.headers.get("User-Agent") == "SigmaRAG/1.0" + client.close() + + def test_merges_custom_headers(self) -> None: + client = create_client(headers={"X-Custom": "value"}) + assert client.headers.get("User-Agent") == "SigmaRAG/1.0" + assert client.headers.get("X-Custom") == "value" + client.close() + + def test_follow_redirects_default_true(self) -> None: + client = create_client() + assert client.follow_redirects is True + client.close() + + def test_follow_redirects_false(self) -> None: + client = create_client(follow_redirects=False) + assert client.follow_redirects is False + client.close() + + def test_timeout_float(self) -> None: + client = create_client(timeout=15.0) + assert client.timeout == httpx.Timeout(15.0) + client.close() + + +class TestHeadUrl: + def test_returns_content_type_size_url(self) -> None: + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.headers = { + "content-type": "text/html; charset=utf-8", + "content-length": "1234", + } + mock_resp.url = httpx.URL("https://example.com/doc") + mock_client = MagicMock(spec=httpx.Client) + mock_client.head.return_value = mock_resp + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + ctype, size, final_url = head_url("https://example.com/doc") + + assert ctype == "text/html" + assert size == 1234 + assert final_url == "https://example.com/doc" + + def test_returns_none_on_http_error(self) -> None: + mock_client = MagicMock(spec=httpx.Client) + mock_client.head.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=MagicMock() + ) + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + result = head_url("https://example.com/404") + + assert result == (None, None, None) + + def test_returns_none_on_connection_error(self) -> None: + mock_client = MagicMock(spec=httpx.Client) + mock_client.head.side_effect = httpx.ConnectError("connection refused") + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + result = head_url("https://example.com/down") + + assert result == (None, None, None) + + def test_returns_none_for_private_url(self) -> None: + with patch("src.shared.http.is_private_url", return_value=True): + result = head_url("http://localhost:8080/secret") + + assert result == (None, None, None) + + def test_returns_none_for_empty_content_type(self) -> None: + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.headers = {} + mock_resp.url = httpx.URL("https://example.com/doc") + mock_client = MagicMock(spec=httpx.Client) + mock_client.head.return_value = mock_resp + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + ctype, size, final_url = head_url("https://example.com/doc") + + assert ctype is None + assert size == 0 + + def test_private_url_skip_by_default(self) -> None: + with patch("src.shared.http.is_private_url", return_value=True): + result = head_url("http://localhost:8080") + + assert result == (None, None, None) + + def test_private_url_not_skipped_when_check_ssrf_false(self) -> None: + mock_client = MagicMock(spec=httpx.Client) + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.headers = {"content-type": "text/plain"} + mock_resp.url = httpx.URL("http://localhost:8080/doc") + mock_client.head.return_value = mock_resp + + with ( + patch("src.shared.http.get_pooled_client", return_value=mock_client), + patch("src.shared.http.is_private_url", return_value=True), + ): + ctype, size, url = head_url("http://localhost:8080/doc", check_ssrf=False) + + assert ctype == "text/plain" + + +class TestDownloadFile: + def test_successful_download(self, tmp_path: Path) -> None: + output = tmp_path / "doc.md" + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.content = b"hello world" + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.return_value = mock_resp + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + ok, status = download_file("https://example.com/doc", output) + + assert ok is True + assert status is None + assert output.read_bytes() == b"hello world" + + def test_skips_private_url(self, tmp_path: Path) -> None: + output = tmp_path / "secret.md" + with patch("src.shared.http.is_private_url", return_value=True): + ok, status = download_file("http://localhost:8080/secret", output) + + assert ok is False + assert status is None + assert not output.exists() + + def test_retries_on_http_429(self, tmp_path: Path) -> None: + output = tmp_path / "retry.md" + mock_resp_429 = MagicMock(spec=httpx.Response) + mock_resp_429.status_code = 429 + mock_resp_429.headers = {} + + mock_resp_ok = MagicMock(spec=httpx.Response) + mock_resp_ok.content = b"content" + + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.side_effect = [ + httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), + mock_resp_ok, + ] + + with ( + patch("src.shared.http.get_pooled_client", return_value=mock_client), + patch("time.sleep"), + ): + ok, status = download_file("https://example.com/doc", output) + + assert ok is True + + def test_retries_on_network_error(self, tmp_path: Path) -> None: + output = tmp_path / "retry_net.md" + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.side_effect = [ + httpx.ConnectError("timeout"), + MagicMock(content=b"ok"), + ] + + with ( + patch("src.shared.http.get_pooled_client", return_value=mock_client), + patch("time.sleep"), + ): + ok, status = download_file("https://example.com/doc", output) + + assert ok is True + + def test_gives_up_after_max_retries(self, tmp_path: Path) -> None: + output = tmp_path / "fail.md" + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.side_effect = httpx.ConnectError("always fails") + + with ( + patch("src.shared.http.get_pooled_client", return_value=mock_client), + patch("time.sleep"), + ): + ok, status = download_file("https://example.com/doc", output, max_retries=2) + + assert ok is False + assert status is None + + def test_non_retryable_http_status(self, tmp_path: Path) -> None: + output = tmp_path / "forbidden.md" + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.status_code = 403 + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.side_effect = httpx.HTTPStatusError( + "403", request=MagicMock(), response=mock_resp + ) + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + ok, status = download_file("https://example.com/forbidden", output) + + assert ok is False + assert status == 403 + + def test_filesystem_error(self, tmp_path: Path) -> None: + output = Path("/nonexistent/path/doc.md") + mock_resp = MagicMock(spec=httpx.Response) + mock_resp.content = b"data" + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.return_value = mock_resp + + with patch("src.shared.http.get_pooled_client", return_value=mock_client): + ok, status = download_file("https://example.com/doc", output) + + assert ok is False + + def test_respects_retry_after_header(self, tmp_path: Path) -> None: + output = tmp_path / "retry_after.md" + mock_resp_429 = MagicMock(spec=httpx.Response) + mock_resp_429.status_code = 429 + mock_resp_429.headers = {"Retry-After": "5"} + + mock_resp_ok = MagicMock(spec=httpx.Response) + mock_resp_ok.content = b"content" + + mock_client = MagicMock(spec=httpx.Client) + mock_client.get.side_effect = [ + httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), + mock_resp_ok, + ] + + with ( + patch("src.shared.http.get_pooled_client", return_value=mock_client), + patch("time.sleep") as mock_sleep, + ): + ok, status = download_file("https://example.com/doc", output) + + assert ok is True + mock_sleep.assert_called_once_with(5.0) + + +class TestBackoffDelay: + def test_first_attempt(self) -> None: + assert _backoff_delay(1) == 1.0 + + def test_second_attempt(self) -> None: + assert _backoff_delay(2) == 4.0 + + def test_third_attempt(self) -> None: + assert _backoff_delay(3) == 9.0 + + def test_beyond_list_falls_back(self) -> None: + assert _backoff_delay(10) == 9.0 + + +class TestGetRetryAfter: + def test_valid_header(self) -> None: + resp = MagicMock(spec=httpx.Response) + resp.headers = {"Retry-After": "30"} + assert _get_retry_after(resp) == 30 + + def test_missing_header(self) -> None: + resp = MagicMock(spec=httpx.Response) + resp.headers = {} + assert _get_retry_after(resp) is None + + def test_invalid_value(self) -> None: + resp = MagicMock(spec=httpx.Response) + resp.headers = {"Retry-After": "invalid"} + assert _get_retry_after(resp) is None + + +class TestPooledClient: + def test_returns_same_client_for_same_params(self) -> None: + c1 = get_pooled_client(timeout=30.0, follow_redirects=True) + c2 = get_pooled_client(timeout=30.0, follow_redirects=True) + assert c1 is c2 + + def test_returns_different_client_for_different_timeout(self) -> None: + c1 = get_pooled_client(timeout=30.0) + c2 = get_pooled_client(timeout=45.0) + assert c1 is not c2 + + def test_returns_different_client_for_different_redirects(self) -> None: + c1 = get_pooled_client(follow_redirects=True) + c2 = get_pooled_client(follow_redirects=False) + assert c1 is not c2 + + def test_close_all_clears_pool(self) -> None: + c1 = get_pooled_client(timeout=30.0) + close_all_pooled_clients() + c2 = get_pooled_client(timeout=30.0) + assert c1 is not c2 diff --git a/tests/unit/shared/utils/test_identify_file_type.py b/tests/unit/shared/utils/test_identify_file_type.py index c886a702..5fef697d 100644 --- a/tests/unit/shared/utils/test_identify_file_type.py +++ b/tests/unit/shared/utils/test_identify_file_type.py @@ -295,9 +295,12 @@ def test_values_are_filetype(self) -> None: class TestSupportedReferenceDocTypes: def test_contains_relevant_types(self) -> None: assert "markdown" in SUPPORTED_REFERENCE_DOC_TYPES - assert "pdf" in SUPPORTED_REFERENCE_DOC_TYPES - assert "plain_text" in SUPPORTED_REFERENCE_DOC_TYPES - assert "office_document" in SUPPORTED_REFERENCE_DOC_TYPES + assert "sigma_rule" in SUPPORTED_REFERENCE_DOC_TYPES + assert "html" not in SUPPORTED_REFERENCE_DOC_TYPES + assert "yaml" not in SUPPORTED_REFERENCE_DOC_TYPES + assert "pdf" not in SUPPORTED_REFERENCE_DOC_TYPES + assert "plain_text" not in SUPPORTED_REFERENCE_DOC_TYPES + assert "office_document" not in SUPPORTED_REFERENCE_DOC_TYPES def test_excludes_media_types(self) -> None: excluded = {"image", "audio", "video", "archive", "executable"} diff --git a/tests/unit/shared/utils/test_registry_utils.py b/tests/unit/shared/utils/test_registry_utils.py new file mode 100644 index 00000000..05ae0bf2 --- /dev/null +++ b/tests/unit/shared/utils/test_registry_utils.py @@ -0,0 +1,102 @@ +"""Tests for registry entry builder.""" + +from __future__ import annotations + +from src.shared.utils.registry_utils import build_registry_entry + + +class TestBuildRegistryEntry: + def test_minimal_entry(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="markdown", + rule_id="rule-001", + title="Test Rule", + ) + assert entry["normalized_url"] == "https://example.com/doc.md" + assert entry["content_type"] == "markdown" + assert entry["rule_id"] == "rule-001" + assert entry["title"] == "Test Rule" + assert entry["org"] == "sigmaref" + assert entry["repo"] == "references" + assert entry["embed_status"] == "discovery" + assert entry["file_name"] == "doc.md" + assert entry["file_size"] == 0 + assert entry["url_hash"] != "" + + def test_with_url_hash(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="html", + rule_id="rule-002", + title="Another Rule", + url_hash="abc123", + ) + assert entry["url_hash"] == "abc123" + + def test_head_verified_status(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="html", + rule_id="rule-003", + title="Head Verified", + embed_status="head_verified", + file_size=1234, + ) + assert entry["embed_status"] == "head_verified" + assert entry["file_size"] == 1234 + + def test_downloaded_entry(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="pdf", + rule_id="rule-004", + title="Downloaded Doc", + content_sha256="abcdef1234567890", + file_name="abc123.pdf", + file_size=5678, + ) + assert entry["content_sha256"] == "abcdef1234567890" + assert entry["file_name"] == "abc123.pdf" + assert entry["file_size"] == 5678 + assert entry["embed_status"] == "discovery" + + def test_original_url_falls_back(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="markdown", + rule_id="rule-005", + title="Fallback", + ) + assert entry["original_url"] == "https://example.com/doc.md" + + def test_custom_original_url(self) -> None: + entry = build_registry_entry( + normalized_url="https://raw.githubusercontent.com/user/repo/doc.md", + content_type="markdown", + rule_id="rule-006", + title="Custom Original", + original_url="https://github.com/user/repo/blob/main/doc.md", + ) + assert entry["original_url"] == "https://github.com/user/repo/blob/main/doc.md" + assert entry["normalized_url"] == "https://raw.githubusercontent.com/user/repo/doc.md" + + def test_timestamp_is_set(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/doc.md", + content_type="markdown", + rule_id="rule-007", + title="Timestamp test", + ) + assert entry["timestamp"] is not None + assert entry["last_seen"] is not None + assert entry["timestamp"] == entry["last_seen"] + + def test_file_name_from_url_when_not_given(self) -> None: + entry = build_registry_entry( + normalized_url="https://example.com/path/to/document.pdf", + content_type="pdf", + rule_id="rule-008", + title="File name from URL", + ) + assert entry["file_name"] == "document.pdf" diff --git a/tests/unit/workers/conftest.py b/tests/unit/workers/conftest.py index fe0754e0..9b1456cf 100644 --- a/tests/unit/workers/conftest.py +++ b/tests/unit/workers/conftest.py @@ -47,6 +47,10 @@ def mock_db() -> MagicMock: db.get_local_file_count = MagicMock(return_value=0) db.get_repos_with_selected_dirs = MagicMock(return_value=[]) db.get_selected_dirs = MagicMock(return_value=[]) + db.get_doc_registry_url_hashes_by_repo = MagicMock(return_value=[]) + db.delete_doc_registry_by_url_hashes = MagicMock() + db.get_rule_id_by_url_hash = MagicMock(return_value=None) + db.get_rule_reference_paths = MagicMock(return_value=[]) return db diff --git a/tests/unit/workers/test_discovery_workers.py b/tests/unit/workers/test_discovery_workers.py index f9ead339..4036cd72 100644 --- a/tests/unit/workers/test_discovery_workers.py +++ b/tests/unit/workers/test_discovery_workers.py @@ -280,3 +280,146 @@ def test_process_reports_source_type(self, mock_db: MagicMock, tmp_path: Path) - mock_db.batch_upsert_doc_registry.assert_called_once() entries = mock_db.batch_upsert_doc_registry.call_args[0][0] assert all(e["org"] == "local" for e in entries) + + +class TestGarbageCollection: + def test_garbage_collect_clears_stale_entries_across_repos( + self, mock_db: MagicMock, tmp_path: Path + ) -> None: + """GC should clean stale entries from ALL repos with active selections, not just the scanned one.""" + repo_a_dir = tmp_path / "org-a" / "repo-a" + repo_a_dir.mkdir(parents=True) + (repo_a_dir / "stale.md").write_text("# Stale") + + repo_b_dir = tmp_path / "org-b" / "repo-b" + repo_b_dir.mkdir(parents=True) + (repo_b_dir / "active.md").write_text("# Active") + + mock_db.get_repos_with_selected_dirs.return_value = ["org-a/repo-a"] + + stale_hash = "stale-hash-123" + active_hash = "active-hash-456" + + active_entry = {"url_hash": active_hash, "org": "org-b", "repo": "repo-b"} + + mock_db.get_doc_registry_url_hashes_by_repo.side_effect = lambda org, repo: ( + [stale_hash] if org == "org-a" and repo == "repo-a" else [] + ) + mock_db.get_rule_id_by_url_hash.return_value = "00000000-0000-0000-0000-000000000000" + mock_db.get_rule_reference_paths.return_value = [] + + worker = GenericDiscoveryWorker( + db=mock_db, source_type=SourceType.GITHUB, base_dir=tmp_path + ) + + # Call GC directly with stale entries from repo A, active entries from repo B + repo_items = [("org-a", "repo-a"), ("org-b", "repo-b")] + current_entries = [active_entry] # Only active entries in current scan + + worker._garbage_collect_github(repo_items, current_entries) + + # Verify: stale entry from repo A was deleted + mock_db.delete_doc_registry_by_url_hashes.assert_called_once_with([stale_hash]) + + def test_garbage_collect_preserves_active_entries( + self, mock_db: MagicMock, tmp_path: Path + ) -> None: + """GC should NOT delete entries that are still in the current scan.""" + active_hash = "active-hash-789" + active_entry = {"url_hash": active_hash, "org": "org", "repo": "repo"} + + mock_db.get_doc_registry_url_hashes_by_repo.return_value = [active_hash] + mock_db.get_rule_id_by_url_hash.return_value = "00000000-0000-0000-0000-000000000000" + mock_db.get_rule_reference_paths.return_value = [] + + worker = GenericDiscoveryWorker( + db=mock_db, source_type=SourceType.GITHUB, base_dir=tmp_path + ) + + repo_items = [("org", "repo")] + current_entries = [active_entry] # Active entry is in current scan + + worker._garbage_collect_github(repo_items, current_entries) + + # Verify: delete was NOT called (no stale entries) + mock_db.delete_doc_registry_by_url_hashes.assert_not_called() + + def test_garbage_collect_cleanup_rule_references( + self, mock_db: MagicMock, tmp_path: Path + ) -> None: + """GC should delete reference files before deleting doc_registry entries.""" + ref_path = tmp_path / "ref.pdf" + ref_path.write_bytes(b"fake pdf") + + stale_hash = "stale-ref-hash" + rule_id = "test-rule-id-123" + + mock_db.get_doc_registry_url_hashes_by_repo.return_value = [stale_hash] + mock_db.get_rule_id_by_url_hash.return_value = rule_id + mock_db.get_rule_reference_paths.return_value = [ref_path] + + worker = GenericDiscoveryWorker( + db=mock_db, source_type=SourceType.GITHUB, base_dir=tmp_path + ) + + repo_items = [("org", "repo")] + # Pass a dummy active entry so GC doesn't return early + active_entry = {"url_hash": "active-dummy", "org": "org", "repo": "repo"} + current_entries: list[dict] = [active_entry] + + worker._garbage_collect_github(repo_items, current_entries) + + # Verify: reference file was deleted + assert not ref_path.exists() + # Verify: doc_registry entry was deleted + mock_db.delete_doc_registry_by_url_hashes.assert_called_once_with([stale_hash]) + + def test_garbage_collect_skips_null_rule_id(self, mock_db: MagicMock, tmp_path: Path) -> None: + """GC should skip entries with null or empty rule_id.""" + stale_hash = "stale-no-rule" + + mock_db.get_doc_registry_url_hashes_by_repo.return_value = [stale_hash] + mock_db.get_rule_id_by_url_hash.return_value = None # No rule association + + worker = GenericDiscoveryWorker( + db=mock_db, source_type=SourceType.GITHUB, base_dir=tmp_path + ) + + repo_items = [("org", "repo")] + # Pass a dummy active entry so GC doesn't return early + active_entry = {"url_hash": "active-dummy", "org": "org", "repo": "repo"} + current_entries: list[dict] = [active_entry] + + worker._garbage_collect_github(repo_items, current_entries) + + # Verify: no cleanup attempted (rule_id is None) + mock_db.get_rule_reference_paths.assert_not_called() + mock_db.delete_doc_registry_by_url_hashes.assert_called_once_with([stale_hash]) + + def test_garbage_collect_handles_missing_files( + self, mock_db: MagicMock, tmp_path: Path + ) -> None: + """GC should not fail if reference files are already missing.""" + ref_path = tmp_path / "missing.pdf" + # Don't create the file + + stale_hash = "stale-missing" + rule_id = "test-rule-id" + + mock_db.get_doc_registry_url_hashes_by_repo.return_value = [stale_hash] + mock_db.get_rule_id_by_url_hash.return_value = rule_id + mock_db.get_rule_reference_paths.return_value = [ref_path] # Path exists but file doesn't + + worker = GenericDiscoveryWorker( + db=mock_db, source_type=SourceType.GITHUB, base_dir=tmp_path + ) + + repo_items = [("org", "repo")] + # Pass a dummy active entry so GC doesn't return early + active_entry = {"url_hash": "active-dummy", "org": "org", "repo": "repo"} + current_entries: list[dict] = [active_entry] + + worker._garbage_collect_github(repo_items, current_entries) + + # Verify: no error, deletion is idempotent (missing_ok=True) + mock_db.delete_doc_registry_by_url_hashes.assert_called_once_with([stale_hash]) diff --git a/tests/unit/workers/test_gc_worker.py b/tests/unit/workers/test_gc_worker.py new file mode 100644 index 00000000..19c27e32 --- /dev/null +++ b/tests/unit/workers/test_gc_worker.py @@ -0,0 +1,264 @@ +"""Tests for DocGCWorker and orphaned sigmaref file cleaning.""" + +from __future__ import annotations + +import threading +from pathlib import Path +from unittest.mock import MagicMock, patch + +import duckdb + +from src.infrastructure.database import DatabaseService +from src.workers.document.gc_worker import DocGCWorker, _is_orphan_candidate + +_VALID_HASH = "ab" * 32 # 64-char hex +_ORPHAN_HASH = "ef" * 32 + + +class TestIsOrphanCandidate: + def test_accepts_unknown_hash(self) -> None: + assert _is_orphan_candidate(_VALID_HASH, set()) is True + + def test_rejects_known_hash(self) -> None: + assert _is_orphan_candidate(_VALID_HASH, {_VALID_HASH}) is False + + def test_rejects_short_stem(self) -> None: + assert _is_orphan_candidate("too-short", set()) is False + + def test_rejects_non_hex_stem(self) -> None: + assert _is_orphan_candidate("z" + _VALID_HASH[1:], set()) is False + + def test_rejects_dotfile(self) -> None: + assert _is_orphan_candidate(".keep", set()) is False + + +def _make_db(known_hashes: set[str] | None = None) -> MagicMock: + db = MagicMock() + lock = MagicMock() + db._lock = lock + db._writer_conn.execute.return_value.fetchall.return_value = [ + (h,) for h in (known_hashes or {_VALID_HASH}) + ] + return db + + +def _make_worker(known_hashes: set[str] | None = None) -> DocGCWorker: + """Build a DocGCWorker with a mocked DB that returns *known_hashes*.""" + db = _make_db(known_hashes) + worker = DocGCWorker(db=db) + worker.dispatcher = MagicMock() + return worker + + +def _touch(d: Path, name: str) -> Path: + p = d / name + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text("test") + return p + + +class TestGcOrphanedSigmarefFiles: + def test_removes_orphan_in_flat_layout(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + _touch(base, _VALID_HASH + ".md") # known + orphan = _touch(base, _ORPHAN_HASH + ".md") # orphan + + worker = _make_worker() + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + result = worker._gc_orphaned_sigmaref_files() + + assert result == 1 + assert not orphan.exists() + known_file = _VALID_HASH + ".md" + assert (base / known_file).exists() + + def test_removes_orphan_in_subdir(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + orphan = _touch(base / "markdown", _ORPHAN_HASH + ".md") + + worker = _make_worker() + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + result = worker._gc_orphaned_sigmaref_files() + + assert result == 1 + assert not orphan.exists() + + def test_skips_known_file(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + known = _touch(base, _VALID_HASH + ".md") + + worker = _make_worker() + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + result = worker._gc_orphaned_sigmaref_files() + + assert result == 0 + assert known.exists() + + def test_skips_dotfiles(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + dotfile = _touch(base, ".keep") + + worker = _make_worker(known_hashes=set()) + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + result = worker._gc_orphaned_sigmaref_files() + + assert result == 0 + assert dotfile.exists() + + def test_removes_empty_subdir(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + sub = base / "pdf" + sub.mkdir() + + worker = _make_worker(known_hashes=set()) + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + worker._gc_orphaned_sigmaref_files() + + assert not sub.exists() + + def test_files_go_to_trash(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + orphan = _touch(base, _ORPHAN_HASH + ".md") + + worker = _make_worker(known_hashes=set()) + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + worker._gc_orphaned_sigmaref_files() + + assert (base / ".trash" / (_ORPHAN_HASH + ".md")).exists() + assert not orphan.exists() + + def test_noop_when_base_missing(self) -> None: + worker = _make_worker() + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = "/nonexistent" + result = worker._gc_orphaned_sigmaref_files() + + assert result == 0 + + def test_db_error_returns_zero(self, tmp_path: Path) -> None: + base = tmp_path / "sigmaref" + base.mkdir() + + db = _make_db(known_hashes=set()) + db._writer_conn.execute.side_effect = RuntimeError("db down") + worker = DocGCWorker(db=db) + + with patch("src.workers.document.gc_worker.get_config") as mock_cfg: + mock_cfg.return_value.sigmaref_documents_path = str(base) + result = worker._gc_orphaned_sigmaref_files() + + assert result == 0 + + +class TestDeleteUnreferencedEntries: + def _make_db_service(self) -> DatabaseService: + """Create a DatabaseService with an in-memory DuckDB and required tables.""" + db = DatabaseService.__new__(DatabaseService) + db._lock = threading.Lock() + db._writer_conn = duckdb.connect(":memory:") + db._writer_conn.execute(""" + CREATE TABLE doc_registry ( + url_hash TEXT PRIMARY KEY, + org TEXT, + repo TEXT, + content_type TEXT, + file_name TEXT, + content_sha256 TEXT, + file_size BIGINT, + original_url TEXT NOT NULL, + normalized_url TEXT, + rule_id TEXT, + title TEXT, + timestamp TEXT, + last_seen TEXT, + embed_status TEXT DEFAULT 'discovery' + ) + """) + db._writer_conn.execute(""" + CREATE TABLE rule_references ( + rule_id TEXT NOT NULL, + url_hash TEXT NOT NULL, + ref_url TEXT NOT NULL, + created TEXT + ) + """) + db._writer_conn.execute(""" + CREATE TABLE doc_error ( + url_hash TEXT PRIMARY KEY, + original_url TEXT NOT NULL, + normalized_url TEXT NOT NULL, + error_code INTEGER, + error_message TEXT, + org TEXT, + repo TEXT, + timestamp TEXT + ) + """) + db._safe_query = lambda sql, params=None: db._writer_conn.execute( + sql, params or [] + ).fetchall() + return db + + def _insert_doc(self, db: DatabaseService, url_hash: str, org: str = "sigmaref") -> None: + db._writer_conn.execute( + "INSERT INTO doc_registry (url_hash, org, original_url, normalized_url) VALUES (?, ?, ?, ?)", + [url_hash, org, f"https://example.com/{url_hash}", f"https://example.com/{url_hash}"], + ) + + def _insert_ref(self, db: DatabaseService, url_hash: str) -> None: + db._writer_conn.execute( + "INSERT INTO rule_references (rule_id, url_hash, ref_url) VALUES (?, ?, ?)", + ["rule-1", url_hash, f"https://example.com/{url_hash}"], + ) + + def test_deletes_orphan_entries(self) -> None: + db = self._make_db_service() + self._insert_doc(db, "orphan1") + self._insert_doc(db, "orphan2") + self._insert_doc(db, "referenced1") + self._insert_ref(db, "referenced1") + + result = db.delete_unreferenced_entries() + + assert result == 2 + remaining = db._writer_conn.execute( + "SELECT url_hash FROM doc_registry ORDER BY url_hash" + ).fetchall() + assert remaining == [("referenced1",)] + + def test_noop_when_no_orphans(self) -> None: + db = self._make_db_service() + self._insert_doc(db, "ref1") + self._insert_ref(db, "ref1") + + result = db.delete_unreferenced_entries() + + assert result == 0 + remaining = db._writer_conn.execute("SELECT COUNT(*) FROM doc_registry").fetchone()[0] + assert remaining == 1 + + def test_skips_local_and_github_entries(self) -> None: + db = self._make_db_service() + self._insert_doc(db, "sigmaref_orphan", org="sigmaref") + self._insert_doc(db, "local_file", org="local") + self._insert_doc(db, "github_file", org="sigmahq") + + result = db.delete_unreferenced_entries() + + assert result == 1 # only sigmaref_orphan deleted + remaining = db._writer_conn.execute( + "SELECT url_hash FROM doc_registry ORDER BY url_hash" + ).fetchall() + assert remaining == [("github_file",), ("local_file",)] diff --git a/uv.lock b/uv.lock index cf17aef5..bddb06f0 100644 --- a/uv.lock +++ b/uv.lock @@ -301,14 +301,14 @@ wheels = [ [[package]] name = "click" -version = "8.4.1" +version = "8.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/98/518d8e5081007684232226f475082b30087d0f585e8457db087298259f49/click-8.4.1.tar.gz", hash = "sha256:918b5633eddf6b41c32d4f454bf0de810065c74e3f7dbf8ee5452f8be88d3e96", size = 353007, upload-time = "2026-05-22T04:08:37.769Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, ] [[package]] @@ -499,36 +499,36 @@ wheels = [ [[package]] name = "duckdb" -version = "1.5.3" +version = "1.5.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/69/00/d579dcb2a536b6ea3a2563cdad6844f77d81a9b2d4b22a858097f2468acf/duckdb-1.5.3.tar.gz", hash = "sha256:df39428eb130faa35ae96fd35245bdeae6ecf43936250b116b5fead568eb9f16", size = 18026640, upload-time = "2026-05-20T11:55:31.901Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/c4/2e34929b16c8d544ef664fad8f7f3a2a9db05746aae1e7c8c4ee3a8b23e4/duckdb-1.5.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ff11a457258148337ef9a392148a8cdbd1069b6c27c21958816c7b67fe6c542d", size = 32626494, upload-time = "2026-05-20T11:54:33.738Z" }, - { url = "https://files.pythonhosted.org/packages/3a/53/3af681793d03771365ae3e2215331151c196a3ac8193f613344840694671/duckdb-1.5.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5fd25f533cb1b6b2c84cc767a9a9bab7769bb1aa44571a2a0bfc91ac3e4a38ac", size = 17301121, upload-time = "2026-05-20T11:54:36.928Z" }, - { url = "https://files.pythonhosted.org/packages/15/e2/c80af1eac2ab5d35fc2c372ef0a84668842e549fbbf7799277b3fccf3e39/duckdb-1.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:10960400ed60cdf0fe05bab2086fa8eb733889cb0ceca18d07ff9a00c0e0be7b", size = 15449283, upload-time = "2026-05-20T11:54:39.777Z" }, - { url = "https://files.pythonhosted.org/packages/2d/9a/c63af233c9f761bf5178a5210437e1bc6bcb30fa8a9073de6398cfb12c03/duckdb-1.5.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5f18e7561403054433706c187589e86629a7af09a7efc23a06a8b308e6acc68", size = 19332762, upload-time = "2026-05-20T11:54:42.51Z" }, - { url = "https://files.pythonhosted.org/packages/21/cc/2d77af4fff86012f334ef82e6d54a995a86c8745e58074f1218ed7d25171/duckdb-1.5.3-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9fb7516255a8764545e30f7efacea408cc847764a3027b3b0b3e7d1a7bebbc5c", size = 21453290, upload-time = "2026-05-20T11:54:45.272Z" }, - { url = "https://files.pythonhosted.org/packages/8d/5e/9bc4817a98feb4dab83e56f2245cd3a30d00ee646d4dec7926464e2b3f28/duckdb-1.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:8001eccbc28be244dfd04d708526f34ddd6460b47a8aeb5d0e39d6f7f9e3fe15", size = 13118308, upload-time = "2026-05-20T11:54:48.058Z" }, - { url = "https://files.pythonhosted.org/packages/81/35/e3f32e4e53e2450ddb1db8312a17d1ce455d60cc4941b6ad2cfc908794b0/duckdb-1.5.3-cp312-cp312-win_arm64.whl", hash = "sha256:6d2835e39bb6af73891f73c0f8d4324f98afe00d0b00c6d34b2a582c2256cbb0", size = 13927187, upload-time = "2026-05-20T11:54:50.584Z" }, - { url = "https://files.pythonhosted.org/packages/cc/9c/a528eb09d8be51954c485864bd06753e616939a080cbc3dd4417e8c94a57/duckdb-1.5.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:e75a6122c12579a99848517f6f00a4e342aebda3590c30fe9b5cc5f39d5e6afc", size = 32626254, upload-time = "2026-05-20T11:54:53.65Z" }, - { url = "https://files.pythonhosted.org/packages/ec/3c/1534c0a6db347c05eb7d0f6ecfb7aefbe74cbff398e4892a8fd1903a20e8/duckdb-1.5.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fd3963c1cb9d9567777f4a898a9dbe388a2fe9724681801b1e7d6d93eecf1b76", size = 17300917, upload-time = "2026-05-20T11:54:56.628Z" }, - { url = "https://files.pythonhosted.org/packages/23/fa/beafb91e6e152d2161c4a9cbc472334c87607eb61ad7104b5a7fa8d8d7b1/duckdb-1.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3d5db8c0b55e072cf437948ebb5d7e23d7b9d03d905fa5f9145583e65aa447f7", size = 15449411, upload-time = "2026-05-20T11:54:59.089Z" }, - { url = "https://files.pythonhosted.org/packages/50/0a/49b6fe04e2fcd63729eb607dadd44818dde77342a4f5ce086c6c92f1dd4d/duckdb-1.5.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ce80aed7a538422129a57eaca9141e3afb51f8bf562b1908b1576c9725b5b22", size = 19333120, upload-time = "2026-05-20T11:55:01.727Z" }, - { url = "https://files.pythonhosted.org/packages/63/4c/0907c3f76adb9dd90e67610b31e0304a35814e65c4c41a354a262c09b885/duckdb-1.5.3-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:787df63824f07bf18022dbc3b8ca4b2bfab0ebe616464f55c6e8cd0f59ea762e", size = 21453266, upload-time = "2026-05-20T11:55:04.5Z" }, - { url = "https://files.pythonhosted.org/packages/6d/9c/d2f23a7803ddbbd9413f7572ecf66a15120ed5ced7ce5c73e698c1406b76/duckdb-1.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:bb5bb5dcdd09d62ee60f0ddbbef918e71cce304ffe28428b1131949d39ffaabf", size = 13118640, upload-time = "2026-05-20T11:55:07.389Z" }, - { url = "https://files.pythonhosted.org/packages/27/d5/7ba2316415bcdab6edd765bbbe35c2ca8a3800f2fe695cd70e3cdb997f09/duckdb-1.5.3-cp313-cp313-win_arm64.whl", hash = "sha256:2fa17ecdd5d3db122836cb71bb93601c2106a3be883c17dffddc02fbf3fa7888", size = 13926409, upload-time = "2026-05-20T11:55:10.166Z" }, - { url = "https://files.pythonhosted.org/packages/a5/c2/d4b6f8a5e4d3bc25773be6da76a99d9661ebbf3552c007c460d2dd59dbf8/duckdb-1.5.3-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:4bfa9a4dadf71e83e2c4eaca2f9421c82a54defecc1b0b4c0be95e2389dec4fe", size = 32636685, upload-time = "2026-05-20T11:55:13.158Z" }, - { url = "https://files.pythonhosted.org/packages/42/58/e835c8298979d29db7a62cb5acc29e9b57aeaca7cdde2fcd3ac980f5cb18/duckdb-1.5.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:aea7baf67ad7e1829ac76f67d7dcbd7fb1f57c3eb179d55ac30952df4709ae30", size = 17308134, upload-time = "2026-05-20T11:55:16.194Z" }, - { url = "https://files.pythonhosted.org/packages/c9/46/617b51363f5613418c8b224b3cce16b58e6dde80904566bec232579c1d4e/duckdb-1.5.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:0b0b4f088a65d77e1217ce5d7eff889e63fedc44281200d899ff47c84d8ff836", size = 15449891, upload-time = "2026-05-20T11:55:18.687Z" }, - { url = "https://files.pythonhosted.org/packages/b3/72/354146656e8d9ba3853d3a5ee80a481b8c5f70edfc3d5ae80a8c4479c967/duckdb-1.5.3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe8d0c1f6a120aa03fa6e0d03897c71a1842e6cf7afd31d181348391f7108fe1", size = 19338499, upload-time = "2026-05-20T11:55:21.34Z" }, - { url = "https://files.pythonhosted.org/packages/56/8f/65fc623b51448f2bfba1a9ec6ab3debb4664c0876c0113a5e782600b53ac/duckdb-1.5.3-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0405eae18ec6e8210a471c97dbfe87a7e4d605274b7fe572a1f276e92158f13", size = 21455828, upload-time = "2026-05-20T11:55:23.847Z" }, - { url = "https://files.pythonhosted.org/packages/2b/db/d0274cbe9f5fe219f77c0bdf900ac77103569e83c102a4225ce04cbc607d/duckdb-1.5.3-cp314-cp314-win_amd64.whl", hash = "sha256:33ae08b3e818d7613d8936744b67718c2062c2f530376895bfd89efb51b81538", size = 13640011, upload-time = "2026-05-20T11:55:26.276Z" }, - { url = "https://files.pythonhosted.org/packages/07/5d/8f1899b8bef291caf953992fcd6c24df9f29387a35645e58c2504a5ca473/duckdb-1.5.3-cp314-cp314-win_arm64.whl", hash = "sha256:746433e49bbc667b4df283153415fbe37e9083e0eff6c3cd6e54de7536869cd4", size = 14411554, upload-time = "2026-05-20T11:55:29.037Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/31/29/9bad86ed7aa812d8c822a27c15c355b6d5423b991feeec86ed18027b6daa/duckdb-1.5.4.tar.gz", hash = "sha256:f9e32f1cdd106793d79d190186bed9e75289d51e68bd9174e47c04bffedeab6f", size = 18046634, upload-time = "2026-06-17T10:48:52.499Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/f2/e2f4b477ae3a3b40e8b5f429832e48edb62ed9da99807cc4902e157e5646/duckdb-1.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:291a9e7502551170af989ff63139a7a49e99d68edbc5ef5017ac27541fe54c65", size = 32708876, upload-time = "2026-06-17T10:48:01.527Z" }, + { url = "https://files.pythonhosted.org/packages/2e/2b/b698d82a5e1e30b6a05748d72045f672994c6b22f4f0f8423523608b991f/duckdb-1.5.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:83e8c089bbb756ca4471d8b05943b80a106058697cf00615e70423106bb783bc", size = 17346125, upload-time = "2026-06-17T10:48:04.035Z" }, + { url = "https://files.pythonhosted.org/packages/71/75/37e13f39268eaf34864453b3a039c4a1ff0b088d3eae45a4289b41c98c1b/duckdb-1.5.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ff96d2a342b200e1ec6f1f19986c77f4ac16a49b6112f71c5b763989203a9d60", size = 15488133, upload-time = "2026-06-17T10:48:06.312Z" }, + { url = "https://files.pythonhosted.org/packages/cc/59/2d082af578f689231798245b54562c61416e49049b0bda81a06c56a4b53e/duckdb-1.5.4-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8f935ef210ab00bc94bb1e3052697adaa36bb0ce7bdfeda8b0f34e2ff1643870", size = 19367895, upload-time = "2026-06-17T10:48:08.59Z" }, + { url = "https://files.pythonhosted.org/packages/52/2b/55c34d2863a76ca824ef8274691e84240b4ff1acde3d231709e82557c240/duckdb-1.5.4-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cda263d8c20addb8d4f95464787cbe0af1144f7ab7e21db3709fb826ee01725", size = 21486499, upload-time = "2026-06-17T10:48:10.963Z" }, + { url = "https://files.pythonhosted.org/packages/cf/30/ade5952b8182fac86fab43b95ebe3836e66381d0ad64eb1e54bd8207c988/duckdb-1.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:266c7c909558ce7377f57d082cee408aadebdd9111be017558ca54e44a031037", size = 13147934, upload-time = "2026-06-17T10:48:13.061Z" }, + { url = "https://files.pythonhosted.org/packages/f5/00/278f0f70e25b9911afe2fd227b9460f2e6d76177f0dcc03f7f1454afefa5/duckdb-1.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:f14e79a006341f29ee5a2692a24dac5114e77533d579c57ec39124adf0135033", size = 13965235, upload-time = "2026-06-17T10:48:15.782Z" }, + { url = "https://files.pythonhosted.org/packages/da/69/3fcb34e523a9bad1f0557a6c7691a71ba66c43a05e5be9ee96a9a841ed65/duckdb-1.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:42a612e67d64450b446eb69695290d460713eef46e0f64467ab9dfe96264ee05", size = 32708366, upload-time = "2026-06-17T10:48:18.084Z" }, + { url = "https://files.pythonhosted.org/packages/f5/5f/bff5054c2c1d65decab36aa6296621e51a2a575a9f250db7ab9b83a325d6/duckdb-1.5.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3fb6f07d54ecf4d0d3c5179a2361fdddfafa14de4fc42696de4632479b703421", size = 17345735, upload-time = "2026-06-17T10:48:20.67Z" }, + { url = "https://files.pythonhosted.org/packages/93/12/d1b2b344e9699246aada6f9de5156e708fb476e2780e5bff9b5d95fe11d9/duckdb-1.5.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:0f32ad7e0286c1c29ab6c73b29118c86101f8eee46aae54f54d0b50916f542f6", size = 15488568, upload-time = "2026-06-17T10:48:23.038Z" }, + { url = "https://files.pythonhosted.org/packages/c1/d1/ac56c6096e3e95da60b2c5dd5a0f0eb5540a80622e2e4f8faab893ec4e96/duckdb-1.5.4-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:698ec90bd5d5538bd5f6d212a4b61af443d240703cf45f134738535026556ea5", size = 19368184, upload-time = "2026-06-17T10:48:25.601Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0b/2ae4c3e157a19d9b4ac1f09a5dea6f93012334cc2db09f1e0c71eb99693d/duckdb-1.5.4-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:136cea7f886b78caf4035485b4b1e766e8b309e999f9e83a966f81ebb8122844", size = 21486523, upload-time = "2026-06-17T10:48:27.817Z" }, + { url = "https://files.pythonhosted.org/packages/64/7b/c3d8d21e0d0db8faa81eeeb3a55b9932f5a0a16466cb968dc713a653d701/duckdb-1.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:bd6777e8ddd74fb603a6d09766bfcff28638189f8aaa61fc0dffd9e9a4baa8e5", size = 13147807, upload-time = "2026-06-17T10:48:30.017Z" }, + { url = "https://files.pythonhosted.org/packages/44/48/ddf8d3740e3d28582944f70d84e720b5dc28c10ec22b668a0e0bd965f2f2/duckdb-1.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:73f4878a3012283024a64a1909e440aac12091ef336f671fc142f7e87449ce0c", size = 13965189, upload-time = "2026-06-17T10:48:32.251Z" }, + { url = "https://files.pythonhosted.org/packages/62/01/67ac4cbc8e552a1e14c029b5c443d828e68f94d5d913c574f577e1db277e/duckdb-1.5.4-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:4647968629d0677bbcc2416c7aeda8685eb84e4ca15a6dbd4f82a66cfc91a532", size = 32714364, upload-time = "2026-06-17T10:48:34.724Z" }, + { url = "https://files.pythonhosted.org/packages/4e/0e/eb44d983fa56b175f971eea251bde284a36d26cbb93fcb68035061f54078/duckdb-1.5.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:e8fcef301cf68d3951ea1eb8ac4d76cea0a6f6a08f4c78fe4026fc96d217bebc", size = 17349820, upload-time = "2026-06-17T10:48:37.126Z" }, + { url = "https://files.pythonhosted.org/packages/10/b2/b9dc7624b105d414585b8530451c1162c0b4750c0be9be2e497bb47a8a9b/duckdb-1.5.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:f6f39cd0dc6948dee17fd130aec55114f97a8ef6e1db519b9774087962bc5c8c", size = 15498160, upload-time = "2026-06-17T10:48:40.032Z" }, + { url = "https://files.pythonhosted.org/packages/b7/57/61356444f6a8c62dec3c3d129abfc53f428de1d484093d1bb381db441231/duckdb-1.5.4-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:262f068158beb5943f2c618f4e54b46db8306b959f90dce956f90a89f613673d", size = 19374183, upload-time = "2026-06-17T10:48:42.698Z" }, + { url = "https://files.pythonhosted.org/packages/b0/f4/d5d633dd7c5138d8f7c434e6ac2553c831b7fb658494efa8d0bc73df8623/duckdb-1.5.4-cp314-cp314-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d2307a76d199077b0055b354e90e857479461a0d875437535dd4833172c8b6d", size = 21487202, upload-time = "2026-06-17T10:48:45.407Z" }, + { url = "https://files.pythonhosted.org/packages/c0/26/5be13bbd5c3421dccfc1ad4ca9da4b97c5a3ddd73f66542092f3167ec52c/duckdb-1.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:6dcbb81a1276bc48deb4d562bce4f8895e4fc6348750a096e30052345c6d6552", size = 13666989, upload-time = "2026-06-17T10:48:47.764Z" }, + { url = "https://files.pythonhosted.org/packages/dc/82/4d52f3f9f9703a226b26b80bdae3f6905aeefe5221bf1815fc93ff02ca25/duckdb-1.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:0f8722346024e5d9f02b58bf7b0491a629f97fdc8a04a10e432940f471ee387a", size = 14449863, upload-time = "2026-06-17T10:48:50.18Z" }, ] [[package]] name = "fastapi" -version = "0.136.3" +version = "0.139.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-doc" }, @@ -537,9 +537,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/81/2d/ff8d91d7b564d464629a0fd50a4489c97fcb836ac230bf3a7269232a9b1f/fastapi-0.136.3.tar.gz", hash = "sha256:e487fae93ad408e6f47641ee4dfe389864fd7bec92e547ea8498fc13f43e83ab", size = 396410, upload-time = "2026-05-23T18:53:15.192Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d3/af/a5f50ccfa659ec1802cb4ca842c23f06d906a8cc9aef6016a2caeea3d4ed/fastapi-0.139.0.tar.gz", hash = "sha256:99ab7b2d92223c76d6cf10757ab3f89d45b38267fc20b2a136cf02f6beac3145", size = 423016, upload-time = "2026-07-01T16:35:33.436Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e0/82/45359b62a067409bd929ae8a56b8ed13e5a8c8a61194b3c236920999ab83/fastapi-0.136.3-py3-none-any.whl", hash = "sha256:3d2a69bdf04b7e9f3afa292c3bc7a98816bbfafa10bc9b45f3f3700d2f761620", size = 117481, upload-time = "2026-05-23T18:53:16.924Z" }, + { url = "https://files.pythonhosted.org/packages/9e/7c/8e3c6ad324ea5cb36604fc3f968554887891c316d9dfde57761611d907ad/fastapi-0.139.0-py3-none-any.whl", hash = "sha256:cf15e1e9e667ddb0ad63811e60bd11390d1aac838ca4a7a23f421807b2308189", size = 130339, upload-time = "2026-07-01T16:35:32.19Z" }, ] [[package]] @@ -829,34 +829,34 @@ wheels = [ [[package]] name = "hf-xet" -version = "1.5.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/d8/5c06fc76461418326a7decf8367480c35be11a41fd938633929c60a9ec6b/hf_xet-1.5.0.tar.gz", hash = "sha256:e0fb0a34d9f406eed88233e829a67ec016bec5af19e480eac65a233ea289a948", size = 837196, upload-time = "2026-05-06T06:18:15.583Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/68/9b/6912c99070915a4f28119e3c5b52a9abd1eec0ad5cb293b8c967a0c6f5a2/hf_xet-1.5.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:7d70fe2ce97b9db73b9c9b9c81fe3693640aec83416a966c446afea54acfae3c", size = 4023383, upload-time = "2026-05-06T06:17:53.947Z" }, - { url = "https://files.pythonhosted.org/packages/0f/6d/9563cfde59b5d8128a9c7ec972a087f4c782e4f7bac5a85234edfd5d5e49/hf_xet-1.5.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:73a0dae8c71de3b0633a45c73f4a4a5ed09e94b43441d82981a781d4f12baa42", size = 3792751, upload-time = "2026-05-06T06:17:51.791Z" }, - { url = "https://files.pythonhosted.org/packages/07/a5/ed5a0cf35b49a0571af5a8f53416dad1877a718c021c9937c3a53cb45781/hf_xet-1.5.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a60290ec57e9b71767fba7c3645ddafdd0759974b540441510c629c6db6db24a", size = 4456058, upload-time = "2026-05-06T06:17:40.735Z" }, - { url = "https://files.pythonhosted.org/packages/60/fb/3ae8bf2a7a37a4197d0195d7247fd25b3952e15cb8a599e285dfaa6f52b3/hf_xet-1.5.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e5de0f6deada0dada870bb376a11bcd1f08abf3a968a6d118f33e72d1b1eb480", size = 4250783, upload-time = "2026-05-06T06:17:38.412Z" }, - { url = "https://files.pythonhosted.org/packages/a2/9b/8bae40d4d91525085137196e84eb0ed49cf65b5e96e5c3ecdadd8bd0fac2/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c799d49f1a5544a0ef7591c0ee75e0d6b93d6f56dc7a4979f59f7518d2872216", size = 4445594, upload-time = "2026-05-06T06:18:04.219Z" }, - { url = "https://files.pythonhosted.org/packages/13/59/c74efbbd4e8728172b2cc72a2bc014d2947a4b7bdced932fbd3f5da1a4e5/hf_xet-1.5.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2baea1b0b989e5c152fe81425f7745ddc8901280ba3d97c98d8cdece7b706c60", size = 4663995, upload-time = "2026-05-06T06:18:06.1Z" }, - { url = "https://files.pythonhosted.org/packages/73/32/8e1e0410af64cda9b139d1dcebdc993a8ff9c8c7c0e2696ae356d75ccc0d/hf_xet-1.5.0-cp313-cp313t-win_amd64.whl", hash = "sha256:526345b3ed45f374f6317349df489167606736c876241ba984105afe7fd4839d", size = 3966608, upload-time = "2026-05-06T06:18:19.74Z" }, - { url = "https://files.pythonhosted.org/packages/fc/34/a8febc8f4edbea8b3e21b02ebc8b628679b84ba7e45cde624a7736b51500/hf_xet-1.5.0-cp313-cp313t-win_arm64.whl", hash = "sha256:786d28e2eb8315d5035544b9d137b4a842d600c434bb91bf7d0d953cce906ad4", size = 3796946, upload-time = "2026-05-06T06:18:17.568Z" }, - { url = "https://files.pythonhosted.org/packages/2a/20/8fc8996afe5815fa1a6be8e9e5c02f24500f409d599e905800d498a4e14d/hf_xet-1.5.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:872d5601e6deea30d15865ede55d29eac6daf5a534ab417b99b6ef6b076dd96c", size = 4023495, upload-time = "2026-05-06T06:18:01.94Z" }, - { url = "https://files.pythonhosted.org/packages/32/6a/93d84463c00cecb561a7508aa6303e35ee2894294eac14245526924415fe/hf_xet-1.5.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:9929561f5abf4581c8ea79587881dfef6b8abb2a0d8a51915936fc2a614f4e73", size = 3792731, upload-time = "2026-05-06T06:18:00.021Z" }, - { url = "https://files.pythonhosted.org/packages/9d/5a/8ec8e0c863b382d00b3c2e2af6ded6b06371be617144a625903a6d562f4b/hf_xet-1.5.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f7b7bbae318e583a86fb21e5a4a175d6721d628a2874f4bd022d0e660c32a682", size = 4456738, upload-time = "2026-05-06T06:17:49.574Z" }, - { url = "https://files.pythonhosted.org/packages/c5/ca/f7effa1a67717da2bcc6b6c28f71c6ca648c77acaec4e2c32f40cbe16d85/hf_xet-1.5.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:cf7b2dc6f31a4ea754bb50f74cde482dcf5d366d184076d8530b9872787f3761", size = 4251622, upload-time = "2026-05-06T06:17:47.096Z" }, - { url = "https://files.pythonhosted.org/packages/65/f2/19247dba3e231cf77dec59ddfb878f00057635ff773d099c9b59d37812c3/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8dbcbab554c9ef158ef2c991545c3e970ddd8cc7acdcd0a78c5a41095dab4ded", size = 4445667, upload-time = "2026-05-06T06:18:11.983Z" }, - { url = "https://files.pythonhosted.org/packages/7f/64/6f116801a3bcfb6f59f5c251f48cadc47ea54026441c4a385079286a94fa/hf_xet-1.5.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5906bf7718d3636dc13402914736abe723492cb730f744834f5f5b67d3a12702", size = 4664619, upload-time = "2026-05-06T06:18:13.771Z" }, - { url = "https://files.pythonhosted.org/packages/5c/e8/069542d37946ed08669b127e1496fa99e78196d71de8d41eda5e9f1b7a58/hf_xet-1.5.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5f3dc2248fc01cc0a00cd392ab497f1ca373fcbc7e3f2da1f452480b384e839e", size = 3966802, upload-time = "2026-05-06T06:18:28.162Z" }, - { url = "https://files.pythonhosted.org/packages/f9/91/fc6fdec27b14d04e88c386ac0a0129732b53fa23f7c4a78f4b83a039c567/hf_xet-1.5.0-cp314-cp314t-win_arm64.whl", hash = "sha256:b285cea1b5bab46b758772716ba8d6854a1a0310fed1c249d678a8b38601e5a0", size = 3797168, upload-time = "2026-05-06T06:18:26.287Z" }, - { url = "https://files.pythonhosted.org/packages/3d/fb/69ff198a82cae7eb1a69fb84d93b3a3e4816564d76817fe541ddc96874eb/hf_xet-1.5.0-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:dad0dc84e941b8ba3c860659fe1fdc35c049d47cce293f003287757e971a8f56", size = 4030814, upload-time = "2026-05-06T06:17:57.933Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ff/edcc2b40162bef3ff78e14ab637e5f3b89243d6aee72f5949d3bb6a5af83/hf_xet-1.5.0-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:fd6e5a9b0fdac4ed03ed45ef79254a655b1aaab514a02202617fbf643f5fdf7a", size = 3798444, upload-time = "2026-05-06T06:17:55.79Z" }, - { url = "https://files.pythonhosted.org/packages/49/4d/103f76b04310e5e57656696cc184690d20c466af0bca3ca88f8c8ea5d4f3/hf_xet-1.5.0-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3531b1823a0e6d77d80f9ed15ca0e00f0d115094f8ac033d5cae88f4564cc949", size = 4465986, upload-time = "2026-05-06T06:17:44.886Z" }, - { url = "https://files.pythonhosted.org/packages/c4/a2/546f47f464737b3edbab6f8ddb57f2599b93d2cbb66f06abb475ccb48651/hf_xet-1.5.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9a0ee58cd18d5ea799f7ed11290bbccbe56bdd8b1d97ca74b9cc49a3945d7a3b", size = 4259865, upload-time = "2026-05-06T06:17:42.639Z" }, - { url = "https://files.pythonhosted.org/packages/95/7f/1be593c1f28613be2e196473481cd81bfc5910795e30a34e8f744f6cac4f/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1e60df5a42e9bed8628b6416af2cba4cba57ae9f02de226a06b020d98e1aab18", size = 4459835, upload-time = "2026-05-06T06:18:08.026Z" }, - { url = "https://files.pythonhosted.org/packages/aa/b2/703569fc881f3284487e68cda7b42179978480da3c438042a6bbbb4a671c/hf_xet-1.5.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4b35549ce62601b84da4ff9b24d970032ace3d4430f52d91bcbb26c901d6c690", size = 4672414, upload-time = "2026-05-06T06:18:09.864Z" }, - { url = "https://files.pythonhosted.org/packages/af/37/1b6def445c567286b50aa3b33828158e135b1be44938dde59f11382a500c/hf_xet-1.5.0-cp37-abi3-win_amd64.whl", hash = "sha256:2806c7c17b4d23f8d88f7c4814f838c3b6150773fe339c20af23e1cfaf2797e4", size = 3977238, upload-time = "2026-05-06T06:18:23.621Z" }, - { url = "https://files.pythonhosted.org/packages/62/94/3b66b148778ee100dcfd69c2ca22b57b41b44d3063ceec934f209e9184ce/hf_xet-1.5.0-cp37-abi3-win_arm64.whl", hash = "sha256:b6c9df403040248c76d808d3e047d64db2d923bae593eb244c41e425cf6cd7be", size = 3806916, upload-time = "2026-05-06T06:18:21.7Z" }, +version = "1.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4b/2d/57fd21d84d93efb4bd0b962383790e19dd1bc053501b4264c97903b4e83e/hf_xet-1.5.1.tar.gz", hash = "sha256:51ef4500dab3764b41135ee1381a4b62ce56fc54d4c92b719b59e597d6df5bf6", size = 876636, upload-time = "2026-06-08T23:02:53.897Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/ee/dd9ba7beae1005e54131b7d45263cc74c8a066d47d354e6d58ae9445a388/hf_xet-1.5.1-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:dbf48c0d02cf0b2e568944330c60d9120c272dabe013bd892d48e25bc6797577", size = 4069485, upload-time = "2026-06-08T23:02:13.193Z" }, + { url = "https://files.pythonhosted.org/packages/b6/bc/9cae6cfeb4e03070874e73e5c97c66eb90369d3206b6a2b1ef5f96520888/hf_xet-1.5.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:e78e4e5192ad2b674c2e1160b651cb9134db974f8ae1835bdfbfb0166b894a43", size = 3838493, upload-time = "2026-06-08T23:02:15.282Z" }, + { url = "https://files.pythonhosted.org/packages/ba/b4/d5c01e0eb6d9f2ca2dacd84d0d1b71e6cfbb2ef3208c968528e010e9b3d7/hf_xet-1.5.1-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6f7a04a8ad962422e225bc49fbbac99dc1806764b1f3e54dbd154bffa7593947", size = 4505658, upload-time = "2026-06-08T23:02:17.196Z" }, + { url = "https://files.pythonhosted.org/packages/76/c5/29a7598c0c6383c523dc22186d577f4e04267a626cd95ae60f67c00bfe66/hf_xet-1.5.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d48199c2bf4f8df0adc55d31d1368b6ec0e4d4f45bc86b08038089c23db0bed8", size = 4292822, upload-time = "2026-06-08T23:02:18.608Z" }, + { url = "https://files.pythonhosted.org/packages/04/9a/dceaf6ca69390126b86ea825fb354b93d01163199070b7bd849225de9468/hf_xet-1.5.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:97f212a88d14bbf573619a74b7fecb238de77d08fc702e54dec6f78276ca3283", size = 4491255, upload-time = "2026-06-08T23:02:20.124Z" }, + { url = "https://files.pythonhosted.org/packages/48/a7/e5a7afaacf6c1791fdbeeac42951fb81c3d2bc482992b115dedcc86d963e/hf_xet-1.5.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:f61e3665892a6c8c5e765395838b8ddf36185da835253d4bc4509a81e49fb342", size = 4711062, upload-time = "2026-06-08T23:02:21.863Z" }, + { url = "https://files.pythonhosted.org/packages/53/49/2802f8433c9742ce281bddc1e65c02c32268ca3098d66828b05e12e45ee2/hf_xet-1.5.1-cp313-cp313t-win_amd64.whl", hash = "sha256:f4ad3ebd4c32dd2b27099d69dc7b2df821e30767e46fb6ee6a0713778243b8ff", size = 4017205, upload-time = "2026-06-08T23:02:23.495Z" }, + { url = "https://files.pythonhosted.org/packages/9e/5a/50c71195b9fb883659f596e7252faf4c18c58e753a9013bdbf9bac5d2250/hf_xet-1.5.1-cp313-cp313t-win_arm64.whl", hash = "sha256:8298485c1e36e7e67cbd01eeb1376619b7af43d4f1ec245caae306f890a8a32d", size = 3845426, upload-time = "2026-06-08T23:02:25.124Z" }, + { url = "https://files.pythonhosted.org/packages/05/24/5e0c28f80371c17d49fed004597d9d132cb75c1f6f53db2cb95f459d2312/hf_xet-1.5.1-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:3474760d10e3bb6f92ff3f024fcb00c0b3e4001e9b035c7483e49a5dd17aa70f", size = 4069676, upload-time = "2026-06-08T23:02:26.759Z" }, + { url = "https://files.pythonhosted.org/packages/d2/17/261ba565b6a4d960fb478f61fdf919c0be5824645aaf1c319eca660c1611/hf_xet-1.5.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6762d89b9e3267dfd502b29b2a327b4525f33b17e7b509a78d94e2151a30ce30", size = 3838509, upload-time = "2026-06-08T23:02:28.573Z" }, + { url = "https://files.pythonhosted.org/packages/4e/44/7ffdc2e184b0d41fc0f683ba3936ef669ab63cf242cf36ef50e57d683668/hf_xet-1.5.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bf67e6ed10260cef62e852789dc91ebb03f382d5bdc4b1dbeb64763ea275e7d6", size = 4505881, upload-time = "2026-06-08T23:02:30.257Z" }, + { url = "https://files.pythonhosted.org/packages/63/b6/788060d5aa4d5e671f1a31bf69624c314eb2d8babab3aa562f9e5d53444e/hf_xet-1.5.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:c6b6cd08ca095058780b50b8ce4d6cbf6787bcf27841705d58a9d32246e3e47a", size = 4292995, upload-time = "2026-06-08T23:02:31.993Z" }, + { url = "https://files.pythonhosted.org/packages/22/93/c5540cbd6b55529b7dc42f6734e88cebee21aefbea34128b66229df56c57/hf_xet-1.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e1af0de8ca6f190d4294a28b88023db64a1e2d1d719cab044baf75bec569e7a9", size = 4491570, upload-time = "2026-06-08T23:02:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/03/f3/9d8ceab30f44f36c1679b1b8683054c71a0dadc787dbf07421891742d3ca/hf_xet-1.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:4f561cbbb92f80960772059864b7fb07eae879adde1b2e781ec6f86f6ac26c59", size = 4711565, upload-time = "2026-06-08T23:02:35.454Z" }, + { url = "https://files.pythonhosted.org/packages/cd/54/27ed9a5e2cc583b4df82f75a03a4df8dbf55f5a9fa1f47f1fadfb20dbeac/hf_xet-1.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:e7dbb40617410f432182d918e37c12303fe6700fd6aa6c5964e30a535a4461d6", size = 4017343, upload-time = "2026-06-08T23:02:37.14Z" }, + { url = "https://files.pythonhosted.org/packages/ae/12/ecb2fc8d45e767580e3a37faa97cb895608b614965567efb4f18cff67e27/hf_xet-1.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:6071d5ccb4d8d2cbd5fea5cc798da4f0ba3f44e25369591c4e89a4987050e61d", size = 3845716, upload-time = "2026-06-08T23:02:39.073Z" }, + { url = "https://files.pythonhosted.org/packages/7a/d8/5e54cf37434759d1f4f2ba9b66077ff9d4c4e1f37b6bd7975da5c40d94ab/hf_xet-1.5.1-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:6abd35c3221eff63836618ddfb954dcf84798603f71d8e33e3ed7b04acfdbe6e", size = 4077794, upload-time = "2026-06-08T23:02:40.656Z" }, + { url = "https://files.pythonhosted.org/packages/35/94/4b2ecfbad8f8b04701a23aefb62f540b9137d058b7e1dbef16a32676f0e9/hf_xet-1.5.1-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:94e761bbd266bf4c03cee73753916062665ce8365aa40ed321f45afcb934b41e", size = 3845354, upload-time = "2026-06-08T23:02:42.702Z" }, + { url = "https://files.pythonhosted.org/packages/de/cc/f99f4bc7295023d7bd9ebbfd51f75cc530ca262c1227666268b8208f4b77/hf_xet-1.5.1-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:892e3a3a3aecc12aded8b93cf4f9cd059282c7de0732f7d55026f3abdf474350", size = 4514864, upload-time = "2026-06-08T23:02:44.497Z" }, + { url = "https://files.pythonhosted.org/packages/cd/6e/21f7e5a2381278bd3b7b7a5a4d90038518bb6308a0c1daf5d9f8268bb178/hf_xet-1.5.1-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a93df2039190502835b1db8cd7e178b0b7b889fe9ab51299d5ced26e0dd879a4", size = 4303784, upload-time = "2026-06-08T23:02:46.203Z" }, + { url = "https://files.pythonhosted.org/packages/35/0e/f992bb6927ac1cb30ef74e62268f551f338bc32b2191f7c96a44c6f7283e/hf_xet-1.5.1-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0c97106032ef70467b4f6bc2d0ccc266d7613ee076afc56516c502f87ce1c4a6", size = 4500703, upload-time = "2026-06-08T23:02:47.628Z" }, + { url = "https://files.pythonhosted.org/packages/fb/d1/90a498d05447980b977b1669246eeeeae4cfb0ea3e7a286eaba627f91bf9/hf_xet-1.5.1-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6208adb15d192b90e4c2ad2a27ed864359b2cb0f2494eb6d7c7f3699ac02e2bf", size = 4719498, upload-time = "2026-06-08T23:02:49.268Z" }, + { url = "https://files.pythonhosted.org/packages/6d/b6/20f99cfe97cc663a711f7b33cc21d4793e51968e9a26125b4afcd77315ba/hf_xet-1.5.1-cp37-abi3-win_amd64.whl", hash = "sha256:f7b3002f95d1c13e24bcb4537baa8f0eb3838957067c91bb4959bc004a6435f5", size = 4026419, upload-time = "2026-06-08T23:02:50.829Z" }, + { url = "https://files.pythonhosted.org/packages/f9/fa/77453694888f03e5a8c8852d1514a0894d8e81c622d39edbaf308ea0dcf4/hf_xet-1.5.1-cp37-abi3-win_arm64.whl", hash = "sha256:93d090b57b211133f6c0dab0205ef5cb6d89162979ba75a74845045cc3063b8e", size = 3855178, upload-time = "2026-06-08T23:02:52.452Z" }, ] [[package]] @@ -903,7 +903,7 @@ http2 = [ [[package]] name = "huggingface-hub" -version = "1.17.0" +version = "1.22.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -914,12 +914,11 @@ dependencies = [ { name = "packaging" }, { name = "pyyaml" }, { name = "tqdm" }, - { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bd/65/9826515abb600b5722bcf53f8b4a2fb58340b1f8bfcaee19f83561c13a44/huggingface_hub-1.17.0.tar.gz", hash = "sha256:fad842b6763ef70ebc3919665b1b9273645203185400a7d6c5eddc2323cc3435", size = 797082, upload-time = "2026-05-28T15:12:13.347Z" } +sdist = { url = "https://files.pythonhosted.org/packages/77/ea/dc54b4dda5841cb3a7812a178695be776e7c15c597887c2ed892f17d015a/huggingface_hub-1.22.0.tar.gz", hash = "sha256:e2dfe5fe1ec3b87ba2709aa34555b23e3f3f6ad4d7255238e13ddb8348e6bbfa", size = 914232, upload-time = "2026-07-03T09:46:44.685Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/02/28/d7cef5e477b855c25d415b8f57e5bc7347c7a90cad3acf1725d0c92ca294/huggingface_hub-1.17.0-py3-none-any.whl", hash = "sha256:3b8156d23118e87f6a587648bfbc04f04a12a757ccb4ed298b35c4ae638bf24c", size = 671546, upload-time = "2026-05-28T15:12:11.441Z" }, + { url = "https://files.pythonhosted.org/packages/64/9c/a1a377265abd8b823a2c661c665028ccb6b9fba1ca9d08e52ff679c20ecd/huggingface_hub-1.22.0-py3-none-any.whl", hash = "sha256:b09e19309ae09ee0a71892701c4fe70af39ab4e00817321dc62f2289a977249b", size = 765085, upload-time = "2026-07-03T09:46:42.832Z" }, ] [[package]] @@ -1113,7 +1112,7 @@ wheels = [ [[package]] name = "llama-index" -version = "0.14.22" +version = "0.14.23" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "llama-index-core" }, @@ -1121,14 +1120,14 @@ dependencies = [ { name = "llama-index-llms-openai" }, { name = "nltk" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/28/89/3b6f3318ea2249158daab3ff22777ef5ffa87a63c011659a6cfc55e54c35/llama_index-0.14.22.tar.gz", hash = "sha256:c2c9b31f50d2815abdc191085db4acaf96b7c01851ac66b2e4cc82be8cde589e", size = 8565, upload-time = "2026-05-14T20:22:21.006Z" } +sdist = { url = "https://files.pythonhosted.org/packages/cb/5c/1f9334c3c0edd880367d0734939256f8536ca94165c6a5fdb455a7bcf180/llama_index-0.14.23.tar.gz", hash = "sha256:eac2049816a7410ff4568490cce4bdff99cda3ab99d59f52f6227dad22cda44b", size = 8566, upload-time = "2026-06-24T19:36:38.241Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/59/fd/f0837c4ce049d8ece7525bbf64564e93e3f16333856c2a0b47fecb58f317/llama_index-0.14.22-py3-none-any.whl", hash = "sha256:14b4bdd799112062e38288eab6aa16643f29d7532505ab174b0b6d5b0817fe94", size = 7115, upload-time = "2026-05-14T20:22:19.611Z" }, + { url = "https://files.pythonhosted.org/packages/e0/92/da2a737bf712fe31c1aaec00f51745797d4570ecdf7e4f5d2b4b93b3d337/llama_index-0.14.23-py3-none-any.whl", hash = "sha256:c205de2442a7186b8e05096f0771f96fa6bdc9603fcad2e07eab1bc96dcf086a", size = 7114, upload-time = "2026-06-24T19:36:37.21Z" }, ] [[package]] name = "llama-index-core" -version = "0.14.22" +version = "0.14.23" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -1160,9 +1159,9 @@ dependencies = [ { name = "typing-inspect" }, { name = "wrapt" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/96/7f/94a4b940ef0d069840df0fd6d361a2aa832a2dd73b4cecdf86e8f8c353c8/llama_index_core-0.14.22.tar.gz", hash = "sha256:1384410f89bdbd32349aab444ef4f5c828c338787bc65bd1ffd8e86dfb44ac41", size = 11584786, upload-time = "2026-05-14T20:21:37.271Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/ac/f885ae14317af43a026c909ea4d2083fcee2f0d014f90426b5b9aa1f9912/llama_index_core-0.14.23.tar.gz", hash = "sha256:c4baf2f2ab4f84e95090fe7941e0c87d6c514304f7bd2a749b8fa22164c1822b", size = 11588373, upload-time = "2026-06-24T19:35:55.43Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/39/15/e1a26d8d56aa55fa07587a3e9c7e85294d2df5af6c2229193019bc549ef6/llama_index_core-0.14.22-py3-none-any.whl", hash = "sha256:9cfffde46fd5b7937101e1c0c9bb5c21bd7ff8c8a56937810b87ba3542f31225", size = 11920774, upload-time = "2026-05-14T20:21:40.409Z" }, + { url = "https://files.pythonhosted.org/packages/ec/d5/05d61f34c01c6578fb758d0a3ddef58d36c6ffa9a9f84a5c9a16262ad94d/llama_index_core-0.14.23-py3-none-any.whl", hash = "sha256:6a54d267826732a8507f81df40785b107f7592af20f451a39a59005147caf84c", size = 11924908, upload-time = "2026-06-24T19:35:52.833Z" }, ] [[package]] @@ -1250,16 +1249,16 @@ wheels = [ [[package]] name = "llama-index-vector-stores-qdrant" -version = "0.10.1" +version = "0.10.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, { name = "llama-index-core" }, { name = "qdrant-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ec/9c/7df54b10e6f7d7b3fde3853d25f0c0ec6e701e9dea35359d316e5ea6d942/llama_index_vector_stores_qdrant-0.10.1.tar.gz", hash = "sha256:fef4ca8411c3e33636aabcf883941fde7a9d6deaa45254ee80627f2d0ffbf551", size = 14731, upload-time = "2026-05-04T15:17:50.368Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/9f/6b335df0161d46e83894513d4f07f0617a45e0a4418ebecdbd53dd5477b9/llama_index_vector_stores_qdrant-0.10.2.tar.gz", hash = "sha256:3122b644901c7b58e616fd9e7ed4fd1ec2604c63f1b85c5d6ad44820af329be2", size = 14721, upload-time = "2026-07-02T15:02:33.807Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ab/9d/32ac29c621db5e635cd248ef824cf49ea934c52085d19c409d6aea4d4167/llama_index_vector_stores_qdrant-0.10.1-py3-none-any.whl", hash = "sha256:90e8c0c0cd96309a9dd84a29cfe77aa6b3ef2f2596f0a0360eaaa8e3a004f1ff", size = 15000, upload-time = "2026-05-04T15:17:51.117Z" }, + { url = "https://files.pythonhosted.org/packages/75/d6/d50b4b2cfddacec12e1eb29c6884ff86c3017eab29ac7f8c13a4f5cbb3f3/llama_index_vector_stores_qdrant-0.10.2-py3-none-any.whl", hash = "sha256:51070d47d3374860e8072bfa7f5b079222a58a6f1e1d78443e2babc8cab6d0c9", size = 14998, upload-time = "2026-07-02T15:02:32.799Z" }, ] [[package]] @@ -2254,18 +2253,19 @@ wheels = [ [[package]] name = "pymupdf" -version = "1.27.2.3" +version = "1.28.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/22/32/708bedc9dde7b328d45abbc076091769d44f2f24ad151ad92d56a6ec142b/pymupdf-1.27.2.3.tar.gz", hash = "sha256:7a92faa25129e8bbec5e50eeb9214f187665428c31b05c4ef6e36c58c0b1c6d2", size = 85759618, upload-time = "2026-04-24T14:13:14.42Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8e/e9/6d6c5d6c0a3551bffd47681a6240caf941727f195b45593cf20ab36f018f/pymupdf-1.28.0.tar.gz", hash = "sha256:e53f3567403a92da15caa9e7ae0164327fff48817e9f40175367fb9de524258d", size = 87637751, upload-time = "2026-06-29T09:08:47.547Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/09/ddbdfa7ee91fbabd6f63d7d744884cbdfe3e7ff9b8604749fb38bddf5c5d/pymupdf-1.27.2.3-cp310-abi3-macosx_10_9_x86_64.whl", hash = "sha256:fc1bc3cae6e9e150b0dbb0a9221bdfd411d65f0db2fe359eaa22467d7cc2a05f", size = 24002636, upload-time = "2026-04-24T14:09:17.459Z" }, - { url = "https://files.pythonhosted.org/packages/01/89/3f8edd6c4f50ca370e2a2f2a3011face36f3760728ffe76dffec91c0fca0/pymupdf-1.27.2.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:660d93cb6da5bbddf11d3982ae27745dd3a9902d9f24cdb69adab83962294b5a", size = 23278238, upload-time = "2026-04-24T14:09:32.882Z" }, - { url = "https://files.pythonhosted.org/packages/c3/26/b7e5a70eb83bd189f8b5df87ec442746b992f2f632662839b288170d357d/pymupdf-1.27.2.3-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:1dd460a3ae4597a755f00a3bd9771f5ebf1531dc111f6a36bf05dd00a6b84425", size = 24333923, upload-time = "2026-04-24T14:09:47.341Z" }, - { url = "https://files.pythonhosted.org/packages/e4/a0/aa1ee2240f29481a04a827c313333b4ecd8a14d6ac3e15d3f41a30574781/pymupdf-1.27.2.3-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:857842b4888827bd6155a1131341b2822a7ebe9a8c15a975fd7d490d7a64a30c", size = 24963198, upload-time = "2026-04-24T14:10:07.408Z" }, - { url = "https://files.pythonhosted.org/packages/69/49/4f742451f980840829fc00ba158bebb25d389c846d8f4f8c65936ee55de8/pymupdf-1.27.2.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:580983849c64a08d08344ca3d1580e87c01f046a8392421797bc850efd72a5b6", size = 25184609, upload-time = "2026-04-24T14:10:22.911Z" }, - { url = "https://files.pythonhosted.org/packages/f6/3f/3853d6608f394faf6eec2bd4e8ea9f6a00beea329b071abdb29f4164cc3d/pymupdf-1.27.2.3-cp310-abi3-win32.whl", hash = "sha256:a5c1088a87189891a4946ab314a14b7934ac4c5b6077f7e74ebee956f8906d0e", size = 18019286, upload-time = "2026-04-24T14:10:34.239Z" }, - { url = "https://files.pythonhosted.org/packages/44/47/5fb10fe73f96b31253a41647c362ea9e0380920bddf16028414a051247fc/pymupdf-1.27.2.3-cp310-abi3-win_amd64.whl", hash = "sha256:d20f68ef15195e073071dbc4ae7455257c7889af7584e39df490c0a92728526e", size = 19249102, upload-time = "2026-04-24T14:10:46.72Z" }, - { url = "https://files.pythonhosted.org/packages/53/a4/b9e91aac82293f9c954654c85581ee8212b5b05efadc534b581141241e6f/pymupdf-1.27.2.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:77691604c5d1d0233827139bbcdea61fd57879c84712b8e49b1f45520f7ab9c2", size = 25000393, upload-time = "2026-04-24T14:11:01.669Z" }, + { url = "https://files.pythonhosted.org/packages/c8/b7/88043e38cc7529de070f0c9bd267fa258035cca0b4ad5260536b994594a7/pymupdf-1.28.0-cp310-abi3-macosx_10_15_x86_64.whl", hash = "sha256:892b89ba88e8f98b53133b62877a9dc9b5e7dc6a4aeb837b612db56a8d2e03ac", size = 24597385, upload-time = "2026-06-29T09:03:30.608Z" }, + { url = "https://files.pythonhosted.org/packages/33/f4/23775bbda0781b61fc398cc75079a2b0e64696d8fcf93271748883e9627e/pymupdf-1.28.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:4d692dcf44d3566ae96bc6f6346c6ad432274a29ba617bf7a9fe18009e24adb4", size = 23828292, upload-time = "2026-06-29T09:03:46.129Z" }, + { url = "https://files.pythonhosted.org/packages/1c/f5/bf75fc7a415722f8b33662054f82d88520c0cbfd4c36d0e08aeaec605e49/pymupdf-1.28.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:47a5c29ed4eb0744de9c4e37bb49b1259b18d4d75fcc8a7c130f7c9fa15956f6", size = 25045507, upload-time = "2026-06-29T09:04:03.86Z" }, + { url = "https://files.pythonhosted.org/packages/58/69/5d12c9f1f2d76f28383d6110a069c79fbfced5a4f97bb1ee6e8354f52bb7/pymupdf-1.28.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:44f0973f5e5edbaec95bc34b64e71d1959d4ee90b1328de1b4f4f5b4fa78673f", size = 25716599, upload-time = "2026-06-29T09:04:19.367Z" }, + { url = "https://files.pythonhosted.org/packages/4d/b4/ec0e017bc42857cc86bd651441dbc41cc18be48d4698ecd27aac491e0c9a/pymupdf-1.28.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4d61ec323a706e153a12e262e51febfb43eeaa20977785ace135d18d48bcdc83", size = 25940489, upload-time = "2026-06-29T09:04:36.624Z" }, + { url = "https://files.pythonhosted.org/packages/06/86/f831fef09013f33b3c9c09fb3923f2ff53e1e437f6ace14b8ae46392f558/pymupdf-1.28.0-cp310-abi3-win32.whl", hash = "sha256:caea2b3b67347fd79e5d15ed7929b0e886aac594ea228073b6d39de0078189da", size = 18489703, upload-time = "2026-06-29T20:50:30.599Z" }, + { url = "https://files.pythonhosted.org/packages/2e/5d/1a03f53eb0449900469335fcfc742ca28e3ba159b7d650e0921d50b8b308/pymupdf-1.28.0-cp310-abi3-win_amd64.whl", hash = "sha256:e01e90fd86abfeb37ceb921eddb951f988a11d45ff6ce6b7664f2039849068ec", size = 19773102, upload-time = "2026-06-29T09:04:49.773Z" }, + { url = "https://files.pythonhosted.org/packages/72/f6/1e52ce243ca792254f6223b4017c5667194c146ce9b88baf37bc5eb3d1c9/pymupdf-1.28.0-cp313-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:74c6d00ba2a9aad3a635db73b07c15db462b480741d831a34a75a56535ebc22b", size = 18357011, upload-time = "2026-06-29T20:50:50.353Z" }, + { url = "https://files.pythonhosted.org/packages/62/b1/46b5b3d8ef3cc71114667cf10c4d8b33f39af97253af32e9a0986775b638/pymupdf-1.28.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:b3e1399c7a64c6914239116a369efcdaac4cfb9e838bde2656d7accc4a85c72d", size = 25753599, upload-time = "2026-06-29T09:05:09.398Z" }, ] [[package]] @@ -2347,11 +2347,11 @@ wheels = [ [[package]] name = "python-multipart" -version = "0.0.30" +version = "0.0.32" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4b/82/c8cd43a6e0719bf5a3b034f6726dd701f75829c08944c83d4b95d02ed0e8/python_multipart-0.0.30.tar.gz", hash = "sha256:0edfe0475c1f46ddd3ff7785a626f6118af32bdcf359bb21260367313bb32118", size = 46316, upload-time = "2026-05-31T19:24:55.198Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5b/42/55c32bb9b12693c092ad250a0e82edb5b31ddeda6eb772de5f308b3804ad/python_multipart-0.0.32.tar.gz", hash = "sha256:be54b7f3fa167bb83e4fcd936b887b708f4e57fe75911c02aebf53efaf8d938e", size = 46881, upload-time = "2026-06-04T16:18:58.647Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1c/fd/0318007beb234790993d3ec5afd051d1dbceb733e81e3afe2b981ece3f37/python_multipart-0.0.30-py3-none-any.whl", hash = "sha256:830964def8c90607ac5daa00514e3987815865713ade8d20febc9177ac0c3c5b", size = 29730, upload-time = "2026-05-31T19:24:53.814Z" }, + { url = "https://files.pythonhosted.org/packages/e1/04/e8135ebd1ad02c56ec633277529b2602ff99ff634be76cdba5744cf554fd/python_multipart-0.0.32-py3-none-any.whl", hash = "sha256:ff6d3f776f16878c894e52e107296ffc890e913c611b1a4ec6c44e2821fe2e23", size = 30042, upload-time = "2026-06-04T16:18:57.319Z" }, ] [[package]] @@ -2442,7 +2442,7 @@ wheels = [ [[package]] name = "qdrant-client" -version = "1.17.1" +version = "1.18.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "grpcio" }, @@ -2453,9 +2453,9 @@ dependencies = [ { name = "pydantic" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/30/dd/f8a8261b83946af3cd65943c93c4f83e044f01184e8525404989d22a81a5/qdrant_client-1.17.1.tar.gz", hash = "sha256:22f990bbd63485ed97ba551a4c498181fcb723f71dcab5d6e4e43fe1050a2bc0", size = 344979, upload-time = "2026-03-13T17:13:44.678Z" } +sdist = { url = "https://files.pythonhosted.org/packages/65/45/5b1bdd15a3c7730eefb9c113600829e20d689b82b5a23f9e07d107094004/qdrant_client-1.18.0.tar.gz", hash = "sha256:52e8ece1a7d40519801bf0b70713bfa0f6b7ae28c7275bbe0b0286fbed7f6db4", size = 352580, upload-time = "2026-05-11T14:12:38.702Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/68/69/77d1a971c4b933e8c79403e99bcbb790463da5e48333cc4fd5d412c63c98/qdrant_client-1.17.1-py3-none-any.whl", hash = "sha256:6cda4064adfeaf211c751f3fbc00edbbdb499850918c7aff4855a9a759d56cbd", size = 389947, upload-time = "2026-03-13T17:13:43.156Z" }, + { url = "https://files.pythonhosted.org/packages/d6/10/c437bd2ac41ef30d3019063e6ce537dc111e9214473b337ee88f7fa6359a/qdrant_client-1.18.0-py3-none-any.whl", hash = "sha256:093aa8cf8a420ee3ad2a68b007e1378d7992b2600e0b53c193fc172674f659cd", size = 398126, upload-time = "2026-05-11T14:12:36.998Z" }, ] [[package]] @@ -2786,6 +2786,7 @@ dependencies = [ { name = "python-multipart" }, { name = "python-pptx" }, { name = "pyyaml" }, + { name = "qdrant-client" }, { name = "rich" }, { name = "tomli-w" }, { name = "uvicorn" }, @@ -2805,26 +2806,27 @@ dev = [ [package.metadata] requires-dist = [ { name = "docx2txt", specifier = ">=0.9" }, - { name = "duckdb", specifier = ">=1.5.3" }, - { name = "fastapi", specifier = ">=0.136.3" }, + { name = "duckdb", specifier = ">=1.5.4" }, + { name = "fastapi", specifier = ">=0.139.0" }, { name = "gitpython", specifier = ">=3.1.50" }, - { name = "hf-xet", specifier = ">=1.5.0" }, + { name = "hf-xet", specifier = ">=1.5.1" }, { name = "httpx", specifier = ">=0.28.1" }, - { name = "huggingface-hub", specifier = ">=1.17.0" }, + { name = "huggingface-hub", specifier = ">=1.22.0" }, { name = "jinja2", specifier = ">=3.1.6" }, - { name = "llama-index", specifier = ">=0.14.22" }, + { name = "llama-index", specifier = ">=0.14.23" }, { name = "llama-index-embeddings-huggingface", specifier = ">=0.7.0" }, { name = "llama-index-llms-openai-like", specifier = ">=0.7.2" }, { name = "llama-index-readers-file", specifier = ">=0.6.0" }, - { name = "llama-index-vector-stores-qdrant", specifier = ">=0.10.1" }, + { name = "llama-index-vector-stores-qdrant", specifier = ">=0.10.2" }, { name = "puremagic", specifier = ">=2.2.0" }, - { name = "pymupdf", specifier = ">=1.27.2.3" }, - { name = "python-multipart", specifier = ">=0.0.30" }, + { name = "pymupdf", specifier = ">=1.28.0" }, + { name = "python-multipart", specifier = ">=0.0.32" }, { name = "python-pptx", specifier = ">=1.0.2" }, { name = "pyyaml", specifier = ">=6.0.3" }, - { name = "rich", specifier = ">=14.0.0" }, + { name = "qdrant-client", specifier = ">=1.18.0" }, + { name = "rich", specifier = ">=15.0.0" }, { name = "tomli-w", specifier = ">=1.2.0" }, - { name = "uvicorn", specifier = ">=0.48.0" }, + { name = "uvicorn", specifier = ">=0.49.0" }, ] [package.metadata.requires-dev] @@ -3250,15 +3252,15 @@ wheels = [ [[package]] name = "uvicorn" -version = "0.48.0" +version = "0.49.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "h11" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e6/bf/f6544ba992ddb9a6077343a576f9844f7f8f06ab819aefd00206e9255f18/uvicorn-0.48.0.tar.gz", hash = "sha256:a5504207195d08c2511bf9125ede5ac4a4b71725d519e758d01dcf0bc2d31c37", size = 91074, upload-time = "2026-05-24T12:08:41.925Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/1f/fa18009dea8469069cca78a4e877a008ab78f08b064bfc9ab891579077ff/uvicorn-0.49.0.tar.gz", hash = "sha256:ebf4271aa580d9de97f93192d4595176df6e91f9aae919ca73e4fc07df1e66a3", size = 91284, upload-time = "2026-06-03T22:01:30.448Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/01/be/72532be3da7acc5fdfbccdb95215cd04f995a0886532a5b423f929cda4cc/uvicorn-0.48.0-py3-none-any.whl", hash = "sha256:48097851328b87ec36117d3d575234519eb58c2b22d79666e9bbc6c49a761dad", size = 71410, upload-time = "2026-05-24T12:08:40.258Z" }, + { url = "https://files.pythonhosted.org/packages/88/fa/e1388bbcf24ef3274f45c0c1c7b501fd14971037c1b6ee23610553307497/uvicorn-0.49.0-py3-none-any.whl", hash = "sha256:ba3d14c3ee7e41c6c654c46c9eb489d33213cdd30aa1696eab1374337c13f68f", size = 71376, upload-time = "2026-06-03T22:01:29.037Z" }, ] [[package]]