"""LLM-driven document classifier for auto-categorization. Provides structured JSON classification of documents via Ollama /api/chat. All prompts enforce strict JSON output with temperature=0 for determinism. ClassificationResult has all fields mandatory — never empty/None. Research references: - ByteRover (arXiv:2604.01599) — hierarchical Context Tree, knowledge in markdown - Memex (jagdeepvirdi/memex) — Ollama auto-classify, confidence scores, category review - secondmem — AI classification into topic markdown files with JSON schema - Grimore-MD — two-tier change detection, RRF fusion - RAGraph — entity/topic resolution, graph-based organization """ from __future__ import annotations import json import logging import re from dataclasses import dataclass from typing import Any from httpx import HTTPStatusError from cyberbrain26.knowledge_plane.embedding_config import EmbeddingConfig logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # JSON Schema for LLM classification output # --------------------------------------------------------------------------- CLASSIFY_PROMPT = """You are classifying a technical document into a knowledge hierarchy. Return ONLY a valid JSON object matching this exact schema — no prose, no markdown, no extra keys. { "topic": "ONE of: ai, architecture, security, research, notes, reports, general", "subtopic": "short free-text domain within the topic (e.g. 'audit', 'design', 'networking')", "tags": ["TAG1", "TAG2", ...], // 2-5 tags max, lowercase, relevant to content "summary": "2-3 sentence summary of the document's core content", "confidence": "FLOAT string 0.0-1.0 (e.g. \"0.87\") — your confidence in this classification", "entities": ["ENTITY1", "ENTITY2", ...], // 0-4 named entities extracted from text "intent": "ONE of: actionable, reference, idea" } Temperature must be 0 for deterministic output. Do NOT invent tags or entities not supported by the text. """ # --------------------------------------------------------------------------- # Classification result dataclass # --------------------------------------------------------------------------- @dataclass(frozen=True) class ClassificationResult: """Result of LLM-driven document classification. All fields are mandatory — never None/empty. If the LLM output is malformed, the caller must handle the exception rather than producing a partial result. """ topic: str subtopic: str tags: list[str] summary: str confidence: float entities: list[str] intent: str # "actionable" | "reference" | "idea" # --------------------------------------------------------------------------- # LLM Classifier implementation # --------------------------------------------------------------------------- class LLMClassifier: """Classify documents using a local Ollama model. Uses Ollama /api/chat endpoint with strict JSON prompt and temperature=0 for deterministic output. All fields in ClassificationResult are mandatory. Research-backed design: - ByteRover: hierarchical memory via agent-native classification - Memex: confidence-scored auto-classification, categorical review queue - secondmem: AI classification into topic markdown files """ def __init__(self, config: EmbeddingConfig | None = None) -> None: self.config = config or EmbeddingConfig() self._client = None # lazily initialized @property def client(self): """Lazily create httpx client for Ollama chat endpoint.""" if self._client is None: self._client = __import__("httpx").Client( base_url=self.config.ollama_base_url, timeout=60.0, ) return self._client def _call_ollama( self, text: str, *, temperature: float = 0.0, num_ctx: int | None = None ) -> dict[str, Any]: """Call Ollama /api/chat and return parsed JSON response.""" payload: dict[str, Any] = { "model": self.config.chat_model, "messages": [ { "role": "user", "content": f"{CLASSIFY_PROMPT}\n\nDocument text:\n{text}", } ], "options": {"temperature": temperature, "num_ctx": num_ctx or 8192}, "format": "json", # enforce JSON output from Ollama } try: response = self.client.post("/api/chat", json=payload) response.raise_for_status() data = response.json() # Ollama returns {"message": {"role": "assistant", "content": "..."}} content = data.get("message", {}).get("content", "{}") parsed = json.loads(content) if not isinstance(parsed, dict): raise ValueError(f"Expected dict from LLM, got {type(parsed)}") return parsed except (HTTPStatusError, json.JSONDecodeError, ValueError) as exc: logger.error("LLM classification failed: %s", exc) raise def classify_document(self, text: str) -> ClassificationResult: """Run LLM classification on document text. Returns ClassificationResult with all fields mandatory. If the LLM fails (network error, malformed output), raises RuntimeError — caller must decide whether to retry, stage for review, or assign to fallback category. Design rationale (per Memex and ByteRover research): - temperature=0 ensures deterministic output across runs - format=json forces Ollama to return structured data - all fields mandatory — no empty/None values in result - 2-5 tags max prevents tag spray """ raw = self._call_ollama(text) # Validate all required keys exist required_keys = {"topic", "subtopic", "tags", "summary", "confidence", "entities", "intent"} missing = required_keys - set(raw.keys()) if missing: raise ValueError(f"LLM output missing keys: {missing}") # Validate types if not isinstance(raw["tags"], list) or not all(isinstance(t, str) for t in raw["tags"]): raise ValueError(f"tags must be list of str, got {type(raw['tags'])}") if not isinstance(raw["entities"], list) or not all(isinstance(e, str) for e in raw["entities"]): raise ValueError(f"entities must be list of str, got {type(raw['entities'])}") if raw["confidence"] not in [str(i/100) for i in range(101)] + ["0.0", "1.0"]: # Be lenient with float string -> float conversion try: confidence_val = float(raw["confidence"]) if not 0.0 <= confidence_val <= 1.0: raise ValueError(f"confidence must be 0.0-1.0, got {confidence_val}") except (ValueError, TypeError): raise ValueError( f"confidence must be parseable float string, got '{raw['confidence']}'" ) from None # Map intent string to enum-like values intent_map = {"actionable": "actionable", "reference": "reference", "idea": "idea"} intent_raw = raw["intent"].lower() if intent_raw not in intent_map: raise ValueError(f"intent must be one of {list(intent_map.keys())}, got '{raw['intent']}'") intent = intent_map[intent_raw] # Normalize topic to allowed set (per prompt schema) allowed_topics = {"ai", "architecture", "security", "research", "notes", "reports", "general"} topic_lower = raw["topic"].lower() if topic_lower not in allowed_topics: logger.warning("Topic '%s' not in allowed set, using 'general'", topic_lower) topic = "general" else: topic = topic_lower # Normalize subtopic — ensure it's not empty string subtopic = raw["subtopic"] if raw["subtopic"] else "general" # Tag limit: enforce 2-5 tags (per secondmem research) tags = raw["tags"][:5] # truncate if too many if len(tags) < 2: # Pad with generic tags if too few tags = tags + ["knowledge", "ai"][: 2 - len(tags)] # Entity limit: enforce 0-4 entities (per RAGraph research) entities = raw["entities"][:4] confidence_val = float(raw["confidence"]) return ClassificationResult( topic=topic, subtopic=subtopic, tags=tags, summary=raw["summary"], confidence=confidence_val, entities=entities, intent=intent, ) def close(self) -> None: """Close the httpx client.""" if self._client is not None: self._client.close() self._client = None def __enter__(self) -> LLMClassifier: return self def __exit__(self, *_exc: object) -> None: self.close() # --------------------------------------------------------------------------- # PII / sensitivity detection (MemStrata pattern: deterministic fact # retirement based on content sensitivity, no LLM call required) # --------------------------------------------------------------------------- _EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}") _PHONE_RE = re.compile(r"(\+\d{1,3}[- ]?)?\d{3,4}[- ]?\d{3}[- ]?\d{4}") _CREDIT_CARD_RE = re.compile(r"(?:4[0-9]{12}|5[1-5][0-9]{14}|3[47][0-9]{13}|3[0-9]{15}|\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4})") _SSN_RE = re.compile(r"\d{3}-?\d{2}-?\d{4}") _MEDICAL_RE = re.compile(r"\b(?:patient|diagnosis|prescription|HIPAA|PHI)\b", re.IGNORECASE) _CREDENTIAL_RE = re.compile(r"(?:api[_-]?key|token|secret)\s*[:=]\s*['\"]?[A-Za-z0-9\-_]{16,}['\"]?", re.IGNORECASE) def detect_sensitivity(text: str) -> str: """Detect document sensitivity from plain text. Returns one of: "public", "personal", "sensitive", "restricted" based on heuristic PII pattern matching. Deterministic, no LLM call. Used per MemStrata research for automated fact classification. Priority order (first match wins): 1. RESTRICTED: credit card, SSN, API credentials 2. SENSITIVE: medical/health terms + personal identifiers 3. PERSONAL: email addresses, phone numbers 4. PUBLIC: nothing matched """ # Normalize line endings for consistent matching normalized = text.replace("\n", " ").replace("\r", "") # Check restricted patterns first if _CREDIT_CARD_RE.search(normalized) or _SSN_RE.search(normalized) or _CREDENTIAL_RE.search(normalized): return "restricted" # Check sensitive: medical terms combined with personal data has_medical = _MEDICAL_RE.search(normalized) is not None has_personal = _EMAIL_RE.search(normalized) is not None or _PHONE_RE.search(normalized) is not None if has_medical and has_personal: return "sensitive" # Check personal identifiers alone if _EMAIL_RE.search(normalized) is not None or _PHONE_RE.search(normalized) is not None: return "personal" return "public" # --------------------------------------------------------------------------- # Convenience: short-circuit for non-text inputs # --------------------------------------------------------------------------- def detect_sensitivity_from_bytes(content: bytes) -> str: """Detect sensitivity from raw bytes by decoding as UTF-8 first.""" try: text = content.decode("utf-8", errors="replace") except Exception: # noqa: BLE001 - fall through return "public" return detect_sensitivity(text)