Building an Internal Document Search Agent with Claude API — Hybrid RAG, Role-Based Access Control, and Audit Logging in Production

import anthropic
import psycopg2
from psycopg2.extras import execute_values
import hashlib
from pathlib import Path
from typing import Optional
import pypdf
import tiktoken
 
client = anthropic.Anthropic()
enc = tiktoken.encoding_for_model("gpt-4o")  # Approximate token counting
 
def extract_text_from_pdf(pdf_path: str) -> str:
    """Extract text content from a PDF file."""
    text_parts = []
    with open(pdf_path, "rb") as f:
        reader = pypdf.PdfReader(f)
        for page in reader.pages:
            text_parts.append(page.extract_text())
    return "\n".join(text_parts)
 
def chunk_document(
    text: str,
    chunk_size: int = 512,
    overlap: int = 64
) -> list[dict]:
    """
    Split document into overlapping chunks.
    - chunk_size: target size in tokens
    - overlap: tokens of context shared between adjacent chunks
    
    Why 512/64? This lets us fit 8–10 chunks into a Claude context window
    while preserving enough overlap to avoid losing meaning at boundaries.
    """
    tokens = enc.encode(text)
    chunks = []
    start = 0
    
    while start < len(tokens):
        end = min(start + chunk_size, len(tokens))
        chunk_text = enc.decode(tokens[start:end])
        
        # Skip fragments shorter than 50 tokens (end of document)
        if len(enc.encode(chunk_text)) < 50:
            break
        
        chunks.append({
            "text": chunk_text,
            "token_start": start,
            "token_end": end
        })
        start += chunk_size - overlap
    
    return chunks
 
def ingest_document(
    file_path: str,
    department: str,  # "hr", "engineering", "finance", "all"
    doc_title: str,
    db_conn: psycopg2.extensions.connection
) -> int:
    """
    Ingest a document into the vector store.
    Returns the number of chunks created (0 if already ingested).
    """
    path = Path(file_path)
    
    if path.suffix == ".pdf":
        text = extract_text_from_pdf(file_path)
    elif path.suffix in [".md", ".txt"]:
        text = path.read_text(encoding="utf-8")
    else:
        raise ValueError(f"Unsupported file type: {path.suffix}")
    
    # Deduplication via content hash
    doc_hash = hashlib.sha256(text.encode()).hexdigest()
    
    with db_conn.cursor() as cur:
        cur.execute(
            "SELECT id FROM documents WHERE content_hash = %s",
            (doc_hash,)
        )
        if cur.fetchone():
            print(f"Skipping {doc_title} — already ingested")
            return 0
    
    chunks = chunk_document(text)
    
    with db_conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO documents (title, department, content_hash, file_path)
            VALUES (%s, %s, %s, %s) RETURNING id
            """,
            (doc_title, department, doc_hash, file_path)
        )
        doc_id = cur.fetchone()[0]
    
    chunk_data = [
        (doc_id, i, chunk["text"], department)
        for i, chunk in enumerate(chunks)
    ]
    
    with db_conn.cursor() as cur:
        execute_values(
            cur,
            """
            INSERT INTO document_chunks (doc_id, chunk_index, content, department)
            VALUES %s
            """,
            chunk_data
        )
    
    db_conn.commit()
    print(f"✅ Ingested: {doc_title} ({len(chunks)} chunks, dept: {department})")
    return len(chunks)

def hybrid_search(
    query: str,
    user_department: str,
    query_embedding: list[float],
    db_conn: psycopg2.extensions.connection,
    top_k: int = 10,
    vector_weight: float = 0.6,
    bm25_weight: float = 0.4
) -> list[dict]:
    """
    Hybrid search using RRF (Reciprocal Rank Fusion).
    
    Critical: This function must be called with a DB connection that has
    RLS configured for user_department. The function itself does NOT
    apply department filtering — the database handles that automatically.
    Never add application-layer department filtering on top of this.
    """
    rrf_k = 60  # Standard RRF parameter from literature
    
    with db_conn.cursor() as cur:
        # Vector search — RLS automatically filters by department
        cur.execute(
            """
            SELECT id, content, doc_id, chunk_index,
                   1 - (embedding <=> %s::vector) AS score
            FROM document_chunks
            ORDER BY embedding <=> %s::vector
            LIMIT %s
            """,
            (query_embedding, query_embedding, top_k * 2)
        )
        vector_results = {row[0]: (i, row) for i, row in enumerate(cur.fetchall())}
        
        # BM25 full-text search — also subject to RLS
        cur.execute(
            """
            SELECT id, content, doc_id, chunk_index,
                   ts_rank(to_tsvector('english', content),
                           plainto_tsquery('english', %s)) AS score
            FROM document_chunks
            WHERE to_tsvector('english', content) @@ plainto_tsquery('english', %s)
            ORDER BY score DESC
            LIMIT %s
            """,
            (query, query, top_k * 2)
        )
        bm25_results = {row[0]: (i, row) for i, row in enumerate(cur.fetchall())}
    
    # Merge with RRF
    all_ids = set(vector_results.keys()) | set(bm25_results.keys())
    scored_results = []
    
    for chunk_id in all_ids:
        rrf_score = 0.0
        row_data = None
        
        if chunk_id in vector_results:
            rank, row = vector_results[chunk_id]
            rrf_score += vector_weight / (rrf_k + rank + 1)
            row_data = row
        
        if chunk_id in bm25_results:
            rank, row = bm25_results[chunk_id]
            rrf_score += bm25_weight / (rrf_k + rank + 1)
            if row_data is None:
                row_data = row
        
        scored_results.append({
            "id": chunk_id,
            "content": row_data[1],
            "doc_id": row_data[2],
            "chunk_index": row_data[3],
            "rrf_score": rrf_score
        })
    
    scored_results.sort(key=lambda x: x["rrf_score"], reverse=True)
    return scored_results[:top_k]

-- Table definition
CREATE TABLE document_chunks (
    id          BIGSERIAL PRIMARY KEY,
    doc_id      BIGINT NOT NULL REFERENCES documents(id),
    chunk_index INT NOT NULL,
    content     TEXT NOT NULL,
    department  TEXT NOT NULL,  -- "hr", "engineering", "finance", "all"
    embedding   vector(1024)    -- voyage-3-large dimensionality
);
 
-- Enable RLS (FORCE applies it to table owners too)
ALTER TABLE document_chunks ENABLE ROW LEVEL SECURITY;
ALTER TABLE document_chunks FORCE ROW LEVEL SECURITY;
 
-- Standard user policy: see your own department + "all"
CREATE POLICY dept_access ON document_chunks
    USING (
        department = current_setting('app.user_department')
        OR department = 'all'
    );
 
-- Admin bypass (separate role)
CREATE POLICY admin_access ON document_chunks
    TO admin_role
    USING (true);
 
-- App role for regular users
CREATE ROLE app_user;
GRANT SELECT ON document_chunks TO app_user;

from contextlib import contextmanager
 
@contextmanager
def get_db_connection_for_user(user_department: str):
    """
    Return a DB connection scoped to user_department via RLS.
    
    The third argument to set_config is crucial: 'true' means the setting
    is only valid for the current transaction. 'false' would make it
    session-scoped — dangerous with connection pooling because the setting
    could persist to a different user's request.
    """
    conn = psycopg2.connect(
        dsn=os.environ["DATABASE_URL"],
        options="-c role=app_user"
    )
    try:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT set_config('app.user_department', %s, true)",
                (user_department,)
            )
        yield conn
    finally:
        conn.close()

SYSTEM_PROMPT = """You are an internal document search assistant.
 
Rules:
1. Answer only based on the document chunks provided
2. If information is not in the chunks, say explicitly: "This information is not in the provided documents"
3. Always cite which document chunks support your answer
4. Flag if your answer might involve sensitive or confidential information
"""
 
def answer_query(
    query: str,
    chunks: list[dict],
    user_info: dict
) -> dict:
    """
    Generate an answer from filtered search results using Claude.
    
    Note on chunk count: claude-sonnet-4-6 has a 200K token window,
    but we limit to 8 chunks to keep costs reasonable. Each 512-token
    chunk + context overhead costs roughly 0.015 USD per query with sonnet.
    """
    if not chunks:
        return {
            "answer": "No matching documents found. Try different keywords "
                      "or contact your admin to check access permissions.",
            "cited_chunks": [],
            "model": "none",
            "input_tokens": 0,
            "output_tokens": 0
        }
    
    # Limit chunks to control cost
    display_chunks = chunks[:8]
    
    context_parts = []
    for i, chunk in enumerate(display_chunks):
        context_parts.append(
            f"[Chunk {i+1}] (Doc ID: {chunk['doc_id']}, "
            f"Position: {chunk['chunk_index']})\n{chunk['content']}"
        )
    context = "\n\n---\n\n".join(context_parts)
    
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        system=SYSTEM_PROMPT,
        messages=[
            {
                "role": "user",
                "content": (
                    f"Document chunks:\n\n{context}\n\n"
                    f"Question: {query}"
                )
            }
        ]
    )
    
    return {
        "answer": response.content[0].text,
        "cited_chunks": [c["id"] for c in display_chunks],
        "model": response.model,
        "input_tokens": response.usage.input_tokens,
        "output_tokens": response.usage.output_tokens
    }

import uuid
from datetime import datetime, timezone
 
def log_search_event(
    db_conn: psycopg2.extensions.connection,
    user_id: str,
    user_department: str,
    query: str,
    cited_chunk_ids: list[int],
    answer_summary: str,
    model: str,
    input_tokens: int,
    output_tokens: int,
    session_id: Optional[str] = None
) -> str:
    """
    Write a search event to the audit log.
    
    Design note: store only the first 200 characters of the answer.
    Storing full responses significantly increases storage costs and
    may create compliance issues if responses contain PII.
    Returns the event UUID for downstream correlation.
    """
    event_id = str(uuid.uuid4())
    
    with db_conn.cursor() as cur:
        cur.execute(
            """
            INSERT INTO search_audit_log (
                event_id, user_id, user_department, query_text,
                cited_chunk_ids, answer_summary, model_used,
                input_tokens, output_tokens, session_id, created_at
            ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """,
            (
                event_id, user_id, user_department, query,
                cited_chunk_ids, answer_summary[:200], model,
                input_tokens, output_tokens, session_id,
                datetime.now(timezone.utc)
            )
        )
        db_conn.commit()
    
    return event_id
 
# Monthly cost report query
MONTHLY_COST_QUERY = """
SELECT
    DATE_TRUNC('day', created_at) AS day,
    user_department,
    COUNT(*) AS query_count,
    SUM(input_tokens) AS total_input_tokens,
    SUM(output_tokens) AS total_output_tokens,
    -- claude-sonnet-4-6: ~$3/M input, ~$15/M output (approx. May 2026)
    ROUND(
        (SUM(input_tokens) * 3.0 + SUM(output_tokens) * 15.0) / 1000000,
        4
    ) AS estimated_cost_usd
FROM search_audit_log
WHERE created_at >= DATE_TRUNC('month', CURRENT_DATE)
GROUP BY 1, 2
ORDER BY 1, 2;
"""

# Don't do this
chunk_size = 2048  # Reduces chunk count per query → misses information
 
# This works better in practice
chunk_size = 512   # Fits 8–10 chunks comfortably, preserves context
overlap = 64       # ~12.5% overlap prevents boundary artifacts

-- This is insufficient — table owner bypasses it
ALTER TABLE document_chunks ENABLE ROW LEVEL SECURITY;
 
-- This enforces it for everyone, including the table owner
ALTER TABLE document_chunks ENABLE ROW LEVEL SECURITY;
ALTER TABLE document_chunks FORCE ROW LEVEL SECURITY;

def update_document(doc_id: int, new_text: str, db_conn):
    """Always regenerate embeddings when updating document text."""
    new_hash = hashlib.sha256(new_text.encode()).hexdigest()
    chunks = chunk_document(new_text)
    
    with db_conn.cursor() as cur:
        # Remove all old chunks
        cur.execute("DELETE FROM document_chunks WHERE doc_id = %s", (doc_id,))
        
        # Re-insert with fresh embeddings
        for i, chunk in enumerate(chunks):
            # embedding = generate_embedding(chunk["text"])  # Required
            cur.execute(
                """
                INSERT INTO document_chunks (doc_id, chunk_index, content, department)
                VALUES (%s, %s, %s, (SELECT department FROM documents WHERE id = %s))
                """,
                (doc_id, i, chunk["text"], doc_id)
            )
        
        cur.execute(
            "UPDATE documents SET content_hash = %s, updated_at = NOW() WHERE id = %s",
            (new_hash, doc_id)
        )
    
    db_conn.commit()

def test_department_isolation():
    # Engineering user should not see HR documents
    with get_db_connection_for_user("engineering") as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT content FROM document_chunks WHERE department = 'hr'")
            rows = cur.fetchall()
            assert len(rows) == 0, "Engineering user must not see HR documents"
 
    # 'all' department documents should be visible to everyone
    with get_db_connection_for_user("finance") as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT content FROM document_chunks WHERE department = 'all'")
            rows = cur.fetchall()
            assert len(rows) > 0, "All users should see 'all' department documents"
 
def test_rls_not_bypassable_with_explicit_filter():
    # Even with explicit WHERE department = 'hr', RLS blocks access for non-HR users
    with get_db_connection_for_user("engineering") as conn:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT COUNT(*) FROM document_chunks WHERE department = 'hr'"
            )
            count = cur.fetchone()[0]
            # RLS enforces zero rows regardless of what the WHERE clause says
            assert count == 0, "RLS must block access regardless of explicit filters"

-- IVFFlat: faster queries, slight recall trade-off
-- Set 'lists' to roughly sqrt(row count)
CREATE INDEX idx_chunks_embedding_ivfflat
    ON document_chunks
    USING ivfflat (embedding vector_cosine_ops)
    WITH (lists = 100);
 
-- HNSW: better recall than IVFFlat, higher memory usage
CREATE INDEX idx_chunks_embedding_hnsw
    ON document_chunks
    USING hnsw (embedding vector_cosine_ops)
    WITH (m = 16, ef_construction = 64);

import redis
import json
import hashlib
 
redis_client = redis.Redis.from_url(os.environ["REDIS_URL"])
 
def get_or_generate_embedding(text: str) -> list[float]:
    cache_key = f"emb:{hashlib.sha256(text.encode()).hexdigest()}"
    cached = redis_client.get(cache_key)
    if cached:
        return json.loads(cached)
    
    embedding = generate_embedding(text)  # Your embedding API call
    # Embeddings are deterministic — a long TTL (1-24h) is safe
    redis_client.setex(cache_key, 3600, json.dumps(embedding))
    return embedding

import asyncio
from concurrent.futures import ThreadPoolExecutor
 
async def ingest_document_batch(
    file_paths: list[tuple[str, str, str]],
    db_conn,
    max_workers: int = 4
) -> dict[str, int]:
    """Ingest multiple documents in parallel. Returns {title: chunk_count}."""
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        loop = asyncio.get_event_loop()
        tasks = [
            loop.run_in_executor(
                executor,
                lambda p=path, d=dept, t=title: ingest_document(p, d, t, db_conn)
            )
            for path, dept, title in file_paths
        ]
        chunk_counts = await asyncio.gather(*tasks)
    return {title: count for (_, _, title), count in zip(file_paths, chunk_counts)}

import voyageai
import os
 
voyage_client = voyageai.Client(api_key=os.environ["VOYAGE_API_KEY"])
 
def generate_document_embedding(text: str) -> list[float]:
    result = voyage_client.embed(
        texts=[text],
        model="voyage-3-large",
        input_type="document"  # Use this for ingestion
    )
    return result.embeddings[0]
 
def generate_query_embedding(query: str) -> list[float]:
    result = voyage_client.embed(
        texts=[query],
        model="voyage-3-large",
        input_type="query"  # Use this for search queries — different from "document"
    )
    return result.embeddings[0]