DEV Community

shashank ms
shashank ms

Posted on

Leveraging LLM for Information Retrieval

I have built a small retrieval agent that answers questions over an internal wiki without spinning up a full search engine. In this tutorial, we will embed a handful of text chunks using Oxlo.ai embeddings, find the most relevant ones with vector similarity, and synthesize a concise answer with Llama 3.3 70B. The whole pipeline fits in a single Python file and runs on flat per-request pricing, so long documents do not inflate your bill.

What you'll need

  • Python 3.10 or higher
  • An Oxlo.ai API key from https://portal.oxlo.ai
  • The OpenAI SDK and NumPy: pip install openai numpy

Step 1: Configure the Oxlo.ai client

I start by initializing the OpenAI SDK to point at Oxlo.ai. Because the API is fully compatible, I only need to change the base_url and pass my Oxlo.ai key.

from openai import OpenAI

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

Step 2: Prepare the knowledge base

I scraped six passages from a fictional data-platform wiki. In a real deployment you would read these from Markdown or Confluence, but a list of strings keeps the example self-contained.

CORPUS = [
    "ETL pipelines should be idempotent. Running the same job twice with the same input must produce identical output without duplicating rows in the warehouse.",
    "Schema evolution is handled through additive changes only. Never drop a column that downstream dashboards depend on. Use a deprecation window of 30 days.",
    "Airflow DAGs are versioned in Git and deployed via CI. Each task runs in an isolated Kubernetes pod with resource limits set to prevent noisy neighbor issues.",
    "Data retention for user events is 90 days in hot storage and 2 years in cold S3 Glacier. PII must be tokenized before cross-region replication.",
    "Monitoring relies on Prometheus metrics emitted from each pipeline. Alertmanager routes critical alerts to PagerDuty during business hours and Slack otherwise.",
    "The feature store serves pre-computed embeddings for model inference. It is backed by Redis for online serving and S3 for offline batch exports.",
]

Step 3: Embed the corpus with BGE-Large

We batch the six chunks to Oxlo.ai's BGE-Large embedding endpoint. I then L2-normalize the vectors so that cosine similarity reduces to a simple dot product later.

import numpy as np

def get_embeddings(texts):
    response = client.embeddings.create(
        model="bge-large",
        input=texts,
    )
    return [np.array(d.embedding) for d in response.data]

corpus_embeddings = get_embeddings(CORPUS)
corpus_embeddings = [v / np.linalg.norm(v) for v in corpus_embeddings]

Step 4: Build the retriever

When a user asks a question, we embed it with the same model and score every document with np.dot. The two highest scores are our retrieval context.

def retrieve(query, top_k=2):
    q_emb = get_embeddings([query])[0]
    q_emb = q_emb / np.linalg.norm(q_emb)

    scores = [float(np.dot(q_emb, doc_emb)) for doc_emb in corpus_embeddings]
    ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)

    results = []
    for idx, score in ranked[:top_k]:
        results.append((score, CORPUS[idx]))
    return results

Step 5: Define the system prompt and agent

Here is the system prompt I use to keep the model grounded in the retrieved text:

SYSTEM_PROMPT = """You are a precise technical assistant.
Answer the user's question using ONLY the context provided below.
If the answer is not in the context, say: "I don't have that information."
Keep your answer under three sentences and cite the relevant source numbers."""

The ask_agent function formats the retrieved chunks and calls Llama 3.3 70B through Oxlo.ai.

def ask_agent(question):
    retrieved = retrieve(question, top_k=2)
    context_blocks = "\n\n".join(
        f"[Source {i+1}] {text}" for i, (_, text) in enumerate(retrieved)
    )

    user_content = f"Context:\n{context_blocks}\n\nQuestion: {question}"

    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
    )
    return response.choices[0].message.content

Run it

I run the script with a question about schema policy. The agent fetches the relevant wiki snippets and returns a synthesized answer.

if __name__ == "__main__":
    question = "What is the policy for changing database schemas?"
    answer = ask_agent(question)
    print(f"Q: {question}\nA: {answer}")

Example output:

Q: What is the policy for changing database schemas?
A: Schema changes must be additive only, and columns should not be dropped if downstream dashboards depend on them. A 30-day deprecation window is required before any removal. [Source 1]

Wrap-up

Swap the static list for a directory of Markdown files and you have a minimal internal search tool. If you need to scale beyond a few hundred chunks, move the embeddings into a vector database like pgvector or Milvus, but keep the same Oxlo.ai client for both retrieval and generation.

Top comments (0)