DEV Community

shashank ms
shashank ms

Posted on

Introduction to Knowledge Graph Embedding in LLM

We will build a pipeline that extracts a knowledge graph from raw text, embeds the entities for semantic retrieval, and answers questions using those structured triples. This gives you explainable, grounded responses that pure vector RAG often obscures. We will run everything against Oxlo.ai's OpenAI-compatible API so you can swap models without changing code.

What you'll need

I use Oxlo.ai here because its flat per-request pricing keeps costs predictable even when we feed long source documents into the extraction prompt. You can explore details at https://oxlo.ai/pricing.

Step 1: Initialize the client and extraction prompt

First, point the OpenAI SDK at Oxlo.ai. Then define the system prompt that forces the model to emit clean JSON triples.

from openai import OpenAI

client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
SYSTEM_PROMPT = (
    "You are a precise knowledge extraction engine. "
    "Read the user's text and return a JSON list of triples. "
    "Each triple must have exactly three keys: subject, predicate, object. "
    "Include only facts explicitly stated in the text. "
    "Output raw JSON only, with no markdown formatting."
)

Step 2: Extract triples from raw text

I pass a sample paragraph to Llama 3.3 70B and parse the JSON it returns. This becomes our knowledge graph.

from openai import OpenAI
import json

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

text = (
    "Oxlo.ai is a developer-first AI inference platform. "
    "It offers flat per-request pricing for open-source LLMs. "
    "Unlike token-based providers, Oxlo.ai does not scale cost with input length. "
    "The platform supports over 45 models including Llama 3.3 70B and DeepSeek R1."
)

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

triples = json.loads(response.choices[0].message.content)
print(json.dumps(triples, indent=2))

Step 3: Embed entities via Oxlo.ai

To enable semantic retrieval, I embed every unique entity using Oxlo.ai's BGE-Large model through the same OpenAI-compatible client.

from openai import OpenAI
import numpy as np

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

# Gather unique entities from subjects and objects
entities = list({ent for t in triples for ent in [t["subject"], t["object"]]})

embeddings = {}
for ent in entities:
    res = client.embeddings.create(
        model="bge-large",
        input=ent,
    )
    embeddings[ent] = np.array(res.data[0].embedding)

print(f"Embedded {len(embeddings)} entities.")

Step 4: Build a similarity index

Stack the vectors into a matrix and define a retriever that embeds a query and returns the top-k entities by cosine similarity.

from openai import OpenAI
import numpy as np

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

entity_names = list(embeddings.keys())
entity_matrix = np.stack([embeddings[e] for e in entity_names])

def retrieve_entities(query, top_k=3):
    q_res = client.embeddings.create(model="bge-large", input=query)
    q_vec = np.array(q_res.data[0].embedding)

    # Cosine similarity
    q_vec = q_vec / np.linalg.norm(q_vec)
    mat = entity_matrix / np.linalg.norm(entity_matrix, axis=1, keepdims=True)
    scores = mat @ q_vec

    idx = np.argsort(scores)[::-1][:top_k]
    return [entity_names[i] for i in idx]

# Quick test
print(retrieve_entities("pricing model"))

Step 5: Query the graph with natural language

Finally, I retrieve relevant triples and feed them into Kimi K2.6 for a grounded answer.

from openai import OpenAI

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

def answer_question(question):
    relevant = retrieve_entities(question, top_k=3)

    context = [
        t for t in triples
        if t["subject"] in relevant or t["object"] in relevant
    ]

    context_str = "\n".join([
        f"- {t['subject']} {t['predicate']} {t['object']}" for t in context
    ])

    user_msg = f"Context triples:\n{context_str}\n\nQuestion: {question}"

    response = client.chat.completions.create(
        model="kimi-k2.6",
        messages=[
            {"role": "system", "content": "Answer the question using only the provided knowledge graph triples."},
            {"role": "user", "content": user_msg},
        ],
    )
    return response.choices[0].message.content

print(answer_question("How does Oxlo.ai charge users?"))

Run it

Putting the pieces together and executing the script produces the following output.

Extracted triples:
[
  {
    "subject": "Oxlo.ai",
    "predicate": "is",
    "object": "developer-first AI inference platform"
  },
  {
    "subject": "Oxlo.ai",
    "predicate": "offers",
    "object": "flat per-request pricing for open-source LLMs"
  },
  {
    "subject": "Oxlo.ai",
    "predicate": "does not scale",
    "object": "cost with input length"
  },
  {
    "subject": "Oxlo.ai",
    "predicate": "supports",
    "object": "over 45 models"
  }
]

Top entities for 'pricing model':
['flat per-request pricing for open-source LLMs', 'Oxlo.ai', 'cost with input length']

Answer:
Oxlo.ai charges users a flat rate per API request. The cost does not increase with prompt length.

Wrap-up and next steps

You now have a working knowledge graph embedding pipeline powered entirely by Oxlo.ai. The flat per-request pricing is especially useful here because extraction prompts grow as your source documents grow, yet your cost stays constant per call.

Two concrete next steps: persist the triples and vectors in Neo4j or a vector database to scale beyond memory, or swap the extraction model to Qwen 3 32B if you start processing multilingual documents. Both changes require only a single line of code thanks to the OpenAI-compatible API.

Top comments (0)