@cohortapp/agent-sdk 2.3.0

package/scripts/user-context-search.py

@@ -0,0 +1,659 @@
+#!/usr/bin/env python3
+"""
+user-context-search.py — Per-user ring-fenced context search (Phase 1 + Phase 2)
+
+Searches the agent repo's files scoped to a user's access level.
+Supports two backends:
+  - grep (default): file-based grep with user-scoped directory paths (Phase 1)
+  - sqlite: FTS5 full-text search over pre-indexed SQLite database (Phase 2)
+
+No external dependencies — stdlib only (PyYAML used if available, else manual parsing).
+
+Usage:
+    python3 scripts/user-context-search.py --user alex --query "the regulator"
+    python3 scripts/user-context-search.py --user jordan --query "portal" --max-results 5
+    python3 scripts/user-context-search.py --user unknown --query "northwind" --format text
+    python3 scripts/user-context-search.py --user alex --query "flight" --backend sqlite
+
+Author: Northwind
+Date: 2026-04-03
+"""
+
+import argparse
+import json
+import os
+import re
+import sqlite3
+import subprocess
+import sys
+from pathlib import Path
+
+# ── Constants ──────────────────────────────────────────────────────────────────
+
+BASE_DIR = Path(os.environ.get("AGENT_DIR", str(Path(__file__).resolve().parent.parent)))  # agent repo root
+CONFIG_PATH = BASE_DIR / "config" / "caller-id-map.yaml"
+DEFAULT_DB_PATH = BASE_DIR / "state" / "rag" / "search.db"
+DEFAULT_MAX_RESULTS = 10
+CONTEXT_LINES = 2  # lines before/after match to include
+
+
+# ── YAML Parsing (no PyYAML dependency) ────────────────────────────────────────
+
+def _try_import_yaml():
+    """Try to import PyYAML; return None if unavailable."""
+    try:
+        import yaml
+        return yaml
+    except ImportError:
+        return None
+
+
+def _parse_yaml_manually(filepath):
+    """
+    Minimal YAML parser for the caller-id-map.yaml structure.
+    Handles the specific nested structure we expect — not a general-purpose parser.
+    Returns a dict matching the YAML structure.
+    """
+    users = {}
+    current_user = None
+    current_key = None
+
+    with open(filepath, "r") as f:
+        for raw_line in f:
+            line = raw_line.rstrip("\n")
+
+            # Skip empty lines and comments
+            stripped = line.strip()
+            if not stripped or stripped.startswith("#"):
+                continue
+
+            # Top-level "users:" key
+            if line == "users:":
+                continue
+
+            # User slug line (2-space indent, ends with colon)
+            match_user = re.match(r"^  (\w[\w-]*):\s*$", line)
+            if match_user:
+                current_user = match_user.group(1)
+                users[current_user] = {}
+                current_key = None
+                continue
+
+            # Key-value pair (4-space indent)
+            match_kv = re.match(r'^    (\w[\w_]*):\s*(.*)$', line)
+            if match_kv and current_user:
+                key = match_kv.group(1)
+                value = match_kv.group(2).strip()
+
+                # Handle list on same line: ["+971..."]
+                if value.startswith("[") and value.endswith("]"):
+                    inner = value[1:-1].strip()
+                    if inner:
+                        items = [s.strip().strip('"').strip("'") for s in inner.split(",")]
+                        users[current_user][key] = items
+                    else:
+                        users[current_user][key] = []
+                    current_key = key
+                # Handle quoted string
+                elif value.startswith('"') and value.endswith('"'):
+                    users[current_user][key] = value.strip('"')
+                elif value.startswith("'") and value.endswith("'"):
+                    users[current_user][key] = value.strip("'")
+                # Handle empty list start
+                elif value == "[]":
+                    users[current_user][key] = []
+                    current_key = key
+                else:
+                    users[current_user][key] = value
+                current_key = key
+                continue
+
+            # List item (6-space indent, starts with -)
+            match_list = re.match(r'^      - (.*)$', line)
+            if match_list and current_user and current_key:
+                val = match_list.group(1).strip().strip('"').strip("'")
+                if not isinstance(users[current_user].get(current_key), list):
+                    users[current_user][current_key] = []
+                users[current_user][current_key].append(val)
+                continue
+
+    return {"users": users}
+
+
+def load_config(filepath):
+    """Load the caller-id-map.yaml config file."""
+    if not filepath.exists():
+        print(f"Error: Config file not found: {filepath}", file=sys.stderr)
+        sys.exit(1)
+
+    yaml_mod = _try_import_yaml()
+    if yaml_mod:
+        with open(filepath, "r") as f:
+            return yaml_mod.safe_load(f)
+    else:
+        return _parse_yaml_manually(filepath)
+
+
+# ── User Resolution ────────────────────────────────────────────────────────────
+
+def resolve_user(config, user_id):
+    """
+    Resolve a user_id to their config entry.
+    user_id can be: slug (alex), email (alex@example.com),
+    Slack ID (U099N1JGKRQ), or phone (+971585291799).
+    Returns (slug, user_dict) or (None, None) if not found.
+    """
+    users = config.get("users", {})
+
+    # Direct slug match
+    if user_id in users:
+        return user_id, users[user_id]
+
+    # Search by email, slack_id, or phone
+    for slug, data in users.items():
+        if data.get("email", "") == user_id:
+            return slug, data
+        if data.get("slack_id", "") == user_id:
+            return slug, data
+        phones = data.get("phone", [])
+        if isinstance(phones, list) and user_id in phones:
+            return slug, data
+        elif isinstance(phones, str) and phones == user_id:
+            return slug, data
+
+    return None, None
+
+
+# ── Access Scoping ─────────────────────────────────────────────────────────────
+
+# Mapping from user slug to their personal interaction directories
+USER_INTERACTION_DIRS = {
+    "alex": {
+        "slack_dm": ["memory/interactions/slack/dm-alex-chen/", "memory/interactions/slack/alex-chen/"],
+        "email": ["memory/interactions/email/alex-chen/"],
+        "phone": ["logs/phone/"],
+    },
+    "jordan": {
+        "slack_dm": [],  # No known DM channel yet
+        "email": ["memory/interactions/email/sam-rivera/"],
+        "phone": [],
+    },
+    "nima": {
+        "slack_dm": [],
+        "email": ["memory/interactions/email/jordan-blake/"],
+        "phone": [],
+    },
+    "bronwyn": {
+        "slack_dm": [],
+        "email": [],
+        "phone": [],
+    },
+    "shayan": {
+        "slack_dm": ["memory/interactions/slack/dm-casey-morgan/"],
+        "email": [],
+        "phone": [],
+    },
+}
+
+# Shared channel directories accessible to leadership+
+SHARED_CHANNELS = [
+    "memory/interactions/slack/ceo-office/",
+    "memory/interactions/slack/channels/",
+]
+
+
+def get_search_paths(access_level, user_slug):
+    """
+    Return list of directory paths (relative to BASE_DIR) that this user can search.
+
+    Access levels:
+      ceo        — Everything
+      leadership — Company knowledge/docs/research/briefs + shared channels + own interactions
+      partner    — Public knowledge/sources + research + own interactions
+      default    — Public knowledge/sources only
+    """
+    paths = []
+
+    if access_level == "ceo":
+        # CEO sees everything
+        paths = [
+            "memory/",
+            "knowledge/",
+            "state/",
+            "outputs/",
+            "docs/",
+            "logs/",
+        ]
+    elif access_level == "leadership":
+        # Shared knowledge and docs
+        paths = [
+            "knowledge/",
+            "docs/",
+            "outputs/research/",
+            "outputs/briefs/",
+            "outputs/tasks/",
+        ]
+        # Shared channels
+        for ch in SHARED_CHANNELS:
+            paths.append(ch)
+        # Own interaction logs
+        user_dirs = USER_INTERACTION_DIRS.get(user_slug, {})
+        for dir_list in user_dirs.values():
+            paths.extend(dir_list)
+    elif access_level == "partner":
+        # Public knowledge only + own interactions
+        paths = [
+            "knowledge/sources/",
+            "outputs/research/",
+        ]
+        # Own interaction logs
+        user_dirs = USER_INTERACTION_DIRS.get(user_slug, {})
+        for dir_list in user_dirs.values():
+            paths.extend(dir_list)
+    else:
+        # Default / unknown — minimal access
+        paths = [
+            "knowledge/sources/",
+        ]
+
+    return paths
+
+
+# ── Search Execution ───────────────────────────────────────────────────────────
+
+def run_grep_search(query, search_paths, max_results, base_dir):
+    """
+    Run grep across allowed paths and return matches.
+    Uses grep -rli for file discovery, then extracts context from matching files.
+    """
+    matches = []
+    files_found = set()
+
+    for rel_path in search_paths:
+        abs_path = base_dir / rel_path
+        if not abs_path.exists():
+            continue
+
+        try:
+            # First pass: find matching files (case-insensitive)
+            result = subprocess.run(
+                ["grep", "-rli", "--include=*.yaml", "--include=*.yml",
+                 "--include=*.md", "--include=*.jsonl", "--include=*.json",
+                 "--include=*.txt", "--include=*.html",
+                 query, str(abs_path)],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+            if result.returncode == 0:
+                for filepath in result.stdout.strip().split("\n"):
+                    filepath = filepath.strip()
+                    if filepath and filepath not in files_found:
+                        files_found.add(filepath)
+        except subprocess.TimeoutExpired:
+            continue
+        except Exception as e:
+            print(f"Warning: grep failed on {abs_path}: {e}", file=sys.stderr)
+            continue
+
+    # Second pass: extract context from matching files
+    for filepath in sorted(files_found):
+        if len(matches) >= max_results:
+            break
+
+        try:
+            result = subprocess.run(
+                ["grep", "-in", "-C", str(CONTEXT_LINES), query, filepath],
+                capture_output=True,
+                text=True,
+                timeout=15,
+            )
+            if result.returncode == 0:
+                # Parse grep output into individual matches
+                context_blocks = result.stdout.strip().split("--\n")
+                for block in context_blocks:
+                    if len(matches) >= max_results:
+                        break
+
+                    lines = block.strip().split("\n")
+                    if not lines or not lines[0]:
+                        continue
+
+                    # Extract line number from first matching line
+                    line_num = None
+                    match_line = ""
+                    context_text = block.strip()
+
+                    for l in lines:
+                        # grep -n output: "123:matched text" or "123-context text"
+                        m = re.match(r'^(\d+)[:-](.*)$', l)
+                        if m:
+                            if line_num is None:
+                                line_num = int(m.group(1))
+                            # The actual match line (with :)
+                            if ":" in l[:l.index(m.group(2))]:
+                                match_line = m.group(2).strip()
+
+                    rel_filepath = os.path.relpath(filepath, base_dir)
+                    matches.append({
+                        "file": rel_filepath,
+                        "line": line_num,
+                        "match": match_line[:500],  # truncate long lines
+                        "context": context_text[:1000],  # truncate very long contexts
+                    })
+
+        except subprocess.TimeoutExpired:
+            continue
+        except Exception:
+            continue
+
+    return matches
+
+
+# ── SQLite FTS5 Search (Phase 2) ──────────────────────────────────────────────
+
+def _build_scope_filter(access_level, user_slug):
+    """
+    Build a list of allowed user_scope values for the given access level.
+
+    Access level mapping:
+      ceo        — sees everything (all scopes)
+      leadership — "public" + "leadership" + own user slug
+      partner    — "public" + own user slug
+      default    — "public" only
+    """
+    if access_level == "ceo":
+        return None  # No filter — CEO sees all
+    elif access_level == "leadership":
+        scopes = ["public", "leadership"]
+        if user_slug:
+            scopes.append(user_slug)
+        return scopes
+    elif access_level == "partner":
+        scopes = ["public"]
+        if user_slug:
+            scopes.append(user_slug)
+        return scopes
+    else:
+        return ["public"]
+
+
+def run_sqlite_search(query, access_level, user_slug, max_results, db_path=None):
+    """
+    Run FTS5 search against the SQLite index with access control filtering.
+    Returns matches in the same format as run_grep_search.
+    Falls back to None if database doesn't exist.
+    """
+    if db_path is None:
+        db_path = DEFAULT_DB_PATH
+
+    if not db_path.exists():
+        return None  # Signal caller to fall back to grep
+
+    try:
+        conn = sqlite3.connect(str(db_path))
+        conn.row_factory = sqlite3.Row
+    except Exception:
+        return None
+
+    try:
+        # Build the FTS5 query — escape special characters for FTS5 syntax
+        # FTS5 uses double quotes for phrase queries; we just do a simple term search
+        fts_query = _sanitize_fts_query(query)
+
+        # Build scope filter
+        allowed_scopes = _build_scope_filter(access_level, user_slug)
+
+        if allowed_scopes is None:
+            # CEO — no scope filter
+            sql = """
+                SELECT d.id, d.source_path, d.source_type, d.user_scope,
+                       d.content, d.title, d.timestamp, d.metadata_json,
+                       rank
+                FROM documents_fts fts
+                JOIN documents d ON d.rowid = fts.rowid
+                WHERE documents_fts MATCH ?
+                ORDER BY rank
+                LIMIT ?
+            """
+            params = (fts_query, max_results)
+        else:
+            # Build placeholders for IN clause
+            placeholders = ",".join(["?"] * len(allowed_scopes))
+            sql = f"""
+                SELECT d.id, d.source_path, d.source_type, d.user_scope,
+                       d.content, d.title, d.timestamp, d.metadata_json,
+                       rank
+                FROM documents_fts fts
+                JOIN documents d ON d.rowid = fts.rowid
+                WHERE documents_fts MATCH ?
+                  AND d.user_scope IN ({placeholders})
+                ORDER BY rank
+                LIMIT ?
+            """
+            params = (fts_query, *allowed_scopes, max_results)
+
+        rows = conn.execute(sql, params).fetchall()
+
+        matches = []
+        for row in rows:
+            content = row["content"] or ""
+            # Extract a snippet around the match (first occurrence)
+            match_line = _extract_snippet(content, query)
+            context_text = _extract_context(content, query)
+
+            matches.append({
+                "file": row["source_path"],
+                "line": None,
+                "match": match_line[:500],
+                "context": context_text[:1000],
+                "score": abs(row["rank"]) if row["rank"] else 0,
+                "source_type": row["source_type"],
+                "user_scope": row["user_scope"],
+                "timestamp": row["timestamp"] or "",
+            })
+
+        conn.close()
+        return matches
+
+    except Exception as e:
+        print(f"Warning: SQLite search failed: {e}", file=sys.stderr)
+        conn.close()
+        return None
+
+
+def _sanitize_fts_query(query):
+    """
+    Sanitize a query string for FTS5 syntax.
+    Wraps each word in quotes to handle special characters.
+    """
+    # Remove FTS5 special operators that could cause syntax errors
+    # Keep it simple: split into words and quote each one
+    words = query.strip().split()
+    if not words:
+        return '""'
+
+    # If it's a single word, just return it (FTS5 handles simple terms fine)
+    # For multi-word: join with spaces (implicit AND in FTS5)
+    sanitized = []
+    for word in words:
+        # Remove characters that break FTS5 syntax
+        clean = re.sub(r'[^\w\-@.]', '', word)
+        if clean:
+            sanitized.append(f'"{clean}"')
+
+    return " ".join(sanitized) if sanitized else '""'
+
+
+def _extract_snippet(content, query, max_len=200):
+    """Extract the line containing the first match of query in content."""
+    query_lower = query.lower()
+    content_lower = content.lower()
+    idx = content_lower.find(query_lower)
+    if idx == -1:
+        # Try individual words
+        for word in query.split():
+            idx = content_lower.find(word.lower())
+            if idx != -1:
+                break
+    if idx == -1:
+        return content[:max_len]
+
+    # Find the line containing the match
+    start = content.rfind("\n", 0, idx)
+    start = start + 1 if start != -1 else 0
+    end = content.find("\n", idx)
+    end = end if end != -1 else len(content)
+
+    line = content[start:end].strip()
+    return line[:max_len]
+
+
+def _extract_context(content, query, context_chars=400):
+    """Extract a window of context around the match."""
+    query_lower = query.lower()
+    content_lower = content.lower()
+    idx = content_lower.find(query_lower)
+    if idx == -1:
+        for word in query.split():
+            idx = content_lower.find(word.lower())
+            if idx != -1:
+                break
+    if idx == -1:
+        return content[:context_chars]
+
+    start = max(0, idx - context_chars // 2)
+    end = min(len(content), idx + context_chars // 2)
+    return content[start:end].strip()
+
+
+# ── Output Formatting ──────────────────────────────────────────────────────────
+
+def format_json(user_slug, user_data, access_level, paths_searched, matches, query):
+    """Format results as JSON."""
+    output = {
+        "query": query,
+        "user": user_slug or "unknown",
+        "user_name": user_data.get("name", "Unknown") if user_data else "Unknown",
+        "access_level": access_level,
+        "paths_searched": paths_searched,
+        "result_count": len(matches),
+        "matches": matches,
+    }
+    return json.dumps(output, indent=2)
+
+
+def format_text(user_slug, user_data, access_level, paths_searched, matches, query):
+    """Format results as human-readable text."""
+    lines = []
+    name = user_data.get("name", "Unknown") if user_data else "Unknown"
+    lines.append(f"Search: \"{query}\"")
+    lines.append(f"User: {name} ({user_slug or 'unknown'}) | Access: {access_level}")
+    lines.append(f"Paths searched: {len(paths_searched)}")
+    lines.append(f"Results: {len(matches)}")
+    lines.append("-" * 60)
+
+    if not matches:
+        lines.append("No matches found.")
+    else:
+        for i, m in enumerate(matches, 1):
+            lines.append(f"\n[{i}] {m['file']}:{m.get('line', '?')}")
+            lines.append(f"    {m['match']}")
+
+    return "\n".join(lines)
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Per-user ring-fenced context search (Phase 1 + Phase 2). "
+                    "Searches the agent repo's files scoped to a user's access level.",
+        epilog="Examples:\n"
+               "  %(prog)s --user alex --query \"the regulator\"\n"
+               "  %(prog)s --user jordan --query \"portal\" --max-results 5\n"
+               "  %(prog)s --user U099N1JGKRQ --query \"flight\" --format text\n"
+               "  %(prog)s --user unknown --query \"northwind\"\n"
+               "  %(prog)s --user alex --query \"flight\" --backend sqlite\n",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--user", required=True,
+        help="User identifier: slug (alex), email (alex@example.com), "
+             "Slack ID (U099N1JGKRQ), or phone (+971585291799)"
+    )
+    parser.add_argument(
+        "--query", required=True,
+        help="Search term (case-insensitive grep pattern)"
+    )
+    parser.add_argument(
+        "--max-results", type=int, default=DEFAULT_MAX_RESULTS,
+        help=f"Maximum number of results to return (default: {DEFAULT_MAX_RESULTS})"
+    )
+    parser.add_argument(
+        "--format", choices=["json", "text"], default="json",
+        help="Output format: json (default, machine-readable) or text (human-readable)"
+    )
+    parser.add_argument(
+        "--backend", choices=["grep", "sqlite"], default="grep",
+        help="Search backend: grep (default, Phase 1) or sqlite (Phase 2 FTS5)"
+    )
+    parser.add_argument(
+        "--db", type=str, default=None,
+        help="Path to SQLite database (only used with --backend sqlite)"
+    )
+
+    args = parser.parse_args()
+
+    # Load config
+    config = load_config(CONFIG_PATH)
+
+    # Resolve user
+    user_slug, user_data = resolve_user(config, args.user)
+
+    if user_slug:
+        access_level = user_data.get("access_level", "default")
+    else:
+        # Unknown user — default (minimal) access
+        user_slug = None
+        user_data = None
+        access_level = "default"
+
+    # Determine backend and run search
+    backend_used = args.backend
+    matches = None
+
+    if args.backend == "sqlite":
+        db_path = Path(args.db) if args.db else DEFAULT_DB_PATH
+        matches = run_sqlite_search(
+            args.query, access_level, user_slug,
+            args.max_results, db_path
+        )
+        if matches is None:
+            # Fall back to grep if SQLite DB doesn't exist
+            print("Warning: SQLite database not found, falling back to grep backend",
+                  file=sys.stderr)
+            backend_used = "grep"
+
+    if matches is None:
+        # Use grep backend (default or fallback)
+        search_paths = get_search_paths(access_level, user_slug)
+        valid_paths = [p for p in search_paths if (BASE_DIR / p).exists()]
+        matches = run_grep_search(args.query, valid_paths, args.max_results, BASE_DIR)
+    else:
+        # For sqlite backend, build valid_paths for the output
+        valid_paths = ["(sqlite index)"]
+
+    # Output
+    if args.format == "json":
+        output = format_json(user_slug, user_data, access_level, valid_paths, matches, args.query)
+        # Add backend info to JSON output
+        output_dict = json.loads(output)
+        output_dict["backend"] = backend_used
+        print(json.dumps(output_dict, indent=2))
+    else:
+        print(format_text(user_slug, user_data, access_level, valid_paths, matches, args.query))
+
+
+if __name__ == "__main__":
+    main()