Real-Time RAG Pipeline

-- Every edit is a new version. A deletion is a version with deleted = true.
-- The topic is append-only, which is what the vector index downstream requires.
CREATE SOURCE doc_versions (
  doc_id       BIGINT,
  version      INT,
  title        VARCHAR,
  body         VARCHAR,
  product_id   BIGINT,           -- NULL for policies and general articles
  region       VARCHAR,          -- 'us' | 'eu' | 'all'
  deleted      BOOLEAN,
  published_at TIMESTAMPTZ
) WITH (
  connector = 'kafka', topic = 'doc_versions',
  properties.bootstrap.server = 'broker:9092'
) FORMAT PLAIN ENCODE JSON;

-- Live metadata the answer must respect comes from the product database as change streams.
CREATE SOURCE shop_pg WITH (
  connector = 'postgres-cdc', hostname = 'pg', port = '5432',
  username = 'rw', password = '...', database.name = 'shop', schema.name = 'public'
);

CREATE TABLE products (
  product_id BIGINT PRIMARY KEY, name VARCHAR, price NUMERIC,
  in_stock BOOLEAN, discontinued BOOLEAN
) FROM shop_pg TABLE 'public.products';

-- The embedding model and key are defined once, in a SQL function used for both indexing and querying.
CREATE FUNCTION embed(VARCHAR) RETURNS VECTOR(1536) LANGUAGE SQL AS $$
  SELECT openai_embedding(
    '{"model": "text-embedding-3-small", "api_key": "sk-..."}'::jsonb, $1
  )::vector(1536)
$$;

-- Chunk in SQL: split on blank lines, drop fragments. unnest preserves the append-only property.
CREATE MATERIALIZED VIEW doc_chunks AS
SELECT d.doc_id, d.version, d.title, d.product_id, d.region, d.published_at,
       c.chunk
FROM doc_versions d,
     unnest(regexp_split_to_array(d.body, E'\n\s*\n')) AS c(chunk)
WHERE d.deleted = false
  AND length(c.chunk) > 40;

-- Embed in-stream. One API call per chunk per version, made once, when the version arrives.
CREATE MATERIALIZED VIEW chunk_embeddings AS
SELECT doc_id, version, title, product_id, region, published_at, chunk,
       embed(title || E'\n' || chunk) AS embedding
FROM doc_chunks;

-- The index. INCLUDE carries the columns the retrieval query returns and filters on.
CREATE INDEX chunk_embeddings_hnsw ON chunk_embeddings
USING HNSW (embedding)
INCLUDE (doc_id, version, title, product_id, region, published_at, chunk)
WITH (distance_type = 'cosine', m = 16, ef_construction = 64);

-- Which version of each document is current? Group Top-N over all versions, including deletions.
CREATE MATERIALIZED VIEW doc_current AS
SELECT doc_id, version, deleted, region, product_id, published_at
FROM (
  SELECT *, ROW_NUMBER() OVER (PARTITION BY doc_id ORDER BY version DESC) AS rn
  FROM doc_versions
) t
WHERE rn = 1;

-- Feedback the application writes after each answer: which chunks were shown, and did they help.
CREATE TABLE retrieval_feedback (
  feedback_id VARCHAR PRIMARY KEY,
  doc_id      BIGINT,
  version     INT,
  helpful     BOOLEAN,
  created_at  TIMESTAMPTZ
);

CREATE MATERIALIZED VIEW doc_feedback AS
SELECT doc_id, version,
       COUNT(*)                          AS shown,
       COUNT(*) FILTER (WHERE helpful)   AS helpful
FROM retrieval_feedback
GROUP BY doc_id, version;

-- retrieve.sql: hybrid retrieval in one statement.
--   $1 question text, $2 caller's region
-- Vector distance from the index, current-version and deletion check, live stock filter,
-- and a small rerank from feedback. Raise ef_search when filters are selective.
SET batch_hnsw_ef_search = 200;

SELECT e.doc_id, e.version, e.title, e.chunk,
       p.name AS product, p.price, p.in_stock,
       e.published_at,
       (e.embedding <=> embed($1)) AS distance
FROM chunk_embeddings e
JOIN doc_current d
  ON d.doc_id = e.doc_id AND d.version = e.version          -- current version only
LEFT JOIN products p ON p.product_id = e.product_id
LEFT JOIN doc_feedback f ON f.doc_id = e.doc_id AND f.version = e.version
WHERE d.deleted = false
  AND (e.region = 'all' OR e.region = $2)
  AND (p.product_id IS NULL OR (p.in_stock AND NOT p.discontinued))
ORDER BY (e.embedding <=> embed($1))
         - 0.05 * COALESCE(f.helpful::float / NULLIF(f.shown, 0), 0)   -- helpful chunks rank higher
LIMIT 8;

from typing import List
import re, psycopg
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

RETRIEVE_SQL = open("retrieve.sql").read()

def bind(sql: str) -> str:
    """The SQL files use Postgres $n placeholders; psycopg wants %(name)s."""
    return re.sub(r"\$(\d+)", r"%(p\1)s", sql)

class RisingWaveRetriever(BaseRetriever):
    """Hybrid retrieval over RisingWave: vector distance plus live metadata, in one query."""
    dsn: str = "postgresql://rag_ro@rw:4566/dev"
    region: str = "us"

    def _get_relevant_documents(self, query: str) -> List[Document]:
        with psycopg.connect(self.dsn, autocommit=True) as conn, conn.cursor() as cur:
            cur.execute(bind(RETRIEVE_SQL), {"p1": query, "p2": self.region})
            rows = cur.fetchall()
        return [
            Document(
                page_content=chunk,
                metadata={"doc_id": doc_id, "version": version, "title": title,
                          "product": product, "price": price, "in_stock": in_stock,
                          "published_at": str(published_at), "distance": distance},
            )
            for (doc_id, version, title, chunk, product, price, in_stock,
                 published_at, distance) in rows
        ]

# Use it like any retriever. Each Document carries the live price and stock,
# so the prompt can say "as of now" and mean it.
retriever = RisingWaveRetriever(region="eu")
docs = retriever.invoke("Can I return a sale item after 30 days?")

import json, time, uuid, psycopg
from kafka import KafkaProducer

# Freshness probe: publish a version with a nonce, time until it is retrievable.
nonce = uuid.uuid4().hex
producer = KafkaProducer(bootstrap_servers="broker:9092",
                         value_serializer=lambda v: json.dumps(v).encode())
producer.send("doc_versions", {
    "doc_id": 999_999_999, "version": int(time.time()), "title": "freshness probe",
    "body": f"Probe document {nonce}. It exists only to measure indexing latency.",
    "product_id": None, "region": "all", "deleted": False,
    "published_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
})
producer.flush()
t0 = time.time()

rw = psycopg.connect("postgresql://rag_ro@rw:4566/dev", autocommit=True)
while True:
    row = rw.execute(
        "SELECT 1 FROM chunk_embeddings e JOIN doc_current d USING (doc_id, version) "
        "WHERE e.chunk LIKE %s LIMIT 1", (f"%{nonce}%",)).fetchone()
    if row:
        break
    time.sleep(0.2)
print(f"indexed and retrievable after {time.time() - t0:.1f} s")   # typically 2 to 5 s

-- Optional: pre-retrieve for every incoming question as a stream (RisingWave 2.7 and later).
-- Each question row gets its top 5 chunk ids attached, so the agent reads one row per question.
CREATE SOURCE questions (
  question_id VARCHAR, text VARCHAR, region VARCHAR, asked_at TIMESTAMPTZ
) WITH (connector = 'kafka', topic = 'questions',
        properties.bootstrap.server = 'broker:9092') FORMAT PLAIN ENCODE JSON;

CREATE MATERIALIZED VIEW question_candidates AS
SELECT q.question_id, q.text, q.asked_at,
       array(
         SELECT row(e.doc_id, e.version, e.title)
         FROM chunk_embeddings e FOR SYSTEM_TIME AS OF PROCTIME()
         ORDER BY embed(q.text) <=> e.embedding
         LIMIT 5
       ) AS candidates
FROM questions q;