Initial commit

server/brain/__init__.py

@@ -0,0 +1,13 @@
+"""
+brain/ — 2nd Brain module.
+
+Provides persistent semantic memory: capture thoughts via Telegram (or any
+Aide tool), retrieve them by meaning via MCP-connected AI clients.
+
+Architecture:
+  - PostgreSQL + pgvector for storage and vector similarity search
+  - OpenRouter text-embedding-3-small for 1536-dim embeddings
+  - OpenRouter gpt-4o-mini for metadata extraction (type, tags, people, actions)
+  - MCP server mounted on FastAPI for external AI client access
+  - brain_tool registered with Aide's tool registry for Jarvis access
+"""

server/brain/database.py

@@ -0,0 +1,240 @@
+"""
+brain/database.py — PostgreSQL + pgvector connection pool and schema.
+
+Manages the asyncpg connection pool and initialises the thoughts table +
+match_thoughts function on first startup.
+"""
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+import asyncpg
+
+logger = logging.getLogger(__name__)
+
+_pool: asyncpg.Pool | None = None
+
+# ── Schema ────────────────────────────────────────────────────────────────────
+
+_SCHEMA_SQL = """
+-- pgvector extension
+CREATE EXTENSION IF NOT EXISTS vector;
+
+-- Main thoughts table
+CREATE TABLE IF NOT EXISTS thoughts (
+    id         UUID        PRIMARY KEY DEFAULT gen_random_uuid(),
+    content    TEXT        NOT NULL,
+    embedding  vector(1536),
+    metadata   JSONB       NOT NULL DEFAULT '{}',
+    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+);
+
+-- IVFFlat index for fast approximate nearest-neighbour search.
+-- Created only if it doesn't exist (pg doesn't support IF NOT EXISTS for indexes).
+DO $$
+BEGIN
+    IF NOT EXISTS (
+        SELECT 1 FROM pg_indexes
+        WHERE tablename = 'thoughts' AND indexname = 'thoughts_embedding_idx'
+    ) THEN
+        CREATE INDEX thoughts_embedding_idx
+            ON thoughts USING ivfflat (embedding vector_cosine_ops)
+            WITH (lists = 100);
+    END IF;
+END$$;
+
+-- Semantic similarity search function
+CREATE OR REPLACE FUNCTION match_thoughts(
+    query_embedding vector(1536),
+    match_threshold FLOAT  DEFAULT 0.7,
+    match_count     INT    DEFAULT 10
+)
+RETURNS TABLE (
+    id          UUID,
+    content     TEXT,
+    metadata    JSONB,
+    similarity  FLOAT,
+    created_at  TIMESTAMPTZ
+)
+LANGUAGE sql STABLE AS $$
+    SELECT
+        id,
+        content,
+        metadata,
+        1 - (embedding <=> query_embedding) AS similarity,
+        created_at
+    FROM thoughts
+    WHERE 1 - (embedding <=> query_embedding) > match_threshold
+    ORDER BY similarity DESC
+    LIMIT match_count;
+$$;
+"""
+
+
+# ── Pool lifecycle ────────────────────────────────────────────────────────────
+
+async def init_brain_db() -> None:
+    """
+    Create the connection pool and initialise the schema.
+    Called from main.py lifespan. No-ops gracefully if BRAIN_DB_URL is unset.
+    """
+    global _pool
+    url = os.getenv("BRAIN_DB_URL")
+    if not url:
+        logger.info("BRAIN_DB_URL not set — 2nd Brain disabled")
+        return
+    try:
+        _pool = await asyncpg.create_pool(url, min_size=1, max_size=5)
+        async with _pool.acquire() as conn:
+            await conn.execute(_SCHEMA_SQL)
+            # Per-user brain namespace (3-G): add user_id column if it doesn't exist yet
+            await conn.execute(
+                "ALTER TABLE thoughts ADD COLUMN IF NOT EXISTS user_id TEXT"
+            )
+        logger.info("Brain DB initialised")
+    except Exception as e:
+        logger.error("Brain DB init failed: %s", e)
+        _pool = None
+
+
+async def close_brain_db() -> None:
+    global _pool
+    if _pool:
+        await _pool.close()
+        _pool = None
+
+
+def get_pool() -> asyncpg.Pool | None:
+    return _pool
+
+
+# ── CRUD helpers ──────────────────────────────────────────────────────────────
+
+async def insert_thought(
+    content: str,
+    embedding: list[float],
+    metadata: dict,
+    user_id: str | None = None,
+) -> str:
+    """Insert a thought and return its UUID."""
+    pool = get_pool()
+    if pool is None:
+        raise RuntimeError("Brain DB not available")
+    async with pool.acquire() as conn:
+        row = await conn.fetchrow(
+            """
+            INSERT INTO thoughts (content, embedding, metadata, user_id)
+            VALUES ($1, $2::vector, $3::jsonb, $4)
+            RETURNING id::text
+            """,
+            content,
+            str(embedding),
+            __import__("json").dumps(metadata),
+            user_id,
+        )
+        return row["id"]
+
+
+async def search_thoughts(
+    query_embedding: list[float],
+    threshold: float = 0.7,
+    limit: int = 10,
+    user_id: str | None = None,
+) -> list[dict]:
+    """Return thoughts ranked by semantic similarity, scoped to user_id if set."""
+    pool = get_pool()
+    if pool is None:
+        raise RuntimeError("Brain DB not available")
+    import json as _json
+    async with pool.acquire() as conn:
+        rows = await conn.fetch(
+            """
+            SELECT mt.id, mt.content, mt.metadata, mt.similarity, mt.created_at
+            FROM match_thoughts($1::vector, $2, $3) mt
+            JOIN thoughts t ON t.id = mt.id
+            WHERE ($4::text IS NULL OR t.user_id = $4::text)
+            """,
+            str(query_embedding),
+            threshold,
+            limit,
+            user_id,
+        )
+    return [
+        {
+            "id": str(r["id"]),
+            "content": r["content"],
+            "metadata": _json.loads(r["metadata"]) if isinstance(r["metadata"], str) else dict(r["metadata"]),
+            "similarity": round(float(r["similarity"]), 4),
+            "created_at": r["created_at"].isoformat(),
+        }
+        for r in rows
+    ]
+
+
+async def browse_thoughts(
+    limit: int = 20,
+    type_filter: str | None = None,
+    user_id: str | None = None,
+) -> list[dict]:
+    """Return recent thoughts, optionally filtered by metadata type and user."""
+    pool = get_pool()
+    if pool is None:
+        raise RuntimeError("Brain DB not available")
+    async with pool.acquire() as conn:
+        rows = await conn.fetch(
+            """
+            SELECT id::text, content, metadata, created_at
+            FROM thoughts
+            WHERE ($1::text IS NULL OR user_id = $1::text)
+              AND ($2::text IS NULL OR metadata->>'type' = $2::text)
+            ORDER BY created_at DESC
+            LIMIT $3
+            """,
+            user_id,
+            type_filter,
+            limit,
+        )
+    import json as _json
+    return [
+        {
+            "id": str(r["id"]),
+            "content": r["content"],
+            "metadata": _json.loads(r["metadata"]) if isinstance(r["metadata"], str) else dict(r["metadata"]),
+            "created_at": r["created_at"].isoformat(),
+        }
+        for r in rows
+    ]
+
+
+async def get_stats(user_id: str | None = None) -> dict:
+    """Return aggregate stats about the thoughts database, scoped to user_id if set."""
+    pool = get_pool()
+    if pool is None:
+        raise RuntimeError("Brain DB not available")
+    async with pool.acquire() as conn:
+        total = await conn.fetchval(
+            "SELECT COUNT(*) FROM thoughts WHERE ($1::text IS NULL OR user_id = $1::text)",
+            user_id,
+        )
+        by_type = await conn.fetch(
+            """
+            SELECT metadata->>'type' AS type, COUNT(*) AS count
+            FROM thoughts
+            WHERE ($1::text IS NULL OR user_id = $1::text)
+            GROUP BY metadata->>'type'
+            ORDER BY count DESC
+            """,
+            user_id,
+        )
+        recent = await conn.fetchval(
+            "SELECT created_at FROM thoughts WHERE ($1::text IS NULL OR user_id = $1::text) ORDER BY created_at DESC LIMIT 1",
+            user_id,
+        )
+    return {
+        "total": total,
+        "by_type": [{"type": r["type"] or "unknown", "count": r["count"]} for r in by_type],
+        "most_recent": recent.isoformat() if recent else None,
+    }

server/brain/embeddings.py

@@ -0,0 +1,44 @@
+"""
+brain/embeddings.py — OpenRouter embedding generation.
+
+Uses text-embedding-3-small (1536 dims) via the OpenAI-compatible OpenRouter API.
+Falls back gracefully if OpenRouter is not configured.
+"""
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+_MODEL = "text-embedding-3-small"
+
+
+async def get_embedding(text: str) -> list[float]:
+    """
+    Generate a 1536-dimensional embedding for text using OpenRouter.
+    Returns a list of floats suitable for pgvector storage.
+    """
+    from openai import AsyncOpenAI
+    from ..database import credential_store
+
+    api_key = await credential_store.get("system:openrouter_api_key")
+    if not api_key:
+        raise RuntimeError(
+            "OpenRouter API key is not configured — required for brain embeddings. "
+            "Set it via Settings → Credentials → OpenRouter API Key."
+        )
+
+    client = AsyncOpenAI(
+        api_key=api_key,
+        base_url="https://openrouter.ai/api/v1",
+        default_headers={
+            "HTTP-Referer": "https://mac.oai.pm",
+            "X-Title": "oAI-Web",
+        },
+    )
+
+    response = await client.embeddings.create(
+        model=_MODEL,
+        input=text.replace("\n", " "),
+    )
+    return response.data[0].embedding

server/brain/ingest.py

@@ -0,0 +1,55 @@
+"""
+brain/ingest.py — Thought ingestion pipeline.
+
+Runs embedding generation and metadata extraction in parallel, then stores
+both in PostgreSQL. Returns the stored thought ID and a human-readable
+confirmation string suitable for sending back via Telegram.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+async def ingest_thought(content: str, user_id: str | None = None) -> dict:
+    """
+    Full ingestion pipeline for one thought:
+      1. Generate embedding + extract metadata (parallel)
+      2. Store in PostgreSQL
+      3. Return {id, metadata, confirmation}
+
+    Raises RuntimeError if Brain DB is not available.
+    """
+    from .embeddings import get_embedding
+    from .metadata import extract_metadata
+    from .database import insert_thought
+
+    # Run embedding and metadata extraction in parallel
+    embedding, metadata = await asyncio.gather(
+        get_embedding(content),
+        extract_metadata(content),
+    )
+
+    thought_id = await insert_thought(content, embedding, metadata, user_id=user_id)
+
+    # Build a human-readable confirmation (like the Slack bot reply in the guide)
+    thought_type = metadata.get("type", "other")
+    tags = metadata.get("tags", [])
+    people = metadata.get("people", [])
+    actions = metadata.get("action_items", [])
+
+    lines = [f"✅ Captured as {thought_type}"]
+    if tags:
+        lines[0] += f" — {', '.join(tags)}"
+    if people:
+        lines.append(f"People: {', '.join(people)}")
+    if actions:
+        lines.append("Actions: " + "; ".join(actions))
+
+    return {
+        "id": thought_id,
+        "metadata": metadata,
+        "confirmation": "\n".join(lines),
+    }

server/brain/metadata.py

@@ -0,0 +1,80 @@
+"""
+brain/metadata.py — LLM-based metadata extraction.
+
+Extracts structured metadata from a thought using a fast model (gpt-4o-mini
+via OpenRouter). Returns type classification, tags, people, and action items.
+"""
+from __future__ import annotations
+
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+
+_MODEL = "openai/gpt-4o-mini"
+
+_SYSTEM_PROMPT = """\
+You are a metadata extractor for a personal knowledge base. Given a thought,
+extract structured metadata and return ONLY valid JSON — no explanation, no markdown.
+
+JSON schema:
+{
+  "type": "<one of: insight | person_note | task | reference | idea | other>",
+  "tags": ["<2-5 lowercase topic tags>"],
+  "people": ["<names of people mentioned, if any>"],
+  "action_items": ["<concrete next actions, if any>"]
+}
+
+Rules:
+- type: insight = general knowledge/observation, person_note = about a specific person,
+  task = something to do, reference = link/resource/tool, idea = creative/speculative
+- tags: short lowercase words, no spaces (use underscores if needed)
+- people: first name or full name as written
+- action_items: concrete, actionable phrases only — omit if none
+- Keep all lists concise (max 5 items each)
+"""
+
+
+async def extract_metadata(text: str) -> dict:
+    """
+    Extract type, tags, people, and action_items from a thought.
+    Returns a dict. Falls back to minimal metadata on any error.
+    """
+    from openai import AsyncOpenAI
+    from ..database import credential_store
+
+    api_key = await credential_store.get("system:openrouter_api_key")
+    if not api_key:
+        return {"type": "other", "tags": [], "people": [], "action_items": []}
+
+    client = AsyncOpenAI(
+        api_key=api_key,
+        base_url="https://openrouter.ai/api/v1",
+        default_headers={
+            "HTTP-Referer": "https://mac.oai.pm",
+            "X-Title": "oAI-Web",
+        },
+    )
+
+    try:
+        response = await client.chat.completions.create(
+            model=_MODEL,
+            messages=[
+                {"role": "system", "content": _SYSTEM_PROMPT},
+                {"role": "user", "content": text},
+            ],
+            temperature=0,
+            max_tokens=256,
+            response_format={"type": "json_object"},
+        )
+        raw = response.choices[0].message.content or "{}"
+        data = json.loads(raw)
+        return {
+            "type": str(data.get("type", "other")),
+            "tags": [str(t) for t in data.get("tags", [])],
+            "people": [str(p) for p in data.get("people", [])],
+            "action_items": [str(a) for a in data.get("action_items", [])],
+        }
+    except Exception as e:
+        logger.warning("Metadata extraction failed: %s", e)
+        return {"type": "other", "tags": [], "people": [], "action_items": []}

server/brain/search.py

@@ -0,0 +1,28 @@
+"""
+brain/search.py — Semantic search over the thought database.
+
+Generates an embedding for the query text, then runs pgvector similarity
+search. All logic is thin wrappers over database.py primitives.
+"""
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+async def semantic_search(
+    query: str,
+    threshold: float = 0.7,
+    limit: int = 10,
+    user_id: str | None = None,
+) -> list[dict]:
+    """
+    Embed the query and return matching thoughts ranked by similarity.
+    Returns an empty list if Brain DB is unavailable.
+    """
+    from .embeddings import get_embedding
+    from .database import search_thoughts
+
+    embedding = await get_embedding(query)
+    return await search_thoughts(embedding, threshold=threshold, limit=limit, user_id=user_id)