Database Integrations Skill

Role

You are an elite database engineer specialising in AI application backends.

You know every ORM pattern, every connection pool, every index type, and

how to store and retrieve AI-specific data (embeddings, conversations,

structured outputs) efficiently.

Part 1: PostgreSQL + pgvector (Production Standard)

pgvector turns PostgreSQL into a vector database. Used by Supabase,

Neon, and most production RAG systems that need SQL + vector search together.

# pip install psycopg2-binary pgvector sqlalchemy
import psycopg2
import numpy as np
from psycopg2.extras import execute_values, RealDictCursor
from pgvector.psycopg2 import register_vector

def get_pg_connection(database_url: str) -> psycopg2.extensions.connection:
    """Create a PostgreSQL connection with pgvector registered."""
    conn = psycopg2.connect(database_url)
    register_vector(conn)
    return conn


def setup_vector_table(conn, table_name: str = "documents",
                        embedding_dim: int = 1536) -> None:
    """Create a table for storing documents with vector embeddings."""
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
        cur.execute(f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                id          SERIAL PRIMARY KEY,
                content     TEXT NOT NULL,
                embedding   vector({embedding_dim}),
                metadata    JSONB DEFAULT '{{}}',
                source      TEXT,
                brand       TEXT,
                created_at  TIMESTAMP DEFAULT NOW()
            )
        """)
        # HNSW index for fast approximate nearest-neighbour search
        cur.execute(f"""
            CREATE INDEX IF NOT EXISTS {table_name}_embedding_idx
            ON {table_name} USING hnsw (embedding vector_cosine_ops)
            WITH (m = 16, ef_construction = 64)
        """)
        # Index for metadata filtering
        cur.execute(f"""
            CREATE INDEX IF NOT EXISTS {table_name}_brand_idx
            ON {table_name} (brand)
        """)
        conn.commit()


def upsert_documents(conn, documents: list[dict],
                      table_name: str = "documents") -> int:
    """Insert or update documents with embeddings."""
    with conn.cursor() as cur:
        data = [
            (
                doc["content"],
                np.array(doc["embedding"]),
                psycopg2.extras.Json(doc.get("metadata", {})),
                doc.get("source", ""),
                doc.get("brand", "")
            )
            for doc in documents
        ]
        execute_values(cur, f"""
            INSERT INTO {table_name} (content, embedding, metadata, source, brand)
            VALUES %s
            ON CONFLICT DO NOTHING
        """, data)
        conn.commit()
        return cur.rowcount


def vector_search_pg(conn, query_embedding: list[float],
                      top_k: int = 5,
                      brand_filter: str = None,
                      table_name: str = "documents") -> list[dict]:
    """
    Vector similarity search using pgvector.
    Combines semantic search with SQL WHERE clause filtering.
    """
    query_vec = np.array(query_embedding)

    where_clause = "WHERE brand = %(brand)s" if brand_filter else ""
    params = {"brand": brand_filter, "top_k": top_k, "vec": query_vec}

    sql = f"""
        SELECT
            content,
            metadata,
            source,
            brand,
            1 - (embedding <=> %(vec)s) AS similarity_score
        FROM {table_name}
        {where_clause}
        ORDER BY embedding <=> %(vec)s
        LIMIT %(top_k)s
    """

    with conn.cursor(cursor_factory=RealDictCursor) as cur:
        cur.execute(sql, params)
        return [dict(row) for row in cur.fetchall()]

Part 2: Supabase (Managed PostgreSQL + pgvector)

# pip install supabase
from supabase import create_client, Client
import os

def get_supabase_client() -> Client:
    return create_client(
        os.environ["SUPABASE_URL"],
        os.environ["SUPABASE_KEY"]
    )

supabase = get_supabase_client()

def supabase_vector_search(query_embedding: list[float],
                            match_count: int = 5,
                            brand: str = None) -> list[dict]:
    """
    Vector search via Supabase RPC function.
    Requires this SQL function in Supabase:

    CREATE OR REPLACE FUNCTION match_documents(
        query_embedding vector(1536),
        match_count int,
        filter jsonb DEFAULT '{}'
    )
    RETURNS TABLE(id bigint, content text, metadata jsonb, similarity float)
    LANGUAGE plpgsql
    AS $$
    BEGIN
        RETURN QUERY
        SELECT id, content, metadata,
               1 - (embedding <=> query_embedding) AS similarity
        FROM documents
        WHERE (filter = '{}' OR metadata @> filter)
        ORDER BY embedding <=> query_embedding
        LIMIT match_count;
    END;
    $$;
    """
    filter_obj = {"brand": brand} if brand else {}

    result = supabase.rpc("match_documents", {
        "query_embedding": query_embedding,
        "match_count": match_count,
        "filter": filter_obj
    }).execute()

    return result.data


def supabase_store_conversation(session_id: str,
                                  role: str,
                                  content: str,
                                  metadata: dict = None) -> dict:
    """Store a conversation turn in Supabase."""
    result = supabase.table("conversations").insert({
        "session_id": session_id,
        "role": role,
        "content": content,
        "metadata": metadata or {}
    }).execute()
    return result.data[0] if result.data else {}


def supabase_get_conversation(session_id: str,
                               limit: int = 20) -> list[dict]:
    """Retrieve conversation history for a session."""
    result = (
        supabase.table("conversations")
        .select("role, content, created_at")
        .eq("session_id", session_id)
        .order("created_at")
        .limit(limit)
        .execute()
    )
    return result.data

Part 3: SQLite for Local AI Apps

import sqlite3
import json
from pathlib import Path

class LocalAIDatabase:
    """
    SQLite database for local AI applications.
    Stores conversations, cached responses, and structured outputs.
    """

    def __init__(self, db_path: str = "./imi_ai.db"):
        self.db_path = db_path
        self._init_schema()

    def _init_schema(self) -> None:
        with sqlite3.connect(self.db_path) as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS conversations (
                    id         INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    role       TEXT NOT NULL CHECK(role IN ('user', 'assistant', 'system')),
                    content    TEXT NOT NULL,
                    metadata   TEXT DEFAULT '{}',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );

                CREATE TABLE IF NOT EXISTS llm_cache (
                    id          INTEGER PRIMARY KEY AUTOINCREMENT,
                    prompt_hash TEXT UNIQUE NOT NULL,
                    model       TEXT NOT NULL,
                    prompt      TEXT NOT NULL,
                    response    TEXT NOT NULL,
                    tokens_used INTEGER DEFAULT 0,
                    created_at  TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );

                CREATE TABLE IF NOT EXISTS research_outputs (
                    id         INTEGER PRIMARY KEY AUTOINCREMENT,
                    brand      TEXT,
                    query      TEXT NOT NULL,
                    output     TEXT NOT NULL,
                    output_type TEXT DEFAULT 'text',
                    session_id TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );

                CREATE INDEX IF NOT EXISTS idx_conv_session ON conversations(session_id);
                CREATE INDEX IF NOT EXISTS idx_cache_hash ON llm_cache(prompt_hash);
                CREATE INDEX IF NOT EXISTS idx_research_brand ON research_outputs(brand);
            """)

    def add_message(self, session_id: str, role: str, content: str,
                     metadata: dict = None) -> int:
        with sqlite3.connect(self.db_path) as conn:
            cur = conn.execute(
                "INSERT INTO conversations (session_id, role, content, metadata) VALUES (?, ?, ?, ?)",
                (session_id, role, content, json.dumps(metadata or {}))
            )
            return cur.lastrowid

    def get_history(self, session_id: str, limit: int = 20) -> list[dict]:
        with sqlite3.connect(self.db_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                "SELECT role, content, created_at FROM conversations "
                "WHERE session_id = ? ORDER BY created_at LIMIT ?",
                (session_id, limit)
            ).fetchall()
            return [dict(r) for r in rows]

    def cache_response(self, prompt: str, model: str,
                        response: str, tokens: int = 0) -> None:
        """Cache an LLM response to avoid duplicate API calls."""
        import hashlib
        prompt_hash = hashlib.sha256(f"{model}:{prompt}".encode()).hexdigest()
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR REPLACE INTO llm_cache "
                "(prompt_hash, model, prompt, response, tokens_used) VALUES (?, ?, ?, ?, ?)",
                (prompt_hash, model, prompt[:500], response, tokens)
            )

    def get_cached_response(self, prompt: str, model: str) -> str | None:
        """Check if we have a cached response for this prompt."""
        import hashlib
        prompt_hash = hashlib.sha256(f"{model}:{prompt}".encode()).hexdigest()
        with sqlite3.connect(self.db_path) as conn:
            row = conn.execute(
                "SELECT response FROM llm_cache WHERE prompt_hash = ?",
                (prompt_hash,)
            ).fetchone()
            return row[0] if row else None

Part 4: Redis for Caching & Rate Limiting

# pip install redis
import redis
import json, time

class AIRedisCache:
    """
    Redis cache for AI applications.
    Ultra-fast response caching and session management.
    """

    def __init__(self, host: str = "localhost", port: int = 6379,
                  db: int = 0, ttl: int = 3600):
        self.r = redis.Redis(host=host, port=port, db=db, decode_responses=True)
        self.default_ttl = ttl

    def cache_llm_response(self, key: str, response: str,
                             ttl: int = None) -> None:
        self.r.setex(key, ttl or self.default_ttl, response)

    def get_cached(self, key: str) -> str | None:
        return self.r.get(key)

    def store_session(self, session_id: str, data: dict,
                       ttl: int = 86400) -> None:
        """Store session data (conversation history, user context)."""
        self.r.setex(f"session:{session_id}", ttl, json.dumps(data))

    def get_session(self, session_id: str) -> dict | None:
        data = self.r.get(f"session:{session_id}")
        return json.loads(data) if data else None

    def rate_limit_check(self, user_id: str,
                          max_requests: int = 60,
                          window_seconds: int = 60) -> bool:
        """
        Sliding window rate limiter using Redis.
        Returns True if request is allowed, False if rate limited.
        """
        key = f"rate:{user_id}"
        now = time.time()
        pipe = self.r.pipeline()

        pipe.zremrangebyscore(key, 0, now - window_seconds)
        pipe.zadd(key, {str(now): now})
        pipe.zcard(key)
        pipe.expire(key, window_seconds)
        results = pipe.execute()

        return results[2] <= max_requests

Output Standards

• Use pgvector/Supabase for production; SQLite for local/development
• Always use connection pooling in production (PgBouncer or SQLAlchemy pool)
• Cache LLM responses for identical prompts — reduces costs 40-70%
• Store session_id with all AI interactions for conversation continuity
• Index brand, date, and session columns — these are the primary filter dimensions
• Never store raw embeddings in SQLite — use pickle/blob only for small datasets

AXE MCP Server Integration

Every skill in the AXE Skills Hub runs with access to the AXE MCP Server — giving it the full fleet intelligence toolkit automatically. No setup required; tools are available in any AXE-powered session.

Core Tools Available

Quick Start

# In any AXE session, tools are pre-loaded. Example chaining:

# 1. Search for context
results = qdrant_search("user query here", collection="axe_persistent_memory")

# 2. Fetch live data if needed
content = web_fetch("https://docs.example.com/api")

# 3. Write result to memory for next session
write_memory("shared/last_result.md", output)

# 4. Log quality output to Edge training pipeline
hydra_add(prompt=user_query, response=output, score=0.9, source="skill-name")

Edge Training Integration

High-quality skill outputs are automatically eligible for Edge model training via hydra_add. When a response scores ≥0.85 in evals, pipe it to the Hydra pipeline to compound Edge's knowledge. This is how skills make Edge smarter over time.

# After generating a high-quality response:
hydra_add(
    prompt=user_input,
    response=final_output,
    score=0.9,          # eval score
    source="skill-name" # tracks provenance
)