Skip to main content

Recipe: Retrieval with embeddings + reranking

A two-stage retrieval pipeline: embed your documents and query for a fast vector shortlist, then rerank the shortlist for precision. Both stages use your prx- key.

Prerequisites​

1. Embed documents and the query​

from openai import OpenAI
import numpy as np

client = OpenAI(base_url="https://gateway.primexos.com/v1", api_key="prx-YOUR_KEY")

docs = [
    "Failed requests aren't billed on Primex.",
    "VM credentials arrive after provisioning completes.",
    "Primex bills per request from your credit balance.",
]
emb = client.embeddings.create(model="YOUR_EMBEDDING_MODEL", input=docs + ["How is billing handled?"])
vectors = [np.array(e.embedding) for e in emb.data]
*doc_vecs, query_vec = vectors

def cosine(a, b): return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))
shortlist = sorted(range(len(docs)), key=lambda i: cosine(doc_vecs[i], query_vec), reverse=True)[:5]

2. Rerank the shortlist​

import requests

resp = requests.post(
    "https://gateway.primexos.com/v1/rerank",
    headers={"Authorization": "Bearer prx-YOUR_KEY"},
    json={
        "model": "YOUR_RERANK_MODEL",
        "query": "How is billing handled?",
        "documents": [docs[i] for i in shortlist],
        "top_n": 1,
    },
).json()
best = resp  # inspect the returned indexes/scores; feed the top document into your prompt

3. Answer with the top document​

Pass the best-ranked document as context to a chat model. See Chat Completions.

Notes​

  • Embeddings bill per token; rerank bills per query. See Billing & Credits.
  • Cache your document embeddings; only the query needs embedding at request time.
  • See Embeddings and Reranking.