Vector Search & RAG¶

import re
from typing import List

def chunk_text(text: str, chunk_size: int = 512, overlap: int = 64) -> List[str]:
    """Split text into overlapping chunks of approximately chunk_size tokens."""
    # Approximate tokenization by splitting on whitespace
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk_words = words[i : i + chunk_size]
        chunks.append(" ".join(chunk_words))
        i += chunk_size - overlap
    return chunks

# Load and chunk your document
with open("my_document.txt") as f:
    raw_text = f.read()

chunks = chunk_text(raw_text, chunk_size=256, overlap=32)
print(f"Created {len(chunks)} chunks")

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("all-MiniLM-L6-v2")   # 384-dimensional

embeddings = model.encode(
    chunks,
    batch_size=64,
    show_progress_bar=True,
    normalize_embeddings=True,  # L2-normalize for cosine similarity
)

print(f"Embedding shape: {embeddings.shape}")  # (n_chunks, 384)

import requests
import json

INDENTIA_URL = "http://localhost:7001"

# Create SurrealQL records with embeddings
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
    record = {
        "id":         f"chunk:{i}",
        "content":    chunk,
        "source":     "my_document.txt",
        "chunk_idx":  i,
        "embedding":  embedding.tolist(),
    }

    response = requests.post(
        f"{INDENTIA_URL}/sql",
        headers={"Content-Type": "application/json", "Accept": "application/json"},
        json={"query": f"""
            CREATE chunks CONTENT {{
                content:   $content,
                source:    $source,
                chunk_idx: $chunk_idx,
                embedding: $embedding
            }}
        """, "vars": {
            "content":   chunk,
            "source":    "my_document.txt",
            "chunk_idx": i,
            "embedding": embedding.tolist(),
        }}
    )
    response.raise_for_status()

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

def retrieve_chunks(query: str, k: int = 5) -> list[dict]:
    """Retrieve the k most relevant chunks for a query."""
    query_embedding = model.encode([query], normalize_embeddings=True)[0]

    response = requests.post(
        f"{INDENTIA_URL}/sql",
        headers={"Content-Type": "application/json", "Accept": "application/json"},
        json={
            "query": """
                SELECT id, content, source, chunk_idx,
                       vector::similarity::cosine(embedding, $query) AS score
                FROM chunks
                WHERE embedding <|$k, 40|> $query
                ORDER BY score DESC
                LIMIT $k
            """,
            "vars": {
                "query": query_embedding.tolist(),
                "k":     k,
            }
        }
    )
    response.raise_for_status()
    results = response.json()[0]["result"]
    return results

chunks_retrieved = retrieve_chunks("How does IndentiaDB handle vector indexing?", k=5)
for c in chunks_retrieved:
    print(f"[{c['score']:.3f}] {c['content'][:100]}...")

import openai   # or any other LLM client

def answer_question(question: str, k: int = 5) -> str:
    """Retrieve relevant context and answer a question using an LLM."""
    relevant_chunks = retrieve_chunks(question, k=k)

    # Build context string
    context_parts = []
    for i, chunk in enumerate(relevant_chunks, 1):
        context_parts.append(
            f"[{i}] (relevance: {chunk['score']:.2f})\n"
            f"Source: {chunk['source']}, chunk {chunk['chunk_idx']}\n"
            f"{chunk['content']}"
        )
    context = "\n\n---\n\n".join(context_parts)

    prompt = f"""You are a helpful assistant. Answer the question using ONLY the provided context.
If the context does not contain enough information, say so explicitly.

Context:
{context}

Question: {question}

Answer:"""

    client = openai.OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
        max_tokens=512,
    )

    return response.choices[0].message.content

answer = answer_question("What HNSW parameters should I use for a 10M document corpus?")
print(answer)

import requests

SPARQL_ENDPOINT = "http://localhost:7001/sparql"
CONFIDENCE_THRESHOLD = 0.7

def fetch_kg_context(topic_uri: str) -> list[dict]:
    """Fetch high-confidence RDF-star facts about a topic for RAG context."""
    sparql_query = f"""
        PREFIX ig:   <http://indentiadb.nl/ontology/>
        PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>

        SELECT ?subject ?predName ?object ?source ?confidence WHERE {{
            << ?subject ?predicate ?object >>
                ig:source     ?source ;
                ig:confidence ?confidence .

            OPTIONAL {{ ?predicate rdfs:label ?predName . }}
            FILTER (?confidence > {CONFIDENCE_THRESHOLD})
            FILTER (STRSTARTS(STR(?subject), STR(<{topic_uri}>)))
        }}
        ORDER BY DESC(?confidence)
        LIMIT 20
    """

    response = requests.post(
        SPARQL_ENDPOINT,
        headers={
            "Content-Type": "application/sparql-query",
            "Accept":       "application/sparql-results+json",
        },
        data=sparql_query,
    )
    response.raise_for_status()
    bindings = response.json()["results"]["bindings"]

    facts = []
    for b in bindings:
        facts.append({
            "subject":    b["subject"]["value"],
            "predicate":  b.get("predName", b["predicate"])["value"],
            "object":     b["object"]["value"],
            "source":     b["source"]["value"],
            "confidence": float(b["confidence"]["value"]),
        })
    return facts


def kg_rag_answer(question: str, topic_uri: str) -> str:
    """Answer a question using KG facts + vector-retrieved chunks."""
    import openai

    # Fetch structured KG facts
    facts = fetch_kg_context(topic_uri)

    # Fetch unstructured relevant chunks
    query_embedding = model.encode([question], normalize_embeddings=True)[0]
    chunks = retrieve_chunks(question, k=3)

    # Build context
    fact_lines = "\n".join(
        f"  - [{f['confidence']:.0%} confidence] {f['subject']} {f['predicate']} "
        f"{f['object']} (source: {f['source']})"
        for f in facts
    )
    chunk_lines = "\n\n".join(c["content"] for c in chunks)

    context = f"""Structured Knowledge Graph Facts (confidence-weighted):
{fact_lines}

Unstructured Document Excerpts:
{chunk_lines}"""

    client = openai.OpenAI()
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{
            "role": "user",
            "content": (
                f"Answer the question using the knowledge graph facts and document "
                f"excerpts below. Cite confidence levels where relevant.\n\n"
                f"{context}\n\nQuestion: {question}\nAnswer:"
            )
        }],
        temperature=0.1,
    )
    return response.choices[0].message.content


answer = kg_rag_answer(
    question="What is the HNSW index based on?",
    topic_uri="http://example.org/HNSWIndex"
)
print(answer)