intersystems-community/fhir-ai-hackathon-kit/iris-vector-search

CREATE TABLE Diabetes.VectorStore (
    Source    VARCHAR(100),
    Text      VARCHAR(10000),
    Embedding VECTOR(DOUBLE, 1536)
)

import pymupdf
from langchain_text_splitters import RecursiveCharacterTextSplitter
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=100,
    separators=[". ", " ", ""],   # prefer sentence boundaries, then words
    keep_separator="end",
)
def get_chunks(pdf_path: str) -> list[str]:
    with pymupdf.open(pdf_path) as doc:
        text = " ".join(page.get_text("text") for page in doc)
    return splitter.split_text(" ".join(text.split()))   # collapse whitespace first

from openai import OpenAI
client = OpenAI()
def embed(texts: list[str]) -> list[list[float]]:
    resp = client.embeddings.create(input=texts, model="text-embedding-3-small")   # batch: one call, many texts
    return [d.embedding for d in resp.data]
insert = "INSERT INTO Diabetes.VectorStore (Source, Text, Embedding) VALUES (?, ?, TO_VECTOR(?))"
for n, (source, chunks) in enumerate(chunks_by_doc.items(), 1):
    vectors = embed(chunks)
    cursor.executemany(insert, [[source, c, str(v)] for c, v in zip(chunks, vectors)])
    print(f"[{n}/{len(chunks_by_doc)}] {source}: {len(chunks)} chunks")

def vector_search(prompt: str, top_k: int = 3) -> list:
    """Return (Source, Text) rows most similar to prompt."""
    query_vec = embed([prompt])[0]
    conn = iris.connect(**connection_args)
    cursor = conn.cursor()
    try:
        cursor.execute(
            f"SELECT TOP {int(top_k)} Source, Text FROM Diabetes.VectorStore "
            "ORDER BY VECTOR_COSINE(Embedding, TO_VECTOR(?, DOUBLE)) DESC",
            [str(query_vec)],
        )
        return cursor.fetchall()
    finally:
        cursor.close()
        conn.close()