RAG: embed + rerank + chat

import os, httpx, numpy as np
from openai import OpenAI

EK = os.environ["EPITHRE_KEY"]
client = OpenAI(api_key=EK, base_url="https://api.epithre.com/v1")

# ============================================================
# 1) ONE-TIME: embed your corpus
# ============================================================
corpus = [
    "UU 41/1999 pasal 50 - barangsiapa merusak hutan akan dipidana penjara paling lama 10 tahun.",
    "Permen LHK No. 92/2018 tentang pengelolaan satwa liar dilindungi.",
    "PP No. 23/2021 tentang penyelenggaraan kehutanan dan rehabilitasi lahan.",
    # ... thousands of docs in production
]

e = client.embeddings.create(
    model="epithre-embed",
    input=corpus,
    extra_body={"instruction": "Represent this document for retrieval:"},
)
corpus_vecs = np.array([row.embedding for row in e.data])   # (N, 4000)
# Store corpus_vecs + corpus in your vector DB (pgvector, Qdrant, etc.)

# ============================================================
# 2) AT QUERY TIME: embed the question, find top-K
# ============================================================
question = "Apa hukuman untuk perusakan hutan lindung?"

qe = client.embeddings.create(
    model="epithre-embed",
    input=[question],
    extra_body={"instruction": "Represent this query for retrieving relevant documents:"},
)
qv = np.array(qe.data[0].embedding)

# Cosine sim = dot product (vectors are L2-normalized)
scores = corpus_vecs @ qv
top_k_idx = np.argsort(-scores)[:10]
candidates = [corpus[i] for i in top_k_idx]

# ============================================================
# 3) Rerank to narrow to top 3
# ============================================================
r = httpx.post("https://api.epithre.com/v1/rerank",
    headers={"Authorization": f"Bearer {EK}"},
    json={
        "model": "epithre-rerank",
        "query": question,
        "documents": candidates,
        "top_n": 3,
        "return_documents": True,
    },
).json()
context_blocks = [item["document"]["text"] for item in r["results"]]
context = "\n\n".join(f"[{i+1}] {t}" for i, t in enumerate(context_blocks))

# ============================================================
# 4) Generate answer
# ============================================================
resp = client.chat.completions.create(
    model="epithre-omni",
    messages=[
        {"role": "system", "content": [
            {"type": "text",
             "text": ("Kamu asisten hukum Indonesia. Jawab berdasarkan konteks "
                      "yang diberikan. Selalu sebutkan nomor pasal/peraturan yang "
                      "relevan. Kalau konteks tidak cukup, bilang terus terang."),
             "cache_control": {"type": "ephemeral"}}
        ]},
        {"role": "user", "content": f"Konteks:\n{context}\n\nPertanyaan: {question}"},
    ],
)
print(resp.choices[0].message.content)

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    text TEXT NOT NULL,
    embedding halfvec(4000) NOT NULL,
    created_at TIMESTAMPTZ DEFAULT now()
);

CREATE INDEX ON documents USING hnsw (embedding halfvec_cosine_ops);

import psycopg2

for text, vec in zip(corpus, corpus_vecs):
    vec_str = "[" + ",".join(f"{v:.7f}" for v in vec) + "]"
    cur.execute(
        "INSERT INTO documents (text, embedding) VALUES (%s, %s::halfvec(4000))",
        (text, vec_str),
    )

qvec_str = "[" + ",".join(f"{v:.7f}" for v in qv) + "]"
cur.execute("""
    SELECT id, text, 1 - (embedding <=> %s::halfvec(4000)) AS score
    FROM documents
    ORDER BY embedding <=> %s::halfvec(4000)
    LIMIT 10
""", (qvec_str, qvec_str))
candidates = [row[1] for row in cur.fetchall()]

# Upload once
client.files.create(file=open("regulasi.pdf", "rb"), purpose="knowledge")

# Query later (handles embed + cosine search server-side)
hits = httpx.post(".../v1/retrieval", json={
    "query": question, "top_k": 10,
}, headers={"Authorization": f"Bearer {EK}"}).json()["results"]

# Then rerank (top 10 -> top 3) and chat as before