Build a Multi-Tenant Fraud Investigation Assistant with YugabyteDB pg_dist_rag

				
					#!/usr/bin/env python3

import json
import os
import sys

import psycopg
from openai import OpenAI


TENANTS = {
    "bank-a": "11111111-1111-4111-8111-111111111111",
    "bank-b": "22222222-2222-4222-8222-222222222222",
}


if len(sys.argv) < 3:
    print(
        "Usage: python fraud_search.py "
        "<bank-a|bank-b> <question>"
    )
    sys.exit(1)


tenant_name = sys.argv[1].lower()
question = " ".join(sys.argv[2:])


if tenant_name not in TENANTS:
    raise ValueError(f"Unknown tenant: {tenant_name}")


tenant_id = TENANTS[tenant_name]

embedding_model = os.getenv(
    "EMBEDDING_MODEL",
    "text-embedding-3-small",
)

client = OpenAI()


response = client.embeddings.create(
    model=embedding_model,
    input=question,
)

embedding = response.data[0].embedding

vector_literal = (
    "["
    + ",".join(str(value) for value in embedding)
    + "]"
)


sql = """
SELECT
    chunk_text,
    metadata_filters,
    embeddings <=> %s::vector AS distance
FROM public.fraud_kb
WHERE tenant_id = %s::uuid
ORDER BY embeddings <=> %s::vector
LIMIT 5
"""


with psycopg.connect(
    os.environ["YUGABYTEDB_CONNECTION_STRING"]
) as conn:

    with conn.cursor() as cur:

        cur.execute(
            sql,
            (
                vector_literal,
                tenant_id,
                vector_literal,
            ),
        )

        rows = cur.fetchall()


print()
print(f"Tenant:   {tenant_name}")
print(f"Question: {question}")
print()


for number, row in enumerate(rows, start=1):

    chunk_text = row[0]
    metadata = row[1]
    distance = row[2]

    print(f"--- Result {number} ---")
    print(f"Distance: {distance}")
    print(
        "Type:     "
        f"{metadata.get('knowledge_type')}"
    )
    print(chunk_text)
    print()