contextos-memory-runtime 1.0.0rc2__py3-none-any.whl

contextos/services/privacy.py

@@ -0,0 +1,331 @@
+"""Two-stage local privacy gate for raw input and extracted candidates."""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import re
+import secrets
+from dataclasses import dataclass
+from typing import Any
+
+from contextos.core.enums import (
+    PrivacyClassification,
+    PrivacyDecision,
+    PrivacySeverity,
+    SecretDetectionMode,
+    SecretType,
+    SourceRole,
+    SourceTrust,
+)
+from contextos.core.models import (
+    CandidateMemory,
+    PrivacyAssessment,
+    PrivacyFinding,
+    SecretMatch,
+)
+from contextos.core.protocols import SecretScanner
+
+
+_CRITICAL_TYPES = {
+    SecretType.PRIVATE_KEY,
+    SecretType.SSH_PRIVATE_KEY,
+    SecretType.PASSWORD,
+    SecretType.CONNECTION_STRING,
+    SecretType.AUTHORIZATION_HEADER,
+    SecretType.SESSION_COOKIE,
+    SecretType.ACCESS_TOKEN,
+}
+_CREDENTIAL_LANGUAGE = re.compile(
+    r"\b(?:api key|password|passwd|otp|pin|token|authorization|cookie|session|"
+    r"private key|secret key|connection string)\b",
+    re.I,
+)
+_UNTRUSTED_INSTRUCTION = re.compile(
+    r"\b(?:ignore (?:all |the )?(?:previous|prior) instructions|system instruction|"
+    r"developer instruction|save .{0,40}(?:password|token|secret).{0,20}(?:memory|forever)|"
+    r"permanently remember)\b",
+    re.I,
+)
+_FINGERPRINT_KEY = secrets.token_bytes(32)
+
+
+@dataclass(frozen=True)
+class GatedInput:
+    assessment: PrivacyAssessment
+    content: str
+    source_type: str
+    source_uri: str | None
+    tags: list[str]
+
+
+def _hash(value: str) -> str:
+    return hmac.new(_FINGERPRINT_KEY, value.encode("utf-8"), hashlib.sha256).hexdigest()
+
+
+def classify_source(source_type: str, source_role: SourceRole) -> SourceTrust:
+    """Classify source trust as metadata, never as proof that content is true."""
+    normalized = source_type.casefold().replace("-", "_")
+    if source_role in {SourceRole.ASSISTANT, SourceRole.SYSTEM}:
+        return SourceTrust.MODEL_OUTPUT
+    if any(marker in normalized for marker in ("web", "browser", "url")):
+        return SourceTrust.EXTERNAL_WEBPAGE
+    if any(marker in normalized for marker in ("document", "file", "import")):
+        return SourceTrust.IMPORTED_DOCUMENT
+    if "tool" in normalized:
+        return SourceTrust.TOOL_OUTPUT
+    if "connector" in normalized:
+        return SourceTrust.LOCAL_TRUSTED_CONNECTOR
+    return SourceTrust.DIRECT_USER
+
+
+class PrivacyGate:
+    """Deterministic pre-ingest and post-extraction security policy."""
+
+    def __init__(self, scanner: SecretScanner) -> None:
+        self._scanner = scanner
+
+    def gate_input(
+        self,
+        content: str,
+        *,
+        source_type: str,
+        source_uri: str | None,
+        tags: list[str],
+        source_role: SourceRole,
+        mode: SecretDetectionMode,
+    ) -> GatedInput:
+        safe_content, findings = self._scan_text(content, "content")
+        safe_source_type, source_type_findings = self._scan_text(
+            source_type, "source_type"
+        )
+        findings.extend(source_type_findings)
+        safe_source_uri: str | None = None
+        if source_uri is not None:
+            safe_source_uri, source_uri_findings = self._scan_text(
+                source_uri, "source_uri"
+            )
+            findings.extend(source_uri_findings)
+        safe_tags: list[str] = []
+        for index, tag in enumerate(tags):
+            safe_tag, tag_findings = self._scan_text(tag, f"tags[{index}]")
+            safe_tags.append(safe_tag)
+            findings.extend(tag_findings)
+
+        trust = classify_source(safe_source_type, source_role)
+        untrusted_instruction = (
+            trust not in {SourceTrust.DIRECT_USER, SourceTrust.LOCAL_TRUSTED_CONNECTOR}
+            and bool(_UNTRUSTED_INSTRUCTION.search(safe_content))
+        )
+        decision, classification = self._policy(findings, mode)
+        raw_envelope = json.dumps(
+            [content, source_type, source_uri, tags], ensure_ascii=False, separators=(",", ":")
+        )
+        assessment = PrivacyAssessment(
+            decision=decision,
+            classification=classification,
+            source_trust=trust,
+            findings=findings,
+            sanitized_text=safe_content,
+            scanned_length=sum(
+                len(value) for value in [content, source_type, source_uri or "", *tags]
+            ),
+            input_hash=_hash(raw_envelope),
+            untrusted_instruction_detected=untrusted_instruction,
+        )
+        return GatedInput(
+            assessment=assessment,
+            content=safe_content,
+            source_type=safe_source_type,
+            source_uri=safe_source_uri,
+            tags=safe_tags,
+        )
+
+    def assess_input(
+        self,
+        text: str,
+        *,
+        source_type: str,
+        source_role: SourceRole,
+        mode: SecretDetectionMode,
+    ) -> PrivacyAssessment:
+        """Compatibility helper for callers assessing only a content string."""
+        return self.gate_input(
+            text,
+            source_type=source_type,
+            source_uri=None,
+            tags=[],
+            source_role=source_role,
+            mode=mode,
+        ).assessment
+
+    def assess_candidate(
+        self,
+        candidate: CandidateMemory,
+        *,
+        source_trust: SourceTrust,
+    ) -> tuple[CandidateMemory | None, PrivacyAssessment]:
+        content, findings = self._scan_text(candidate.content, "candidate.content")
+        evidence, evidence_findings = self._scan_text(
+            candidate.evidence, "candidate.evidence"
+        )
+        findings.extend(evidence_findings)
+        source_type, source_type_findings = self._scan_text(
+            candidate.source_type, "candidate.source_type"
+        )
+        findings.extend(source_type_findings)
+        source_uri = None
+        if candidate.source_uri is not None:
+            source_uri, source_uri_findings = self._scan_text(
+                candidate.source_uri, "candidate.source_uri"
+            )
+            findings.extend(source_uri_findings)
+        tags, tag_findings = self._sanitize_structure(candidate.tags, "candidate.tags")
+        metadata, metadata_findings = self._sanitize_structure(
+            candidate.metadata, "candidate.metadata"
+        )
+        findings.extend(tag_findings)
+        findings.extend(metadata_findings)
+
+        contains_redaction = "[REDACTED:" in content or "[REDACTED:" in evidence
+        credential_only = contains_redaction and bool(_CREDENTIAL_LANGUAGE.search(content))
+        combined = f"{content}\n{evidence}"
+        untrusted_instruction = (
+            source_trust not in {
+                SourceTrust.DIRECT_USER,
+                SourceTrust.LOCAL_TRUSTED_CONNECTOR,
+            }
+            and bool(_UNTRUSTED_INSTRUCTION.search(combined))
+        )
+        blocked = credential_only or untrusted_instruction
+        assessment = PrivacyAssessment(
+            decision=PrivacyDecision.REJECT if blocked else (
+                PrivacyDecision.REDACT if findings else PrivacyDecision.ALLOW
+            ),
+            classification=PrivacyClassification.BLOCKED if blocked else (
+                PrivacyClassification.SECRET if findings else PrivacyClassification.SAFE
+            ),
+            source_trust=source_trust,
+            findings=findings,
+            sanitized_text=combined,
+            scanned_length=len(combined),
+            input_hash=_hash(json.dumps(
+                candidate.model_dump(mode="json", exclude={"privacy_assessment"}),
+                sort_keys=True,
+                separators=(",", ":"),
+            )),
+            untrusted_instruction_detected=untrusted_instruction,
+        )
+        if blocked:
+            return None, assessment
+        safe_candidate = candidate.model_copy(update={
+            "content": content,
+            "evidence": evidence,
+            "source_type": source_type,
+            "source_uri": source_uri,
+            "tags": tags,
+            "metadata": metadata,
+            "privacy_assessment": assessment,
+        })
+        return safe_candidate, assessment
+
+    def gate_candidates(
+        self,
+        candidates: list[CandidateMemory],
+        *,
+        source_trust: SourceTrust,
+    ) -> tuple[list[CandidateMemory], list[PrivacyAssessment]]:
+        safe: list[CandidateMemory] = []
+        blocked: list[PrivacyAssessment] = []
+        for candidate in candidates:
+            gated, assessment = self.assess_candidate(candidate, source_trust=source_trust)
+            if gated is None:
+                blocked.append(assessment)
+            else:
+                safe.append(gated)
+        return safe, blocked
+
+    def _scan_text(self, text: str, location: str) -> tuple[str, list[PrivacyFinding]]:
+        scan = self._scanner.scan(text)
+        findings = self._safe_findings(text, scan.matches, location)
+        return self._redact_from_findings(text, findings), findings
+
+    def _sanitize_structure(
+        self, value: Any, location: str
+    ) -> tuple[Any, list[PrivacyFinding]]:
+        if isinstance(value, str):
+            return self._scan_text(value, location)
+        if isinstance(value, list):
+            safe_list: list[Any] = []
+            findings: list[PrivacyFinding] = []
+            for index, item in enumerate(value):
+                safe_item, item_findings = self._sanitize_structure(
+                    item, f"{location}[{index}]"
+                )
+                safe_list.append(safe_item)
+                findings.extend(item_findings)
+            return safe_list, findings
+        if isinstance(value, dict):
+            safe_dict: dict[str, Any] = {}
+            findings = []
+            for key, item in value.items():
+                safe_key, key_findings = self._scan_text(str(key), f"{location}.key")
+                safe_item, item_findings = self._sanitize_structure(
+                    item, f"{location}.{safe_key}"
+                )
+                safe_dict[safe_key] = safe_item
+                findings.extend(key_findings)
+                findings.extend(item_findings)
+            return safe_dict, findings
+        return value, []
+
+    @staticmethod
+    def _policy(
+        findings: list[PrivacyFinding], mode: SecretDetectionMode
+    ) -> tuple[PrivacyDecision, PrivacyClassification]:
+        if not findings:
+            return PrivacyDecision.ALLOW, PrivacyClassification.SAFE
+        if mode == SecretDetectionMode.STRICT:
+            return PrivacyDecision.REJECT, PrivacyClassification.BLOCKED
+        if mode == SecretDetectionMode.WARN:
+            return PrivacyDecision.QUARANTINE, PrivacyClassification.BLOCKED
+        return PrivacyDecision.REDACT, PrivacyClassification.SECRET
+
+    @staticmethod
+    def _safe_findings(
+        text: str, matches: list[SecretMatch], location: str
+    ) -> list[PrivacyFinding]:
+        findings: list[PrivacyFinding] = []
+        for match in matches:
+            value = text[match.start:match.end]
+            severity = (
+                PrivacySeverity.CRITICAL
+                if match.secret_type in _CRITICAL_TYPES
+                else PrivacySeverity.HIGH
+                if match.secret_type != SecretType.HIGH_ENTROPY
+                else PrivacySeverity.MEDIUM
+            )
+            findings.append(PrivacyFinding(
+                category=match.secret_type,
+                severity=severity,
+                start=match.start,
+                end=match.end,
+                detector="pattern-secret-scanner",
+                location=location,
+                confidence=match.confidence,
+                safe_preview=f"[REDACTED:{match.secret_type.value}]",
+                fingerprint=_hash(value),
+            ))
+        return findings
+
+    @staticmethod
+    def _redact_from_findings(text: str, findings: list[PrivacyFinding]) -> str:
+        redacted = text
+        for finding in sorted(findings, key=lambda item: item.start, reverse=True):
+            redacted = (
+                redacted[:finding.start]
+                + finding.safe_preview
+                + redacted[finding.end:]
+            )
+        return redacted

contextos/services/retrieval.py

@@ -0,0 +1,302 @@
+"""Side-effect-free lexical, dense, and hybrid memory retrieval."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Sequence
+from uuid import UUID
+
+from contextos.core.enums import (
+    CandidateTemporalStatus,
+    MemoryStatus,
+    RetrievalMode,
+    TemporalScope,
+)
+from contextos.core.models import (
+    LexicalResult,
+    Memory,
+    RetrievalConfig,
+    RetrievalQuery,
+    RetrievalResult,
+    RetrievalTrace,
+    ScoredMemory,
+    StageTrace,
+    VectorResult,
+)
+from contextos.core.protocols import EmbeddingService, LexicalIndex, MemoryRepository, VectorStore
+from contextos.services.retrieval_index import RetrievalIndexSynchronizer
+
+
+class HybridRetrievalEngine:
+    """Retrieve with BM25, cosine similarity, or Reciprocal Rank Fusion.
+
+    RRF combines ranks rather than adding incomparable BM25 and cosine scores.
+    A bounded metadata factor can increase the relevance score by at most 5%.
+    """
+
+    def __init__(
+        self,
+        *,
+        memory_repo: MemoryRepository,
+        vector_store: VectorStore,
+        lexical_index: LexicalIndex,
+        embedding_service: EmbeddingService,
+        index_synchronizer: RetrievalIndexSynchronizer | None = None,
+    ) -> None:
+        self._memory_repo = memory_repo
+        self._vector_store = vector_store
+        self._lexical_index = lexical_index
+        self._embedding_service = embedding_service
+        self._index_synchronizer = index_synchronizer
+
+    async def retrieve(
+        self,
+        query: str | RetrievalQuery,
+        config: RetrievalConfig | None = None,
+    ) -> RetrievalResult:
+        request = self._coerce_query(query, config)
+        started = time.perf_counter()
+        stages: list[StageTrace] = []
+
+        if self._index_synchronizer is not None:
+            sync_started = time.perf_counter()
+            rebuilt = await self._index_synchronizer.ensure_current()
+            stages.append(self._stage(
+                "index_sync", 0, 0, sync_started, {"rebuilt": rebuilt}
+            ))
+
+        corpus_size = max(
+            await self._lexical_index.count(),
+            await self._vector_store.count(),
+        )
+        candidate_limit = max(corpus_size, request.k)
+        lexical: list[LexicalResult] = []
+        dense: list[VectorResult] = []
+
+        if request.mode in {
+            RetrievalMode.LEXICAL,
+            RetrievalMode.HYBRID,
+            RetrievalMode.GRAPH,
+            RetrievalMode.HYBRID_GRAPH,
+        }:
+            stage_started = time.perf_counter()
+            lexical = await self._lexical_index.search(request.text, candidate_limit)
+            stages.append(self._stage(
+                "lexical_search", 1, len(lexical), stage_started, {"top_k": candidate_limit}
+            ))
+
+        if request.mode in {
+            RetrievalMode.DENSE,
+            RetrievalMode.HYBRID,
+            RetrievalMode.GRAPH,
+            RetrievalMode.HYBRID_GRAPH,
+        }:
+            stage_started = time.perf_counter()
+            vector = await self._embedding_service.embed_query(request.text)
+            dense = await self._vector_store.search(vector, candidate_limit)
+            stages.append(self._stage(
+                "dense_search", 1, len(dense), stage_started, {"top_k": candidate_limit}
+            ))
+
+        resolution_started = time.perf_counter()
+        ids = {item.id for item in lexical} | {item.id for item in dense}
+        memories = await self._eligible_memories(ids, request)
+        lexical = [item for item in lexical if item.id in memories]
+        dense = [item for item in dense if item.id in memories]
+        stages.append(self._stage(
+            "eligibility_filter",
+            len(ids),
+            len(memories),
+            resolution_started,
+            {"temporal_scope": request.temporal_scope.value},
+        ))
+
+        rank_started = time.perf_counter()
+        scored = self._rank(request, memories, lexical, dense, rrf_k=(config.rrf_k if config else 60))
+        minimum = config.min_score if config else 0.0
+        pre_limit_scored = [item for item in scored if item.final_score >= minimum]
+        pre_limit_ids = [str(item.memory.id) for item in pre_limit_scored]
+        scored = pre_limit_scored[:request.k]
+        for rank, item in enumerate(scored, 1):
+            item.rank = rank
+        stages.append(self._stage(
+            "fusion_rerank",
+            len(set(item.id for item in lexical) | set(item.id for item in dense)),
+            len(scored),
+            rank_started,
+            {"mode": request.mode.value, "method": "rrf", "metadata_cap": 0.05},
+        ))
+
+        total_latency = (time.perf_counter() - started) * 1000
+        trace = RetrievalTrace(
+            stages=stages if request.include_trace else [],
+            total_latency_ms=total_latency,
+            total_candidates=len(ids),
+            total_results=len(scored),
+            lexical_candidate_ids=[str(item.id) for item in lexical[:100]],
+            dense_candidate_ids=[str(item.id) for item in dense[:100]],
+            pre_limit_candidate_ids=pre_limit_ids[:200],
+            channel_candidates_truncated=len(lexical) > 100 or len(dense) > 100,
+            pre_limit_candidates_truncated=len(pre_limit_ids) > 200,
+        )
+        strategy_results = {
+            name: [item for item in scored if name in item.retrieval_sources]
+            for name in ("lexical", "dense")
+            if any(name in item.retrieval_sources for item in scored)
+        }
+        return RetrievalResult(
+            query=request.text,
+            memories=scored,
+            strategy_results=strategy_results,
+            trace=trace,
+        )
+
+    @staticmethod
+    def _coerce_query(
+        query: str | RetrievalQuery, config: RetrievalConfig | None
+    ) -> RetrievalQuery:
+        if isinstance(query, RetrievalQuery):
+            return query
+        if config is None:
+            return RetrievalQuery(text=query)
+        statuses = {MemoryStatus.ACTIVE}
+        if config.include_superseded:
+            statuses.add(MemoryStatus.SUPERSEDED)
+        if config.include_contradicted:
+            statuses.add(MemoryStatus.CONTRADICTED)
+        if config.include_expired:
+            statuses.add(MemoryStatus.EXPIRED)
+        return RetrievalQuery(text=query, k=config.max_results, allowed_statuses=statuses)
+
+    async def _eligible_memories(
+        self, ids: set[str], request: RetrievalQuery
+    ) -> dict[str, Memory]:
+        get_many = getattr(self._memory_repo, "get_many", None)
+        if callable(get_many):
+            found = await get_many(ids)
+            return {
+                raw_id: memory for raw_id, memory in found.items()
+                if self._eligible(memory, request)
+            }
+        eligible: dict[str, Memory] = {}
+        for raw_id in sorted(ids):
+            try:
+                memory = await self._memory_repo.get(UUID(raw_id))
+            except ValueError:
+                continue
+            if memory is not None and self._eligible(memory, request):
+                eligible[raw_id] = memory
+        return eligible
+
+    @staticmethod
+    def _eligible(memory: Memory, query: RetrievalQuery) -> bool:
+        if memory.status in {MemoryStatus.DELETED, MemoryStatus.PURGED, MemoryStatus.MERGED}:
+            return False
+        if query.allowed_statuses is not None:
+            if memory.status not in query.allowed_statuses:
+                return False
+        else:
+            statuses = {
+                TemporalScope.CURRENT: {MemoryStatus.ACTIVE},
+                TemporalScope.HISTORICAL: {
+                    MemoryStatus.HISTORICAL,
+                    MemoryStatus.SUPERSEDED,
+                },
+                TemporalScope.ALL: {
+                    MemoryStatus.ACTIVE,
+                    MemoryStatus.HISTORICAL,
+                    MemoryStatus.SUPERSEDED,
+                    MemoryStatus.CONTRADICTED,
+                    MemoryStatus.EXPIRED,
+                },
+            }[query.temporal_scope]
+            if memory.status not in statuses:
+                return False
+        if query.allowed_statuses is None:
+            if (
+                query.temporal_scope == TemporalScope.CURRENT
+                and memory.temporal_status in {
+                    CandidateTemporalStatus.FUTURE,
+                    CandidateTemporalStatus.HISTORICAL,
+                }
+            ):
+                return False
+            if (
+                query.temporal_scope == TemporalScope.HISTORICAL
+                and memory.temporal_status == CandidateTemporalStatus.FUTURE
+            ):
+                return False
+        if query.allowed_memory_types and memory.type not in query.allowed_memory_types:
+            return False
+        if query.source_types and memory.source_type not in query.source_types:
+            return False
+        if query.tags and not query.tags.issubset(memory.tags):
+            return False
+        if query.created_after and memory.created_at < query.created_after:
+            return False
+        if query.created_before and memory.created_at > query.created_before:
+            return False
+        if query.min_confidence is not None and memory.confidence < query.min_confidence:
+            return False
+        return not (
+            query.min_importance is not None and memory.importance < query.min_importance
+        )
+
+    @staticmethod
+    def _rank(
+        query: RetrievalQuery,
+        memories: dict[str, Memory],
+        lexical: Sequence[LexicalResult],
+        dense: Sequence[VectorResult],
+        *,
+        rrf_k: int,
+    ) -> list[ScoredMemory]:
+        lexical_ranks = {item.id: rank for rank, item in enumerate(lexical, 1)}
+        dense_ranks = {item.id: rank for rank, item in enumerate(dense, 1)}
+        lexical_scores = {item.id: item.score for item in lexical}
+        dense_scores = {item.id: item.score for item in dense}
+        identifiers = set(lexical_ranks) | set(dense_ranks)
+        results: list[ScoredMemory] = []
+        for identifier in identifiers:
+            sources: list[str] = []
+            base_score = 0.0
+            if identifier in lexical_ranks:
+                sources.append("lexical")
+                base_score += 1.0 / (rrf_k + lexical_ranks[identifier])
+            if identifier in dense_ranks:
+                sources.append("dense")
+                base_score += 1.0 / (rrf_k + dense_ranks[identifier])
+            memory = memories[identifier]
+            adjustment = (
+                base_score * 0.05 * ((memory.confidence + memory.importance) / 2.0)
+                if query.apply_metadata_rerank
+                else 0.0
+            )
+            results.append(ScoredMemory(
+                memory=memory,
+                final_score=base_score + adjustment,
+                vector_score=dense_scores.get(identifier),
+                bm25_score=lexical_scores.get(identifier),
+                lexical_rank=lexical_ranks.get(identifier),
+                dense_rank=dense_ranks.get(identifier),
+                metadata_adjustment=adjustment,
+                retrieval_sources=sources,
+            ))
+        results.sort(key=lambda item: (-item.final_score, str(item.memory.id)))
+        return results
+
+    @staticmethod
+    def _stage(
+        name: str,
+        input_count: int,
+        output_count: int,
+        started: float,
+        metadata: dict[str, object] | None = None,
+    ) -> StageTrace:
+        return StageTrace(
+            stage_name=name,
+            input_count=input_count,
+            output_count=output_count,
+            latency_ms=(time.perf_counter() - started) * 1000,
+            metadata=metadata or {},
+        )