The Role of Knowledge Graph Embedding in LLM Models

What we are building

I am building a Knowledge Graph QA Agent that turns a small set of entity-relationship triples into vector embeddings, retrieves the most relevant facts for a question, and grounds an LLM's answer with those structured facts. This matters when unstructured chunk retrieval is too fuzzy and you need precise, relational reasoning. I will use Oxlo.ai for both the embedding and generation calls.

What you'll need

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

Step 1: Bootstrap the Oxlo.ai client

I start by importing the OpenAI SDK and pointing it at Oxlo.ai's endpoint. Because Oxlo.ai uses flat per-request pricing, I can pack long prompts with retrieved triples without the cost scaling with token count.

from openai import OpenAI

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

Step 2: Define and embed the knowledge graph

I create a handful of real triples about tech acquisitions and founders, then batch embed them with Oxlo.ai's BGE-Large model. I format each triple as a simple sentence so the embedding captures the full relation.

import numpy as np

triples = [
    ("Microsoft", "acquired", "LinkedIn"),
    ("Microsoft", "founded_by", "Bill Gates"),
    ("Microsoft", "founded_by", "Paul Allen"),
    ("LinkedIn", "founded_by", "Reid Hoffman"),
    ("Google", "acquired", "YouTube"),
    ("Google", "founded_by", "Larry Page"),
    ("Google", "founded_by", "Sergey Brin"),
    ("YouTube", "founded_by", "Steve Chen"),
    ("YouTube", "founded_by", "Chad Hurley"),
    ("Amazon", "acquired", "Twitch"),
    ("Amazon", "founded_by", "Jeff Bezos"),
    ("Twitch", "founded_by", "Emmett Shear"),
    ("Twitch", "founded_by", "Justin Kan"),
]

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

triple_texts = [f"{s} {p} {o}" for s, p, o in triples]
triple_embeddings = embed_texts(triple_texts)

Step 3: Index the embeddings for fast retrieval

I stack the embeddings into a matrix and L2-normalize each row so I can compute cosine similarity later with a single dot product.

embedding_matrix = np.vstack(triple_embeddings)

# L2 normalize for cosine similarity via dot product
norms = np.linalg.norm(embedding_matrix, axis=1, keepdims=True)
normalized_matrix = embedding_matrix / norms

Step 4: Retrieve relevant triples for a question

When a question comes in, I embed it using the same model, normalize it, and score it against every triple. I return the top 3 matches.

def retrieve_triples(query, top_k=3):
    q_emb = embed_texts([query])[0]
    q_norm = q_emb / np.linalg.norm(q_emb)
    scores = normalized_matrix @ q_norm
    top_idx = np.argsort(scores)[::-1][:top_k]
    return [triples[i] for i in top_idx]

Step 5: Define the agent system prompt

I keep the system prompt strict. The model must answer only from the retrieved triples and cite them explicitly. This prevents hallucination when the graph is incomplete.

SYSTEM_PROMPT = """You are a knowledge graph agent. Answer the user's question using ONLY the provided triples. Each triple is given as (subject, predicate, object). Cite every triple you use. If the triples do not contain enough information, say so clearly."""

Step 6: Wire retrieval and generation into an agent

I combine the retrieval function with an Oxlo.ai chat completion call. I format the retrieved triples as a context block and pass them to Llama 3.3 70B.

def ask_agent(question):
    retrieved = retrieve_triples(question, top_k=3)
    context = "\n".join([f"- {s} {p} {o}" for s, p, o in retrieved])
    user_message = f"Triples:\n{context}\n\nQuestion: {question}"

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

Run it

I test the agent with a multi-hop question that requires connecting an acquisition to a founder. The retrieval step fetches the relevant triples, and the LLM reasons over them.

if __name__ == "__main__":
    question = "Who founded the company that acquired LinkedIn?"
    answer, sources = ask_agent(question)

    print("Retrieved triples:")
    for s, p, o in sources:
        print(f"  ({s}, {p}, {o})")

    print(f"\nAnswer: {answer}")

Example output:

Retrieved triples:
  (Microsoft, acquired, LinkedIn)
  (Microsoft, founded_by, Bill Gates)
  (Microsoft, founded_by, Paul Allen)

Answer: Microsoft acquired LinkedIn. Microsoft was founded by Bill Gates and Paul Allen.

Wrap-up and next steps

This agent proves that even a tiny knowledge graph can ground an LLM when the triples are embedded and retrieved with semantic similarity. To take it further, swap the in-memory list for a real graph database like Neo4j and rely on Oxlo.ai's flat per-request pricing to keep costs predictable as your context grows. See https://oxlo.ai/pricing for plan details. Another solid next step is to experiment with multi-hop retrieval by walking graph neighbors before embedding the path.