Merge branch 'Memory' into bug-bounty-mode

Made-with: Cursor

Author: Av7danger

pyproject.toml

@@ -23,6 +23,10 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+memory = [
+  "pgvector>=0.4.0",
+  "sentence-transformers>=2.2.0",
+]
 dev = [
   "pytest==8.3.5",
   "pytest-asyncio==0.25.3",

src/secnodeapi/infra/config.py

@@ -2,6 +2,21 @@
 import os
 
 
+def get_semantic_memory_enabled() -> bool:
+    return os.environ.get("SECNODE_SEMANTIC_MEMORY_ENABLED", "false").lower() in (
+        "true",
+        "1",
+        "yes",
+    )
+
+
+def get_embedding_model() -> str:
+    return os.environ.get(
+        "SECNODE_EMBEDDING_MODEL",
+        "sentence-transformers/all-MiniLM-L6-v2",
+    )
+
+
 def get_queue_backend() -> str:
     return os.environ.get("SECNODE_QUEUE_BACKEND", "memory")
 

src/secnodeapi/memory/__init__.py

@@ -1,6 +1,7 @@
 """Memory subsystem."""
 from typing import Optional
 
+from ..infra.config import get_embedding_model, get_postgres_dsn, get_semantic_memory_enabled
 from .history.store import AttackHistoryStore
 from .history.pg_store import PgAttackHistoryStore
 from .metrics.store import SkillMetricsStore
@@ -43,4 +44,26 @@ def build_memory_service(
         raise ValueError(f"Unsupported backend: {backend}. Use 'memory' or 'postgres'.")
 
 
-__all__ = ["MemoryService", "build_memory_service"]
+def build_semantic_memory_service():
+    """Build SemanticMemoryService when enabled via SECNODE_SEMANTIC_MEMORY_ENABLED.
+
+    Returns:
+        SemanticMemoryService if enabled and dependencies available, else None.
+    """
+    if not get_semantic_memory_enabled():
+        return None
+    try:
+        from .semantic import PgVectorSemanticBackend, SemanticMemoryService
+        dsn = get_postgres_dsn()
+        model = get_embedding_model()
+        backend = PgVectorSemanticBackend(dsn=dsn, embedding_model=model)
+        return SemanticMemoryService(backend=backend)
+    except ImportError:
+        return None
+
+
+__all__ = [
+    "MemoryService",
+    "build_memory_service",
+    "build_semantic_memory_service",
+]

src/secnodeapi/memory/semantic/__init__.py

@@ -0,0 +1,31 @@
+"""Semantic memory subsystem for long-term cross-session learning."""
+from .contracts import (
+    AuthFlow,
+    BugBountyProgram,
+    EndpointPattern,
+    ExploitAttempt,
+    ExploitChain,
+    GlobalExploitPattern,
+    SemanticMemoryRecord,
+    SkillPerformance,
+    StackFingerprint,
+    TargetProfile,
+)
+from .backend import PgVectorSemanticBackend, SemanticMemoryBackend
+from .service import SemanticMemoryService
+
+__all__ = [
+    "AuthFlow",
+    "BugBountyProgram",
+    "EndpointPattern",
+    "ExploitAttempt",
+    "ExploitChain",
+    "GlobalExploitPattern",
+    "PgVectorSemanticBackend",
+    "SemanticMemoryBackend",
+    "SemanticMemoryRecord",
+    "SemanticMemoryService",
+    "SkillPerformance",
+    "StackFingerprint",
+    "TargetProfile",
+]

src/secnodeapi/memory/semantic/backend.py

@@ -0,0 +1,210 @@
+"""Semantic memory backend abstraction and pgvector implementation."""
+from __future__ import annotations
+
+import json
+import uuid
+from typing import Any, Dict, List, Optional, Protocol
+
+from .contracts import SemanticMemoryRecord
+
+# Embedding dimension for all-MiniLM-L6-v2
+EMBEDDING_DIM = 384
+
+
+def _get_embedding_function(model_name: str):
+    """Lazy-load sentence-transformers to avoid startup cost when disabled."""
+    try:
+        from sentence_transformers import SentenceTransformer
+        return SentenceTransformer(model_name)
+    except ImportError:
+        return None
+
+
+class SemanticMemoryBackend(Protocol):
+    """Protocol for semantic memory storage and retrieval."""
+
+    async def add(self, record: SemanticMemoryRecord) -> str: ...
+    async def search(
+        self,
+        query: str,
+        memory_types: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 20,
+    ) -> List[SemanticMemoryRecord]: ...
+    async def delete(self, record_id: str) -> bool: ...
+
+
+class PgVectorSemanticBackend:
+    """Postgres + pgvector implementation of semantic memory."""
+
+    def __init__(
+        self,
+        dsn: str,
+        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
+        embedding_dim: int = EMBEDDING_DIM,
+    ) -> None:
+        self._dsn = dsn
+        self._embedding_model_name = embedding_model
+        self._embedding_dim = embedding_dim
+        self._model = None
+        self._pool = None
+
+    def _get_model(self):
+        if self._model is None:
+            self._model = _get_embedding_function(self._embedding_model_name)
+            if self._model is None:
+                raise ImportError(
+                    "sentence-transformers is required for semantic memory. "
+                    "Install with: uv sync --extra memory"
+                )
+        return self._model
+
+    def _embed(self, text: str) -> List[float]:
+        model = self._get_model()
+        return model.encode(text, convert_to_numpy=True).tolist()
+
+    async def _ensure_pool(self):
+        if self._pool is not None:
+            return
+        import asyncpg
+        from pgvector.asyncpg import register_vector
+
+        self._pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=5)
+
+        async with self._pool.acquire() as conn:
+            await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+            await register_vector(conn)
+
+            await conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS semantic_memories (
+                    id TEXT PRIMARY KEY,
+                    memory_type TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    embedding vector(384),
+                    metadata JSONB DEFAULT '{}',
+                    confidence REAL DEFAULT 0.5,
+                    verification_status TEXT DEFAULT 'unverified',
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    expires_at TIMESTAMP
+                )
+                """
+            )
+            await conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS idx_semantic_memories_type
+                ON semantic_memories(memory_type)
+                """
+            )
+            await conn.execute(
+                """
+                CREATE INDEX IF NOT EXISTS idx_semantic_memories_created
+                ON semantic_memories(created_at DESC)
+                """
+            )
+
+    async def add(self, record: SemanticMemoryRecord) -> str:
+        await self._ensure_pool()
+        record_id = record.id or str(uuid.uuid4())
+        embedding = self._embed(record.content)
+
+        async with self._pool.acquire() as conn:
+            from pgvector.asyncpg import register_vector
+            await register_vector(conn)
+
+            await conn.execute(
+                """
+                INSERT INTO semantic_memories
+                (id, memory_type, content, embedding, metadata, confidence,
+                 verification_status, created_at, expires_at)
+                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+                """,
+                record_id,
+                record.memory_type,
+                record.content,
+                embedding,
+                json.dumps(record.metadata),
+                record.confidence,
+                record.verification_status,
+                record.created_at,
+                record.expires_at,
+            )
+        return record_id
+
+    async def search(
+        self,
+        query: str,
+        memory_types: Optional[List[str]] = None,
+        filters: Optional[Dict[str, Any]] = None,
+        limit: int = 20,
+    ) -> List[SemanticMemoryRecord]:
+        await self._ensure_pool()
+        query_embedding = self._embed(query)
+
+        conditions = ["(expires_at IS NULL OR expires_at > NOW())"]
+        params: List[Any] = [query_embedding]
+        idx = 2
+
+        if memory_types:
+            placeholders = ", ".join(f"${idx + i}" for i in range(len(memory_types)))
+            conditions.append(f"memory_type IN ({placeholders})")
+            params.extend(memory_types)
+            idx += len(memory_types)
+
+        if filters:
+            for key, value in filters.items():
+                safe_key = key.replace("'", "''")
+                if isinstance(value, str):
+                    conditions.append(f"metadata->>'{safe_key}' = ${idx}")
+                elif isinstance(value, (int, float)):
+                    conditions.append(f"(metadata->>'{safe_key}')::numeric = ${idx}")
+                params.append(value)
+                idx += 1
+
+        params.append(limit)
+        where_clause = " AND ".join(conditions)
+        limit_param = f"${idx}"
+
+        async with self._pool.acquire() as conn:
+            from pgvector.asyncpg import register_vector
+            await register_vector(conn)
+
+            rows = await conn.fetch(
+                f"""
+                SELECT id, memory_type, content, metadata, confidence,
+                       verification_status, created_at, expires_at
+                FROM semantic_memories
+                WHERE {where_clause}
+                ORDER BY embedding <=> $1
+                LIMIT {limit_param}
+                """,
+                *params,
+            )
+
+        return [
+            SemanticMemoryRecord(
+                id=r["id"],
+                memory_type=r["memory_type"],
+                content=r["content"],
+                metadata=r["metadata"] or {},
+                confidence=r["confidence"],
+                verification_status=r["verification_status"],
+                created_at=r["created_at"],
+                expires_at=r["expires_at"],
+            )
+            for r in rows
+        ]
+
+    async def delete(self, record_id: str) -> bool:
+        await self._ensure_pool()
+        async with self._pool.acquire() as conn:
+            result = await conn.execute(
+                "DELETE FROM semantic_memories WHERE id = $1",
+                record_id,
+            )
+        return result == "DELETE 1"
+
+    async def close(self) -> None:
+        if self._pool:
+            await self._pool.close()
+            self._pool = None