shashank ms

Posted on Aug 10

LLM for Information Retrieval

#learnai #oxlo #ai

We are going to build a minimal retrieval-augmented generation pipeline that embeds a document corpus with Oxlo.ai embeddings, retrieves the top matches with cosine similarity, and generates a grounded answer with an Oxlo.ai LLM. This pattern is useful for anyone adding semantic search to internal docs, support wikis, or product catalogs without managing a separate vector database.

What you'll need

Python 3.10 or newer.
An Oxlo.ai API key from https://portal.oxlo.ai.
The OpenAI SDK and numpy: pip install openai numpy python-dotenv.

Step 1: Install dependencies and configure the Oxlo.ai client

I keep my API key in a .env file so I do not hardcode it. The client setup is identical to the OpenAI SDK, just pointed at Oxlo.ai.

# .env
# OXLO_API_KEY=your_key_here

import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key=os.getenv("OXLO_API_KEY")
)

Step 2: Prepare a toy document corpus

For this demo I am using a handful of sentences about Oxlo.ai. In production this would be your wiki, help center, or research papers.

CORPUS = [
    "Oxlo.ai offers flat per-request pricing for LLM inference. Unlike token-based providers, cost does not scale with input length.",
    "The platform hosts 45+ models including Llama 3.3 70B, Qwen 3 32B, DeepSeek R1, and Kimi K2.6.",
    "Oxlo.ai is fully OpenAI SDK compatible. You can switch your base URL and API key to start using it.",
    "Embeddings available on Oxlo.ai include BGE-Large and E5-Large for vector search applications.",
    "Oxlo.ai provides vision models such as Gemma 3 27B and Kimi VL A3B, plus image generation through Flux.1 and Stable Diffusion 3.5.",
]

Step 3: Embed the corpus with Oxlo.ai embeddings

I am using bge-large from Oxlo.ai's embeddings category. I call the embeddings endpoint through the same OpenAI-compatible client.

import numpy as np

def get_embedding(text):
    text = text.replace("\n", " ")
    resp = client.embeddings.create(
        model="bge-large",
        input=[text],
    )
    return np.array(resp.data[0].embedding, dtype=np.float32)

corpus_embeddings = np.vstack([get_embedding(doc) for doc in CORPUS])
print(f"Embedded {len(CORPUS)} documents into shape {corpus_embeddings.shape}")

Step 4: Build a cosine-similarity retriever

This is a pure numpy implementation. It takes a query embedding, computes cosine similarity against the corpus, and returns the top-k chunks.

def retrieve(query, top_k=2):
    q_emb = get_embedding(query)
    # Cosine similarity
    norms = np.linalg.norm(corpus_embeddings, axis=1) * np.linalg.norm(q_emb)
    similarities = corpus_embeddings.dot(q_emb) / norms
    top_idx = np.argsort(similarities)[::-1][:top_k]
    return [CORPUS[i] for i in top_idx], similarities[top_idx].tolist()

Step 5: Define the system prompt for grounded answers

The system prompt instructs the model to stick to the retrieved context and avoid hallucination. I inject the retrieved documents into the {context} slot at runtime.

SYSTEM_PROMPT = """You are a precise information retrieval assistant.
Answer the user's question using ONLY the provided context below.
If the context does not contain enough information, say you do not know.
Do not make up facts. Cite the relevant sentence in your answer.

Context:
{context}
"""

Step 6: Wire retrieval and generation into a pipeline

This function orchestrates the two stages: retrieve, format the prompt, and call llama-3.3-70b through Oxlo.ai. Because Oxlo.ai uses flat per-request pricing, long context windows from retrieved chunks do not inflate cost the way token-based billing would.

def answer(query):
    docs, scores = retrieve(query, top_k=2)
    context = "\n".join(f"- {d}" for d in docs)
    prompt = SYSTEM_PROMPT.format(context=context)

    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": query},
        ],
        temperature=0.1,
    )
    return {
        "query": query,
        "retrieved": docs,
        "scores": scores,
        "answer": response.choices[0].message.content,
    }

Run it

Here is the driver script and the output I get.

if __name__ == "__main__":
    result = answer("What pricing model does Oxlo.ai use?")
    print("Query:", result["query"])
    print("Retrieved:")
    for d, s in zip(result["retrieved"], result["scores"]):
        print(f"  [{s:.3f}] {d}")
    print("Answer:", result["answer"])

Example output:

Query: What pricing model does Oxlo.ai use?
Retrieved:
  [0.912] Oxlo.ai offers flat per-request pricing for LLM inference. Unlike token-based providers, cost does not scale with input length.
  [0.745] The platform hosts 45+ models including Llama 3.3 70B, Qwen 3 32B, DeepSeek R1, and Kimi K2.6.
Answer: Oxlo.ai uses a flat per-request pricing model for LLM inference. Unlike token-based providers, the cost does not scale with input length.

Next steps

Swap the in-memory numpy index for a real vector database such as Qdrant or Pinecone when you move past a few thousand documents. You can also experiment with Oxlo.ai's request-based pricing for reranking or multi-hop retrieval workflows, where long contexts do not drive up cost. See https://oxlo.ai/pricing for plan details.

DEV Community