Retrieval-augmented generation

import math
from together import Together

client = Together()

EMBEDDING_MODEL = "intfloat/multilingual-e5-large-instruct"
CHAT_MODEL = "openai/gpt-oss-20b"

# A tiny corpus. In a real app, load from your data source and chunk first.
corpus = [
    "Photosynthesis converts sunlight, water, and carbon dioxide into glucose and oxygen, primarily in the chloroplasts of plant leaves.",
    "Mitochondria generate ATP through cellular respiration and are often called the powerhouse of the cell.",
    "Plate tectonics explains the slow movement of Earth's lithospheric plates and accounts for earthquakes and volcanoes.",
    "The water cycle moves water between oceans, atmosphere, and land through evaporation, condensation, and precipitation.",
    "Natural selection favors organisms whose inherited traits improve their chance of surviving and reproducing.",
    "Neural networks are layered computations of weighted sums and nonlinear activations, loosely inspired by biological neurons.",
]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(x * x for x in b))
    return dot / (na * nb) if na and nb else 0.0


# 1. Embed the corpus once.
doc_embeddings = client.embeddings.create(
    model=EMBEDDING_MODEL, input=corpus
).data
index = list(zip(corpus, [d.embedding for d in doc_embeddings]))


def rag(query: str, top_k: int = 3) -> str:
    # 2. Embed the query.
    q_emb = (
        client.embeddings.create(model=EMBEDDING_MODEL, input=query)
        .data[0]
        .embedding
    )

    # 3. Retrieve top_k by cosine similarity.
    ranked = sorted(index, key=lambda d: cosine(q_emb, d[1]), reverse=True)
    context = "\n\n".join(text for text, _ in ranked[:top_k])

    # 4. Generate an answer grounded in the retrieved context.
    response = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Answer the question using only the context below. "
                    "If the context is insufficient, say so.\n\n"
                    f"Context:\n{context}"
                ),
            },
            {"role": "user", "content": query},
        ],
    )
    return response.choices[0].message.content


print(rag("How do plants make their food?"))

RERANK_MODEL = (
    "mixedbread-ai/mxbai-rerank-large-v2"  # requires dedicated endpoint
)


def rag_with_rerank(query: str, retrieve_k: int = 20, top_n: int = 3) -> str:
    q_emb = (
        client.embeddings.create(model=EMBEDDING_MODEL, input=query)
        .data[0]
        .embedding
    )

    # 1. Over-retrieve from the vector store.
    candidates = sorted(
        index, key=lambda d: cosine(q_emb, d[1]), reverse=True
    )[:retrieve_k]
    candidate_texts = [text for text, _ in candidates]

    # 2. Rerank the candidates with a Together reranker.
    reranked = client.rerank.create(
        model=RERANK_MODEL,
        query=query,
        documents=candidate_texts,
        top_n=top_n,
    )
    context = "\n\n".join(candidate_texts[r.index] for r in reranked.results)

    # 3. Generate the final answer.
    response = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=[
            {
                "role": "system",
                "content": f"Answer using only the context below.\n\nContext:\n{context}",
            },
            {"role": "user", "content": query},
        ],
    )
    return response.choices[0].message.content

from pinecone import Pinecone, ServerlessSpec
from together import Together

pc = Pinecone(api_key="<PINECONE_API_KEY>", source_tag="TOGETHER_AI")
client = Together()

pc.create_index(
    name="together-rag",
    dimension=1024,  # match your embedding model's output dimension
    metric="cosine",
    spec=ServerlessSpec(cloud="aws", region="us-west-2"),
)
index = pc.Index("together-rag")

texts = ["Our solar system orbits the Milky Way at about 515,000 mph."]
embeddings = client.embeddings.create(
    model="intfloat/multilingual-e5-large-instruct", input=texts
).data

index.upsert(
    vectors=[
        {
            "id": f"doc_{i}",
            "values": e.embedding,
            "metadata": {"text": texts[i]},
        }
        for i, e in enumerate(embeddings)
    ]
)

from pymongo import MongoClient
from together import Together

mongo = MongoClient("<MONGODB_ATLAS_URI>")
collection = mongo["rag_db"]["documents"]
client = Together()

text = "Our solar system orbits the Milky Way at about 515,000 mph."
embedding = (
    client.embeddings.create(
        model="intfloat/multilingual-e5-large-instruct", input=text
    )
    .data[0]
    .embedding
)

collection.insert_one({"text": text, "embedding": embedding})

import pixeltable as pxt
from pixeltable.functions.together import embeddings

docs = pxt.create_table("rag.documents", {"text": pxt.String})
docs.add_computed_column(
    embedding=embeddings(
        input=docs.text, model="intfloat/multilingual-e5-large-instruct"
    )
)
docs.add_embedding_index(
    "text",
    string_embed=embeddings.using(
        model="intfloat/multilingual-e5-large-instruct"
    ),
)