DEV Community

shashank ms
shashank ms

Posted on

Best Practices for Training LLMs on Company Data

I built an internal knowledge agent that answers questions grounded in our company wiki. It chunks our Markdown docs, embeds them with Oxlo.ai, and retrieves relevant context before calling Llama 3.3 70B. This pattern is cheaper and faster than fine-tuning for most teams, and Oxlo.ai's flat per-request pricing keeps costs predictable even when we stuff long technical pages into the prompt.

What you'll need

  • Python 3.10 or newer
  • openai and numpy (pip install openai numpy)
  • A folder of Markdown files representing your company knowledge base
  • An Oxlo.ai API key from https://portal.oxlo.ai

Step 1: Ingest and chunk company documents

I keep our runbooks and policies in a docs/ folder as Markdown files. The script below reads every .md file and splits it into chunks at level-2 headers. I cap each chunk at roughly 1,500 characters to stay well under token limits.

import os
import re

def load_chunks(docs_dir="docs", max_chars=1500):
    chunks = []
    for filename in os.listdir(docs_dir):
        if not filename.endswith(".md"):
            continue
        filepath = os.path.join(docs_dir, filename)
        with open(filepath, "r", encoding="utf-8") as f:
            text = f.read()
        # Split on ## headers
        sections = re.split(r'\n## ', text)
        for section in sections:
            section = section.strip()
            if not section:
                continue
            if len(section) > max_chars:
                section = section[:max_chars]
            chunks.append({
                "source": filename,
                "text": section
            })
    return chunks

chunks = load_chunks()
print(f"Loaded {len(chunks)} chunks")

Step 2: Embed chunks with Oxlo.ai

Next, embed each chunk. Oxlo.ai hosts BGE-Large, which works well for semantic search on technical documentation. I call the embeddings endpoint in small batches to stay efficient.

from openai import OpenAI
import numpy as np

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

def embed_texts(texts, batch_size=32):
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        response = client.embeddings.create(
            model="bge-large",
            input=batch
        )
        all_embeddings.extend([item.embedding for item in response.data])
    return np.array(all_embeddings, dtype=np.float32)

texts = [c["text"] for c in chunks]
embeddings = embed_texts(texts)

# Normalize for cosine similarity
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
embeddings = embeddings / norms

for i, chunk in enumerate(chunks):
    chunk["embedding"] = embeddings[i]

print(f"Embedded {len(chunks)} chunks")

Step 3: Build a minimal vector store

I skip the heavy vector database and use a simple numpy dot product. For a few thousand chunks, this is instant and removes another dependency.

class SimpleRetriever:
    def __init__(self, chunks):
        self.chunks = chunks
        self.vectors = np.stack([c["embedding"] for c in chunks])

    def search(self, query, top_k=3):
        res = client.embeddings.create(model="bge-large", input=[query])
        q_vec = np.array(res.data[0].embedding, dtype=np.float32)
        q_vec = q_vec / np.linalg.norm(q_vec)

        scores = self.vectors @ q_vec
        top_indices = np.argpartition(scores, -top_k)[-top_k:]
        top_indices = top_indices[np.argsort(-scores[top_indices])]
        return [self.chunks[i] for i in top_indices]

retriever = SimpleRetriever(chunks)

Step 4: Define the system prompt

The system prompt is the contract with the model. It tells Llama 3.3 70B to stay grounded in the retrieved context, cite the source file, and refuse to guess when the answer is not present.

SYSTEM_PROMPT = """You are an internal company knowledge assistant.
Answer the user's question using ONLY the provided context below.
If the context does not contain the answer, say "I don't have that information."
Cite the source filename for every fact you use.

Context:
{context}
"""

Step 5: Assemble the agent

Finally, wire retrieval and generation together. The agent searches for the top 3 chunks, injects them into the system prompt, and calls Oxlo.ai.

def ask_agent(question):
    hits = retriever.search(question, top_k=3)
    context = "\n\n".join(
        f"[Source: {h['source']}]\n{h['text']}" for h in hits
    )
    prompt = SYSTEM_PROMPT.format(context=context)

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

answer, sources = ask_agent("What is our password rotation policy?")
print(answer)

Run it

With a few sample docs in place, running the script produces a grounded answer. Here is what I see after adding our security runbook:

$ python company_agent.py
Loaded 12 chunks
Embedded 12 chunks

Answer:
All production credentials must be rotated every 90 days.
[Source: security-runbook.md]

Because the prompt explicitly restricts the model to the retrieved text, hallucinations drop sharply compared to asking the model raw.

Wrap-up

This agent is already useful, but two upgrades matter. First, add a re-ranking step with a cross-encoder to improve retrieval accuracy once you pass a few hundred documents. Second, if your source material is massive, swap in Kimi K2.6 on Oxlo.ai; its 131K context window lets you retrieve larger chunks or more of them per request without worrying about token costs, since Oxlo.ai charges per request, not per token.

Top comments (0)