From 8d24148dbfe4be3e6fe529800f5516613a0217c6 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Thu, 2 Jul 2026 06:44:20 +0200 Subject: [PATCH 01/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20extract?= =?UTF-8?q?=20shared=20crypto/url=20utils,=20remove=20dead=20code,=20norma?= =?UTF-8?q?lize=20imports=20(Phase=201)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../documents/sigma_ref_downloader.py | 84 +++---------------- .../documents/sigma_ref_processor.py | 43 +++------- src/core/pipeline/__init__.py | 2 - src/core/pipeline/orchestrator.py | 21 +---- src/core/sigma/chunker.py | 14 ---- src/shared/utils/crypto_utils.py | 28 +++++++ src/shared/utils/url_utils.py | 55 ++++++++++++ src/workers/sigma/discovery_base.py | 22 +++-- src/workers/sigma/discovery_worker.py | 3 +- .../documents/test_sigma_ref_downloader.py | 7 +- 10 files changed, 122 insertions(+), 157 deletions(-) create mode 100644 src/shared/utils/crypto_utils.py create mode 100644 src/shared/utils/url_utils.py diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index f832ee3e..b386b3d2 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -2,9 +2,7 @@ from __future__ import annotations -import ipaddress import logging -import re import threading import time import urllib.parse @@ -16,34 +14,17 @@ import httpx import yaml +from src.shared.utils.crypto_utils import compute_sha256_file, compute_sha256_str from src.shared.utils.identify_file_type import ( SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, ) +from src.shared.utils.url_utils import is_private_url, normalize_url, url_ext from src.infrastructure.database import DatabaseService from src.core.sigma.models import is_sigma_rule_dict from src.shared.utils import iso_now 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} @@ -56,25 +37,6 @@ def _is_private_url(url: str) -> bool: _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. @@ -481,7 +443,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 +451,7 @@ def _collect_yaml_files() -> list[Path]: total_refs += 1 normalized = normalize_url(ref) - url_hash = _sha256(normalized) + url_hash = compute_sha256_str(normalized) if url_hash in error_registry: logger.debug("Skipping previously failed URL: %s", normalized) @@ -503,7 +465,10 @@ def _collect_yaml_files() -> list[Path]: output_file = output_path / 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, @@ -524,7 +489,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 = url_ext(normalized) ftype = _detect_url_type(normalized) if ftype is None and url_hash in registry: ct = registry[url_hash].get("content_type") @@ -560,7 +525,7 @@ def _collect_yaml_files() -> list[Path]: output_file = output_path / f"{url_hash}{ext}" 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: @@ -675,7 +640,7 @@ 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 "" + content_hash = compute_sha256_file(output_file) if output_file.exists() else "" registry[item["url_hash"]] = _make_entry( url_hash=item["url_hash"], original_url=item["original_url"], @@ -715,33 +680,6 @@ def _collect_yaml_files() -> list[Path]: 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 _empty_summary() -> dict[str, Any]: """Return an empty summary dict.""" return { diff --git a/src/application/documents/sigma_ref_processor.py b/src/application/documents/sigma_ref_processor.py index 83c66f5d..c2d73bbe 100644 --- a/src/application/documents/sigma_ref_processor.py +++ b/src/application/documents/sigma_ref_processor.py @@ -10,8 +10,10 @@ import httpx +from src.shared.utils.crypto_utils import compute_sha256_bytes from src.shared.utils.identify_file_type import SUPPORTED_REFERENCE_DOC_TYPES from src.shared.utils.sigma_utils import extract_sigma_references +from src.shared.utils.url_utils import is_private_url, normalize_url from src.config.settings import get_config from src.shared.utils import iso_now @@ -19,7 +21,6 @@ DEFAULT_REQUEST_DELAY = 0.5 DEFAULT_MAX_WORKERS = 5 -GITHUB_BLOB_PATTERN: Any = None # imported lazily to avoid circular imports def _build_head_entry( @@ -31,7 +32,7 @@ def _build_head_entry( ) -> dict[str, Any]: now = iso_now() return { - "url_hash": _sha256_bytes(normalized_url.encode()), + "url_hash": compute_sha256_bytes(normalized_url.encode()), "org": "sigmaref", "repo": "references", "content_type": content_type, @@ -60,7 +61,7 @@ def _build_download_entry( """Build a doc_registry entry for a downloaded reference document.""" now = iso_now() return { - "url_hash": _sha256_bytes(normalized_url.encode()), + "url_hash": compute_sha256_bytes(normalized_url.encode()), "org": "sigmaref", "repo": "references", "content_type": content_type, @@ -159,8 +160,8 @@ def process_sigma_refs( if not ref_url_clean: continue - norm_url = _normalize_url(ref_url_clean) - url_hash = _sha256_bytes(norm_url.encode()) + norm_url = normalize_url(ref_url_clean) + url_hash = compute_sha256_bytes(norm_url.encode()) # Deduplicate within a single run (same URL across multiple rules) if url_hash in seen_urls: @@ -218,8 +219,8 @@ def process_sigma_refs( 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()) + norm_url = normalize_url(final_url or url) + url_hash = compute_sha256_bytes(norm_url.encode()) if content_type not in supported_types: logger.info("Reference skipped (unsupported type): %s (%s)", url, content_type) @@ -273,7 +274,7 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, int] | None: file_path = output_path / _sanitize_filename(url) if file_path.exists(): - url_hash = _sha256_bytes(url.encode()) + url_hash = compute_sha256_bytes(url.encode()) existing = db.get_entry(url_hash) if existing and existing.get("content_sha256"): logger.info("Reference already present: %s", url) @@ -288,7 +289,7 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, int] | None: content = resp.content file_path.write_bytes(content) logger.info("Reference downloaded: %s", url) - return ("ok", _sha256_bytes(content), len(content)) + return ("ok", compute_sha256_bytes(content), len(content)) except Exception as e: logger.error("Reference download failed: %s - %s", url, e) return ("fail", "", 0) @@ -369,13 +370,11 @@ def _resolve_rule_path(entry: dict, cfg: Any) -> Path | None: 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.""" + if is_private_url(url): + logger.warning("Skipping private URL: %s", url) + return None, None, None try: with httpx.Client( timeout=httpx.Timeout(15.0), headers={"User-Agent": "SigmaRAG/1.0"} @@ -389,22 +388,6 @@ def _head_request(url: str, delay: float = 0.0) -> tuple[str | None, int | None, 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 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/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/sigma/chunker.py b/src/core/sigma/chunker.py index 8a3657ec..c00b4bb8 100644 --- a/src/core/sigma/chunker.py +++ b/src/core/sigma/chunker.py @@ -500,20 +500,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/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/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/sigma/discovery_base.py b/src/workers/sigma/discovery_base.py index 2de885e1..f38d8788 100644 --- a/src/workers/sigma/discovery_base.py +++ b/src/workers/sigma/discovery_base.py @@ -1,7 +1,7 @@ -import hashlib import logging from pathlib import Path +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 +14,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: @@ -44,7 +42,7 @@ def _make_doc_registry_entry( title: str, rule_id: str = "00000000-0000-0000-0000-000000000000", ) -> 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 +73,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..5a26cca8 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Optional 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 @@ -440,8 +441,6 @@ def _prepare_entry( rule_id = "00000000-0000-0000-0000-000000000000" 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 diff --git a/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 93db3e66..8f0f7b70 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -8,20 +8,19 @@ import httpx 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.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, ) From 24f1f77b8753896cf24c020008da888e5fd8abe1 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Thu, 2 Jul 2026 17:27:29 +0200 Subject: [PATCH 02/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20extract?= =?UTF-8?q?=20shared=20http=20utils=20and=20registry=20entry=20builder=20(?= =?UTF-8?q?P1.2=20+=20P1.3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - src/shared/http.py: create_client, head_url, download_file with retry/backoff - src/shared/utils/registry_utils.py: unified build_registry_entry() - sigma_ref_downloader: use shared http and registry builder (-173 lines) - sigma_ref_processor: use shared http and registry builder (-112 lines) - Tests: 28 tests for http utils, 8 tests for registry builder, 20 regression tests --- .../documents/sigma_ref_downloader.py | 173 +---------- .../documents/sigma_ref_processor.py | 112 +------ src/shared/http.py | 220 +++++++++++++ src/shared/utils/registry_utils.py | 81 +++++ .../documents/test_sigma_ref_downloader.py | 167 +--------- .../test_sigma_validator_regression.py | 184 +++++++++++ tests/unit/shared/test_http.py | 290 ++++++++++++++++++ .../unit/shared/utils/test_registry_utils.py | 102 ++++++ 8 files changed, 919 insertions(+), 410 deletions(-) create mode 100644 src/shared/http.py create mode 100644 src/shared/utils/registry_utils.py create mode 100644 tests/unit/application/services/test_sigma_validator_regression.py create mode 100644 tests/unit/shared/test_http.py create mode 100644 tests/unit/shared/utils/test_registry_utils.py diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index b386b3d2..d7927abb 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -4,16 +4,18 @@ import logging 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.http import RETRY_STATUSES +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_file, compute_sha256_str from src.shared.utils.identify_file_type import ( SUPPORTED_DOC_EXTENSION_MAP, @@ -25,9 +27,6 @@ from src.shared.utils import iso_now logger = logging.getLogger(__name__) -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] = { @@ -80,128 +79,6 @@ def _detect_url_type(url: str, content_type: str | None = None) -> str | None: 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. - - 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. - - 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. @@ -293,35 +170,6 @@ def _maybe_record_error( 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, - } - - def download_references( rules_dir: str, output_dir: str, @@ -542,14 +390,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, @@ -580,7 +427,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): @@ -588,7 +435,7 @@ 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: @@ -627,9 +474,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, item["normalized_url"], item["output_file"], + check_ssrf=False, ) future_map[future] = item @@ -641,14 +489,13 @@ def _collect_yaml_files() -> list[Path]: if success: output_file = item["output_file"] content_hash = compute_sha256_file(output_file) if output_file.exists() else "" - registry[item["url_hash"]] = _make_entry( + 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, diff --git a/src/application/documents/sigma_ref_processor.py b/src/application/documents/sigma_ref_processor.py index c2d73bbe..2d73323e 100644 --- a/src/application/documents/sigma_ref_processor.py +++ b/src/application/documents/sigma_ref_processor.py @@ -8,14 +8,14 @@ from pathlib import Path from typing import Any, cast -import httpx - +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 from src.shared.utils.identify_file_type import SUPPORTED_REFERENCE_DOC_TYPES from src.shared.utils.sigma_utils import extract_sigma_references -from src.shared.utils.url_utils import is_private_url, normalize_url +from src.shared.utils.url_utils import normalize_url from src.config.settings import get_config -from src.shared.utils import iso_now logger = logging.getLogger(__name__) @@ -23,61 +23,6 @@ DEFAULT_MAX_WORKERS = 5 -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": compute_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": compute_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, @@ -203,7 +148,7 @@ def process_sigma_refs( 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"]) + future = head_executor.submit(http_head_url, item["url"], 15.0) head_futures[future] = item for future in as_completed(head_futures): @@ -226,11 +171,12 @@ def process_sigma_refs( logger.info("Reference skipped (unsupported type): %s (%s)", url, content_type) db.batch_upsert_doc_registry( [ - _build_head_entry( + 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", ) ] ) @@ -240,12 +186,13 @@ def process_sigma_refs( # Cache HEAD result immediately db.batch_upsert_doc_registry( [ - _build_head_entry( + 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", ) ] ) @@ -270,8 +217,7 @@ def process_sigma_refs( # 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) + file_path = Path(output_dir) / _sanitize_filename(url) if file_path.exists(): url_hash = compute_sha256_bytes(url.encode()) @@ -280,19 +226,12 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, int] | None: 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) + ok, _ = http_download_file(url, file_path, check_ssrf=False) + if ok: + content = file_path.read_bytes() return ("ok", compute_sha256_bytes(content), len(content)) - except Exception as e: - logger.error("Reference download failed: %s - %s", url, e) - return ("fail", "", 0) + logger.error("Reference download failed: %s", url) + return ("fail", "", 0) with ThreadPoolExecutor(max_workers=max_workers) as executor: futures: dict[Future, dict[str, Any]] = { @@ -308,13 +247,14 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, int] | None: elif result[0] == "ok": final_url = item["final_url"] head_result = cast("tuple[str, str, int]", result) - entry = _build_download_entry( + entry = build_registry_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], + embed_status="discovery", ) db.batch_upsert_doc_registry([entry]) downloaded += 1 @@ -370,24 +310,6 @@ def _resolve_rule_path(entry: dict, cfg: Any) -> Path | None: return None -def _head_request(url: str, delay: float = 0.0) -> tuple[str | None, int | None, str | None]: - """HEAD request to resolve content type and size.""" - if is_private_url(url): - logger.warning("Skipping private URL: %s", url) - return None, None, None - 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 _sanitize_filename(url: str) -> str: import re diff --git a/src/shared/http.py b/src/shared/http.py new file mode 100644 index 00000000..5e88bc7f --- /dev/null +++ b/src/shared/http.py @@ -0,0 +1,220 @@ +"""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 time +from pathlib import Path + +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} + + +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 = {"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: + with create_client(timeout=timeout) as client: + 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) + + for attempt in range(1, max_retries + 1): + try: + with create_client(timeout=timeout) as client: + 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/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/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 8f0f7b70..69565a9d 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -3,9 +3,7 @@ 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 @@ -13,10 +11,7 @@ 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, _load_registry, _registry_lock, _save_registry, @@ -151,108 +146,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() @@ -360,7 +253,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( @@ -424,7 +317,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=[""]) @@ -463,7 +356,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=[""]) @@ -505,7 +398,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=[""]) @@ -533,7 +426,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=[""]) @@ -571,7 +464,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 @@ -613,7 +508,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"), @@ -669,7 +564,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) @@ -678,7 +573,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"), @@ -739,7 +634,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, ), ): @@ -797,7 +692,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, ), ): @@ -835,7 +730,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"), @@ -865,38 +760,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" 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..2aa28d90 --- /dev/null +++ b/tests/unit/application/services/test_sigma_validator_regression.py @@ -0,0 +1,184 @@ +"""Regression tests for SigmaValidator contract. + +These tests capture the current contract of SigmaValidator.validate() +so that when it is changed from returning dict[str, Any] to SigmaRule: +- Every dict access pattern is identified +- Every consumer is accounted for + +Run before and after the refactoring. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +import yaml + +from src.application.sigma.validator import SigmaValidator +from src.shared.exceptions import ValidationError + +_VALID_YAML = b""" +id: regr_test_001 +name: Regression Test Rule +description: A rule for regression testing +detection: + selection: + EventID: 4625 +condition: selection +""" + +_VALID_DICT: dict[str, Any] = { + "id": "regr_test_001", + "name": "Regression Test Rule", + "description": "A rule for regression testing", + "detection": {"selection": {"EventID": 4625}}, + "condition": "selection", +} + + +class TestSigmaValidatorRegression: + """Regression tests: current contract of SigmaValidator.""" + + def setup_method(self) -> None: + self.validator = SigmaValidator() + + # --- Return type contract --- + + def test_validate_returns_dict(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert isinstance(result, dict), "Must return a dict" + + def test_validate_dict_contains_expected_keys(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert "id" in result + assert "name" in result + assert "description" in result + assert "detection" in result + + def test_validate_dict_supports_bracket_access(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert result["id"] == "regr_test_001" + + def test_validate_dict_supports_get_access(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert result.get("condition") == "selection" + + def test_validate_dict_supports_get_missing_key(self) -> None: + result = self.validator.validate(_VALID_YAML) + assert result.get("nonexistent") is None + + def test_validate_dict_yaml_dumpable(self) -> None: + result = self.validator.validate(_VALID_YAML) + dumped = yaml.dump(result, default_flow_style=False) + assert isinstance(dumped, str) + assert "regr_test_001" in dumped + + # --- 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_dict(self) -> None: + yaml_with_deprecated = b""" +id: regr_test_002 +name: 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"] + + # --- 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 validate result. + + ChatService stores the validator result in self._uploaded_rule + (typed as dict[str, Any]) and accesses it via: + - .get("name", "") + - .get("id", "N/A") + - .get("description") + - Passes it to RAGPipeline methods as dict[str, Any] + - Passes it to search_engine.search() as .get("name", "") + """ + + def test_uploaded_rule_get_name(self) -> None: + result = _VALID_DICT + name = result.get("name", "") + assert name == "Regression Test Rule" + + def test_uploaded_rule_get_id(self) -> None: + result = _VALID_DICT + rule_id = result.get("id", "N/A") + assert rule_id == "regr_test_001" + + def test_uploaded_rule_get_description(self) -> None: + result = _VALID_DICT + desc = result.get("description") + assert desc == "A rule for regression testing" + + def test_uploaded_rule_missing_key_returns_default(self) -> None: + result: dict[str, Any] = {} + assert result.get("name", "") == "" + assert result.get("id", "N/A") == "N/A" + assert result.get("description") is None + + def test_rag_pipeline_format_rule_yaml(self) -> None: + dumped = yaml.dump(_VALID_DICT, default_flow_style=False, allow_unicode=True) + assert isinstance(dumped, str) + assert "Regression Test Rule" in dumped + + def test_rag_pipeline_fallback_explanation(self) -> None: + parts = [ + f"**Rule:** {_VALID_DICT.get('name', 'Unknown')}", + f"**ID:** {_VALID_DICT.get('id', 'N/A')}", + ] + if desc := _VALID_DICT.get("description"): + parts.append(f"**Description:** {desc}") + 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: + query = _VALID_DICT.get("name", "") + assert query == "Regression Test Rule" diff --git a/tests/unit/shared/test_http.py b/tests/unit/shared/test_http.py new file mode 100644 index 00000000..91de4a23 --- /dev/null +++ b/tests/unit/shared/test_http.py @@ -0,0 +1,290 @@ +"""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 ( + _backoff_delay, + _get_retry_after, + create_client, + download_file, + head_url, +) + + +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.__enter__.return_value.head.return_value = mock_resp + + with patch("src.shared.http.create_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.__enter__.return_value.head.side_effect = httpx.HTTPStatusError( + "404", request=MagicMock(), response=MagicMock() + ) + + with patch("src.shared.http.create_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.__enter__.return_value.head.side_effect = httpx.ConnectError( + "connection refused" + ) + + with patch("src.shared.http.create_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.__enter__.return_value.head.return_value = mock_resp + + with patch("src.shared.http.create_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.__enter__.return_value.head.return_value = mock_resp + + with ( + patch("src.shared.http.create_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.__enter__.return_value.get.return_value = mock_resp + + with patch("src.shared.http.create_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.__enter__.return_value.get.side_effect = [ + httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), + mock_resp_ok, + ] + + with ( + patch("src.shared.http.create_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.__enter__.return_value.get.side_effect = [ + httpx.ConnectError("timeout"), + MagicMock(content=b"ok"), + ] + + with ( + patch("src.shared.http.create_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.__enter__.return_value.get.side_effect = httpx.ConnectError("always fails") + + with ( + patch("src.shared.http.create_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.__enter__.return_value.get.side_effect = httpx.HTTPStatusError( + "403", request=MagicMock(), response=mock_resp + ) + + with patch("src.shared.http.create_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.__enter__.return_value.get.return_value = mock_resp + + with patch("src.shared.http.create_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.__enter__.return_value.get.side_effect = [ + httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), + mock_resp_ok, + ] + + with ( + patch("src.shared.http.create_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 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" From 9b9883129ba94bb0ef926632a4fd72215151ae74 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Thu, 2 Jul 2026 17:41:11 +0200 Subject: [PATCH 03/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20SigmaVal?= =?UTF-8?q?idator=20returns=20SigmaRule,=20merge=20validate=5Fsigma=5Frule?= =?UTF-8?q?=20(P1.1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - SigmaValidator.validate() returns SigmaRule instead of dict[str, Any] - SigmaRule.name property aliases title for Sigma YAML compatibility - ChatService._uploaded_rule: dict -> SigmaRule | None - RAGPipeline accepts SigmaRule in explain_rule, analyze_coverage, etc. - Merged level/status validation from dead validate_sigma_rule() - Removed validate_sigma_rule(), documents.models ValidationError/Result - Normalized name/title duality in validator - Updated all consumer tests --- src/api/v1/chat/chat.py | 6 +- src/application/chat/rag.py | 52 ++--- src/application/chat/service.py | 17 +- src/application/documents/models.py | 18 -- src/application/documents/validator.py | 83 +------- src/application/sigma/validator.py | 73 +++++-- src/core/sigma/models.py | 8 + .../application/documents/test_documents.py | 81 -------- .../application/services/test_rag_pipeline.py | 176 ++++++++-------- .../services/test_sigma_validator.py | 13 +- .../services/test_sigma_validator_advanced.py | 6 +- .../test_sigma_validator_regression.py | 190 +++++++++++------- 12 files changed, 327 insertions(+), 396 deletions(-) diff --git a/src/api/v1/chat/chat.py b/src/api/v1/chat/chat.py index a702a654..92b05cd4 100644 --- a/src/api/v1/chat/chat.py +++ b/src/api/v1/chat/chat.py @@ -102,11 +102,11 @@ 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) return { - "rule_name": rule_data.get("name", "Unknown"), - "rule_id": rule_data.get("id", ""), + "rule_name": rule.name, + "rule_id": rule.id, "validated": True, } except ValidationError as e: 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..0f48b878 100644 --- a/src/application/chat/service.py +++ b/src/application/chat/service.py @@ -17,6 +17,7 @@ ) from src.application.tools import ToolContext, ToolDispatcher, get_tools from src.core.search.engine import SearchEngine +from src.core.sigma.models import SigmaRule from src.api.v1.chat.schemas import ChatMode logger = logging.getLogger(__name__) @@ -33,7 +34,7 @@ def __init__(self, use_router: bool = True) -> None: self.rag_pipeline = RAGPipeline() self.validator = SigmaValidator() self._history: list[dict[str, str]] = [] - self._uploaded_rule: dict[str, Any] | None = None + self._uploaded_rule: SigmaRule | None = None self._last_citations: list[str] = [] self._current_prompt_id: str = "" @@ -231,7 +232,7 @@ async def _handle_explain(self, message: str) -> str: if not self._uploaded_rule: return "No Sigma rule uploaded. Please upload a .yaml file first." - related = await self.search_engine.search(self._uploaded_rule.get("name", "")) + related = await self.search_engine.search(self._uploaded_rule.name) return await self.rag_pipeline.explain_rule( self._uploaded_rule, related, system_prompt_id=self._current_prompt_id ) @@ -245,7 +246,7 @@ async def _handle_explain_stream( 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(self._uploaded_rule.name) async for token in self.rag_pipeline.explain_rule_stream( self._uploaded_rule, related, system_prompt_id=self._current_prompt_id ): @@ -345,19 +346,19 @@ async def _handle_coverage_stream( 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) -> SigmaRule: """Validate YAML content and store rule data in session. Args: content: Raw YAML file content 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._uploaded_rule = rule self.rag_pipeline.cache.invalidate() - return rule_data + return rule def get_last_citations(self) -> list[str]: """Get citations from the last search response.""" 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/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/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/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/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/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 index 2aa28d90..35c1d912 100644 --- a/tests/unit/application/services/test_sigma_validator_regression.py +++ b/tests/unit/application/services/test_sigma_validator_regression.py @@ -1,26 +1,21 @@ """Regression tests for SigmaValidator contract. -These tests capture the current contract of SigmaValidator.validate() -so that when it is changed from returning dict[str, Any] to SigmaRule: -- Every dict access pattern is identified -- Every consumer is accounted for - -Run before and after the refactoring. +These tests capture the contract of SigmaValidator.validate() +after the refactoring from dict[str, Any] to SigmaRule. """ from __future__ import annotations -from typing import Any - 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 -name: Regression Test Rule +title: Regression Test Rule description: A rule for regression testing detection: selection: @@ -28,52 +23,53 @@ condition: selection """ -_VALID_DICT: dict[str, Any] = { - "id": "regr_test_001", - "name": "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: current contract of SigmaValidator.""" + """Regression tests: updated contract of SigmaValidator.""" def setup_method(self) -> None: self.validator = SigmaValidator() # --- Return type contract --- - def test_validate_returns_dict(self) -> None: - result = self.validator.validate(_VALID_YAML) - assert isinstance(result, dict), "Must return a dict" - - def test_validate_dict_contains_expected_keys(self) -> None: + def test_validate_returns_sigma_rule(self) -> None: result = self.validator.validate(_VALID_YAML) - assert "id" in result - assert "name" in result - assert "description" in result - assert "detection" in result + assert isinstance(result, SigmaRule), "Must return a SigmaRule" - def test_validate_dict_supports_bracket_access(self) -> None: + def test_validate_rule_has_expected_attributes(self) -> None: result = self.validator.validate(_VALID_YAML) - assert result["id"] == "regr_test_001" + 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_dict_supports_get_access(self) -> None: + def test_validate_rule_supports_name_property(self) -> None: result = self.validator.validate(_VALID_YAML) - assert result.get("condition") == "selection" + assert result.name == result.title - def test_validate_dict_supports_get_missing_key(self) -> None: + def test_validate_rule_yaml_dumpable(self) -> None: result = self.validator.validate(_VALID_YAML) - assert result.get("nonexistent") is None - - def test_validate_dict_yaml_dumpable(self) -> None: - result = self.validator.validate(_VALID_YAML) - dumped = yaml.dump(result, default_flow_style=False) + 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: @@ -94,10 +90,10 @@ def test_validate_error_has_message(self) -> None: # --- Deprecated fields (warnings, not errors) --- - def test_deprecated_fields_still_in_dict(self) -> None: + def test_deprecated_fields_still_in_rule(self) -> None: yaml_with_deprecated = b""" -id: regr_test_002 -name: Rule with deprecated +id: regr_test_003 +title: Rule with deprecated description: Contains level and falsepositives level: high falsepositives: @@ -107,8 +103,53 @@ def test_deprecated_fields_still_in_dict(self) -> None: EventID: 4625 """ result = self.validator.validate(yaml_with_deprecated) - assert result["level"] == "high" - assert result["falsepositives"] == ["FP1"] + 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 --- @@ -131,54 +172,55 @@ def test_validate_invalid_yaml_syntax(self) -> None: class TestChatServiceConsumerRegression: - """Regression: ChatService usage patterns on validate result. - - ChatService stores the validator result in self._uploaded_rule - (typed as dict[str, Any]) and accesses it via: - - .get("name", "") - - .get("id", "N/A") - - .get("description") - - Passes it to RAGPipeline methods as dict[str, Any] - - Passes it to search_engine.search() as .get("name", "") + """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 test_uploaded_rule_get_name(self) -> None: - result = _VALID_DICT - name = result.get("name", "") - assert name == "Regression Test Rule" + 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_get_id(self) -> None: - result = _VALID_DICT - rule_id = result.get("id", "N/A") - assert rule_id == "regr_test_001" + def test_uploaded_rule_name(self) -> None: + rule = self._make_rule() + assert rule.name == "Regression Test Rule" - def test_uploaded_rule_get_description(self) -> None: - result = _VALID_DICT - desc = result.get("description") - assert desc == "A rule for regression testing" + def test_uploaded_rule_id(self) -> None: + rule = self._make_rule() + assert rule.id == "regr_test_001" - def test_uploaded_rule_missing_key_returns_default(self) -> None: - result: dict[str, Any] = {} - assert result.get("name", "") == "" - assert result.get("id", "N/A") == "N/A" - assert result.get("description") is None + 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: - dumped = yaml.dump(_VALID_DICT, default_flow_style=False, allow_unicode=True) + rule = self._make_rule() + dumped = yaml.dump(rule.to_dict(), default_flow_style=False, allow_unicode=True) assert isinstance(dumped, str) - assert "Regression Test Rule" in dumped + assert "regr_test_001" in dumped def test_rag_pipeline_fallback_explanation(self) -> None: + rule = self._make_rule() parts = [ - f"**Rule:** {_VALID_DICT.get('name', 'Unknown')}", - f"**ID:** {_VALID_DICT.get('id', 'N/A')}", + f"**Rule:** {rule.name}", + f"**ID:** {rule.id}", ] - if desc := _VALID_DICT.get("description"): - parts.append(f"**Description:** {desc}") + 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: - query = _VALID_DICT.get("name", "") + rule = self._make_rule() + query = rule.name assert query == "Regression Test Rule" From 17fc388fb5047a472a10c64b16481dac4361b4ca Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:52:49 +0200 Subject: [PATCH 04/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20unify=20?= =?UTF-8?q?download=20paths,=20standardize=20{url=5Fhash}{ext}=20naming=20?= =?UTF-8?q?(R1.1=20+=20R1.2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - download_sigma_references(db, output_dir, mode='scan'|'registry') - download_references() delegates to mode='scan' - process_sigma_refs() delegates to mode='registry' (-332 lines) - {url_hash}{ext} naming used consistently in both modes - Removed _sanitize_filename() (registry path naming mismatch) - Contract tests: same URL → same filename in both modes --- .../documents/sigma_ref_downloader.py | 386 +++++++++++++++++- .../documents/sigma_ref_processor.py | 384 +++-------------- .../documents/test_sigma_ref_downloader.py | 142 +++++++ 3 files changed, 579 insertions(+), 333 deletions(-) diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index d7927abb..3f2845a4 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -25,6 +25,8 @@ 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 logger = logging.getLogger(__name__) DEFAULT_REQUEST_DELAY = 0.5 @@ -170,34 +172,42 @@ def _maybe_record_error( logger.warning("Failed to record error for %s", normalized_url) -def download_references( - rules_dir: str, - output_dir: str, +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 all Sigma rule references matching supported document types. + """Download Sigma rule references using the specified mode. - Scans all Sigma rules in the given directory, extracts reference URLs, - filters by supported document types, and downloads matching files. + Two modes are supported: + + - **scan** (default): Scans a local directory of Sigma rule YAML files, + extracts reference URLs, and downloads matching documents. + Requires ``rules_dir``. + + - **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: - 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. + 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 {"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. + 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. - If provided, only files within these directories are processed. - Directories not in this list are excluded from scanning. + selected_dirs: Optional list of relative directory paths to scan + (only used in ``"scan"`` mode). Returns: Dict with summary stats: total_rules, total_refs, downloaded, skipped, failed. @@ -205,6 +215,76 @@ def download_references( if supported_types is None: supported_types = SUPPORTED_REFERENCE_DOC_TYPES + 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( + 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]: + """Download all Sigma rule references matching supported document types. + + Scans all Sigma rules in the given directory, extracts reference URLs, + filters by supported document types, and downloads matching files. + + 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 + rules_path = Path(rules_dir) output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) @@ -536,3 +616,279 @@ 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 _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 _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() + + for rule_entry in sigma_entries: + rule_id: str = rule_entry.get("rule_id", "00000000-0000-0000-0000-000000000000") + 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) + + 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) + + _TYPE_TO_EXT = { + "html": ".html", + "markdown": ".md", + "plain_text": ".txt", + "pdf": ".pdf", + "office_document": ".docx", + } + + def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: + url = item["final_url"] + content_type = item.get("content_type", "") + ext = _TYPE_TO_EXT.get(content_type, ".md") + url_hash = item.get("url_hash") or compute_sha256_str(normalize_url(url)) + file_path = output_path / f"{url_hash}{ext}" + + if file_path.exists(): + existing_entry = registry.get(url_hash) + if existing_entry and existing_entry.get("content_sha256"): + return None + + ok, _ = http_download_file(url, file_path, check_ssrf=False) + if ok: + content = file_path.read_bytes() + content_hash = compute_sha256_str(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[0] == "ok": + _, url_hash, content_hash, size = 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}{_TYPE_TO_EXT.get(item['content_type'], '.md')}", + 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", + ) + + return { + "total_rules": total_rules, + "total_refs": total_refs, + "downloaded": downloaded, + "skipped": skipped, + "failed": failed, + } diff --git a/src/application/documents/sigma_ref_processor.py b/src/application/documents/sigma_ref_processor.py index 2d73323e..30c0ec71 100644 --- a/src/application/documents/sigma_ref_processor.py +++ b/src/application/documents/sigma_ref_processor.py @@ -1,318 +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 - -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 -from src.shared.utils.identify_file_type import SUPPORTED_REFERENCE_DOC_TYPES -from src.shared.utils.sigma_utils import extract_sigma_references -from src.shared.utils.url_utils import normalize_url -from src.config.settings import get_config - -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. - - 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 = compute_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(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_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_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 - - # Cache HEAD result immediately - 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 - - 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"] - file_path = Path(output_dir) / _sanitize_filename(url) - - if file_path.exists(): - url_hash = compute_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 - - ok, _ = http_download_file(url, file_path, check_ssrf=False) - if ok: - content = file_path.read_bytes() - return ("ok", compute_sha256_bytes(content), 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[0] == "ok": - final_url = item["final_url"] - head_result = cast("tuple[str, str, int]", result) - entry = build_registry_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], - 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", - ) - - 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 _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/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 69565a9d..8bcce646 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -774,3 +774,145 @@ 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 / 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 / expected_filename).exists() From c111d2654cbe6f1b0243e51b56d1e26cd4f6dfa7 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 06:31:34 +0200 Subject: [PATCH 05/44] =?UTF-8?q?=E2=9C=A8=20feat:=20Phase=206+7=20?= =?UTF-8?q?=E2=80=94=20quantization,=20batch=20tuning,=20worker=20unificat?= =?UTF-8?q?ion,=20HNSW=20config?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P1.4 — Unify _scan_all_github + _scan_all_spec into _scan_all(prepare_fn), unify _write_entries + _write_spec_entries with batch_upsert_fn param, extract _collect_repo_files() for shared rglob+selected_dirs logic P2.5 — Add NULL_UUID constant in src/shared/constants.py, migrate 6 files away from hardcoded null UUID string Q1.1 — Add ScalarQuantization INT8 to all 3 collection creation paths (collections.py, storage.py, qdrant.py), controlled via enable_quantization Q1.2 — Bump embed_batch_size 8->64, num_workers 0->4 in ingestion.py Q2.3 — Add per-collection HNSW config via collection_hnsw_config(): sigma_rules in-RAM ef_construct=200, sigma_docs/sigma_spec on-disk R1.4 follow-up — delete_unreferenced_entries() + DocGCWorker integration (delete_head_verified_orphans + unreferenced cleanup in GC cycle) ruff clean, 146+ relevant tests green --- src/api/v1/documents/files.py | 4 +- src/api/v1/infrastructure/qdrant.py | 7 + .../documents/sigma_ref_downloader.py | 15 +- src/core/pipeline/ingestion.py | 4 +- src/infrastructure/database/core.py | 1 + src/infrastructure/database/doc_ops.py | 141 ++++++++++++- src/infrastructure/database/initdb.sql | 11 + src/infrastructure/vectorstore/__init__.py | 10 +- src/infrastructure/vectorstore/collections.py | 76 ++++++- src/infrastructure/vectorstore/storage.py | 10 + src/shared/constants.py | 3 + src/workers/document/gc_worker.py | 8 +- src/workers/sigma/discovery_base.py | 3 +- src/workers/sigma/discovery_worker.py | 199 +++++++++--------- .../documents/test_sigma_ref_downloader.py | 127 +++++++++++ 15 files changed, 499 insertions(+), 120 deletions(-) create mode 100644 src/shared/constants.py 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/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index af95d936..7d052b78 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -66,6 +66,13 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N sparse_vectors_config={ "text-sparse": models.SparseVectorParams(index=models.SparseIndexParams()) }, + quantization_config=models.ScalarQuantization( + scalar=models.ScalarQuantizationConfig( + type=models.ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) + ), ) diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index 3f2845a4..1b118dfc 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -12,6 +12,7 @@ import yaml +from src.shared.constants import NULL_UUID from src.shared.http import RETRY_STATUSES from src.shared.http import download_file as http_download_file from src.shared.http import head_url as http_head_url @@ -339,6 +340,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): @@ -381,6 +383,8 @@ def _collect_yaml_files() -> list[Path]: 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) skipped += 1 @@ -596,6 +600,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, @@ -684,9 +691,10 @@ def _download_registry_mode( 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", "00000000-0000-0000-0000-000000000000") + rule_id: str = rule_entry.get("rule_id", NULL_UUID) original_url = rule_entry.get("original_url", "") file_name = rule_entry.get("file_name", "") @@ -719,6 +727,8 @@ def _download_registry_mode( 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"): @@ -885,6 +895,9 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: "downloading", ) + if rule_refs: + db.batch_upsert_rule_references(rule_refs) + return { "total_rules": total_rules, "total_refs": total_refs, diff --git a/src/core/pipeline/ingestion.py b/src/core/pipeline/ingestion.py index 37766885..9c6a18a2 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. diff --git a/src/infrastructure/database/core.py b/src/infrastructure/database/core.py index 324e61d6..ab576f98 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", diff --git a/src/infrastructure/database/doc_ops.py b/src/infrastructure/database/doc_ops.py index ebd787b8..8e7ae56b 100644 --- a/src/infrastructure/database/doc_ops.py +++ b/src/infrastructure/database/doc_ops.py @@ -6,6 +6,8 @@ import logging import os from pathlib import Path + +from src.shared.constants import NULL_UUID from typing import TYPE_CHECKING, Any if TYPE_CHECKING: @@ -110,7 +112,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"), @@ -153,7 +155,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"), @@ -564,3 +566,138 @@ 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() + + # ------------------------------------------------------------------ + # 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. + """ + 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 < strftime('%Y-%m-%dT%H:%M:%fZ', 'now', ?))", + [f"-{grace_days} days"], + ).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..0c34817f 100644 --- a/src/infrastructure/database/initdb.sql +++ b/src/infrastructure/database/initdb.sql @@ -89,6 +89,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', 'now')), + 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, 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/collections.py b/src/infrastructure/vectorstore/collections.py index f798deab..2b3ee38c 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_kb=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_kb=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) @@ -104,9 +139,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( @@ -120,8 +170,24 @@ async def create_collection( index=qdrant_client.models.SparseIndexParams() ) } + 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/storage.py b/src/infrastructure/vectorstore/storage.py index f3866a6f..964c9c3e 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -106,6 +106,9 @@ async def store_embeddings( if collection_name not in existing_collections: from qdrant_client.models import ( Distance, + ScalarQuantization, + ScalarQuantizationConfig, + ScalarType, SparseIndexParams, SparseVectorParams, VectorParams, @@ -117,6 +120,13 @@ async def store_embeddings( sparse_vectors_config={ "text-sparse": SparseVectorParams(index=SparseIndexParams()), }, + quantization_config=ScalarQuantization( + scalar=ScalarQuantizationConfig( + type=ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) + ), ) # Delete old points for each source before upserting 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/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 9253eb63..62ab4610 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -51,10 +51,16 @@ def process(self, task: dict) -> None: ) 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() + 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}" ) except Exception as e: diff --git a/src/workers/sigma/discovery_base.py b/src/workers/sigma/discovery_base.py index f38d8788..766ebca5 100644 --- a/src/workers/sigma/discovery_base.py +++ b/src/workers/sigma/discovery_base.py @@ -1,6 +1,7 @@ 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 @@ -40,7 +41,7 @@ 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 = compute_sha256_str(normalized_url) now = iso_now() diff --git a/src/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index 5a26cca8..14367a4d 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -5,8 +5,10 @@ import logging from enum import Enum from pathlib import Path +from collections.abc import Callable from typing import TYPE_CHECKING, Optional +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 @@ -134,40 +136,15 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: logger.info("[GenericDiscoveryWorker] No repos with selected dirs") return - all_files: list[tuple[Path, Path, str, str]] = [] + repo_items: list[tuple[str, str]] = [] for repo_key in repo_keys: parts = repo_key.split("/") if len(parts) != 2: logger.warning(f"[GenericDiscoveryWorker] Invalid repo key: {repo_key}") continue + repo_items.append((parts[0], parts[1])) - 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 - ): - 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)) + 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") @@ -177,7 +154,11 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: ) 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 ) @@ -206,47 +187,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 - - 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) + if org and repo_name: + repo_items.append((org, repo_name)) - 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") @@ -256,15 +210,23 @@ 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, ) - def _scan_all_spec( + 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 @@ -272,7 +234,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 @@ -328,19 +290,24 @@ 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 + def _write_entries( + 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_sigma_spec(entries) + upsert_fn(entries) except Exception as e: - logger.error( - f"[GenericDiscoveryWorker] Batch upsert sigma_spec failed: {e}", exc_info=True - ) + logger.error(f"[GenericDiscoveryWorker] Batch upsert failed: {e}", exc_info=True) if total > 0 and self.dispatcher: pct = int((processed + skipped) / total * 100) @@ -377,33 +344,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 # ------------------------------------------------------------------ @@ -439,7 +379,7 @@ 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": rid = get_sigma_rule_id(file_path) if rid: @@ -461,6 +401,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 # ------------------------------------------------------------------ diff --git a/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 8bcce646..16e4fe9d 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -916,3 +916,130 @@ def _fake_download( # File should be written with {url_hash}{ext} naming expected_filename = f"{expected_hash}.md" assert (output_dir / 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"] From 73c94d4e95d902265ff7f4f33a398fb7c341f3fb Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 06:52:47 +0200 Subject: [PATCH 06/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20external?= =?UTF-8?q?ise=20ChatService=20session=20state=20into=20SessionStore=20(P2?= =?UTF-8?q?.3)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Nouveau module src/shared/session.py — SessionStore thread-safe - ChatService délègue rule/tool_state à SessionStore - _handle_explain/_handle_search/_handle_coverage et leurs streams acceptent (message, prompt_id, sid) - Routes API lisent X-Session-ID header avec fallback '_default' - P2.1: _chunk_rule() éclatée en 13 méthodes spécialisées - P2.6: FILETYPE_TO_EXT dans identify_file_type.py, utilisé dans sigma_ref_downloader.py - Tests mis à jour pour nouvelles signatures --- src/api/v1/chat/chat.py | 50 +- src/application/chat/service.py | 160 ++++-- .../documents/sigma_ref_downloader.py | 31 +- src/core/sigma/chunker.py | 495 +++++++++--------- src/shared/session.py | 45 ++ src/shared/utils/identify_file_type.py | 8 + tests/integration/test_chat_flow.py | 280 +++++----- .../services/test_chat_service_cache.py | 8 +- 8 files changed, 578 insertions(+), 499 deletions(-) create mode 100644 src/shared/session.py diff --git a/src/api/v1/chat/chat.py b/src/api/v1/chat/chat.py index 92b05cd4..d19b8b71 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,7 +89,10 @@ async def send_chat_message(req: ChatMessageRequest) -> ChatMessageResponse: @router.post("/upload") -async def upload_sigma_rule(file: UploadFile) -> dict: +async def upload_sigma_rule( + file: UploadFile, + session_id: str | None = Depends(_get_session_id), +) -> dict: """Upload and validate a Sigma rule YAML file.""" if not file.filename: raise HTTPException( @@ -102,7 +121,7 @@ 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 = 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.name, @@ -123,7 +142,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 +157,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/application/chat/service.py b/src/application/chat/service.py index 0f48b878..c7d93712 100644 --- a/src/application/chat/service.py +++ b/src/application/chat/service.py @@ -19,24 +19,33 @@ from src.core.search.engine import SearchEngine 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: SigmaRule | None = None - self._last_citations: list[str] = [] - self._current_prompt_id: str = "" # Tool-calling setup self._tool_executor = ToolDispatcher(get_tools()) @@ -46,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() @@ -56,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. @@ -64,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 @@ -71,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( @@ -94,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: @@ -109,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: @@ -126,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. @@ -209,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, @@ -227,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.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.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) @@ -263,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}") @@ -272,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. @@ -296,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 @@ -315,7 +354,7 @@ async def _handle_search_stream( 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, results, system_prompt_id=prompt_id ): yield token found = True @@ -327,9 +366,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 @@ -340,51 +382,57 @@ async def _handle_coverage_stream( try: async for token in self.rag_pipeline.analyze_coverage_stream( - self._uploaded_rule, results, system_prompt_id=self._current_prompt_id + rule, 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) -> SigmaRule: + 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 SigmaRule """ rule = self.validator.validate(content) - self._uploaded_rule = rule + self._set_rule(session_id, rule) self.rag_pipeline.cache.invalidate() 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/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index 1b118dfc..7995a72f 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -19,6 +19,7 @@ from src.shared.utils.registry_utils import build_registry_entry from src.shared.utils.crypto_utils import compute_sha256_file, compute_sha256_str from src.shared.utils.identify_file_type import ( + FILETYPE_TO_EXT, SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, ) @@ -442,14 +443,7 @@ 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_TO_EXT.get(ftype, ".md") if not ext: ext = ".md" @@ -523,14 +517,7 @@ def _collect_yaml_files() -> list[Path]: 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_TO_EXT.get(ftype, ".md") if not ext: ext = ".md" output_file = output_path / f"{item['url_hash']}{ext}" @@ -827,18 +814,10 @@ def _download_registry_mode( all_to_download = head_pending + download_ready total_to_download = len(all_to_download) - _TYPE_TO_EXT = { - "html": ".html", - "markdown": ".md", - "plain_text": ".txt", - "pdf": ".pdf", - "office_document": ".docx", - } - def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: url = item["final_url"] content_type = item.get("content_type", "") - ext = _TYPE_TO_EXT.get(content_type, ".md") + ext = FILETYPE_TO_EXT.get(content_type, ".md") url_hash = item.get("url_hash") or compute_sha256_str(normalize_url(url)) file_path = output_path / f"{url_hash}{ext}" @@ -875,7 +854,7 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: rule_id=item["rule_id"], title=item["rule_title"], content_sha256=content_hash, - file_name=f"{url_hash}{_TYPE_TO_EXT.get(item['content_type'], '.md')}", + file_name=f"{url_hash}{FILETYPE_TO_EXT.get(item['content_type'], '.md')}", file_size=size, embed_status="discovery", ) diff --git a/src/core/sigma/chunker.py b/src/core/sigma/chunker.py index c00b4bb8..18ddcd6c 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] = [] - - # 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 _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"), + } + + 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,157 @@ 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']}?", + ], + ) + + 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: list[dict] = [ + self._build_executive_summary(rule, f), + self._build_metadata_lifecycle(rule, f), + self._build_logsource_context(rule, f), + ] + attack = self._build_mitre_attack_mapping(rule, f) + if attack 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), + ] ) - # Enrich all chunks with LLM-generated summaries and keywords if llm_client is not None: - enriched_chunks: 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_chunks.append(chunk) - chunks = enriched_chunks + chunks = self._enrich_chunks(chunks, llm_client) return chunks 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/identify_file_type.py b/src/shared/utils/identify_file_type.py index c1dadf39..f72ee063 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -56,6 +56,14 @@ class FileType(Enum): ".odp": FileType.OFFICE_DOCUMENT, } +FILETYPE_TO_EXT: dict[str, str] = { + FileType.MARKDOWN.value: ".md", + FileType.PDF.value: ".pdf", + FileType.PLAIN_TEXT.value: ".txt", + FileType.HTML.value: ".html", + FileType.OFFICE_DOCUMENT.value: ".docx", +} + SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, FileType.PDF.value, 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/application/services/test_chat_service_cache.py b/tests/unit/application/services/test_chat_service_cache.py index 1df81ae6..b43cc9f3 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() From 3a2b7e6681c15c34e5aadb087e40ac598dd1d260 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 06:55:23 +0200 Subject: [PATCH 07/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20DuckDB=20strftime?= =?UTF-8?q?=20compat,=20=5FEXPECTED=5FTABLES,=20tests=20cross-platform?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - initdb.sql: strftime('%Y-%m-%dT%H:%M:%fZ', current_timestamp) (DuckDB 1.5+ ne tolère plus deux string literals) - doc_ops.py: cutoff calculé en Python pour éviter strftime à 3 args - duckdb.py: ajout rule_references à _EXPECTED_TABLES - test_duckdb.py: assert tables_missing 10→11, clean no-op passe - test_feedback.py: C:\absolute\path.db → /tmp/absolute/path.db (Linux) - test_datadir.py: count 14→15 (specification dir ajouté) --- src/application/system/duckdb.py | 1 + src/infrastructure/database/doc_ops.py | 8 ++++++-- src/infrastructure/database/initdb.sql | 2 +- tests/unit/application/feedback/test_feedback.py | 4 ++-- tests/unit/application/system/test_datadir.py | 2 +- tests/unit/application/system/test_duckdb.py | 2 +- 6 files changed, 12 insertions(+), 7 deletions(-) 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/infrastructure/database/doc_ops.py b/src/infrastructure/database/doc_ops.py index 8e7ae56b..521dfdc2 100644 --- a/src/infrastructure/database/doc_ops.py +++ b/src/infrastructure/database/doc_ops.py @@ -5,6 +5,7 @@ import hashlib import logging import os +from datetime import datetime, timezone, timedelta from pathlib import Path from src.shared.constants import NULL_UUID @@ -643,6 +644,9 @@ def delete_head_verified_orphans(self, grace_days: int = 7) -> int: 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( @@ -650,8 +654,8 @@ def delete_head_verified_orphans(self, grace_days: int = 7) -> int: "WHERE embed_status = 'head_verified' " "AND (content_sha256 IS NULL OR content_sha256 = '') " "AND (last_seen IS NULL " - " OR last_seen < strftime('%Y-%m-%dT%H:%M:%fZ', 'now', ?))", - [f"-{grace_days} days"], + " OR last_seen < ?)", + [cutoff], ).fetchall() orphan_hashes = [row[0] for row in orphans] diff --git a/src/infrastructure/database/initdb.sql b/src/infrastructure/database/initdb.sql index 0c34817f..67d65138 100644 --- a/src/infrastructure/database/initdb.sql +++ b/src/infrastructure/database/initdb.sql @@ -94,7 +94,7 @@ 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', 'now')), + created TEXT DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', current_timestamp)), PRIMARY KEY (rule_id, url_hash) ); 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/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: From 4955cfd00c86082643b2d3ad42ca55eadacc180b Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 07:02:16 +0200 Subject: [PATCH 08/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20R2.1=20?= =?UTF-8?q?=E2=80=94=20sous-r=C3=A9pertoires=20par=20type=20dans=20sigmare?= =?UTF-8?q?f/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Nouveau FILETYPE_TO_SUBDIR dans identify_file_type.py (markdown/, html/, pdf/, plain_text/, office/) - sigma_ref_downloader.py: _sigmaref_write_path / _sigmaref_resolve_path avec fallback flat pour backward compat - indexer.py: résolution sigmaref → subdir puis flat - gc_worker.py: _file_exists_locally vérifie subdir + flat - Fix bug: compute_sha256_bytes() au lieu de compute_sha256_str() sur contenu binaire --- .../documents/sigma_ref_downloader.py | 66 +++++++++++++++---- src/core/pipeline/indexer.py | 8 ++- src/shared/utils/identify_file_type.py | 8 +++ src/workers/document/gc_worker.py | 11 +++- .../documents/test_sigma_ref_downloader.py | 5 +- 5 files changed, 80 insertions(+), 18 deletions(-) diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index 7995a72f..bd52b14c 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -17,9 +17,14 @@ 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_file, compute_sha256_str +from src.shared.utils.crypto_utils import ( + compute_sha256_bytes, + compute_sha256_file, + compute_sha256_str, +) from src.shared.utils.identify_file_type import ( FILETYPE_TO_EXT, + FILETYPE_TO_SUBDIR, SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, ) @@ -40,6 +45,34 @@ _registry_lock = threading.Lock() +def _subdir_for(content_type: str | None) -> str: + """Return the subdirectory name for a given content type.""" + return FILETYPE_TO_SUBDIR.get(content_type or "", "misc") + + +def _sigmaref_write_path(output_path: Path, content_type: str | None, file_name: str) -> Path: + """Return the subdir path for writing a sigmaref file, creating the subdir as needed.""" + subdir = _subdir_for(content_type) + path = output_path / subdir + path.mkdir(parents=True, exist_ok=True) + return path / file_name + + +def _sigmaref_resolve_path(output_path: Path, content_type: str | None, file_name: str) -> Path: + """Resolve the path to an existing sigmaref file, with flat-layout fallback. + + Checks the subdir first, then the old flat layout for backward compatibility + with files downloaded before R2.1. Returns the subdir path if neither exists. + """ + candidate = output_path / _subdir_for(content_type) / file_name + if candidate.exists(): + return candidate + flat = output_path / file_name + if flat.exists(): + return flat + return candidate + + def _detect_url_type(url: str, content_type: str | None = None) -> str | None: """Detect the document type of a reference URL. @@ -395,7 +428,10 @@ 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 ( @@ -407,7 +443,9 @@ 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, content_type_for_subdir, fname + ), "content_type": registry[url_hash].get( "content_type", "markdown" ), @@ -448,7 +486,8 @@ def _collect_yaml_files() -> list[Path]: 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 = compute_sha256_file(output_file) @@ -460,7 +499,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, @@ -491,7 +530,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, @@ -520,7 +559,8 @@ def _collect_yaml_files() -> list[Path]: ext = FILETYPE_TO_EXT.get(ftype, ".md") 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 @@ -529,7 +569,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"], @@ -819,17 +859,19 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: content_type = item.get("content_type", "") ext = FILETYPE_TO_EXT.get(content_type, ".md") url_hash = item.get("url_hash") or compute_sha256_str(normalize_url(url)) - file_path = output_path / f"{url_hash}{ext}" - - if file_path.exists(): + 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_str(content) + content_hash = compute_sha256_bytes(content) return ("ok", url_hash, content_hash, len(content)) logger.error("Reference download failed: %s", url) return ("fail", "", "", 0) diff --git a/src/core/pipeline/indexer.py b/src/core/pipeline/indexer.py index 7340d808..4961bcc3 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_TO_SUBDIR logger = logging.getLogger(__name__) @@ -175,7 +176,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_TO_SUBDIR.get(row.get("content_type", ""), "misc") + 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/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index f72ee063..836ad092 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -64,6 +64,14 @@ class FileType(Enum): FileType.OFFICE_DOCUMENT.value: ".docx", } +FILETYPE_TO_SUBDIR: dict[str, str] = { + FileType.MARKDOWN.value: "markdown", + FileType.PDF.value: "pdf", + FileType.PLAIN_TEXT.value: "plain_text", + FileType.HTML.value: "html", + FileType.OFFICE_DOCUMENT.value: "office", +} + SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, FileType.PDF.value, diff --git a/src/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 62ab4610..6512631e 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -4,6 +4,7 @@ from pathlib import Path from src.config.settings import get_config +from src.shared.utils.identify_file_type import FILETYPE_TO_SUBDIR from src.workers.base import BaseWorker from src.workers.enums import WorkerName, WorkerStatus @@ -100,6 +101,7 @@ def _gc_entries( org or "", repo or "", file_name or "", + content_type or "", local_base, github_base, ) @@ -124,6 +126,7 @@ def _file_exists_locally( org: str, repo: str, file_name: str, + content_type: str, local_base: Path, github_base: Path, ) -> bool: @@ -133,19 +136,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_TO_SUBDIR.get(content_type, "misc") + 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) diff --git a/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 16e4fe9d..32633983 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -8,6 +8,7 @@ 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_TO_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 ( @@ -831,7 +832,7 @@ def _fake_download( ) assert result["downloaded"] == 1 - assert (output_dir / expected_filename).exists() + assert (output_dir / FILETYPE_TO_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 @@ -915,7 +916,7 @@ def _fake_download( assert result["total_refs"] == 1 # File should be written with {url_hash}{ext} naming expected_filename = f"{expected_hash}.md" - assert (output_dir / expected_filename).exists() + assert (output_dir / FILETYPE_TO_SUBDIR["markdown"] / expected_filename).exists() class TestRuleReferences: From 640efcd172b3169e9cfbfaf2117e8ba6cc2d0ddf Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 07:05:29 +0200 Subject: [PATCH 09/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20R2.2=20?= =?UTF-8?q?=E2=80=94=20GC=20fichiers=20orphelins=20sigmaref/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _gc_orphaned_sigmaref_files() : scanne sigmaref/ (subdirs + flat) pour les fichiers dont le SHA256 stem n'est pas dans doc_registry et les déplace vers .trash/ - _is_orphan_candidate() : vérifie format 64-char hex + appartenance - Supprime les sous-répertoires vides après nettoyage - Intégré dans DocGCWorker.process() avec compteur dans le log - 13 tests unitaires (format, flat, subdir, trash, edge cases) --- src/workers/document/gc_worker.py | 70 +++++++++++- tests/unit/workers/test_gc_worker.py | 158 +++++++++++++++++++++++++++ 2 files changed, 227 insertions(+), 1 deletion(-) create mode 100644 tests/unit/workers/test_gc_worker.py diff --git a/src/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 6512631e..8e0ab2f8 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -11,6 +11,15 @@ 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. @@ -55,13 +64,14 @@ def process(self, task: dict) -> None: 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() scanned = deleted logger.info( f"[DocGCWorker] Complete: scanned={scanned}, removed={deleted}, " f"reappears={skipped_found}, head_verified={removed_head}, " - f"unreferenced={removed_unref}" + f"unreferenced={removed_unref}, orphaned_files={removed_orphan_files}" ) except Exception as e: @@ -167,6 +177,64 @@ 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 _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/tests/unit/workers/test_gc_worker.py b/tests/unit/workers/test_gc_worker.py new file mode 100644 index 00000000..cab6cd66 --- /dev/null +++ b/tests/unit/workers/test_gc_worker.py @@ -0,0 +1,158 @@ +"""Tests for DocGCWorker and orphaned sigmaref file cleaning.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +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 From 445c7c4d29215d30ab433ae59b6296c64e6cf280 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 07:08:10 +0200 Subject: [PATCH 10/44] =?UTF-8?q?=E2=9C=A8=20feat:=20Q2.1=20=E2=80=94=20al?= =?UTF-8?q?pha=20hybride=20par=20collection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ALPHA_BY_COLLECTION globale avec valeurs tunées : sigma_rules=0.5, sigma_docs=0.7, sigma_spec=0.3 - SearchEngine accepte alpha_by_collection en override - alpha global devient fallback (default 0.3) - Rétrocompatible : tous les appelants existants inchangés --- src/core/search/engine.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/src/core/search/engine.py b/src/core/search/engine.py index 37227fc6..021bdd99 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -95,6 +95,14 @@ 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 +} + def _node_to_result(node: Any) -> dict[str, Any]: """Convert a LlamaIndex NodeWithScore to the legacy dict format. @@ -298,7 +306,8 @@ 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, ) -> None: """Initialize search engine. @@ -311,14 +320,21 @@ 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. """ 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 async def search( self, @@ -364,7 +380,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: From c03b9f8e2f9e4c8c539ec7cf639fde935fd0ec30 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 07:11:08 +0200 Subject: [PATCH 11/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20Q2.2=20?= =?UTF-8?q?=E2=80=94=20RRF=20param=C3=A9trable=20+=20Q3.1=20=E2=80=94=20sp?= =?UTF-8?q?arse=20vectors=20disk?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Q2.2: - rrf_k configurable (default 60) dans SearchEngine - rrf_weights par collection (weighted RRF) - col_idx aligné sur collections pour lookup weight Q3.1: - on_disk=True ajouté aux 3 chemins de création collection (qdrant.py, storage.py, collections.py) pour sparse vectors --- src/api/v1/infrastructure/qdrant.py | 5 ++++- src/core/search/engine.py | 19 +++++++++++++++++-- src/infrastructure/vectorstore/collections.py | 3 ++- src/infrastructure/vectorstore/storage.py | 5 ++++- 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/src/api/v1/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index 7d052b78..78657c86 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -64,7 +64,10 @@ 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(), + on_disk=True, + ) }, quantization_config=models.ScalarQuantization( scalar=models.ScalarQuantizationConfig( diff --git a/src/core/search/engine.py b/src/core/search/engine.py index 021bdd99..aca1f025 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -103,6 +103,8 @@ def _get_search_embed_model() -> Any: "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. @@ -308,6 +310,8 @@ def __init__( llm_client: LlamaClient | None = None, 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. @@ -324,6 +328,10 @@ def __init__( 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 @@ -335,6 +343,8 @@ def __init__( 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, @@ -444,16 +454,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/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index 2b3ee38c..2893f2cf 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -167,7 +167,8 @@ 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(), + on_disk=True, ) } quantization_config = None diff --git a/src/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index 964c9c3e..3f3297e8 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -118,7 +118,10 @@ 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(), + on_disk=True, + ), }, quantization_config=ScalarQuantization( scalar=ScalarQuantizationConfig( From 7a3353b773d5ea4b0410e2794dd61303d0a37be3 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 07:15:31 +0200 Subject: [PATCH 12/44] =?UTF-8?q?=F0=9F=8E=A8=20style:=20ruff=20format=20?= =?UTF-8?q?=E2=80=94=20reformat=20apr=C3=A8s=20commit=20hook?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scripts/benchmark_sparse_encoder.py | 272 ++++++++++++++++++ src/api/v1/infrastructure/qdrant.py | 1 + src/core/search/sparse_encoder.py | 103 ++++++- src/infrastructure/vectorstore/collections.py | 1 + src/infrastructure/vectorstore/storage.py | 3 +- tests/unit/core/search/test_sparse_encoder.py | 164 +++++++++++ 6 files changed, 538 insertions(+), 6 deletions(-) create mode 100644 scripts/benchmark_sparse_encoder.py create mode 100644 tests/unit/core/search/test_sparse_encoder.py 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/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index 78657c86..6efc6d7c 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -66,6 +66,7 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N sparse_vectors_config={ "text-sparse": models.SparseVectorParams( index=models.SparseIndexParams(), + modifier=models.Modifier.IDF, on_disk=True, ) }, 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/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index 2893f2cf..293db6c4 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -168,6 +168,7 @@ async def create_collection( sparse_vectors_config = { "text-sparse": qdrant_client.models.SparseVectorParams( index=qdrant_client.models.SparseIndexParams(), + modifier=qdrant_client.models.Modifier.IDF, on_disk=True, ) } diff --git a/src/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index 3f3297e8..41809153 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 @@ -120,6 +120,7 @@ async def store_embeddings( sparse_vectors_config={ "text-sparse": SparseVectorParams( index=SparseIndexParams(), + modifier=Modifier.IDF, on_disk=True, ), }, 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] From 88b348c480ef1f9e3acfe273e92c25e6e2660269 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:24:52 +0200 Subject: [PATCH 13/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20Phase=20?= =?UTF-8?q?8=20-=20code=20quality=20+=20storage=20organization?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - P2.6: Deduplicate FILETYPE_TO_EXT/SUBDIR → FILETYPE_INFO + helpers - P2.2: Share RAGPipeline singleton in translate endpoint - P2.1: Extract _assemble_chunks() from _chunk_rule() - R2.2: Add trash cleanup to DocGCWorker (max_age_days=7) --- src/api/v1/infrastructure/qdrant.py | 4 +- src/api/v1/sigma/translate.py | 13 +++++- .../documents/sigma_ref_downloader.py | 14 +++---- src/core/pipeline/indexer.py | 21 ++++++++-- src/core/sigma/chunker.py | 25 ++++++++--- src/infrastructure/database/doc_ops.py | 21 ++++++++-- src/infrastructure/vectorstore/collections.py | 12 +++--- src/infrastructure/vectorstore/storage.py | 8 ++-- src/shared/utils/identify_file_type.py | 29 +++++++------ src/workers/document/gc_worker.py | 42 +++++++++++++++++-- .../documents/test_sigma_ref_downloader.py | 6 +-- .../services/test_chat_service_cache.py | 4 +- 12 files changed, 150 insertions(+), 49 deletions(-) diff --git a/src/api/v1/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index 6efc6d7c..3f5d563e 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -75,7 +75,9 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N type=models.ScalarType.INT8, always_ram=True, quantile=0.5, - ) + ), + rescore=True, + oversampling=2.0, ), ) 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/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index bd52b14c..046fafa8 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -23,10 +23,10 @@ compute_sha256_str, ) from src.shared.utils.identify_file_type import ( - FILETYPE_TO_EXT, - FILETYPE_TO_SUBDIR, SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, + filetype_ext, + filetype_subdir, ) from src.shared.utils.url_utils import is_private_url, normalize_url, url_ext from src.infrastructure.database import DatabaseService @@ -47,7 +47,7 @@ def _subdir_for(content_type: str | None) -> str: """Return the subdirectory name for a given content type.""" - return FILETYPE_TO_SUBDIR.get(content_type or "", "misc") + return filetype_subdir(content_type or "") def _sigmaref_write_path(output_path: Path, content_type: str | None, file_name: str) -> Path: @@ -481,7 +481,7 @@ def _collect_yaml_files() -> list[Path]: continue if not ext and ftype is not None: - ext = FILETYPE_TO_EXT.get(ftype, ".md") + ext = filetype_ext(ftype) if not ext: ext = ".md" @@ -556,7 +556,7 @@ def _collect_yaml_files() -> list[Path]: ftype = _detect_url_type(item["normalized"], content_type=head_ct) ext = item["ext"] if not ext and ftype is not None: - ext = FILETYPE_TO_EXT.get(ftype, ".md") + ext = filetype_ext(ftype) if not ext: ext = ".md" fname = f"{item['url_hash']}{ext}" @@ -857,7 +857,7 @@ def _download_registry_mode( 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_TO_EXT.get(content_type, ".md") + 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) @@ -896,7 +896,7 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: rule_id=item["rule_id"], title=item["rule_title"], content_sha256=content_hash, - file_name=f"{url_hash}{FILETYPE_TO_EXT.get(item['content_type'], '.md')}", + file_name=f"{url_hash}{filetype_ext(item['content_type'])}", file_size=size, embed_status="discovery", ) diff --git a/src/core/pipeline/indexer.py b/src/core/pipeline/indexer.py index 4961bcc3..f9513a91 100644 --- a/src/core/pipeline/indexer.py +++ b/src/core/pipeline/indexer.py @@ -16,7 +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_TO_SUBDIR +from src.shared.utils.identify_file_type import filetype_subdir logger = logging.getLogger(__name__) @@ -139,7 +139,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": @@ -152,6 +157,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": @@ -177,7 +192,7 @@ def _resolve_path(self, table_name: str, row: dict) -> Path | None: if org == "sigmaref": base = Path(cfg.sigmaref_documents_path).resolve() - subdir = FILETYPE_TO_SUBDIR.get(row.get("content_type", ""), "misc") + subdir = filetype_subdir(row.get("content_type", "")) candidate = base / subdir / file_name if candidate.exists(): return candidate diff --git a/src/core/sigma/chunker.py b/src/core/sigma/chunker.py index 18ddcd6c..c9db8525 100644 --- a/src/core/sigma/chunker.py +++ b/src/core/sigma/chunker.py @@ -386,14 +386,32 @@ def _enrich_chunks(self, chunks: list[dict], llm_client: LLMClientLike) -> list[ 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: + 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), ] - attack = self._build_mitre_attack_mapping(rule, f) - if attack is not None: + + 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) @@ -409,9 +427,6 @@ def _chunk_rule(self, rule: dict, llm_client: LLMClientLike | None = None) -> li ] ) - if llm_client is not None: - chunks = self._enrich_chunks(chunks, llm_client) - return chunks def _dict_to_document(self, chunk_data: dict, source_file: str | None = None) -> Document: diff --git a/src/infrastructure/database/doc_ops.py b/src/infrastructure/database/doc_ops.py index 521dfdc2..760a7303 100644 --- a/src/infrastructure/database/doc_ops.py +++ b/src/infrastructure/database/doc_ops.py @@ -102,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"), @@ -144,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"), @@ -255,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"), diff --git a/src/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index 293db6c4..df66d881 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -25,14 +25,14 @@ def collection_hnsw_config( use on-disk storage to save RAM. """ if collection_name == "sigma_rules": - return qdrant_client.models.HnswConfigDiff( + return qdrant_client.models.HnswConfigDiff( # type: ignore[call-arg] m=16, ef_construct=200, full_scan_threshold_kb=10000, on_disk=False, ) if collection_name in ("sigma_docs", "sigma_spec"): - return qdrant_client.models.HnswConfigDiff( + return qdrant_client.models.HnswConfigDiff( # type: ignore[call-arg] m=16, ef_construct=100, full_scan_threshold_kb=10000, @@ -166,7 +166,7 @@ async def create_collection( sparse_vectors_config: dict[str, Any] | None = None if enable_hybrid: sparse_vectors_config = { - "text-sparse": qdrant_client.models.SparseVectorParams( + "text-sparse": qdrant_client.models.SparseVectorParams( # type: ignore[call-arg] index=qdrant_client.models.SparseIndexParams(), modifier=qdrant_client.models.Modifier.IDF, on_disk=True, @@ -174,12 +174,14 @@ async def create_collection( } quantization_config = None if enable_quantization: - quantization_config = qdrant_client.models.ScalarQuantization( + quantization_config = qdrant_client.models.ScalarQuantization( # type: ignore[call-arg] scalar=qdrant_client.models.ScalarQuantizationConfig( type=qdrant_client.models.ScalarType.INT8, always_ram=True, quantile=0.5, - ) + ), + rescore=True, + oversampling=2.0, ) hnsw_config = collection_hnsw_config(collection_name) await asyncio.to_thread( diff --git a/src/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index 41809153..a121eb00 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -118,18 +118,20 @@ async def store_embeddings( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), sparse_vectors_config={ - "text-sparse": SparseVectorParams( + "text-sparse": SparseVectorParams( # type: ignore[call-arg] index=SparseIndexParams(), modifier=Modifier.IDF, on_disk=True, ), }, - quantization_config=ScalarQuantization( + quantization_config=ScalarQuantization( # type: ignore[call-arg] scalar=ScalarQuantizationConfig( type=ScalarType.INT8, always_ram=True, quantile=0.5, - ) + ), + rescore=True, + oversampling=2.0, ), ) diff --git a/src/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index 836ad092..76371d0b 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -56,21 +56,24 @@ class FileType(Enum): ".odp": FileType.OFFICE_DOCUMENT, } -FILETYPE_TO_EXT: dict[str, str] = { - FileType.MARKDOWN.value: ".md", - FileType.PDF.value: ".pdf", - FileType.PLAIN_TEXT.value: ".txt", - FileType.HTML.value: ".html", - FileType.OFFICE_DOCUMENT.value: ".docx", +FILETYPE_INFO: dict[str, dict[str, str]] = { + FileType.MARKDOWN.value: {"ext": ".md", "subdir": "markdown"}, + FileType.PDF.value: {"ext": ".pdf", "subdir": "pdf"}, + FileType.PLAIN_TEXT.value: {"ext": ".txt", "subdir": "plain_text"}, + FileType.HTML.value: {"ext": ".html", "subdir": "html"}, + FileType.OFFICE_DOCUMENT.value: {"ext": ".docx", "subdir": "office"}, } -FILETYPE_TO_SUBDIR: dict[str, str] = { - FileType.MARKDOWN.value: "markdown", - FileType.PDF.value: "pdf", - FileType.PLAIN_TEXT.value: "plain_text", - FileType.HTML.value: "html", - FileType.OFFICE_DOCUMENT.value: "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, diff --git a/src/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 8e0ab2f8..51258551 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -4,7 +4,7 @@ from pathlib import Path from src.config.settings import get_config -from src.shared.utils.identify_file_type import FILETYPE_TO_SUBDIR +from src.shared.utils.identify_file_type import filetype_subdir from src.workers.base import BaseWorker from src.workers.enums import WorkerName, WorkerStatus @@ -65,13 +65,15 @@ def process(self, task: dict) -> None: 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}, " f"reappears={skipped_found}, head_verified={removed_head}, " - f"unreferenced={removed_unref}, orphaned_files={removed_orphan_files}" + f"unreferenced={removed_unref}, orphaned_files={removed_orphan_files}, " + f"trash_cleaned={removed_trash}" ) except Exception as e: @@ -159,7 +161,7 @@ def _file_exists_locally( from src.config.settings import get_config sigmaref_base = get_config().sigmaref_documents_path - subdir = FILETYPE_TO_SUBDIR.get(content_type, "misc") + 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("."): @@ -225,6 +227,40 @@ def _gc_orphaned_sigmaref_files(self) -> int: 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: diff --git a/tests/unit/application/documents/test_sigma_ref_downloader.py b/tests/unit/application/documents/test_sigma_ref_downloader.py index 32633983..f355f59a 100644 --- a/tests/unit/application/documents/test_sigma_ref_downloader.py +++ b/tests/unit/application/documents/test_sigma_ref_downloader.py @@ -8,7 +8,7 @@ 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_TO_SUBDIR +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 ( @@ -832,7 +832,7 @@ def _fake_download( ) assert result["downloaded"] == 1 - assert (output_dir / FILETYPE_TO_SUBDIR["markdown"] / expected_filename).exists() + 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 @@ -916,7 +916,7 @@ def _fake_download( assert result["total_refs"] == 1 # File should be written with {url_hash}{ext} naming expected_filename = f"{expected_hash}.md" - assert (output_dir / FILETYPE_TO_SUBDIR["markdown"] / expected_filename).exists() + assert (output_dir / filetype_subdir("markdown") / expected_filename).exists() class TestRuleReferences: diff --git a/tests/unit/application/services/test_chat_service_cache.py b/tests/unit/application/services/test_chat_service_cache.py index b43cc9f3..372971de 100644 --- a/tests/unit/application/services/test_chat_service_cache.py +++ b/tests/unit/application/services/test_chat_service_cache.py @@ -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 From 54ee4c0ff7cb07c95e9dd83a322e2cbdb8b71650 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:27:07 +0200 Subject: [PATCH 14/44] =?UTF-8?q?=E2=9C=A8=20feat:=20Phase=206=20-=20golde?= =?UTF-8?q?n=20set=20evaluation=20infrastructure?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add golden set loader/save (JSON) - Implement retrieval metrics (recall@k, precision@k, MRR, AP, context_precision/recall) - Add SearchEvaluator runner with per-query and aggregated results - Add exact vs approximate search comparison baseline - 64 unit tests for eval modules --- src/core/search/eval/__init__.py | 28 +++ src/core/search/eval/exact_search.py | 229 ++++++++++++++++++ src/core/search/eval/golden_set.py | 182 ++++++++++++++ src/core/search/eval/metrics.py | 178 ++++++++++++++ src/core/search/eval/runner.py | 179 ++++++++++++++ .../core/search/eval/test_exact_search.py | 131 ++++++++++ .../unit/core/search/eval/test_golden_set.py | 124 ++++++++++ tests/unit/core/search/eval/test_metrics.py | 141 +++++++++++ tests/unit/core/search/eval/test_runner.py | 121 +++++++++ 9 files changed, 1313 insertions(+) create mode 100644 src/core/search/eval/__init__.py create mode 100644 src/core/search/eval/exact_search.py create mode 100644 src/core/search/eval/golden_set.py create mode 100644 src/core/search/eval/metrics.py create mode 100644 src/core/search/eval/runner.py create mode 100644 tests/unit/core/search/eval/test_exact_search.py create mode 100644 tests/unit/core/search/eval/test_golden_set.py create mode 100644 tests/unit/core/search/eval/test_metrics.py create mode 100644 tests/unit/core/search/eval/test_runner.py 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/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"] From 32397d139c418b1866e5785e2a51dc566918ca97 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:36:29 +0200 Subject: [PATCH 15/44] =?UTF-8?q?=F0=9F=93=9D=20docs:=20mark=20Phase=209?= =?UTF-8?q?=20as=20DONE=20in=20refractor.md?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- refractor.md | 700 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 700 insertions(+) create mode 100644 refractor.md diff --git a/refractor.md b/refractor.md new file mode 100644 index 00000000..211c33f1 --- /dev/null +++ b/refractor.md @@ -0,0 +1,700 @@ +# Refractor — Sigma Rule Management + +Fichier de travail pour le refactoring de la gestion des règles Sigma dans l'application. + +## 1. Problématiques Identifiées + +### 1.1 Chemins de Validation Dupliqués + +Deux chemins de validation distincts avec des règles et des retours incompatibles : + +``` +Chemin A (Production) : + POST /api/v1/chat/upload → ChatService.validate_and_store_yaml() + → SigmaValidator.validate(bytes) + → YAML parse, required: id/name/description/detection + → detection non-empty, condition syntax check + → deprecated field warnings (level, falsepositives) + → Retourne: dict[str, Any] brut + → Stocké en mémoire dans ChatService._uploaded_rule + +Chemin B (Tests uniquement — mort) : + validate_sigma_rule(SigmaRule) + → title/condition/detection non-empty + → level enum, status enum + → Retourne: ValidationResult (Pydantic model) + → Aucun appelant en production +``` + +| Aspect | SigmaValidator (A) | validate_sigma_rule (B) | parser.py | +|--------|-------------------|------------------------|-----------| +| **Taille fichier** | 1MB max | — | 1MB max | +| **Parse YAML** | Oui | — (opère sur modèle) | Oui | +| **Champs requis** | `id`, `name`, `description`, `detection` | `title`, `detection`, `condition` | `title`, `detection` | +| **Validation detection** | Dict non-vide | Dict non-vide | Présence | +| **Validation condition** | Références syntaxe | Non-vide | — | +| **Level** | Warning deprecated | Validation enum | — | +| **Status** | Non vérifié | Validation enum | — | +| **Retour** | `dict` brut | `ValidationResult` | `SigmaRule \| None` | +| **Erreur** | Exception (422) | Pydantic model | `None` / log | + +### 1.2 Incohérence de Nommage des Champs + +- `SigmaValidator` vérifie `name` (norme Sigma v2) +- `SigmaRule` (modèle canonique) utilise `title` +- `parser.py` mappe YAML `title` → `SigmaRule.title` +- Conséquence : si une règle a `name:` dans le YAML, `SigmaValidator` l'accepte mais `SigmaRule.title` sera vide ; inversement si `title:`, le parseur fonctionne mais le validateur rejette. + +### 1.3 Deux Types `ValidationError` Incompatibles + +- `src/shared/exceptions.py:ValidationError(Exception)` — HTTP 422, catchable +- `src/application/documents/models.py:ValidationError(BaseModel)` — Pydantic model, non-catchable +- Même nom, même sémantique (field + message), API incompatible + +### 1.4 Duplication Downloader / Processor + +`src/application/documents/sigma_ref_downloader.py` (753 lignes) et `sigma_ref_processor.py` (413 lignes) partagent le même domaine (téléchargement de documents de référence) avec duplication lourde : + +| Fonction | downloader.py | processor.py | +|----------|--------------|-------------| +| **SHA256 fichier** | `_sha256_file()` (10 lignes) | `_sha256_file()` (7 lignes, quasi-identique) | +| **SHA256 string** | `_sha256()` (5 lignes) | `_sha256_bytes()` (4 lignes) | +| **HEAD request** | `_head_content_type()` | `_head_request()` (signatures différentes) | +| **Download HTTP** | `_download_file()` (3 retries + backoff) | `_download_one()` (lambda, 0 retry) | +| **Entry builder** | `_make_entry()` | `_build_head_entry()` / `_build_download_entry()` | +| **Normalisation URL** | `normalize_url()` (complète) | `_normalize_url()` (triviale — juste strip+rstrip) | +| **Constantes** | `DEFAULT_REQUEST_DELAY`, `DEFAULT_MAX_WORKERS` (lignes 50-51) | Mêmes constantes (lignes 20-21) | +| `_TYPE_TO_EXT` | Défini deux fois dans la même fonction (lignes 548-555, 630-637) | — | +| `httpx.Client` | Nouvelle instance par requête | Nouvelle instance par requête | +| `import hashlib` | Dans le corps des fonctions | Dans le corps des fonctions | + +Risque concret : la normalisation URL différente (downloader supprime fragments, processor non) peut causer des `url_hash` différents entre les deux chemins pour la même URL. + +### 1.5 Duplication Discovery Worker + +`src/workers/sigma/discovery_worker.py` (502 lignes) : + +- `_process_github` / `_process_spec` : structure quasi-identique (~65 lignes chacun), seule la stratégie d'énumération des repos change +- `_scan_all_github` / `_scan_all_spec` : quasi-identiques +- `_write_entries` / `_write_spec_entries` : quasi-identiques (seul le nom de méthode DB change) +- `get_sigma_rule_id` importé à l'intérieur d'une boucle (perte de perf) + +### 1.6 Code Mort + +| Fichier | Code | Statut | +|---------|------|--------| +| `chunker.py` | `_generate_eval_questions()` (lignes 503-515) | Jamais appelé. Le rich chunking utilise des listes inline. | +| `chunker.py` | `chunk_sigma_rules_rich()` (lignes 527-534) | Marqué "backwards-compatible", rien ne l'appelle. `SigmaChunker.process()` est le seul chemin actif. | +| `chunker.py` | `post_process()` — le `if enable_eval_questions` | N'a aucun effet observable. `return documents` est exécuté dans tous les cas. | +| `orchestrator.py` | `class RAGPipeline` (22 lignes) | Thin wrapper autour de `SearchEngine`, probablement mort. | +| `documents/validator.py` | `validate_sigma_rule()` | Importé et testé uniquement dans les tests. Aucun appelant en production. | +| `processor.py` | `GITHUB_BLOB_PATTERN: Any = None` (ligne 22) | Placeholder jamais rempli. | +| `downloader.py` | Paramètre `path` de `_load_registry` / `_save_registry` | Jamais utilisé, seul `db` est utilisé. | +| `downloader.py` | Paramètre `request_delay` de `download_references()` | Déclaré, jamais référencé dans le corps. | +| `processor.py` | Paramètre `request_delay` de `process_sigma_refs()` | Même chose. | + +### 1.7 Complexité du Chunker + +`SigmaChunker._chunk_rule()` fait ~346 lignes (lignes 89-434) et concentre trop de responsabilités : +- Extraction des champs depuis le dict brut +- Construction de templates textes pour 12+ types de chunks +- Itération sur les blocs de détection avec boucles imbriquées (field/operator groups + atomic indicators) +- Enrichissement LLM inline (lignes 414-432) +- Génération de questions d'évaluation inline + +### 1.8 Collision de Noms `RAGPipeline` + +Deux classes portent le même nom dans des packages différents : +- `src/application/chat/rag.py:RAGPipeline` — orchestrateur de génération LLM (actif) +- `src/core/pipeline/orchestrator.py:RAGPipeline` — thin wrapper SearchEngine (probablement mort) + +### 1.9 État Sessionnel ChatService + +`ChatService` maintient un état mutable en mémoire (session) : +- `_history`, `_uploaded_rule`, `_last_citations`, `_current_prompt_id` +- Fonctionne pour un usage mono-utilisateur local +- Empêche toute mise à l'échelle horizontale +- L'état devrait être externalisé (cache distribué, DB, etc.) + +### 1.10 Absence de Module HTTP Partagé + +- `httpx.Client` créé frais à chaque requête dans 4 fichiers : downloader, processor, translate (?), chat +- Aucun pool de connexions HTTP +- Logique de retry/backoff dupliquée (downloader l'a, processor ne l'a pas) +- HEAD requests dupliquées avec signatures différentes + +### 1.11 Double Parsing YAML + +`sigma_utils.py` parse le même fichier YAML deux fois si un appelant veut à la fois les références et le rule_id. Appelé depuis `discovery_worker.py` en boucle par fichier — l'impact est multiplié. + +### 1.12 Deux Chemins d'Indexation ≠ + +- **Chemin A (Upload interactif)** : validation uniquement → stocké en mémoire → PAS indexé dans Qdrant +- **Chemin B (Discovery fond)** : parse → chunk → index dans Qdrant (sigma_rules) +- Écart : les règles uploadées interactivement ne sont PAS présentes dans la recherche vectorielle + +### 1.13 Aucun Tuning Qdrant Spécifique aux Collections + +Les 3 collections Qdrant (`sigma_rules`, `sigma_docs`, `sigma_spec`) sont créées avec la même configuration minimale : + +```python +VectorParams(size=384, distance=Distance.COSINE) +SparseVectorParams(index=SparseIndexParams()) # "text-sparse" +``` + +| Paramètre manquant | Impact | +|-------------------|--------| +| `hnsw_config` | Utilise les defaults serveur (m=16, ef_construct=100) — probablement ok mais non évalué | +| `quantization_config` | Aucune quantification — les vecteurs dense sont stockés en float32 intégral | +| `optimizers_config` | Aucun tuning — indexing_threshold_kb par défaut (20 MB) | +| Per-collection tuning | Les 3 collections ont la même config, alors que leurs usages diffèrent (recherche fréquente vs froide) | + +### 1.14 Custom BM25 Sparse Encoder Non Standard + +`src/core/search/sparse_encoder.py` implémente un BM25 custom avec : +- Token IDs basés sur MD5 modulo 2^24 +- Poids : `1.0 + log(term_frequency)` (pas de IDF, pas de normalisation par doc length) +- Stop words anglais codés en dur (107 mots) +- **Pas de stemming, pas de configuration par langue** + +Points d'attention : +- L'absence d'IDF signifie que les tokens fréquents ne sont pas pénalisés +- Le hash MD5 peut causer des collisions (2^24 ≈ 16M IDs, acceptable mais non déterministe) +- Aucune configuration pour le français alors que le modèle d'embedding est `multilingual-e5-small` +- Pas de `avg_len` configuré pour BM25 — Qdrant ne peut pas normaliser par longueur de document + +### 1.15 Pas d'Évaluation de la Qualité de Recherche + +- Aucun golden set / ground truth dataset pour mesurer recall@k +- Aucune évaluation comparative : exact search vs approximate, flat vs rich chunks +- Le paramètre `alpha=0.3` du hybrid search n'a jamais été tuné sur des données réelles +- Fusion RRF avec `k=60` — valeur par défaut, jamais ajustée + +### 1.16 Pipeline d'Indexation Synchrone avec Batch Size Faible + +- `IngestionPipeline.run()` utilise `num_workers=0` (séquentiel) +- `embed_batch_size=8` — très conservateur pour un modèle 384-dim +- Pas de désactivation HNSW pendant le bulk load (reconstruit à chaque insertion) +- Stratégie delete-and-reindex : pas d'upsert incrémental, reconstruction complète + +### 1.17 Deux Chemins Concurrents pour le Téléchargement des Références + +Le téléchargement des documents de référence des règles Sigma est implémenté par **deux fichiers distincts** qui font la même chose différemment : + +| Aspect | `sigma_ref_downloader.py` (753 lignes) | `sigma_ref_processor.py` (413 lignes) | +|--------|----------------------------------------|---------------------------------------| +| **Point d'entrée** | `download_references(rules_dir, ...)` — scanne les fichiers YAML directement | `process_sigma_refs(db, ...)` — lit les entrées du `doc_registry` | +| **Appelé par** | API endpoint `POST /api/v1/documents/index-sigma-ref` | Worker `SigmaRefProcessor` (background) | +| **Normalisation URL** | Complète : GitHub blob→raw, strips fragments, refs/heads | Triviale : `url.strip().rstrip("/")` | +| **Nommage fichiers** | `{url_hash}{ext}` (ex: `abc123.md`) | `_sanitize_filename(url)` (ex: `documentation_page`) | +| **Registry** | Dict mémoire + flush final vers DB | Appels DB unitaires par opération | +| **Retry download** | Oui (3 retries, backoff exponentiel) | Non (0 retry, pas de backoff) | +| **SSRF protection** | Oui (`_is_private_url()`) | **Non** | +| **HEAD request** | `_head_content_type()` — retourne Content-Type seulement | `_head_request()` — retourne Content-Type + size + final_url | +| **hashlib import** | Dans le corps des fonctions | Dans le corps des fonctions | +| **Lock** | `_registry_lock` global (threading.Lock) | Aucun | + +**Conséquence directe** : le même URL peut avoir un `url_hash` différent entre les deux chemins à cause de la normalisation divergente, causant des doublons dans le registry et des téléchargements redondants. + +### 1.18 Problèmes de Stockage Local des Références + +Les fichiers téléchargés sont stockés dans `data/documents/sigmaref/` : + +- **Aucune structure de sous-répertoires** — tous les fichiers sont à plat dans un seul dossier (plusieurs milliers de fichiers potentiellement) +- **Nommage incohérent** — le downloader utilise des hash (déterministe, prévisible), le processor utilise le nom du fichier original (variable, risque de collision) +- **Pas de nettoyage** — les fichiers orphelins (URLs qui ne sont plus référencés par aucune règle) ne sont jamais supprimés +- **Pas de déduplication cross-rule** — si une URL est référencée par 10 règles, elle est téléchargée 1 fois mais son `rule_id`/`title` dans le registry est écrasé à chaque mise à jour (perte de l'information de provenance) + +### 1.19 Cycle de Vie `embed_status` Incohérent + +Le statut d'embedding traverse `discovery` → `head_verified` → `embedded`, mais : + +- `process_sigma_refs()` crée des entrées avec `embed_status = "head_verified"` pour les URLs de type non supporté → **entrées mortes** qui ne seront jamais promues +- `download_references()` utilise toujours `embed_status = "discovery"` après téléchargement +- Aucun mécanisme de nettoyage pour les entrées bloquées en `head_verified` + +### 1.20 Pas de Table de Jonction Règle↔Référence + +La table `doc_registry` stocke `rule_id` et `title` par URL, mais : +- Si 2 règles partagent la même référence, seule la dernière écriture survit +- Impossible de répondre à "quelles règles référencent ce document ?" +- Impossible de savoir si un document est orphelin (plus référencé par aucune règle) + +## 2. Propositions d'Optimisation + +### P0 — Priorité Critique + +#### P0.1 Extraire les fonctions SHA256 partagées + +Créer `src/shared/utils/crypto_utils.py` : +- `compute_sha256_str(data: str) -> str` +- `compute_sha256_file(path: Path) -> str` + +Supprimer les 3 implémentations dupliquées dans : +- `discovery_base.py:15` +- `sigma_ref_downloader.py:725` +- `sigma_ref_processor.py:392` + +#### P0.2 Unifier la normalisation URL + +Une seule `normalize_url()` dans `src/shared/utils/url_utils.py` avec la logique complète du downloader (GitHub blob→raw, fragments, refs/heads). Utiliser partout pour garantir la cohérence des `url_hash`. + +### P1 — Priorité Haute + +#### P1.1 Consolider les validateurs + +- `SigmaValidator.validate()` retourne un `SigmaRule` (via `SigmaRule.from_dict()`) au lieu d'un `dict` brut +- Fusionner les contrôles de `validate_sigma_rule()` (level enum, status enum) dans `SigmaValidator` +- Normaliser les champs : `title` partout, avec alias `name` en lecture si besoin +- Supprimer `src/application/documents/validator.py` après fusion +- Standardiser sur `src/shared/exceptions.ValidationError` (exception) + +#### P1.2 Extraire le download HTTP partagé + +Créer `src/shared/utils/http_utils.py` : +- `head_url(url, timeout=10) -> tuple[str|None, int|None, str|None]` +- `download_file(url, output_path, max_retries=3) -> tuple[bool, int|None]` +- Pool de connexions `httpx.AsyncClient` partagé +- Logique de retry avec backoff harmonisée + +#### P1.3 Factoriser les builders d'entrées registry + +Fonction unique `build_registry_entry()` dans `src/shared/utils/registry_utils.py` avec paramètre `embed_status` par défaut. Remplacer `_make_entry`, `_build_head_entry`, `_build_download_entry`. + +#### P1.4 Factoriser les méthodes du DiscoveryWorker + +- `_scan_all_github` + `_scan_all_spec` → `_scan_all_files(prepare_entry_fn)` +- `_write_entries` + `_write_spec_entries` → `_write_entries(batch_upsert_fn, spec_mode=False)` +- `_process_github` + `_process_spec` → paramétrer la stratégie d'énumération des repos + +#### P1.5 Supprimer le code mort + +- `_generate_eval_questions()` dans chunker.py +- `chunk_sigma_rules_rich()` dans chunker.py +- `orchestrator.py:RAGPipeline` — supprimer la classe, les appelants utilisent `SearchEngine` directement +- `GITHUB_BLOB_PATTERN: Any = None` dans processor.py +- Paramètres morts (`path`, `request_delay`) dans downloader/processor +- `validate_sigma_rule()` dans documents/validator.py (après fusion P1.1) + +### P2 — Priorité Moyenne + +#### P2.1 Décomposer `_chunk_rule()` + +Extraire : +- `_extract_fields(rule: dict) -> dict` (lignes 98-116) +- `_build_executive_summary(rule, fields) -> dict` (chunk type 1) +- `_build_metadata_lifecycle(rule, fields) -> dict` (chunk type 2) +- ... et ainsi de suite pour les 12+ types +- `_enrich_by_llm(chunks) -> chunks` (étape séparée, lignes 414-432) + +#### P2.2 Résoudre la collision RAGPipeline + +Renommer ou supprimer `src/core/pipeline/orchestrator.py:RAGPipeline`. + +#### P2.3 Externaliser l'état ChatService + +- Propager `session_id` depuis l'API +- Stocker `_history`, `_uploaded_rule` dans un store (cache LRU, DuckDB, etc.) +- Permettre la scalabilité horizontale + +#### P2.4 Normaliser les imports + +- `import hashlib` en haut des fichiers, pas dans le corps des fonctions +- `from src.shared.utils.sigma_utils import get_sigma_rule_id` en haut de `discovery_worker.py`, pas dans la boucle + +#### P2.5 Constante NULL_UUID + +Définir `NULL_UUID = "00000000-0000-0000-0000-000000000000"` dans `src/shared/constants.py` et l'utiliser partout. + +#### P2.6 Dédoublonner `_TYPE_TO_EXT` + +Définir une fois comme constante module dans `identify_file_type.py` ou `sigma_ref_downloader.py`. + +### P3 — Priorité Basse + +#### P3.1 Split downloader.py (753 lignes) + +Diviser `sigma_ref_downloader.py` en modules < 400 lignes : +- `sigma_ref_downloader.py` — orchestration +- `sigma_ref_http.py` — HTTP helpers +- `sigma_ref_registry.py` — registry management + +#### P3.2 Connection pooling HTTP + +Réutiliser `httpx.Client` (ou `AsyncClient`) avec un pool de connexions dans toutes les phases de téléchargement batch. + +#### P3.3 Unifier les deux chemins d'indexation + +Permettre aux règles uploadées interactivement d'être optionnellement indexées dans Qdrant (ex. paramètre `index_after_upload`). + +### R1 — Priorité Haute (Téléchargement Références) + +#### R1.1 Unifier `download_references()` et `process_sigma_refs()` + +Une seule fonction `download_sigma_references(source, db, ...)` avec deux modes d'entrée : +- `mode="scan"` : scanne les fichiers YAML directement (ancien downloader) +- `mode="registry"` : lit les entrées du doc_registry (ancien processor) + +La normalisation URL, le nommage des fichiers, le retry, et la détection SSRF doivent être **identiques** dans les deux modes. + +#### R1.2 Standardiser le Nommage des Fichiers + +- Format unique : `{url_hash}{ext}` (comme le downloader actuel) +- Supprimer `_sanitize_filename()` du processor +- Garantir qu'un même URL produit toujours le même chemin disque + +#### R1.3 Ajouter une Table de Jonction `rule_references` + +```sql +CREATE TABLE rule_references ( + rule_id TEXT NOT NULL, + url_hash TEXT NOT NULL, + PRIMARY KEY (rule_id, url_hash) +); +``` + +- Permet de retrouver toutes les références d'une règle +- Permet de détecter les documents orphelins (URLs non référencées) +- Nettoyage possible : supprimer les fichiers dont l'URL_hash n'est plus dans `rule_references` + +#### R1.4 Nettoyer les Entrées Mortes du Registry + +- Supprimer les entrées `head_verified` de plus de X jours +- Supprimer les entrées dont `url_hash` n'est plus dans `rule_references` +- Tâche planifiée (cron) ou déclenchée manuellement + +#### R1.5 Ajouter SSRF Protection dans le Processor + +Copier `_is_private_url()` du downloader et l'utiliser dans `process_sigma_refs()` avant toute HEAD request. Sans ça, le processor est vulnérable aux SSRF. + +### R2 — Priorité Moyenne (Organisation Stockage) + +#### R2.1 Structure de Sous-Répertoires + +``` +data/documents/sigmaref/ + ├── markdown/ # .md files + ├── html/ # .html files + ├── pdf/ # .pdf files + ├── plain_text/ # .txt files + └── office/ # .docx files +``` + +- Préserve le nom `{url_hash}{ext}` dans chaque sous-répertoire +- Évite d'avoir 10 000 fichiers dans le même dossier + +#### R2.2 Garbage Collection des Fichiers Orphelins + +- Scanner `data/documents/sigmaref/` pour les fichiers dont le nom ne correspond à aucun `url_hash` dans `doc_registry` +- Supprimer les fichiers orphelins (avec confirmation) +- Option : déplacer vers `.trash/` avant suppression définitive + +### R3 — Priorité Basse (Traçabilité) + +#### R3.1 Enrichir le Registry avec la Liste des Règles Sources + +Ajouter un champ `referenced_by: list[str]` dans `doc_registry` (ou utiliser la table `rule_references`) pour savoir quelles règles référencent un document donné. Utile pour : +- Debug : "pourquoi ce document a été téléchargé ?" +- Mise à jour : "re-télécharger les docs des règles modifiées" +- Suppression : "ce document n'est plus référencé, on peut le supprimer" + +## 3. Plan d'Exécution Suggéré + +``` +Phase 1 — P0 (nettoyage critique) + ├── P0.1 Extraire crypto_utils.py (SHA256) + ├── P0.2 Unifier normalize_url() + ├── P1.5 Supprimer le code mort évident + └── P2.4 Normaliser les imports hashlib/sigma_utils + +Phase 2 — P1 (consolidation validateurs + HTTP) + ├── P1.1 Consolider SigmaValidator → retourne SigmaRule + ├── P1.2 Extraire http_utils.py (HEAD + download + retry) + ├── P1.3 Factoriser registry entry builders + └── Supprimer doc la mort (validate_sigma_rule, etc.) + +Phase 3 — P1 (factorisation workers) + ├── P1.4 Factoriser _scan_all_* / _write_* / _process_* + └── P2.5 Constante NULL_UUID + +Phase 4 — P2 (qualité de code) + ├── P2.1 Décomposer _chunk_rule() + ├── P2.2 Résoudre collision RAGPipeline + ├── P2.3 Externaliser état ChatService + └── P2.6 Dédoublonner _TYPE_TO_EXT + +Phase 5 — P3 (architecture) + ├── P3.1 Split downloader.py + ├── P3.2 Connection pooling HTTP + └── P3.3 Unifier chemins d'indexation +``` + +## 4. Propositions Qdrant (Basées sur les Skills) + +### Q0 — Priorité Critique (Qualité de Recherche) + +#### Q0.1 Créer un Golden Set pour Évaluer Recall@k + +- Échantillonner 100-200 requêtes réelles avec jugements de pertinence +- Mesurer recall@k avant/après chaque changement de config +- Voir Qdrant Search Quality Diagnosis skill + +#### Q0.2 Tester Exact Search Comme Baseline + +- Avant tout tuning HNSH, comparer exact search vs approximate search +- Si l'écart est > 5%, tuner `ef` et `m` du HNSW +- Permet d'isoler les problèmes de modèle d'embedding vs index + +### Q1 — Priorité Haute (Performance Indexation) + +#### Q1.1 Ajouter `quantization_config` aux Collections + +```python +quantization_config=ScalarQuantization( + scalar=ScalarQuantizationConfig( + type=ScalarType.INT8, + always_ram=True, + quantile=0.5, + ) +) +``` + +- Réduction mémoire 4x pour les vecteurs dense en RAM +- Perte de qualité < 1% recall avec rescore +- Activation immédiate pour les 3 collections + +#### Q1.2 Augmenter `embed_batch_size` et Paralléliser + +- `embed_batch_size` : 8 → **64** (384-dim, CPU, safe) +- `num_workers` : 0 → **2-4** pour ingestion parallèle +- Désactiver HNSW pendant le bulk load : `indexing_threshold_kb = 0` → temporairement très haut, restaurer après + +### Q2 — Priorité Moyenne (Tuning Recherche) + +#### Q2.1 Tuner l'Alpha du Hybrid Search par Collection + +- Remplacer le `alpha=0.3` global par des valeurs par collection : + - `sigma_rules` : α=0.5 (sémantique + lexical équilibré — les règles ont un vocabulaire technique précis) + - `sigma_docs` : α=0.7 (plus lexical — docs de référence, termes exacts) + - `sigma_spec` : α=0.3 (plus sémantique — spécifications, concepts) +- Évaluer avec le golden set (Q0.1) + +#### Q2.2 Tuner la Fusion RRF + +- `k=60` actuel → tester `k=30`, `k=60`, `k=100` avec le golden set +- Poids par collection dans le RRF (weighted RRF) si une collection domine +- Envisager DBSF si les distributions de score entre dense et sparse sont trop différentes + +#### Q2.3 Ajouter `hnsw_config` par Collection + +```python +hnsw_config=HnswConfigDiff( + m=16, # default ok pour 384-dim + ef_construct=200, # 100 → 200 pour meilleure qualité à l'indexation + full_scan_threshold_kb=10000, # 10 MB, identique au serveur + on_disk=False, # sigma_rules en RAM +) +``` + +- `sigma_docs` et `sigma_spec` (collections froides) : `on_disk=True` + `async_scorer` + +### Q3 — Priorité Basse (Architecture Vectorstore) + +#### Q3.1 Stocker les Sparse Vectors sur Disk + +```python +sparse_vectors_config={ + "text-sparse": SparseVectorParams( + index=SparseIndexParams(on_disk=True) + ) +} +``` + +- Les vecteurs sparse BM25 sont rarement tous consultés +- Bon candidat pour le stockage disque (économise RAM) + +#### Q3.2 Revoir le Sparse Encoder Custom + +- Benchmarker `bm25_sparse_encoder` vs le BM25 natif de Qdrant (configuré par language) +- Avantages BM25 natif Qdrant : + - Calcul côté serveur (pas de transfert des vecteurs sparse) + - Tokenization + stemming standardisés + - Support multi-langue (français, etc.) + - IDF calculé automatiquement +- Si le custom encoder est conservé, ajouter au moins le calcul d'IDF et la normalisation par doc length + +#### Q3.3 Oversampling + Rescore avec Quantification + +```python +quantization_config=ScalarQuantization( + scalar=ScalarQuantizationConfig(type=ScalarType.INT8, always_ram=True), + rescore=True, + oversampling=2.0, +) +``` + +- Permet de chercher dans un pool 2x plus large (vitesse ×2 grâce à la quantification) +- Rescore les top_k sur les vecteurs originaux pour préserver la qualité + +#### Q3.4 Pipeline d'Indexation Incrémental + +- Remplacer la stratégie delete-and-reindex par des upserts par `rule_id` +- Nécessite : identifier les règles nouvelles/modifiées vs inchangées +- Avantage : pas de downtime de la collection, indexation plus rapide +- Combinable avec `indexing_threshold_kb` haut initial + baisse progressive + +## 5. Call Graph — Phase 0 Audit + +### SigmaValidator.validate() + +**Production callers (1 fichier, 1 site d'appel direct) :** +| Fichier | Ligne | Usage | Type attente | +|---------|-------|-------|-------------| +| `src/application/chat/service.py` | 357 | `ChatService.validate_and_store_yaml()` → stocke dans `_uploaded_rule` | `dict[str, Any]` | +| `src/application/chat/rag.py` | 99,133,269,305 | `explain_rule()`, `explain_rule_stream()`, `analyze_coverage()`, `analyze_coverage_stream()` | `dict[str, Any]` (param `rule_data`) | +| `src/application/chat/rag.py` | 380 | `_format_rule_yaml()` | `yaml.dump(rule)` attend un dict | +| `src/application/chat/rag.py` | 387 | `_fallback_explanation()` | `.get('name', 'Unknown')`, `.get('id')`, `.get('description')` | + +**Consommateurs indirects de `_uploaded_rule` (dict) via `ChatService` :** +| Méthode | Accès | Ligne | +|---------|-------|-------| +| `_handle_explain()` | `.get("name", "")` | 234 | +| `_handle_explain_stream()` | `.get("name", "")` | 248 | +| `_handle_coverage()` | passé tel quel | 265 | +| `_handle_coverage_stream()` | passé tel quel | 342 | + +**Tests :** +| Fichier | Usage | +|---------|-------| +| `tests/unit/application/services/test_sigma_validator.py` | 19 tests unitaires, `validate()` sur bytes | +| `tests/unit/application/services/test_sigma_validator_advanced.py` | 6 tests avancés | +| `tests/unit/application/services/test_chat_service_cache.py` | `SigmaValidator` mocké (`patch`) | +| `tests/integration/test_chat_flow.py` | `_uploaded_rule` set manuellement comme dict (l.92) | + +### download_references() + +**Production callers (1 fichier, 1 site) :** +| Fichier | Ligne | Contexte | +|---------|-------|----------| +| `src/api/v1/documents/documents.py` | 35 | `POST /api/v1/documents/index-sigma-ref` | + +**Tests :** +| Fichier | Usage | +|---------|-------| +| `tests/unit/application/documents/test_sigma_ref_downloader.py` | 15+ appels, tests unitaires complets | + +### process_sigma_refs() + +**Production callers (1 fichier, 1 site) :** +| Fichier | Ligne | Contexte | +|---------|-------|----------| +| `src/workers/sigma/sigmaref_worker.py` | 41 | `SigmaRefProcessor.process()` worker background | + +**Tests :** +| Fichier | Usage | +|---------|-------| +| `tests/unit/workers/test_discovery_workers.py` | Mocké (l.28, 50) | + +### validate_sigma_rule() — CONFIRMÉ CODE MORT + +**Production callers : AUCUN** + +**Tests uniquement :** +| Fichier | Usage | +|---------|-------| +| `tests/unit/application/documents/test_documents.py` | 5 appels (l.69, 84, 98, 113, 131) | + +### Risques Identifiés pour P1.1 (dict → SigmaRule) + +1. `ChatService._uploaded_rule` typé `dict[str, Any] | None` → doit passer à `SigmaRule | None` +2. Tous les `.get("name", "")` → deviennent `.title` (ou alias `.name` si préservé) +3. `_format_rule_yaml()` utilise `yaml.dump(rule_dict)` → nécessite `rule.model_dump()` +4. `_fallback_explanation()` utilise `.get()` sur dict → nécessite accès attribut +5. Mock dans `test_chat_service_cache.py` patch `SigmaValidator` → pas de changement nécessaire +6. Test intégration `test_chat_flow.py` set `_uploaded_rule` comme dict brut → à migrer vers `SigmaRule` + +### Tests de Régression à Écrire (avant P1.1) DONE by commit 6474065 + +- [x] Capturer le contrat `dict` actuel de `SigmaValidator.validate()` (bracket access, .get(), yaml.dump(), ValidationError) +- [x] Capturer les patterns de `ChatService._uploaded_rule` (`.get("name", "")`, `.get("id", "N/A")`) +- [x] Capturer le contrat `_format_rule_yaml()` et `_fallback_explanation()` + +--- + +## 6. Plan d'Exécution Final + +``` +Phase 1 — P0 + R (nettoyage critique) DONE by commmit 8d24148dbfe4be3e6fe529800f5516613a0217c6 + ├── P0.1 Extraire crypto_utils.py (SHA256) + ├── P0.2 Unifier normalize_url() + ├── R1.5 Ajouter SSRF protection dans le processor + ├── P1.5 Supprimer le code mort évident + └── P2.4 Normaliser les imports hashlib/sigma_utils + +Phase 2 — P1.2 + P1.3 (infra indépendante, 0 risque) DONE by commit 6474065 + ├── P1.2 Extraire src/shared/http.py (HEAD + download + retry + pool httpx) + ├── P1.3 Factoriser build_registry_entry() — pure factory sans IO + └── Tests : mock httpx, retry/backoff, SSRF, golden path + edge cases + +Phase 3 — P1.1 (validator consolidation, protégé par audit Phase 0) DONE by commit 259bccd + ├── Audit préalable : cartographier tous les callers, tests de régression + ├── SigmaValidator.validate() → retourne SigmaRule (pydantic) + ├── Merge checks de validate_sigma_rule() + shared.exceptions.ValidationError + ├── Normalisation title/name (name alias via @property) + ├── Supprimer validate_sigma_rule(), ValidationError/ValidationResult + └── Tests : 52 unit tests verts, ruff/mypy clean + +Phase 4 — R1.1 + R1.2 (download unification, dépend de P1.2) DONE by commit 2f2efe7 + ├── download_sigma_references(source, db, mode="scan"|"registry") + ├── Standardiser nommage {url_hash}{ext} dans les 2 modes + ├── Contract test : même URL → même nom de fichier + └── processor.py passe de 318l à 56l (délégation pure) + +Phase 5 — R1 (traçabilité références) DONE by commit 33e1f70 + follow-up + ├── R1.3 Ajouter table rule_references (junction rule↔reference) + ├── R1.4 Nettoyer entrées mortes head_verified + │ ├── delete_head_verified_orphans() — head_verified sans content_sha256 + │ └── delete_unreferenced_entries() — sigmaref entries sans url_hash dans rule_references + ├── R1.4 integré dans DocGCWorker.process() + └── R3.1 Enrichir registry avec liste des règles sources + +Phase 6 — Q0-Q1 (qualité recherche + perf indexation Qdrant) + ├── Q0.1 Créer golden set pour évaluation recall@k + ├── Q0.2 Tester exact search comme baseline + ├── Q1.1 Ajouter quantization_config aux 3 collections + └── Q1.2 Augmenter embed_batch_size (8→64) + paralléliser (workers 0→4) + +Phase 7 — P1 (factorisation workers) DONE + ├── P1.4 Factoriser _scan_all_github + _scan_all_spec → _scan_all(prepare_fn) + │ ├── _write_entries + _write_spec_entries → _write_entries(batch_upsert_fn) + │ └── _collect_repo_files() extrait du rglob commun github+spec + ├── P2.5 Constante NULL_UUID dans src/shared/constants.py (6 fichiers modifiés) + └── Q2.3 Ajouter hnsw_config par collection via collection_hnsw_config() + ├── sigma_rules → in-RAM, ef_construct=200 + ├── sigma_docs → on-disk + └── sigma_spec → on-disk + +Phase 8 — P2 + R2 (qualité de code + organisation stockage) + ├── P2.1 Décomposer _chunk_rule() + ├── P2.2 Résoudre collision RAGPipeline + ├── P2.3 Externaliser état ChatService + ├── P2.6 Dédoublonner _TYPE_TO_EXT + ├── R2.1 Structure sous-répertoires par type (markdown/, html/, pdf/) + └── R2.2 Garbage collection fichiers orphelins + +Phase 9 — Q2-Q3 (tuning recherche + architecture vectorstore) DONE + ├── Q2.1 Tuner alpha hybrid search par collection + ├── Q2.2 Tuner fusion RRF (k, weighted) + ├── Q3.1 Stocker sparse vectors sur disk + ├── Q3.2 Revoir sparse encoder (benchmark BM25 natif) + └── Q3.4 Pipeline d'indexation incrémental + +Phase 10 — P3 + Q3 (architecture finale) + ├── P3.1 Split downloader.py (si toujours pertinent après R1.1) + ├── P3.2 Connection pooling HTTP + ├── P3.3 Unifier chemins d'indexation (upload→Qdrant) + └── Q3.3 Oversampling + rescore avec quantification +``` + + +# opencode -s ses_0deda6082ffeIa2vHoUdYiWbpY \ No newline at end of file From 833fe40eb0d3de24d6c9bd92e3c3e993a52a7299 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 16:57:05 +0200 Subject: [PATCH 16/44] =?UTF-8?q?=E2=9C=A8=20feat:=20complete=20P3.1=20spl?= =?UTF-8?q?it=20of=20sigma=5Fref=5Fdownloader.py=20into=20focused=20module?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sigma_ref_paths.py: path resolution helpers (subdir_for, sigmaref_write_path, sigmaref_resolve_path) - sigma_ref_url.py: URL type detection (detect_url_type, is_reference_url, resolve_ext) - sigma_ref_registry.py: registry management via DB (load_registry, save_registry, load_error_registry, maybe_record_error) - Refactored sigma_ref_downloader.py to use the new modules - Fixed detect_url_type to properly map extensions to file types - All 51 tests passing, ruff and mypy clean --- refractor.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/refractor.md b/refractor.md index 211c33f1..bae0d8cf 100644 --- a/refractor.md +++ b/refractor.md @@ -690,7 +690,10 @@ Phase 9 — Q2-Q3 (tuning recherche + architecture vectorstore) DONE └── Q3.4 Pipeline d'indexation incrémental Phase 10 — P3 + Q3 (architecture finale) - ├── P3.1 Split downloader.py (si toujours pertinent après R1.1) + ├── P3.1 ✅ Split downloader.py (si toujours pertinent après R1.1) + │ ├── sigma_ref_paths.py — subdir_for(), sigmaref_write_path(), sigmaref_resolve_path() + │ ├── sigma_ref_url.py — detect_url_type(), is_reference_url(), resolve_ext() + │ └── sigma_ref_registry.py — load_registry(), save_registry(), load_error_registry(), maybe_record_error() ├── P3.2 Connection pooling HTTP ├── P3.3 Unifier chemins d'indexation (upload→Qdrant) └── Q3.3 Oversampling + rescore avec quantification From 579bac36fc266a9a0d3d956ebe93ad7d333c9a47 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:06:20 +0200 Subject: [PATCH 17/44] =?UTF-8?q?=E2=9C=A8=20feat:=20complete=20P3.1=20spl?= =?UTF-8?q?it=20of=20sigma=5Fref=5Fdownloader.py=20into=20focused=20module?= =?UTF-8?q?s?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sigma_ref_paths.py: path resolution helpers (subdir_for, sigmaref_write_path, sigmaref_resolve_path) - sigma_ref_url.py: URL type detection (detect_url_type, is_reference_url, resolve_ext) - sigma_ref_registry.py: registry management via DB (load_registry, save_registry, load_error_registry, maybe_record_error) - Refactored sigma_ref_downloader.py to use the new modules - Fixed detect_url_type to properly map extensions to file types - All 51 tests passing, ruff and mypy clean --- .../documents/sigma_ref_downloader.py | 215 ++---------------- src/application/documents/sigma_ref_paths.py | 86 +++++++ .../documents/sigma_ref_registry.py | 135 +++++++++++ src/application/documents/sigma_ref_url.py | 105 +++++++++ 4 files changed, 345 insertions(+), 196 deletions(-) create mode 100644 src/application/documents/sigma_ref_paths.py create mode 100644 src/application/documents/sigma_ref_registry.py create mode 100644 src/application/documents/sigma_ref_url.py diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index 046fafa8..a3011f7a 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -4,7 +4,6 @@ import logging import threading -import urllib.parse from collections.abc import Callable from concurrent.futures import Future, ThreadPoolExecutor, as_completed from pathlib import Path @@ -13,7 +12,6 @@ import yaml from src.shared.constants import NULL_UUID -from src.shared.http import RETRY_STATUSES 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 @@ -26,187 +24,36 @@ SUPPORTED_DOC_EXTENSION_MAP, SUPPORTED_REFERENCE_DOC_TYPES, filetype_ext, - filetype_subdir, ) -from src.shared.utils.url_utils import is_private_url, normalize_url, url_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__) 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 _subdir_for(content_type: str | None) -> str: - """Return the subdirectory name for a given content type.""" - return filetype_subdir(content_type or "") - - -def _sigmaref_write_path(output_path: Path, content_type: str | None, file_name: str) -> Path: - """Return the subdir path for writing a sigmaref file, creating the subdir as needed.""" - subdir = _subdir_for(content_type) - path = output_path / subdir - path.mkdir(parents=True, exist_ok=True) - return path / file_name - - -def _sigmaref_resolve_path(output_path: Path, content_type: str | None, file_name: str) -> Path: - """Resolve the path to an existing sigmaref file, with flat-layout fallback. - - Checks the subdir first, then the old flat layout for backward compatibility - with files downloaded before R2.1. Returns the subdir path if neither exists. - """ - candidate = output_path / _subdir_for(content_type) / file_name - if candidate.exists(): - return candidate - flat = output_path / file_name - if flat.exists(): - return flat - return candidate - - -def _detect_url_type(url: str, content_type: str | None = None) -> str | None: - """Detect the document type of a reference URL. - - Checks URL extension first, then falls back to HTTP Content-Type. - - Args: - url: The reference URL. - content_type: Optional HTTP Content-Type header value. - - 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 _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. - """ - 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) - - -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 download_sigma_references( db: DatabaseService, output_dir: str, @@ -460,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") @@ -585,7 +432,7 @@ def _collect_yaml_files() -> list[Path]: future_map = {} for item in download_queue: future = executor.submit( - http_download_file, + http_download_file, # type: ignore[arg-type] item["normalized_url"], item["output_file"], check_ssrf=False, @@ -658,30 +505,6 @@ def _empty_summary() -> dict[str, Any]: # ------------------------------------------------------------------ -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 _download_registry_mode( output_dir: str, db: DatabaseService, @@ -887,8 +710,8 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: result = future.result() if result is None: skipped += 1 - elif result[0] == "ok": - _, url_hash, content_hash, size = result + 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"]), diff --git a/src/application/documents/sigma_ref_paths.py b/src/application/documents/sigma_ref_paths.py new file mode 100644 index 00000000..68efe488 --- /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 output_path / 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 base / file_name + + if org and repo: + return Path(cfg.sigmaref_documents_path).resolve() / org / repo / file_name + + return None 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..1a4bf667 --- /dev/null +++ b/src/application/documents/sigma_ref_url.py @@ -0,0 +1,105 @@ +"""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" + if "pdf" in ctype: + return "pdf" + if "html" in ctype: + return "html" + if "text" 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 From 709c509e78242d455779f94df8ed67b1c48fec97 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:10:54 +0200 Subject: [PATCH 18/44] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20perf:=20add=20HTTP?= =?UTF-8?q?=20connection=20pooling=20to=20shared=20http=20module?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add get_pooled_client() with httpx.HTTPTransport for connection reuse - Pool keyed by (timeout, follow_redirects) for correct isolation - Limits: 100 max connections, 20 keepalive, 30s expiry - Update head_url() and download_file() to use pooled clients - Add close_all_pooled_clients() for clean shutdown - Update all tests to mock get_pooled_client instead of create_client - Add 4 tests for pool behavior (dedup, isolation, cleanup) --- src/shared/http.py | 92 +++++++++++++++++++++++++++++----- tests/unit/shared/test_http.py | 79 ++++++++++++++++++----------- 2 files changed, 130 insertions(+), 41 deletions(-) diff --git a/src/shared/http.py b/src/shared/http.py index 5e88bc7f..a136ff20 100644 --- a/src/shared/http.py +++ b/src/shared/http.py @@ -8,8 +8,10 @@ from __future__ import annotations import logging +import threading import time from pathlib import Path +from typing import Any import httpx @@ -24,6 +26,70 @@ 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, @@ -44,7 +110,7 @@ def create_client( ------- httpx.Client """ - merged = {"User-Agent": DEFAULT_USER_AGENT} + merged: dict[str, Any] = {"User-Agent": DEFAULT_USER_AGENT} if headers: merged.update(headers) return httpx.Client( @@ -81,15 +147,15 @@ def head_url( return None, None, None try: - with create_client(timeout=timeout) as client: - 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) + 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 @@ -127,12 +193,12 @@ def download_file( return False, None path = Path(output_path) + client = get_pooled_client(timeout=timeout) for attempt in range(1, max_retries + 1): try: - with create_client(timeout=timeout) as client: - resp = client.get(url) - resp.raise_for_status() + resp = client.get(url) + resp.raise_for_status() path.parent.mkdir(parents=True, exist_ok=True) path.write_bytes(resp.content) return True, None diff --git a/tests/unit/shared/test_http.py b/tests/unit/shared/test_http.py index 91de4a23..860ed6f3 100644 --- a/tests/unit/shared/test_http.py +++ b/tests/unit/shared/test_http.py @@ -8,11 +8,13 @@ import httpx from src.shared.http import ( - _backoff_delay, - _get_retry_after, + close_all_pooled_clients, create_client, download_file, + get_pooled_client, head_url, + _backoff_delay, + _get_retry_after, ) @@ -58,9 +60,9 @@ def test_returns_content_type_size_url(self) -> None: } mock_resp.url = httpx.URL("https://example.com/doc") mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.head.return_value = mock_resp + mock_client.head.return_value = mock_resp - with patch("src.shared.http.create_client", return_value=mock_client): + 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" @@ -69,22 +71,20 @@ def test_returns_content_type_size_url(self) -> None: def test_returns_none_on_http_error(self) -> None: mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.head.side_effect = httpx.HTTPStatusError( + mock_client.head.side_effect = httpx.HTTPStatusError( "404", request=MagicMock(), response=MagicMock() ) - with patch("src.shared.http.create_client", return_value=mock_client): + 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.__enter__.return_value.head.side_effect = httpx.ConnectError( - "connection refused" - ) + mock_client.head.side_effect = httpx.ConnectError("connection refused") - with patch("src.shared.http.create_client", return_value=mock_client): + with patch("src.shared.http.get_pooled_client", return_value=mock_client): result = head_url("https://example.com/down") assert result == (None, None, None) @@ -100,9 +100,9 @@ def test_returns_none_for_empty_content_type(self) -> None: mock_resp.headers = {} mock_resp.url = httpx.URL("https://example.com/doc") mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.head.return_value = mock_resp + mock_client.head.return_value = mock_resp - with patch("src.shared.http.create_client", return_value=mock_client): + 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 @@ -119,10 +119,10 @@ def test_private_url_not_skipped_when_check_ssrf_false(self) -> None: mock_resp = MagicMock(spec=httpx.Response) mock_resp.headers = {"content-type": "text/plain"} mock_resp.url = httpx.URL("http://localhost:8080/doc") - mock_client.__enter__.return_value.head.return_value = mock_resp + mock_client.head.return_value = mock_resp with ( - patch("src.shared.http.create_client", return_value=mock_client), + 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) @@ -136,9 +136,9 @@ def test_successful_download(self, tmp_path: Path) -> None: mock_resp = MagicMock(spec=httpx.Response) mock_resp.content = b"hello world" mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.get.return_value = mock_resp + mock_client.get.return_value = mock_resp - with patch("src.shared.http.create_client", return_value=mock_client): + 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 @@ -164,13 +164,13 @@ def test_retries_on_http_429(self, tmp_path: Path) -> None: mock_resp_ok.content = b"content" mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.get.side_effect = [ + mock_client.get.side_effect = [ httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), mock_resp_ok, ] with ( - patch("src.shared.http.create_client", return_value=mock_client), + patch("src.shared.http.get_pooled_client", return_value=mock_client), patch("time.sleep"), ): ok, status = download_file("https://example.com/doc", output) @@ -180,13 +180,13 @@ def test_retries_on_http_429(self, tmp_path: Path) -> None: 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.__enter__.return_value.get.side_effect = [ + mock_client.get.side_effect = [ httpx.ConnectError("timeout"), MagicMock(content=b"ok"), ] with ( - patch("src.shared.http.create_client", return_value=mock_client), + patch("src.shared.http.get_pooled_client", return_value=mock_client), patch("time.sleep"), ): ok, status = download_file("https://example.com/doc", output) @@ -196,10 +196,10 @@ def test_retries_on_network_error(self, tmp_path: Path) -> None: 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.__enter__.return_value.get.side_effect = httpx.ConnectError("always fails") + mock_client.get.side_effect = httpx.ConnectError("always fails") with ( - patch("src.shared.http.create_client", return_value=mock_client), + 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) @@ -212,11 +212,11 @@ def test_non_retryable_http_status(self, tmp_path: Path) -> None: mock_resp = MagicMock(spec=httpx.Response) mock_resp.status_code = 403 mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.get.side_effect = httpx.HTTPStatusError( + mock_client.get.side_effect = httpx.HTTPStatusError( "403", request=MagicMock(), response=mock_resp ) - with patch("src.shared.http.create_client", return_value=mock_client): + 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 @@ -227,9 +227,9 @@ def test_filesystem_error(self, tmp_path: Path) -> None: mock_resp = MagicMock(spec=httpx.Response) mock_resp.content = b"data" mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.get.return_value = mock_resp + mock_client.get.return_value = mock_resp - with patch("src.shared.http.create_client", return_value=mock_client): + 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 @@ -244,13 +244,13 @@ def test_respects_retry_after_header(self, tmp_path: Path) -> None: mock_resp_ok.content = b"content" mock_client = MagicMock(spec=httpx.Client) - mock_client.__enter__.return_value.get.side_effect = [ + mock_client.get.side_effect = [ httpx.HTTPStatusError("429", request=MagicMock(), response=mock_resp_429), mock_resp_ok, ] with ( - patch("src.shared.http.create_client", return_value=mock_client), + 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) @@ -288,3 +288,26 @@ 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 From 6458c23be592fe5c6ff07b62443220e45debc22c Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:14:13 +0200 Subject: [PATCH 19/44] =?UTF-8?q?=F0=9F=94=A7=20chore:=20add=20index=5Faft?= =?UTF-8?q?er=5Fupload=20parameter=20to=20unify=20indexing=20paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add index_after_upload bool param to index_sigma_rules() - When True, skip delete+full-reindex for incremental indexing - Add index_after_upload query param to /api/v1/chat/upload endpoint - Upload response includes indexed/chunks count when index_after_upload=True - Enables immediate Qdrant indexing after file upload without full reindex --- src/api/v1/chat/chat.py | 20 ++++++++++++++++++-- src/application/documents/indexing.py | 25 ++++++++++++++----------- 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/src/api/v1/chat/chat.py b/src/api/v1/chat/chat.py index d19b8b71..ed0d192f 100644 --- a/src/api/v1/chat/chat.py +++ b/src/api/v1/chat/chat.py @@ -92,8 +92,15 @@ async def send_chat_message( 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.""" + """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, @@ -123,11 +130,20 @@ async def upload_sigma_rule( raise HTTPException(status_code=400, detail="File is not valid UTF-8 text") rule = await _get_chat_service().validate_and_store_yaml(content, session_id) - return { + 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, diff --git a/src/application/documents/indexing.py b/src/application/documents/indexing.py index 61b40c77..09bf1f62 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: From 755efdbe697c763767db0ddfc428a8d9325fcf65 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:14:27 +0200 Subject: [PATCH 20/44] =?UTF-8?q?=F0=9F=93=9D=20docs:=20update=20refractor?= =?UTF-8?q?.md=20with=20Phase=2010=20progress=20(P3.1-P3.3=20done)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- refractor.md | 4 +- .../documents/test_sigma_ref_paths.py | 110 ++++++++ .../documents/test_sigma_ref_registry.py | 256 ++++++++++++++++++ .../documents/test_sigma_ref_url.py | 109 ++++++++ 4 files changed, 477 insertions(+), 2 deletions(-) create mode 100644 tests/unit/application/documents/test_sigma_ref_paths.py create mode 100644 tests/unit/application/documents/test_sigma_ref_registry.py create mode 100644 tests/unit/application/documents/test_sigma_ref_url.py diff --git a/refractor.md b/refractor.md index bae0d8cf..f08d87ec 100644 --- a/refractor.md +++ b/refractor.md @@ -694,8 +694,8 @@ Phase 10 — P3 + Q3 (architecture finale) │ ├── sigma_ref_paths.py — subdir_for(), sigmaref_write_path(), sigmaref_resolve_path() │ ├── sigma_ref_url.py — detect_url_type(), is_reference_url(), resolve_ext() │ └── sigma_ref_registry.py — load_registry(), save_registry(), load_error_registry(), maybe_record_error() - ├── P3.2 Connection pooling HTTP - ├── P3.3 Unifier chemins d'indexation (upload→Qdrant) + ├── P3.2 ✅ Connection pooling HTTP (get_pooled_client with httpx.HTTPTransport) + ├── P3.3 ✅ Unifier chemins d'indexation (upload→Qdrant via index_after_upload) └── Q3.3 Oversampling + rescore avec quantification ``` 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 From 1644ceb14d6a1cd8d584a1e8c4f2cfd4f3ffb7e7 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 17:35:39 +0200 Subject: [PATCH 21/44] =?UTF-8?q?=F0=9F=94=A7=20chore:=20remove=20refracto?= =?UTF-8?q?r.md=20(Phase=2010=20complete)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- refractor.md | 703 --------------------------------------------------- 1 file changed, 703 deletions(-) delete mode 100644 refractor.md diff --git a/refractor.md b/refractor.md deleted file mode 100644 index f08d87ec..00000000 --- a/refractor.md +++ /dev/null @@ -1,703 +0,0 @@ -# Refractor — Sigma Rule Management - -Fichier de travail pour le refactoring de la gestion des règles Sigma dans l'application. - -## 1. Problématiques Identifiées - -### 1.1 Chemins de Validation Dupliqués - -Deux chemins de validation distincts avec des règles et des retours incompatibles : - -``` -Chemin A (Production) : - POST /api/v1/chat/upload → ChatService.validate_and_store_yaml() - → SigmaValidator.validate(bytes) - → YAML parse, required: id/name/description/detection - → detection non-empty, condition syntax check - → deprecated field warnings (level, falsepositives) - → Retourne: dict[str, Any] brut - → Stocké en mémoire dans ChatService._uploaded_rule - -Chemin B (Tests uniquement — mort) : - validate_sigma_rule(SigmaRule) - → title/condition/detection non-empty - → level enum, status enum - → Retourne: ValidationResult (Pydantic model) - → Aucun appelant en production -``` - -| Aspect | SigmaValidator (A) | validate_sigma_rule (B) | parser.py | -|--------|-------------------|------------------------|-----------| -| **Taille fichier** | 1MB max | — | 1MB max | -| **Parse YAML** | Oui | — (opère sur modèle) | Oui | -| **Champs requis** | `id`, `name`, `description`, `detection` | `title`, `detection`, `condition` | `title`, `detection` | -| **Validation detection** | Dict non-vide | Dict non-vide | Présence | -| **Validation condition** | Références syntaxe | Non-vide | — | -| **Level** | Warning deprecated | Validation enum | — | -| **Status** | Non vérifié | Validation enum | — | -| **Retour** | `dict` brut | `ValidationResult` | `SigmaRule \| None` | -| **Erreur** | Exception (422) | Pydantic model | `None` / log | - -### 1.2 Incohérence de Nommage des Champs - -- `SigmaValidator` vérifie `name` (norme Sigma v2) -- `SigmaRule` (modèle canonique) utilise `title` -- `parser.py` mappe YAML `title` → `SigmaRule.title` -- Conséquence : si une règle a `name:` dans le YAML, `SigmaValidator` l'accepte mais `SigmaRule.title` sera vide ; inversement si `title:`, le parseur fonctionne mais le validateur rejette. - -### 1.3 Deux Types `ValidationError` Incompatibles - -- `src/shared/exceptions.py:ValidationError(Exception)` — HTTP 422, catchable -- `src/application/documents/models.py:ValidationError(BaseModel)` — Pydantic model, non-catchable -- Même nom, même sémantique (field + message), API incompatible - -### 1.4 Duplication Downloader / Processor - -`src/application/documents/sigma_ref_downloader.py` (753 lignes) et `sigma_ref_processor.py` (413 lignes) partagent le même domaine (téléchargement de documents de référence) avec duplication lourde : - -| Fonction | downloader.py | processor.py | -|----------|--------------|-------------| -| **SHA256 fichier** | `_sha256_file()` (10 lignes) | `_sha256_file()` (7 lignes, quasi-identique) | -| **SHA256 string** | `_sha256()` (5 lignes) | `_sha256_bytes()` (4 lignes) | -| **HEAD request** | `_head_content_type()` | `_head_request()` (signatures différentes) | -| **Download HTTP** | `_download_file()` (3 retries + backoff) | `_download_one()` (lambda, 0 retry) | -| **Entry builder** | `_make_entry()` | `_build_head_entry()` / `_build_download_entry()` | -| **Normalisation URL** | `normalize_url()` (complète) | `_normalize_url()` (triviale — juste strip+rstrip) | -| **Constantes** | `DEFAULT_REQUEST_DELAY`, `DEFAULT_MAX_WORKERS` (lignes 50-51) | Mêmes constantes (lignes 20-21) | -| `_TYPE_TO_EXT` | Défini deux fois dans la même fonction (lignes 548-555, 630-637) | — | -| `httpx.Client` | Nouvelle instance par requête | Nouvelle instance par requête | -| `import hashlib` | Dans le corps des fonctions | Dans le corps des fonctions | - -Risque concret : la normalisation URL différente (downloader supprime fragments, processor non) peut causer des `url_hash` différents entre les deux chemins pour la même URL. - -### 1.5 Duplication Discovery Worker - -`src/workers/sigma/discovery_worker.py` (502 lignes) : - -- `_process_github` / `_process_spec` : structure quasi-identique (~65 lignes chacun), seule la stratégie d'énumération des repos change -- `_scan_all_github` / `_scan_all_spec` : quasi-identiques -- `_write_entries` / `_write_spec_entries` : quasi-identiques (seul le nom de méthode DB change) -- `get_sigma_rule_id` importé à l'intérieur d'une boucle (perte de perf) - -### 1.6 Code Mort - -| Fichier | Code | Statut | -|---------|------|--------| -| `chunker.py` | `_generate_eval_questions()` (lignes 503-515) | Jamais appelé. Le rich chunking utilise des listes inline. | -| `chunker.py` | `chunk_sigma_rules_rich()` (lignes 527-534) | Marqué "backwards-compatible", rien ne l'appelle. `SigmaChunker.process()` est le seul chemin actif. | -| `chunker.py` | `post_process()` — le `if enable_eval_questions` | N'a aucun effet observable. `return documents` est exécuté dans tous les cas. | -| `orchestrator.py` | `class RAGPipeline` (22 lignes) | Thin wrapper autour de `SearchEngine`, probablement mort. | -| `documents/validator.py` | `validate_sigma_rule()` | Importé et testé uniquement dans les tests. Aucun appelant en production. | -| `processor.py` | `GITHUB_BLOB_PATTERN: Any = None` (ligne 22) | Placeholder jamais rempli. | -| `downloader.py` | Paramètre `path` de `_load_registry` / `_save_registry` | Jamais utilisé, seul `db` est utilisé. | -| `downloader.py` | Paramètre `request_delay` de `download_references()` | Déclaré, jamais référencé dans le corps. | -| `processor.py` | Paramètre `request_delay` de `process_sigma_refs()` | Même chose. | - -### 1.7 Complexité du Chunker - -`SigmaChunker._chunk_rule()` fait ~346 lignes (lignes 89-434) et concentre trop de responsabilités : -- Extraction des champs depuis le dict brut -- Construction de templates textes pour 12+ types de chunks -- Itération sur les blocs de détection avec boucles imbriquées (field/operator groups + atomic indicators) -- Enrichissement LLM inline (lignes 414-432) -- Génération de questions d'évaluation inline - -### 1.8 Collision de Noms `RAGPipeline` - -Deux classes portent le même nom dans des packages différents : -- `src/application/chat/rag.py:RAGPipeline` — orchestrateur de génération LLM (actif) -- `src/core/pipeline/orchestrator.py:RAGPipeline` — thin wrapper SearchEngine (probablement mort) - -### 1.9 État Sessionnel ChatService - -`ChatService` maintient un état mutable en mémoire (session) : -- `_history`, `_uploaded_rule`, `_last_citations`, `_current_prompt_id` -- Fonctionne pour un usage mono-utilisateur local -- Empêche toute mise à l'échelle horizontale -- L'état devrait être externalisé (cache distribué, DB, etc.) - -### 1.10 Absence de Module HTTP Partagé - -- `httpx.Client` créé frais à chaque requête dans 4 fichiers : downloader, processor, translate (?), chat -- Aucun pool de connexions HTTP -- Logique de retry/backoff dupliquée (downloader l'a, processor ne l'a pas) -- HEAD requests dupliquées avec signatures différentes - -### 1.11 Double Parsing YAML - -`sigma_utils.py` parse le même fichier YAML deux fois si un appelant veut à la fois les références et le rule_id. Appelé depuis `discovery_worker.py` en boucle par fichier — l'impact est multiplié. - -### 1.12 Deux Chemins d'Indexation ≠ - -- **Chemin A (Upload interactif)** : validation uniquement → stocké en mémoire → PAS indexé dans Qdrant -- **Chemin B (Discovery fond)** : parse → chunk → index dans Qdrant (sigma_rules) -- Écart : les règles uploadées interactivement ne sont PAS présentes dans la recherche vectorielle - -### 1.13 Aucun Tuning Qdrant Spécifique aux Collections - -Les 3 collections Qdrant (`sigma_rules`, `sigma_docs`, `sigma_spec`) sont créées avec la même configuration minimale : - -```python -VectorParams(size=384, distance=Distance.COSINE) -SparseVectorParams(index=SparseIndexParams()) # "text-sparse" -``` - -| Paramètre manquant | Impact | -|-------------------|--------| -| `hnsw_config` | Utilise les defaults serveur (m=16, ef_construct=100) — probablement ok mais non évalué | -| `quantization_config` | Aucune quantification — les vecteurs dense sont stockés en float32 intégral | -| `optimizers_config` | Aucun tuning — indexing_threshold_kb par défaut (20 MB) | -| Per-collection tuning | Les 3 collections ont la même config, alors que leurs usages diffèrent (recherche fréquente vs froide) | - -### 1.14 Custom BM25 Sparse Encoder Non Standard - -`src/core/search/sparse_encoder.py` implémente un BM25 custom avec : -- Token IDs basés sur MD5 modulo 2^24 -- Poids : `1.0 + log(term_frequency)` (pas de IDF, pas de normalisation par doc length) -- Stop words anglais codés en dur (107 mots) -- **Pas de stemming, pas de configuration par langue** - -Points d'attention : -- L'absence d'IDF signifie que les tokens fréquents ne sont pas pénalisés -- Le hash MD5 peut causer des collisions (2^24 ≈ 16M IDs, acceptable mais non déterministe) -- Aucune configuration pour le français alors que le modèle d'embedding est `multilingual-e5-small` -- Pas de `avg_len` configuré pour BM25 — Qdrant ne peut pas normaliser par longueur de document - -### 1.15 Pas d'Évaluation de la Qualité de Recherche - -- Aucun golden set / ground truth dataset pour mesurer recall@k -- Aucune évaluation comparative : exact search vs approximate, flat vs rich chunks -- Le paramètre `alpha=0.3` du hybrid search n'a jamais été tuné sur des données réelles -- Fusion RRF avec `k=60` — valeur par défaut, jamais ajustée - -### 1.16 Pipeline d'Indexation Synchrone avec Batch Size Faible - -- `IngestionPipeline.run()` utilise `num_workers=0` (séquentiel) -- `embed_batch_size=8` — très conservateur pour un modèle 384-dim -- Pas de désactivation HNSW pendant le bulk load (reconstruit à chaque insertion) -- Stratégie delete-and-reindex : pas d'upsert incrémental, reconstruction complète - -### 1.17 Deux Chemins Concurrents pour le Téléchargement des Références - -Le téléchargement des documents de référence des règles Sigma est implémenté par **deux fichiers distincts** qui font la même chose différemment : - -| Aspect | `sigma_ref_downloader.py` (753 lignes) | `sigma_ref_processor.py` (413 lignes) | -|--------|----------------------------------------|---------------------------------------| -| **Point d'entrée** | `download_references(rules_dir, ...)` — scanne les fichiers YAML directement | `process_sigma_refs(db, ...)` — lit les entrées du `doc_registry` | -| **Appelé par** | API endpoint `POST /api/v1/documents/index-sigma-ref` | Worker `SigmaRefProcessor` (background) | -| **Normalisation URL** | Complète : GitHub blob→raw, strips fragments, refs/heads | Triviale : `url.strip().rstrip("/")` | -| **Nommage fichiers** | `{url_hash}{ext}` (ex: `abc123.md`) | `_sanitize_filename(url)` (ex: `documentation_page`) | -| **Registry** | Dict mémoire + flush final vers DB | Appels DB unitaires par opération | -| **Retry download** | Oui (3 retries, backoff exponentiel) | Non (0 retry, pas de backoff) | -| **SSRF protection** | Oui (`_is_private_url()`) | **Non** | -| **HEAD request** | `_head_content_type()` — retourne Content-Type seulement | `_head_request()` — retourne Content-Type + size + final_url | -| **hashlib import** | Dans le corps des fonctions | Dans le corps des fonctions | -| **Lock** | `_registry_lock` global (threading.Lock) | Aucun | - -**Conséquence directe** : le même URL peut avoir un `url_hash` différent entre les deux chemins à cause de la normalisation divergente, causant des doublons dans le registry et des téléchargements redondants. - -### 1.18 Problèmes de Stockage Local des Références - -Les fichiers téléchargés sont stockés dans `data/documents/sigmaref/` : - -- **Aucune structure de sous-répertoires** — tous les fichiers sont à plat dans un seul dossier (plusieurs milliers de fichiers potentiellement) -- **Nommage incohérent** — le downloader utilise des hash (déterministe, prévisible), le processor utilise le nom du fichier original (variable, risque de collision) -- **Pas de nettoyage** — les fichiers orphelins (URLs qui ne sont plus référencés par aucune règle) ne sont jamais supprimés -- **Pas de déduplication cross-rule** — si une URL est référencée par 10 règles, elle est téléchargée 1 fois mais son `rule_id`/`title` dans le registry est écrasé à chaque mise à jour (perte de l'information de provenance) - -### 1.19 Cycle de Vie `embed_status` Incohérent - -Le statut d'embedding traverse `discovery` → `head_verified` → `embedded`, mais : - -- `process_sigma_refs()` crée des entrées avec `embed_status = "head_verified"` pour les URLs de type non supporté → **entrées mortes** qui ne seront jamais promues -- `download_references()` utilise toujours `embed_status = "discovery"` après téléchargement -- Aucun mécanisme de nettoyage pour les entrées bloquées en `head_verified` - -### 1.20 Pas de Table de Jonction Règle↔Référence - -La table `doc_registry` stocke `rule_id` et `title` par URL, mais : -- Si 2 règles partagent la même référence, seule la dernière écriture survit -- Impossible de répondre à "quelles règles référencent ce document ?" -- Impossible de savoir si un document est orphelin (plus référencé par aucune règle) - -## 2. Propositions d'Optimisation - -### P0 — Priorité Critique - -#### P0.1 Extraire les fonctions SHA256 partagées - -Créer `src/shared/utils/crypto_utils.py` : -- `compute_sha256_str(data: str) -> str` -- `compute_sha256_file(path: Path) -> str` - -Supprimer les 3 implémentations dupliquées dans : -- `discovery_base.py:15` -- `sigma_ref_downloader.py:725` -- `sigma_ref_processor.py:392` - -#### P0.2 Unifier la normalisation URL - -Une seule `normalize_url()` dans `src/shared/utils/url_utils.py` avec la logique complète du downloader (GitHub blob→raw, fragments, refs/heads). Utiliser partout pour garantir la cohérence des `url_hash`. - -### P1 — Priorité Haute - -#### P1.1 Consolider les validateurs - -- `SigmaValidator.validate()` retourne un `SigmaRule` (via `SigmaRule.from_dict()`) au lieu d'un `dict` brut -- Fusionner les contrôles de `validate_sigma_rule()` (level enum, status enum) dans `SigmaValidator` -- Normaliser les champs : `title` partout, avec alias `name` en lecture si besoin -- Supprimer `src/application/documents/validator.py` après fusion -- Standardiser sur `src/shared/exceptions.ValidationError` (exception) - -#### P1.2 Extraire le download HTTP partagé - -Créer `src/shared/utils/http_utils.py` : -- `head_url(url, timeout=10) -> tuple[str|None, int|None, str|None]` -- `download_file(url, output_path, max_retries=3) -> tuple[bool, int|None]` -- Pool de connexions `httpx.AsyncClient` partagé -- Logique de retry avec backoff harmonisée - -#### P1.3 Factoriser les builders d'entrées registry - -Fonction unique `build_registry_entry()` dans `src/shared/utils/registry_utils.py` avec paramètre `embed_status` par défaut. Remplacer `_make_entry`, `_build_head_entry`, `_build_download_entry`. - -#### P1.4 Factoriser les méthodes du DiscoveryWorker - -- `_scan_all_github` + `_scan_all_spec` → `_scan_all_files(prepare_entry_fn)` -- `_write_entries` + `_write_spec_entries` → `_write_entries(batch_upsert_fn, spec_mode=False)` -- `_process_github` + `_process_spec` → paramétrer la stratégie d'énumération des repos - -#### P1.5 Supprimer le code mort - -- `_generate_eval_questions()` dans chunker.py -- `chunk_sigma_rules_rich()` dans chunker.py -- `orchestrator.py:RAGPipeline` — supprimer la classe, les appelants utilisent `SearchEngine` directement -- `GITHUB_BLOB_PATTERN: Any = None` dans processor.py -- Paramètres morts (`path`, `request_delay`) dans downloader/processor -- `validate_sigma_rule()` dans documents/validator.py (après fusion P1.1) - -### P2 — Priorité Moyenne - -#### P2.1 Décomposer `_chunk_rule()` - -Extraire : -- `_extract_fields(rule: dict) -> dict` (lignes 98-116) -- `_build_executive_summary(rule, fields) -> dict` (chunk type 1) -- `_build_metadata_lifecycle(rule, fields) -> dict` (chunk type 2) -- ... et ainsi de suite pour les 12+ types -- `_enrich_by_llm(chunks) -> chunks` (étape séparée, lignes 414-432) - -#### P2.2 Résoudre la collision RAGPipeline - -Renommer ou supprimer `src/core/pipeline/orchestrator.py:RAGPipeline`. - -#### P2.3 Externaliser l'état ChatService - -- Propager `session_id` depuis l'API -- Stocker `_history`, `_uploaded_rule` dans un store (cache LRU, DuckDB, etc.) -- Permettre la scalabilité horizontale - -#### P2.4 Normaliser les imports - -- `import hashlib` en haut des fichiers, pas dans le corps des fonctions -- `from src.shared.utils.sigma_utils import get_sigma_rule_id` en haut de `discovery_worker.py`, pas dans la boucle - -#### P2.5 Constante NULL_UUID - -Définir `NULL_UUID = "00000000-0000-0000-0000-000000000000"` dans `src/shared/constants.py` et l'utiliser partout. - -#### P2.6 Dédoublonner `_TYPE_TO_EXT` - -Définir une fois comme constante module dans `identify_file_type.py` ou `sigma_ref_downloader.py`. - -### P3 — Priorité Basse - -#### P3.1 Split downloader.py (753 lignes) - -Diviser `sigma_ref_downloader.py` en modules < 400 lignes : -- `sigma_ref_downloader.py` — orchestration -- `sigma_ref_http.py` — HTTP helpers -- `sigma_ref_registry.py` — registry management - -#### P3.2 Connection pooling HTTP - -Réutiliser `httpx.Client` (ou `AsyncClient`) avec un pool de connexions dans toutes les phases de téléchargement batch. - -#### P3.3 Unifier les deux chemins d'indexation - -Permettre aux règles uploadées interactivement d'être optionnellement indexées dans Qdrant (ex. paramètre `index_after_upload`). - -### R1 — Priorité Haute (Téléchargement Références) - -#### R1.1 Unifier `download_references()` et `process_sigma_refs()` - -Une seule fonction `download_sigma_references(source, db, ...)` avec deux modes d'entrée : -- `mode="scan"` : scanne les fichiers YAML directement (ancien downloader) -- `mode="registry"` : lit les entrées du doc_registry (ancien processor) - -La normalisation URL, le nommage des fichiers, le retry, et la détection SSRF doivent être **identiques** dans les deux modes. - -#### R1.2 Standardiser le Nommage des Fichiers - -- Format unique : `{url_hash}{ext}` (comme le downloader actuel) -- Supprimer `_sanitize_filename()` du processor -- Garantir qu'un même URL produit toujours le même chemin disque - -#### R1.3 Ajouter une Table de Jonction `rule_references` - -```sql -CREATE TABLE rule_references ( - rule_id TEXT NOT NULL, - url_hash TEXT NOT NULL, - PRIMARY KEY (rule_id, url_hash) -); -``` - -- Permet de retrouver toutes les références d'une règle -- Permet de détecter les documents orphelins (URLs non référencées) -- Nettoyage possible : supprimer les fichiers dont l'URL_hash n'est plus dans `rule_references` - -#### R1.4 Nettoyer les Entrées Mortes du Registry - -- Supprimer les entrées `head_verified` de plus de X jours -- Supprimer les entrées dont `url_hash` n'est plus dans `rule_references` -- Tâche planifiée (cron) ou déclenchée manuellement - -#### R1.5 Ajouter SSRF Protection dans le Processor - -Copier `_is_private_url()` du downloader et l'utiliser dans `process_sigma_refs()` avant toute HEAD request. Sans ça, le processor est vulnérable aux SSRF. - -### R2 — Priorité Moyenne (Organisation Stockage) - -#### R2.1 Structure de Sous-Répertoires - -``` -data/documents/sigmaref/ - ├── markdown/ # .md files - ├── html/ # .html files - ├── pdf/ # .pdf files - ├── plain_text/ # .txt files - └── office/ # .docx files -``` - -- Préserve le nom `{url_hash}{ext}` dans chaque sous-répertoire -- Évite d'avoir 10 000 fichiers dans le même dossier - -#### R2.2 Garbage Collection des Fichiers Orphelins - -- Scanner `data/documents/sigmaref/` pour les fichiers dont le nom ne correspond à aucun `url_hash` dans `doc_registry` -- Supprimer les fichiers orphelins (avec confirmation) -- Option : déplacer vers `.trash/` avant suppression définitive - -### R3 — Priorité Basse (Traçabilité) - -#### R3.1 Enrichir le Registry avec la Liste des Règles Sources - -Ajouter un champ `referenced_by: list[str]` dans `doc_registry` (ou utiliser la table `rule_references`) pour savoir quelles règles référencent un document donné. Utile pour : -- Debug : "pourquoi ce document a été téléchargé ?" -- Mise à jour : "re-télécharger les docs des règles modifiées" -- Suppression : "ce document n'est plus référencé, on peut le supprimer" - -## 3. Plan d'Exécution Suggéré - -``` -Phase 1 — P0 (nettoyage critique) - ├── P0.1 Extraire crypto_utils.py (SHA256) - ├── P0.2 Unifier normalize_url() - ├── P1.5 Supprimer le code mort évident - └── P2.4 Normaliser les imports hashlib/sigma_utils - -Phase 2 — P1 (consolidation validateurs + HTTP) - ├── P1.1 Consolider SigmaValidator → retourne SigmaRule - ├── P1.2 Extraire http_utils.py (HEAD + download + retry) - ├── P1.3 Factoriser registry entry builders - └── Supprimer doc la mort (validate_sigma_rule, etc.) - -Phase 3 — P1 (factorisation workers) - ├── P1.4 Factoriser _scan_all_* / _write_* / _process_* - └── P2.5 Constante NULL_UUID - -Phase 4 — P2 (qualité de code) - ├── P2.1 Décomposer _chunk_rule() - ├── P2.2 Résoudre collision RAGPipeline - ├── P2.3 Externaliser état ChatService - └── P2.6 Dédoublonner _TYPE_TO_EXT - -Phase 5 — P3 (architecture) - ├── P3.1 Split downloader.py - ├── P3.2 Connection pooling HTTP - └── P3.3 Unifier chemins d'indexation -``` - -## 4. Propositions Qdrant (Basées sur les Skills) - -### Q0 — Priorité Critique (Qualité de Recherche) - -#### Q0.1 Créer un Golden Set pour Évaluer Recall@k - -- Échantillonner 100-200 requêtes réelles avec jugements de pertinence -- Mesurer recall@k avant/après chaque changement de config -- Voir Qdrant Search Quality Diagnosis skill - -#### Q0.2 Tester Exact Search Comme Baseline - -- Avant tout tuning HNSH, comparer exact search vs approximate search -- Si l'écart est > 5%, tuner `ef` et `m` du HNSW -- Permet d'isoler les problèmes de modèle d'embedding vs index - -### Q1 — Priorité Haute (Performance Indexation) - -#### Q1.1 Ajouter `quantization_config` aux Collections - -```python -quantization_config=ScalarQuantization( - scalar=ScalarQuantizationConfig( - type=ScalarType.INT8, - always_ram=True, - quantile=0.5, - ) -) -``` - -- Réduction mémoire 4x pour les vecteurs dense en RAM -- Perte de qualité < 1% recall avec rescore -- Activation immédiate pour les 3 collections - -#### Q1.2 Augmenter `embed_batch_size` et Paralléliser - -- `embed_batch_size` : 8 → **64** (384-dim, CPU, safe) -- `num_workers` : 0 → **2-4** pour ingestion parallèle -- Désactiver HNSW pendant le bulk load : `indexing_threshold_kb = 0` → temporairement très haut, restaurer après - -### Q2 — Priorité Moyenne (Tuning Recherche) - -#### Q2.1 Tuner l'Alpha du Hybrid Search par Collection - -- Remplacer le `alpha=0.3` global par des valeurs par collection : - - `sigma_rules` : α=0.5 (sémantique + lexical équilibré — les règles ont un vocabulaire technique précis) - - `sigma_docs` : α=0.7 (plus lexical — docs de référence, termes exacts) - - `sigma_spec` : α=0.3 (plus sémantique — spécifications, concepts) -- Évaluer avec le golden set (Q0.1) - -#### Q2.2 Tuner la Fusion RRF - -- `k=60` actuel → tester `k=30`, `k=60`, `k=100` avec le golden set -- Poids par collection dans le RRF (weighted RRF) si une collection domine -- Envisager DBSF si les distributions de score entre dense et sparse sont trop différentes - -#### Q2.3 Ajouter `hnsw_config` par Collection - -```python -hnsw_config=HnswConfigDiff( - m=16, # default ok pour 384-dim - ef_construct=200, # 100 → 200 pour meilleure qualité à l'indexation - full_scan_threshold_kb=10000, # 10 MB, identique au serveur - on_disk=False, # sigma_rules en RAM -) -``` - -- `sigma_docs` et `sigma_spec` (collections froides) : `on_disk=True` + `async_scorer` - -### Q3 — Priorité Basse (Architecture Vectorstore) - -#### Q3.1 Stocker les Sparse Vectors sur Disk - -```python -sparse_vectors_config={ - "text-sparse": SparseVectorParams( - index=SparseIndexParams(on_disk=True) - ) -} -``` - -- Les vecteurs sparse BM25 sont rarement tous consultés -- Bon candidat pour le stockage disque (économise RAM) - -#### Q3.2 Revoir le Sparse Encoder Custom - -- Benchmarker `bm25_sparse_encoder` vs le BM25 natif de Qdrant (configuré par language) -- Avantages BM25 natif Qdrant : - - Calcul côté serveur (pas de transfert des vecteurs sparse) - - Tokenization + stemming standardisés - - Support multi-langue (français, etc.) - - IDF calculé automatiquement -- Si le custom encoder est conservé, ajouter au moins le calcul d'IDF et la normalisation par doc length - -#### Q3.3 Oversampling + Rescore avec Quantification - -```python -quantization_config=ScalarQuantization( - scalar=ScalarQuantizationConfig(type=ScalarType.INT8, always_ram=True), - rescore=True, - oversampling=2.0, -) -``` - -- Permet de chercher dans un pool 2x plus large (vitesse ×2 grâce à la quantification) -- Rescore les top_k sur les vecteurs originaux pour préserver la qualité - -#### Q3.4 Pipeline d'Indexation Incrémental - -- Remplacer la stratégie delete-and-reindex par des upserts par `rule_id` -- Nécessite : identifier les règles nouvelles/modifiées vs inchangées -- Avantage : pas de downtime de la collection, indexation plus rapide -- Combinable avec `indexing_threshold_kb` haut initial + baisse progressive - -## 5. Call Graph — Phase 0 Audit - -### SigmaValidator.validate() - -**Production callers (1 fichier, 1 site d'appel direct) :** -| Fichier | Ligne | Usage | Type attente | -|---------|-------|-------|-------------| -| `src/application/chat/service.py` | 357 | `ChatService.validate_and_store_yaml()` → stocke dans `_uploaded_rule` | `dict[str, Any]` | -| `src/application/chat/rag.py` | 99,133,269,305 | `explain_rule()`, `explain_rule_stream()`, `analyze_coverage()`, `analyze_coverage_stream()` | `dict[str, Any]` (param `rule_data`) | -| `src/application/chat/rag.py` | 380 | `_format_rule_yaml()` | `yaml.dump(rule)` attend un dict | -| `src/application/chat/rag.py` | 387 | `_fallback_explanation()` | `.get('name', 'Unknown')`, `.get('id')`, `.get('description')` | - -**Consommateurs indirects de `_uploaded_rule` (dict) via `ChatService` :** -| Méthode | Accès | Ligne | -|---------|-------|-------| -| `_handle_explain()` | `.get("name", "")` | 234 | -| `_handle_explain_stream()` | `.get("name", "")` | 248 | -| `_handle_coverage()` | passé tel quel | 265 | -| `_handle_coverage_stream()` | passé tel quel | 342 | - -**Tests :** -| Fichier | Usage | -|---------|-------| -| `tests/unit/application/services/test_sigma_validator.py` | 19 tests unitaires, `validate()` sur bytes | -| `tests/unit/application/services/test_sigma_validator_advanced.py` | 6 tests avancés | -| `tests/unit/application/services/test_chat_service_cache.py` | `SigmaValidator` mocké (`patch`) | -| `tests/integration/test_chat_flow.py` | `_uploaded_rule` set manuellement comme dict (l.92) | - -### download_references() - -**Production callers (1 fichier, 1 site) :** -| Fichier | Ligne | Contexte | -|---------|-------|----------| -| `src/api/v1/documents/documents.py` | 35 | `POST /api/v1/documents/index-sigma-ref` | - -**Tests :** -| Fichier | Usage | -|---------|-------| -| `tests/unit/application/documents/test_sigma_ref_downloader.py` | 15+ appels, tests unitaires complets | - -### process_sigma_refs() - -**Production callers (1 fichier, 1 site) :** -| Fichier | Ligne | Contexte | -|---------|-------|----------| -| `src/workers/sigma/sigmaref_worker.py` | 41 | `SigmaRefProcessor.process()` worker background | - -**Tests :** -| Fichier | Usage | -|---------|-------| -| `tests/unit/workers/test_discovery_workers.py` | Mocké (l.28, 50) | - -### validate_sigma_rule() — CONFIRMÉ CODE MORT - -**Production callers : AUCUN** - -**Tests uniquement :** -| Fichier | Usage | -|---------|-------| -| `tests/unit/application/documents/test_documents.py` | 5 appels (l.69, 84, 98, 113, 131) | - -### Risques Identifiés pour P1.1 (dict → SigmaRule) - -1. `ChatService._uploaded_rule` typé `dict[str, Any] | None` → doit passer à `SigmaRule | None` -2. Tous les `.get("name", "")` → deviennent `.title` (ou alias `.name` si préservé) -3. `_format_rule_yaml()` utilise `yaml.dump(rule_dict)` → nécessite `rule.model_dump()` -4. `_fallback_explanation()` utilise `.get()` sur dict → nécessite accès attribut -5. Mock dans `test_chat_service_cache.py` patch `SigmaValidator` → pas de changement nécessaire -6. Test intégration `test_chat_flow.py` set `_uploaded_rule` comme dict brut → à migrer vers `SigmaRule` - -### Tests de Régression à Écrire (avant P1.1) DONE by commit 6474065 - -- [x] Capturer le contrat `dict` actuel de `SigmaValidator.validate()` (bracket access, .get(), yaml.dump(), ValidationError) -- [x] Capturer les patterns de `ChatService._uploaded_rule` (`.get("name", "")`, `.get("id", "N/A")`) -- [x] Capturer le contrat `_format_rule_yaml()` et `_fallback_explanation()` - ---- - -## 6. Plan d'Exécution Final - -``` -Phase 1 — P0 + R (nettoyage critique) DONE by commmit 8d24148dbfe4be3e6fe529800f5516613a0217c6 - ├── P0.1 Extraire crypto_utils.py (SHA256) - ├── P0.2 Unifier normalize_url() - ├── R1.5 Ajouter SSRF protection dans le processor - ├── P1.5 Supprimer le code mort évident - └── P2.4 Normaliser les imports hashlib/sigma_utils - -Phase 2 — P1.2 + P1.3 (infra indépendante, 0 risque) DONE by commit 6474065 - ├── P1.2 Extraire src/shared/http.py (HEAD + download + retry + pool httpx) - ├── P1.3 Factoriser build_registry_entry() — pure factory sans IO - └── Tests : mock httpx, retry/backoff, SSRF, golden path + edge cases - -Phase 3 — P1.1 (validator consolidation, protégé par audit Phase 0) DONE by commit 259bccd - ├── Audit préalable : cartographier tous les callers, tests de régression - ├── SigmaValidator.validate() → retourne SigmaRule (pydantic) - ├── Merge checks de validate_sigma_rule() + shared.exceptions.ValidationError - ├── Normalisation title/name (name alias via @property) - ├── Supprimer validate_sigma_rule(), ValidationError/ValidationResult - └── Tests : 52 unit tests verts, ruff/mypy clean - -Phase 4 — R1.1 + R1.2 (download unification, dépend de P1.2) DONE by commit 2f2efe7 - ├── download_sigma_references(source, db, mode="scan"|"registry") - ├── Standardiser nommage {url_hash}{ext} dans les 2 modes - ├── Contract test : même URL → même nom de fichier - └── processor.py passe de 318l à 56l (délégation pure) - -Phase 5 — R1 (traçabilité références) DONE by commit 33e1f70 + follow-up - ├── R1.3 Ajouter table rule_references (junction rule↔reference) - ├── R1.4 Nettoyer entrées mortes head_verified - │ ├── delete_head_verified_orphans() — head_verified sans content_sha256 - │ └── delete_unreferenced_entries() — sigmaref entries sans url_hash dans rule_references - ├── R1.4 integré dans DocGCWorker.process() - └── R3.1 Enrichir registry avec liste des règles sources - -Phase 6 — Q0-Q1 (qualité recherche + perf indexation Qdrant) - ├── Q0.1 Créer golden set pour évaluation recall@k - ├── Q0.2 Tester exact search comme baseline - ├── Q1.1 Ajouter quantization_config aux 3 collections - └── Q1.2 Augmenter embed_batch_size (8→64) + paralléliser (workers 0→4) - -Phase 7 — P1 (factorisation workers) DONE - ├── P1.4 Factoriser _scan_all_github + _scan_all_spec → _scan_all(prepare_fn) - │ ├── _write_entries + _write_spec_entries → _write_entries(batch_upsert_fn) - │ └── _collect_repo_files() extrait du rglob commun github+spec - ├── P2.5 Constante NULL_UUID dans src/shared/constants.py (6 fichiers modifiés) - └── Q2.3 Ajouter hnsw_config par collection via collection_hnsw_config() - ├── sigma_rules → in-RAM, ef_construct=200 - ├── sigma_docs → on-disk - └── sigma_spec → on-disk - -Phase 8 — P2 + R2 (qualité de code + organisation stockage) - ├── P2.1 Décomposer _chunk_rule() - ├── P2.2 Résoudre collision RAGPipeline - ├── P2.3 Externaliser état ChatService - ├── P2.6 Dédoublonner _TYPE_TO_EXT - ├── R2.1 Structure sous-répertoires par type (markdown/, html/, pdf/) - └── R2.2 Garbage collection fichiers orphelins - -Phase 9 — Q2-Q3 (tuning recherche + architecture vectorstore) DONE - ├── Q2.1 Tuner alpha hybrid search par collection - ├── Q2.2 Tuner fusion RRF (k, weighted) - ├── Q3.1 Stocker sparse vectors sur disk - ├── Q3.2 Revoir sparse encoder (benchmark BM25 natif) - └── Q3.4 Pipeline d'indexation incrémental - -Phase 10 — P3 + Q3 (architecture finale) - ├── P3.1 ✅ Split downloader.py (si toujours pertinent après R1.1) - │ ├── sigma_ref_paths.py — subdir_for(), sigmaref_write_path(), sigmaref_resolve_path() - │ ├── sigma_ref_url.py — detect_url_type(), is_reference_url(), resolve_ext() - │ └── sigma_ref_registry.py — load_registry(), save_registry(), load_error_registry(), maybe_record_error() - ├── P3.2 ✅ Connection pooling HTTP (get_pooled_client with httpx.HTTPTransport) - ├── P3.3 ✅ Unifier chemins d'indexation (upload→Qdrant via index_after_upload) - └── Q3.3 Oversampling + rescore avec quantification -``` - - -# opencode -s ses_0deda6082ffeIa2vHoUdYiWbpY \ No newline at end of file From 3565726ee9609c43a93b27a9ee50ae9e30467200 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 18:10:34 +0200 Subject: [PATCH 22/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20chore:=20update=20qd?= =?UTF-8?q?rant-client=20to=201.18.0=20and=20fix=20deprecated=20API=20fiel?= =?UTF-8?q?ds?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Update qdrant-client dependency to >=1.18.0 - Remove on_disk from SparseVectorParams (removed in 1.18) - Remove rescore/oversampling from ScalarQuantization (removed in 1.18) - Rename full_scan_threshold_kb to full_scan_threshold in HnswConfigDiff - Sort dependencies alphabetically in pyproject.toml - Update direct dependencies to latest versions --- pyproject.toml | 39 ++-- src/api/v1/infrastructure/qdrant.py | 5 +- src/infrastructure/vectorstore/collections.py | 17 +- src/infrastructure/vectorstore/storage.py | 9 +- uv.lock | 202 +++++++++--------- 5 files changed, 133 insertions(+), 139 deletions(-) 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/src/api/v1/infrastructure/qdrant.py b/src/api/v1/infrastructure/qdrant.py index 3f5d563e..637e0ed5 100644 --- a/src/api/v1/infrastructure/qdrant.py +++ b/src/api/v1/infrastructure/qdrant.py @@ -67,7 +67,6 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N "text-sparse": models.SparseVectorParams( index=models.SparseIndexParams(), modifier=models.Modifier.IDF, - on_disk=True, ) }, quantization_config=models.ScalarQuantization( @@ -75,9 +74,7 @@ def _recreate_collection(client: Any, collection_name: str, vector_size: int | N type=models.ScalarType.INT8, always_ram=True, quantile=0.5, - ), - rescore=True, - oversampling=2.0, + ) ), ) diff --git a/src/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index df66d881..0e2d2b9e 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -25,17 +25,17 @@ def collection_hnsw_config( use on-disk storage to save RAM. """ if collection_name == "sigma_rules": - return qdrant_client.models.HnswConfigDiff( # type: ignore[call-arg] + return qdrant_client.models.HnswConfigDiff( m=16, ef_construct=200, - full_scan_threshold_kb=10000, + full_scan_threshold=10000, on_disk=False, ) if collection_name in ("sigma_docs", "sigma_spec"): - return qdrant_client.models.HnswConfigDiff( # type: ignore[call-arg] + return qdrant_client.models.HnswConfigDiff( m=16, ef_construct=100, - full_scan_threshold_kb=10000, + full_scan_threshold=10000, on_disk=True, ) return None @@ -166,22 +166,19 @@ async def create_collection( sparse_vectors_config: dict[str, Any] | None = None if enable_hybrid: sparse_vectors_config = { - "text-sparse": qdrant_client.models.SparseVectorParams( # type: ignore[call-arg] + "text-sparse": qdrant_client.models.SparseVectorParams( index=qdrant_client.models.SparseIndexParams(), modifier=qdrant_client.models.Modifier.IDF, - on_disk=True, ) } quantization_config = None if enable_quantization: - quantization_config = qdrant_client.models.ScalarQuantization( # type: ignore[call-arg] + quantization_config = qdrant_client.models.ScalarQuantization( scalar=qdrant_client.models.ScalarQuantizationConfig( type=qdrant_client.models.ScalarType.INT8, always_ram=True, quantile=0.5, - ), - rescore=True, - oversampling=2.0, + ) ) hnsw_config = collection_hnsw_config(collection_name) await asyncio.to_thread( diff --git a/src/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index a121eb00..e8f9a6a1 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -118,20 +118,17 @@ async def store_embeddings( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), sparse_vectors_config={ - "text-sparse": SparseVectorParams( # type: ignore[call-arg] + "text-sparse": SparseVectorParams( index=SparseIndexParams(), modifier=Modifier.IDF, - on_disk=True, ), }, - quantization_config=ScalarQuantization( # type: ignore[call-arg] + quantization_config=ScalarQuantization( scalar=ScalarQuantizationConfig( type=ScalarType.INT8, always_ram=True, quantile=0.5, - ), - rescore=True, - oversampling=2.0, + ) ), ) 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]] From 888e7c603fd60f3ff2d98ee0536650a0329740a3 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 18:26:14 +0200 Subject: [PATCH 23/44] =?UTF-8?q?=E2=9C=A8=20feat:=20trigger=20discovery?= =?UTF-8?q?=20only=20for=20selected=20repo=20when=20saving=20selections?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Pass repo_key via task params instead of scanning all repos - Add GITHUB_DISCOVERY worker trigger to select-dirs endpoint - Fix parameter naming conflict (request -> request_body) --- src/api/v1/base/repo_router.py | 18 ++++++++--- src/workers/sigma/discovery_worker.py | 43 ++++++++++++++++----------- 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/src/api/v1/base/repo_router.py b/src/api/v1/base/repo_router.py index 45936e87..540d4252 100644 --- a/src/api/v1/base/repo_router.py +++ b/src/api/v1/base/repo_router.py @@ -424,9 +424,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 +442,19 @@ 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) if result.get("success"): + dispatcher = request.app.state.dispatcher + background_tasks.add_task( + dispatcher.ask_for_worker, + WorkerName.GITHUB_DISCOVERY, + task_type=WorkerName.GITHUB_DISCOVERY.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, diff --git a/src/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index 14367a4d..ee0e25c5 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -125,24 +125,33 @@ 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 - - repo_items: list[tuple[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 - repo_items.append((parts[0], parts[1])) + return + else: + 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 + + repo_items: list[tuple[str, str]] = [] + 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])) all_files = self._collect_repo_files(repo_items, gh_base) @@ -150,7 +159,7 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: 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: From 92b77902d2e75e43838f9abad81e2bfd411024b9 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 18:32:56 +0200 Subject: [PATCH 24/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20resolve=20GitHub=20?= =?UTF-8?q?rule=20paths=20to=20data/github=20instead=20of=20sigmaref=20dir?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/application/documents/sigma_ref_paths.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/application/documents/sigma_ref_paths.py b/src/application/documents/sigma_ref_paths.py index 68efe488..ea7b7e4a 100644 --- a/src/application/documents/sigma_ref_paths.py +++ b/src/application/documents/sigma_ref_paths.py @@ -81,6 +81,6 @@ def resolve_rule_path(entry: dict, cfg: Any) -> Path | None: return base / file_name if org and repo: - return Path(cfg.sigmaref_documents_path).resolve() / org / repo / file_name + return Path(cfg.paths_github_dir).resolve() / org / repo / file_name return None From c35a76f9afcbf88a8d25091ca636e8da450bc21f Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Fri, 3 Jul 2026 19:34:07 +0200 Subject: [PATCH 25/44] =?UTF-8?q?=E2=9C=A8=20feat:=20add=20garbage=20colle?= =?UTF-8?q?ction=20for=20stale=20doc=5Fregistry=20entries=20and=20referenc?= =?UTF-8?q?e=20files?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add _garbage_collect_github/local/spec methods to GenericDiscoveryWorker - Add _cleanup_rule_references to delete downloaded reference files - Fix _process_github to GC all repos with active selections + scanned repo - Add delete_doc_registry_by_url_hashes and get_doc_registry_url_hashes_by_repo to DatabaseService - Add get_rule_id_by_url_hash and get_rule_reference_paths for reference cleanup - Fix sigmaref_resolve_path and resolve_rule_path fallbacks to use subdir-aware paths - Add 5 P0 tests for garbage collection scenarios - Add mocks for new DB methods in conftest --- src/application/documents/sigma_ref_paths.py | 4 +- src/infrastructure/database/doc_ops.py | 55 +++++++ src/workers/sigma/discovery_worker.py | 140 ++++++++++++++++++ tests/unit/workers/conftest.py | 4 + tests/unit/workers/test_discovery_workers.py | 143 +++++++++++++++++++ 5 files changed, 344 insertions(+), 2 deletions(-) diff --git a/src/application/documents/sigma_ref_paths.py b/src/application/documents/sigma_ref_paths.py index ea7b7e4a..f2d8b917 100644 --- a/src/application/documents/sigma_ref_paths.py +++ b/src/application/documents/sigma_ref_paths.py @@ -52,7 +52,7 @@ def sigmaref_resolve_path(output_path: Path, content_type: str | None, file_name candidate = sigmaref_write_path(output_path, content_type, file_name) if candidate.exists(): return candidate - return output_path / file_name + return sigmaref_write_path(output_path, content_type, file_name) def resolve_rule_path(entry: dict, cfg: Any) -> Path | None: @@ -78,7 +78,7 @@ def resolve_rule_path(entry: dict, cfg: Any) -> Path | None: candidate = base / subdir / file_name if candidate.exists(): return candidate - return base / file_name + return candidate if org and repo: return Path(cfg.paths_github_dir).resolve() / org / repo / file_name diff --git a/src/infrastructure/database/doc_ops.py b/src/infrastructure/database/doc_ops.py index 760a7303..7b1b1368 100644 --- a/src/infrastructure/database/doc_ops.py +++ b/src/infrastructure/database/doc_ops.py @@ -449,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( @@ -642,6 +661,42 @@ def delete_rule_references_by_url_hash(self, url_hash: str) -> None: 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) # ------------------------------------------------------------------ diff --git a/src/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index ee0e25c5..d4a5f4a7 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -116,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 @@ -153,6 +179,26 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: continue repo_items.append((parts[0], parts[1])) + # 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() + 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 = self._collect_repo_files(repo_items, gh_base) if self.dispatcher: @@ -171,6 +217,63 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: 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 @@ -230,6 +333,43 @@ def _process_spec(self, task: dict, worker_name: WorkerName) -> None: skipped_count, batch_upsert_fn=self.db.batch_upsert_sigma_spec, ) + self._garbage_collect_spec(repos, entries) + + 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, 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]) From d3f8ecfcc857e50ebc8b859b0ef8df9101d4396a Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 07:48:01 +0200 Subject: [PATCH 26/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20correct=20HF=5FHUB?= =?UTF-8?q?=5FOFFLINE=20handling=20and=20reference=20download=20filtering?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix: use huggingface_hub.constants.HF_HUB_OFFLINE instead of os.environ (env var is cached at import, mutations had no effect) - fix: normalize empty HF_TOKEN to None to prevent invalid Bearer header - fix: bypass offline mode for online API calls (search, model info, download) - fix: restrict reference downloads to markdown only - fix: remove catch-all text/plain → markdown mapping in URL type detection - fix: make list_gguf_files/get_model_info async (were blocking event loop) - fix: frontend now reads error body instead of discarding it on HTTP errors - fix: backend returns 400 with actual error message instead of generic 500 - refactor: remove dead os.environ manipulation in ingestion.py - docs: add AirGap / Offline Mode section to README --- README.md | 16 +++++ src/api/v1/models/models_embedding.py | 4 ++ src/api/v1/models/models_llm.py | 4 +- src/api/v1/models/models_sparse.py | 9 +-- src/application/documents/sigma_ref_url.py | 6 -- src/application/models/download.py | 62 +++++++++++++------ src/application/models/embedding.py | 4 +- src/core/pipeline/ingestion.py | 38 ++++-------- src/presentation/static/js/internal/config.js | 24 ++++--- src/shared/utils/identify_file_type.py | 4 -- .../shared/utils/test_identify_file_type.py | 6 +- 11 files changed, 103 insertions(+), 74 deletions(-) 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/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..f5377aac 100644 --- a/src/api/v1/models/models_llm.py +++ b/src/api/v1/models/models_llm.py @@ -91,7 +91,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 +171,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: 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/application/documents/sigma_ref_url.py b/src/application/documents/sigma_ref_url.py index 1a4bf667..1c43609a 100644 --- a/src/application/documents/sigma_ref_url.py +++ b/src/application/documents/sigma_ref_url.py @@ -32,12 +32,6 @@ def detect_url_type(url: str, content_type: str | None = None) -> str | None: ctype = content_type.split(";")[0].strip().lower() if "markdown" in ctype: return "markdown" - if "pdf" in ctype: - return "pdf" - if "html" in ctype: - return "html" - if "text" in ctype: - return "markdown" # Fall back to URL extension ext = url_ext(url) diff --git a/src/application/models/download.py b/src/application/models/download.py index 776b108a..af10d727 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.""" @@ -106,14 +127,17 @@ async def list_models(self, query: str, task: str | None = None) -> list[HFRepo] 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"} + pipeline_tags = [task] if task else ["feature-extraction"] + if task == "feature-extraction": + pipeline_tags = ["feature-extraction", "sentence-similarity"] seen: set[str] = set() results: list[HFRepo] = [] for tag in pipeline_tags: @@ -125,3 +149,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/core/pipeline/ingestion.py b/src/core/pipeline/ingestion.py index 9c6a18a2..4dd9ba0a 100644 --- a/src/core/pipeline/ingestion.py +++ b/src/core/pipeline/ingestion.py @@ -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/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/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index 76371d0b..d1d658b9 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -77,10 +77,6 @@ def filetype_subdir(content_type: str) -> str: SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, - FileType.PDF.value, - FileType.PLAIN_TEXT.value, - FileType.HTML.value, - FileType.OFFICE_DOCUMENT.value, } PUREMAGIC_TYPE_MAP: dict[str, FileType] = { diff --git a/tests/unit/shared/utils/test_identify_file_type.py b/tests/unit/shared/utils/test_identify_file_type.py index c886a702..e9919fa8 100644 --- a/tests/unit/shared/utils/test_identify_file_type.py +++ b/tests/unit/shared/utils/test_identify_file_type.py @@ -295,9 +295,9 @@ 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 "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"} From 6ea6843bae5af8468f32bc23b513b43a06337f33 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 07:54:56 +0200 Subject: [PATCH 27/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20prevent=20Qdrant=20?= =?UTF-8?q?binary=20deletion=20during=20unmocked=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix: move stale file cleanup in download_binary() to after successful download + extraction, so a failed download no longer deletes the existing binary - fix: add missing @patch('start_download') to two test methods (test_admin_download and test_orchestration_download) that were triggering the real download pipeline with default config and deleting the binary --- src/infrastructure/vectorstore/downloader.py | 20 +++++++++---------- tests/unit/api/v1/test_admin_download.py | 7 ++++++- .../api/v1/test_orchestration_download.py | 7 ++++++- 3 files changed, 22 insertions(+), 12 deletions(-) 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/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_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"}, From d6dc20d3f37259d3322bd423e85e90b2a099b438 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 07:58:36 +0200 Subject: [PATCH 28/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20remove=20implicit?= =?UTF-8?q?=20pipeline=5Ftag=3Dfeature-extraction=20default=20in=20list=5F?= =?UTF-8?q?models?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix: list_models(task=None) no longer filters by feature-extraction, allowing LLM model search to find GGUF repos without pipeline_tag - fix: add DownloadError handling to /llm/search endpoint (same as /embeddings/search) to return meaningful error messages --- src/api/v1/models/models_llm.py | 4 ++++ src/application/models/download.py | 19 +++++++++++++++---- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/api/v1/models/models_llm.py b/src/api/v1/models/models_llm.py index f5377aac..583d22c3 100644 --- a/src/api/v1/models/models_llm.py +++ b/src/api/v1/models/models_llm.py @@ -74,10 +74,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"}) diff --git a/src/application/models/download.py b/src/application/models/download.py index af10d727..aefda9d8 100644 --- a/src/application/models/download.py +++ b/src/application/models/download.py @@ -123,6 +123,7 @@ 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. @@ -135,13 +136,23 @@ async def list_models(self, query: str, task: str | None = None) -> list[HFRepo] try: 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": + 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) From 2b4385c41791e46e7277b30bd824157f91505495 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 08:03:38 +0200 Subject: [PATCH 29/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20add=20HF=5FHUB=5FOF?= =?UTF-8?q?FLINE=20bypass=20and=20token=20normalization=20to=20hf=5Fhub=5F?= =?UTF-8?q?download=20in=20LLM=20download?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - add missing os import to models_llm.py - wrap hf_hub_download with save/restore of hc.HF_HUB_OFFLINE to allow online access even when app is in AirGap mode - normalize empty HF_TOKEN to None to avoid invalid Bearer header --- src/api/v1/models/models_llm.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/api/v1/models/models_llm.py b/src/api/v1/models/models_llm.py index 583d22c3..2ba5763f 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 @@ -187,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() From 76eaf8298acb853051443bba1b0cc0e1e57ab592 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 08:30:59 +0200 Subject: [PATCH 30/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20isolate=20service?= =?UTF-8?q?=20auto-starts=20in=20finally=20block=20and=20remove=20prematur?= =?UTF-8?q?e=20started=20flags?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Wrap _start_qd() and _start_llm() in separate try/except so one failure doesn't prevent the other from running - Remove _qdrant_started_by_us / _llamacpp_started_by_us = True that was set before raising ServiceStartError on health check timeout --- src/api/v1/models/models_llm.py | 10 ++++++++-- src/infrastructure/llm/llamacpp/auto_start.py | 1 - src/infrastructure/vectorstore/auto_start.py | 1 - 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/api/v1/models/models_llm.py b/src/api/v1/models/models_llm.py index 2ba5763f..9c8d2d12 100644 --- a/src/api/v1/models/models_llm.py +++ b/src/api/v1/models/models_llm.py @@ -216,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/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/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") From 7261164da90ce38a956559cb296eb6be00221fbd Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sat, 4 Jul 2026 09:00:06 +0200 Subject: [PATCH 31/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20isolate=20github=20?= =?UTF-8?q?and=20spec=20repo=20selected=20dirs=20with=20source=5Ftype=20co?= =?UTF-8?q?lumn?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add source_type TEXT column to git_selected_dirs table - Auto-migration via ALTER TABLE ADD COLUMN IF NOT EXISTS - Backward-compat loading from old DB files (missing source_type) - Thread source_type from routers ('github' / 'spec') to save_selected_dirs() - Filter get_repos_with_selected_dirs(source_type=) in discovery worker - Update schema_version to 20260704 --- src/api/v1/base/repo_router.py | 37 ++++++++++++++++----- src/api/v1/documents/spec.py | 4 +++ src/api/v1/infrastructure/github.py | 1 + src/infrastructure/database/core.py | 22 ++++++++++++ src/infrastructure/database/initdb.sql | 3 +- src/infrastructure/database/table_ops.py | 18 ++++++---- src/infrastructure/github/git.py | 9 +++-- src/workers/sigma/discovery_worker.py | 5 +-- tests/unit/api/v1/documents/test_spec.py | 5 ++- tests/unit/api/v1/test_github.py | 3 ++ tests/unit/back/database/test_service.py | 15 +++++++++ tests/unit/back/github/test_git_advanced.py | 4 ++- 12 files changed, 103 insertions(+), 23 deletions(-) diff --git a/src/api/v1/base/repo_router.py b/src/api/v1/base/repo_router.py index 540d4252..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, @@ -442,13 +452,15 @@ async def select_dirs( ) except Exception as e: return SelectDirsResponse(success=False, error=str(e)) - result = save_selected_dirs(org, name, request_body.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, - WorkerName.GITHUB_DISCOVERY, - task_type=WorkerName.GITHUB_DISCOVERY.value, + select_dirs_worker, + task_type=select_dirs_worker.value, collection_name="all", repo_key=f"{org}/{name}", ) @@ -491,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/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/infrastructure/database/core.py b/src/infrastructure/database/core.py index ab576f98..3dedf7fc 100644 --- a/src/infrastructure/database/core.py +++ b/src/infrastructure/database/core.py @@ -72,12 +72,25 @@ 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 NOT NULL DEFAULT ''", + ] + 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) @@ -89,6 +102,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/initdb.sql b/src/infrastructure/database/initdb.sql index 67d65138..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) ); @@ -150,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/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index d4a5f4a7..cfab8e99 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -160,8 +160,9 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: logger.warning(f"[GenericDiscoveryWorker] Invalid repo key: {repo_key}") return else: + st = "github" if worker_name == WorkerName.GITHUB_DISCOVERY else "" try: - repo_keys = self.db.get_repos_with_selected_dirs() + 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 @@ -184,7 +185,7 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: gc_repos: list[tuple[str, str]] = [] gc_seen: set[str] = set() try: - all_selected = self.db.get_repos_with_selected_dirs() + all_selected = self.db.get_repos_with_selected_dirs(source_type=st) for rk in all_selected: parts = rk.split("/") if len(parts) == 2: 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_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/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() From ec2ee0a8c0a836c6226104650de2d88ec650e7ab Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:32:13 +0200 Subject: [PATCH 32/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20inject=20references?= =?UTF-8?q?=20into=20embedding=20text,=20filters,=20payload=20indexes,=20a?= =?UTF-8?q?nd=20sigma=5Fdocs=20provenance?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/application/documents/indexing.py | 4 ++++ src/core/pipeline/indexer.py | 13 +++++++++++++ src/core/search/engine.py | 3 +++ src/core/sigma/chunker.py | 6 ++++++ src/infrastructure/vectorstore/collections.py | 1 + 5 files changed, 27 insertions(+) diff --git a/src/application/documents/indexing.py b/src/application/documents/indexing.py index 09bf1f62..f8bf5832 100644 --- a/src/application/documents/indexing.py +++ b/src/application/documents/indexing.py @@ -126,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) @@ -140,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/core/pipeline/indexer.py b/src/core/pipeline/indexer.py index f9513a91..12ff9899 100644 --- a/src/core/pipeline/indexer.py +++ b/src/core/pipeline/indexer.py @@ -107,6 +107,19 @@ 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 diff --git a/src/core/search/engine.py b/src/core/search/engine.py index aca1f025..5c0b8c71 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -37,6 +37,7 @@ "chunk_type", "collection", "tags", + "references", } ) @@ -274,6 +275,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 diff --git a/src/core/sigma/chunker.py b/src/core/sigma/chunker.py index c9db8525..f035ba94 100644 --- a/src/core/sigma/chunker.py +++ b/src/core/sigma/chunker.py @@ -427,6 +427,12 @@ def _assemble_chunks(self, rule: dict, f: dict) -> list[dict]: ] ) + refs = f.get("references", []) + if refs: + ref_text = "\n\nReferences:\n" + "\n".join(f" - {r}" for r in refs) + for chunk in chunks: + chunk["text"] += ref_text + return chunks def _dict_to_document(self, chunk_data: dict, source_file: str | None = None) -> Document: diff --git a/src/infrastructure/vectorstore/collections.py b/src/infrastructure/vectorstore/collections.py index 0e2d2b9e..6ad10846 100644 --- a/src/infrastructure/vectorstore/collections.py +++ b/src/infrastructure/vectorstore/collections.py @@ -93,6 +93,7 @@ def _create_collection_sync( ("chunk_type", "keyword"), ("collection", "keyword"), ("tags", "keyword"), + ("references", "keyword"), ] From 9f77552178756f5f23301ee769d7e6be608997fe Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:32:23 +0200 Subject: [PATCH 33/44] =?UTF-8?q?=F0=9F=8E=A8=20style:=20ruff-format=20fix?= =?UTF-8?q?=20for=20long=20line=20in=20indexer.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/core/pipeline/indexer.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/core/pipeline/indexer.py b/src/core/pipeline/indexer.py index 12ff9899..906ea971 100644 --- a/src/core/pipeline/indexer.py +++ b/src/core/pipeline/indexer.py @@ -113,7 +113,14 @@ async def index(self, route: IndexRoute) -> IndexResult: 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") + for k in ( + "rule_id", + "original_url", + "normalized_url", + "url_hash", + "content_type", + "title", + ) if row.get(k) } if source_meta: From 3cfdc94a3e2b2f96af6b908ec19d2b8798092afa Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:33:22 +0200 Subject: [PATCH 34/44] =?UTF-8?q?=E2=9C=A8=20feat:=20support=20html,=20yam?= =?UTF-8?q?l,=20pdf,=20plain=5Ftext=20reference=20doc=20types?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add html, yaml, pdf, and plain_text to SUPPORTED_REFERENCE_DOC_TYPES so the downloader processes them instead of skipping. Add yaml entry to FILETYPE_INFO for path resolution. --- src/shared/utils/identify_file_type.py | 5 +++++ tests/unit/shared/utils/test_identify_file_type.py | 6 ++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index d1d658b9..2ddaf1a8 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -62,6 +62,7 @@ class FileType(Enum): FileType.PLAIN_TEXT.value: {"ext": ".txt", "subdir": "plain_text"}, FileType.HTML.value: {"ext": ".html", "subdir": "html"}, FileType.OFFICE_DOCUMENT.value: {"ext": ".docx", "subdir": "office"}, + FileType.YAML.value: {"ext": ".yml", "subdir": "yaml"}, } @@ -77,6 +78,10 @@ def filetype_subdir(content_type: str) -> str: SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, + FileType.HTML.value, + FileType.YAML.value, + FileType.PLAIN_TEXT.value, + FileType.PDF.value, } PUREMAGIC_TYPE_MAP: dict[str, FileType] = { diff --git a/tests/unit/shared/utils/test_identify_file_type.py b/tests/unit/shared/utils/test_identify_file_type.py index e9919fa8..ba2645a0 100644 --- a/tests/unit/shared/utils/test_identify_file_type.py +++ b/tests/unit/shared/utils/test_identify_file_type.py @@ -295,8 +295,10 @@ 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" not in SUPPORTED_REFERENCE_DOC_TYPES - assert "plain_text" not in SUPPORTED_REFERENCE_DOC_TYPES + assert "html" in SUPPORTED_REFERENCE_DOC_TYPES + assert "yaml" in SUPPORTED_REFERENCE_DOC_TYPES + assert "pdf" in SUPPORTED_REFERENCE_DOC_TYPES + assert "plain_text" in SUPPORTED_REFERENCE_DOC_TYPES assert "office_document" not in SUPPORTED_REFERENCE_DOC_TYPES def test_excludes_media_types(self) -> None: From 081558b8101bd11bc01f35a47ec3f6c05d355221 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:36:28 +0200 Subject: [PATCH 35/44] =?UTF-8?q?=E2=9C=A8=20feat:=20improve=20router=20pr?= =?UTF-8?q?ompt=20for=20sigma=5Fdocs=20refs=20+=20auto-include=20sigma=5Fd?= =?UTF-8?q?ocs=20on=20references:=20filter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Update router prompt to describe sigma_docs as containing external reference documents linked to rules. Auto-include sigma_docs in search when a 'references:' filter is present in the query. --- src/core/search/engine.py | 14 ++++++++++++++ src/core/search/router.py | 6 +++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/core/search/engine.py b/src/core/search/engine.py index 5c0b8c71..db96460a 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -368,6 +368,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: @@ -382,6 +391,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 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. From 6bcd5d4da1fc16514dea82182311782f233b0efd Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:37:51 +0200 Subject: [PATCH 36/44] =?UTF-8?q?=F0=9F=94=A7=20chore:=20add=20broken=20re?= =?UTF-8?q?ference=20monitoring=20=E2=80=94=20failure=20rate=20warning=20a?= =?UTF-8?q?t=205%=20threshold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add _log_download_summary helper to log download results and warn when the failure rate exceeds 5%. Apply to both scan and registry modes. --- .../documents/sigma_ref_downloader.py | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index a3011f7a..02d695b9 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -460,6 +460,10 @@ def _collect_yaml_files() -> list[Path]: ) 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"], @@ -484,10 +488,24 @@ def _collect_yaml_files() -> list[Path]: "skipped": skipped, "failed": failed, } - logger.info("Download complete: %s", summary) + _log_download_summary(summary) return summary +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]: """Return an empty summary dict.""" return { @@ -742,10 +760,12 @@ def _download_one(item: dict[str, Any]) -> tuple[str, str, str, int] | None: if rule_refs: db.batch_upsert_rule_references(rule_refs) - return { + summary = { "total_rules": total_rules, "total_refs": total_refs, "downloaded": downloaded, "skipped": skipped, "failed": failed, } + _log_download_summary(summary) + return summary From b06427d7c783e6156197a9b8edca4422cc719e3f Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:39:22 +0200 Subject: [PATCH 37/44] =?UTF-8?q?=E2=9C=85=20test:=20add=20tests=20for=20d?= =?UTF-8?q?elete=5Funreferenced=5Fentries=20GC=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add three tests covering orphan deletion, noop when all refs are referenced, and scoping to sigmaref org only (skipping local/github). --- tests/unit/workers/test_gc_worker.py | 108 +++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/unit/workers/test_gc_worker.py b/tests/unit/workers/test_gc_worker.py index cab6cd66..0e29f11b 100644 --- a/tests/unit/workers/test_gc_worker.py +++ b/tests/unit/workers/test_gc_worker.py @@ -2,9 +2,13 @@ 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 @@ -156,3 +160,107 @@ def test_db_error_returns_zero(self, tmp_path: Path) -> None: 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",)] From 2ccb9dbc608de6cda1b53c245a006772231956d7 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 07:48:20 +0200 Subject: [PATCH 38/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20remove=20duplicate?= =?UTF-8?q?=20=5Fwrite=5Fentries=20and=20add=20missing=20batch=5Fupsert=5F?= =?UTF-8?q?fn=20param?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The second definition shadowed the first, causing a TypeError when spec discovery called _write_entries with batch_upsert_fn keyword. --- src/workers/sigma/discovery_worker.py | 36 +++++++-------------------- 1 file changed, 9 insertions(+), 27 deletions(-) diff --git a/src/workers/sigma/discovery_worker.py b/src/workers/sigma/discovery_worker.py index cfab8e99..f1b3ab07 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -440,31 +440,6 @@ def _prepare_spec_entry( logger.error(f"[GenericDiscoveryWorker] Cannot prepare spec entry for {file_path}: {e}") return None - def _write_entries( - 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: - upsert_fn(entries) - except Exception as e: - logger.error(f"[GenericDiscoveryWorker] Batch upsert 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 # ------------------------------------------------------------------ @@ -614,14 +589,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) From 673d62537afd89af3760407332da2c2ede03831a Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Sun, 5 Jul 2026 08:02:29 +0200 Subject: [PATCH 39/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20resolve=20code=20re?= =?UTF-8?q?view=20bugs=20=E2=80=94=20unbound=20variable,=20Qdrant=20leak,?= =?UTF-8?q?=20SQL=20injection,=20nested=20lock,=20deprecated=20typing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - BUG 1 (critical): move assignment before if/else in discovery_worker.py - BUG 2 (high): add in block in storage.py - BUG 3 (high): parameterize SQL query with instead of f-string in gc_worker.py - BUG 8: remove redundant inner (RLock but unnecessary) - BUG 9: replace deprecated with ; with - Bonus: migrate to --- src/core/registry.py | 3 ++- src/infrastructure/vectorstore/storage.py | 5 +++-- src/shared/utils/sigma_utils.py | 4 ++-- src/workers/document/gc_worker.py | 24 +++++++++++------------ src/workers/processor.py | 10 +++++----- src/workers/sigma/discovery_worker.py | 19 +++++++++--------- 6 files changed, 33 insertions(+), 32 deletions(-) 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/infrastructure/vectorstore/storage.py b/src/infrastructure/vectorstore/storage.py index e8f9a6a1..6c8079e1 100644 --- a/src/infrastructure/vectorstore/storage.py +++ b/src/infrastructure/vectorstore/storage.py @@ -97,9 +97,8 @@ 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] @@ -150,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/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/workers/document/gc_worker.py b/src/workers/document/gc_worker.py index 51258551..712425f3 100644 --- a/src/workers/document/gc_worker.py +++ b/src/workers/document/gc_worker.py @@ -54,7 +54,6 @@ 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, @@ -89,7 +88,6 @@ def process(self, task: dict) -> None: def _gc_entries( self, - table: str, local_base: Path, github_base: Path, grace_days: int, @@ -100,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: @@ -121,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 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_worker.py b/src/workers/sigma/discovery_worker.py index f1b3ab07..74e7548f 100644 --- a/src/workers/sigma/discovery_worker.py +++ b/src/workers/sigma/discovery_worker.py @@ -6,7 +6,7 @@ from enum import Enum from pathlib import Path from collections.abc import Callable -from typing import TYPE_CHECKING, Optional +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 @@ -54,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 @@ -151,6 +151,8 @@ 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() + st = "github" if worker_name == WorkerName.GITHUB_DISCOVERY else "" + repo_key = task.get("repo_key") if repo_key: parts = repo_key.split("/") @@ -160,7 +162,6 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: logger.warning(f"[GenericDiscoveryWorker] Invalid repo key: {repo_key}") return else: - st = "github" if worker_name == WorkerName.GITHUB_DISCOVERY else "" try: repo_keys = self.db.get_repos_with_selected_dirs(source_type=st) except Exception as e: @@ -172,7 +173,7 @@ def _process_github(self, task: dict, worker_name: WorkerName) -> None: logger.info("[GenericDiscoveryWorker] No repos with selected dirs") return - repo_items: list[tuple[str, str]] = [] + repo_items = [] for rk in repo_keys: parts = rk.split("/") if len(parts) != 2: From 672860fa1e5d91a24b4d8122065131d8d697b543 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:40:53 +0200 Subject: [PATCH 40/44] =?UTF-8?q?=F0=9F=90=9B=20fix:=20make=20source=5Ftyp?= =?UTF-8?q?e=20migration=20nullable-safe,=20close=20repo=20detail=20on=20d?= =?UTF-8?q?elete,=20reorder=20sidebar=20nav?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/infrastructure/database/core.py | 4 +++- src/presentation/static/js/internal/repo-browser.js | 1 + src/presentation/templates/data/base.html.j2 | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/src/infrastructure/database/core.py b/src/infrastructure/database/core.py index 3dedf7fc..c20cff0e 100644 --- a/src/infrastructure/database/core.py +++ b/src/infrastructure/database/core.py @@ -82,7 +82,9 @@ def initialize(self) -> None: def _apply_migrations(self) -> None: """Apply schema migrations idempotently.""" migrations = [ - "ALTER TABLE git_selected_dirs ADD COLUMN IF NOT EXISTS source_type TEXT NOT NULL DEFAULT ''", + "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: 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 From f01327770967fd545f8edd962a756fe384d44182 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:46:44 +0200 Subject: [PATCH 41/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20limit=20?= =?UTF-8?q?sigma=20reference=20types=20to=20markdown=20and=20sigma=5Frule?= =?UTF-8?q?=20only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove html, yaml, pdf, plain_text from SUPPORTED_REFERENCE_DOC_TYPES since only markdown and sigma_rule have dedicated transforms. --- src/shared/utils/identify_file_type.py | 5 +---- tests/unit/shared/utils/test_identify_file_type.py | 9 +++++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index 2ddaf1a8..1f66d963 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -78,10 +78,7 @@ def filetype_subdir(content_type: str) -> str: SUPPORTED_REFERENCE_DOC_TYPES: set[str] = { FileType.MARKDOWN.value, - FileType.HTML.value, - FileType.YAML.value, - FileType.PLAIN_TEXT.value, - FileType.PDF.value, + FileType.SIGMA_RULE.value, } PUREMAGIC_TYPE_MAP: dict[str, FileType] = { diff --git a/tests/unit/shared/utils/test_identify_file_type.py b/tests/unit/shared/utils/test_identify_file_type.py index ba2645a0..5fef697d 100644 --- a/tests/unit/shared/utils/test_identify_file_type.py +++ b/tests/unit/shared/utils/test_identify_file_type.py @@ -295,10 +295,11 @@ def test_values_are_filetype(self) -> None: class TestSupportedReferenceDocTypes: def test_contains_relevant_types(self) -> None: assert "markdown" in SUPPORTED_REFERENCE_DOC_TYPES - assert "html" in SUPPORTED_REFERENCE_DOC_TYPES - assert "yaml" in SUPPORTED_REFERENCE_DOC_TYPES - assert "pdf" in SUPPORTED_REFERENCE_DOC_TYPES - assert "plain_text" 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: From 071052a541f18bb6c0d336bf4030ee32dd12f398 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:55:25 +0200 Subject: [PATCH 42/44] =?UTF-8?q?=F0=9F=8E=A8=20style:=20ruff-format=20tes?= =?UTF-8?q?t=5Fsearch.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/api/v1/chat/search.py | 5 ++-- src/application/chat/service.py | 10 +++++-- src/core/search/engine.py | 3 ++ src/presentation/static/js/internal/chat.js | 32 +++++++++++++++++++++ tests/unit/api/v1/test_search.py | 28 ++++++++++++++++-- 5 files changed, 71 insertions(+), 7 deletions(-) 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/application/chat/service.py b/src/application/chat/service.py index c7d93712..5c653084 100644 --- a/src/application/chat/service.py +++ b/src/application/chat/service.py @@ -16,7 +16,7 @@ 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 @@ -351,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=prompt_id + augmented_message, formatted_results, system_prompt_id=prompt_id ): yield token found = True @@ -380,9 +382,11 @@ 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( - rule, results, system_prompt_id=prompt_id + rule, formatted_results, system_prompt_id=prompt_id ): yield token except Exception as e: diff --git a/src/core/search/engine.py b/src/core/search/engine.py index db96460a..4dbb47d8 100644 --- a/src/core/search/engine.py +++ b/src/core/search/engine.py @@ -255,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": 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/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: From 875beb20cd8f848fd1678e1209d00cf1d289e913 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:56:31 +0200 Subject: [PATCH 43/44] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20align=20?= =?UTF-8?q?FILETYPE=5FINFO=20subdir=20with=20supported=20sigma=20reference?= =?UTF-8?q?=20types?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep only markdown and sigma in FILETYPE_INFO; remove pdf, plain_text, html, yaml subdirs. Add sigma_rule entry so files land under sigma/ subdirectory. --- src/shared/utils/identify_file_type.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/shared/utils/identify_file_type.py b/src/shared/utils/identify_file_type.py index 1f66d963..59ed2ae3 100644 --- a/src/shared/utils/identify_file_type.py +++ b/src/shared/utils/identify_file_type.py @@ -58,11 +58,8 @@ class FileType(Enum): FILETYPE_INFO: dict[str, dict[str, str]] = { FileType.MARKDOWN.value: {"ext": ".md", "subdir": "markdown"}, - FileType.PDF.value: {"ext": ".pdf", "subdir": "pdf"}, - FileType.PLAIN_TEXT.value: {"ext": ".txt", "subdir": "plain_text"}, - FileType.HTML.value: {"ext": ".html", "subdir": "html"}, + FileType.SIGMA_RULE.value: {"ext": ".yml", "subdir": "sigma"}, FileType.OFFICE_DOCUMENT.value: {"ext": ".docx", "subdir": "office"}, - FileType.YAML.value: {"ext": ".yml", "subdir": "yaml"}, } From cf13217055f8d014555f3523a32c82713780c8d0 Mon Sep 17 00:00:00 2001 From: frack113 <62423083+frack113@users.noreply.github.com> Date: Mon, 13 Jul 2026 06:59:05 +0200 Subject: [PATCH 44/44] =?UTF-8?q?=F0=9F=8E=A8=20style:=20ruff-format=20sig?= =?UTF-8?q?ma=5Fref=5Fdownloader=20and=20test=5Fgc=5Fworker?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/application/documents/sigma_ref_downloader.py | 8 ++++++-- tests/unit/workers/test_gc_worker.py | 4 +--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/application/documents/sigma_ref_downloader.py b/src/application/documents/sigma_ref_downloader.py index 02d695b9..2a3720b5 100644 --- a/src/application/documents/sigma_ref_downloader.py +++ b/src/application/documents/sigma_ref_downloader.py @@ -462,7 +462,9 @@ def _collect_yaml_files() -> list[Path]: else: logger.debug( "Reference download failed | url=%s rule_id=%s status=%s", - item["original_url"], item["rule_id"], status_code, + item["original_url"], + item["rule_id"], + status_code, ) _maybe_record_error( db, @@ -502,7 +504,9 @@ def _log_download_summary(summary: dict[str, Any]) -> None: 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, + fail_rate * 100, + failed, + total_refs, ) diff --git a/tests/unit/workers/test_gc_worker.py b/tests/unit/workers/test_gc_worker.py index 0e29f11b..19c27e32 100644 --- a/tests/unit/workers/test_gc_worker.py +++ b/tests/unit/workers/test_gc_worker.py @@ -246,9 +246,7 @@ def test_noop_when_no_orphans(self) -> None: result = db.delete_unreferenced_entries() assert result == 0 - remaining = db._writer_conn.execute( - "SELECT COUNT(*) FROM doc_registry" - ).fetchone()[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: