RAG Knowledge Base Answer

Inputs

Outputs

Source

python

# Input preparation
inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
question = inp["question"]
hugging_face_token = inp["hugging_face_token"]
model_name = inp["model_name"]
milvus_uri = inp["milvus_uri"]
collection_name = inp["collection_name"]
openai_api_key = inp["openai_api_key"]
openai_model = inp.get("openai_model", "gpt-4o")
system_message = inp.get("system_message", "You are a helpful assistant. Answer strictly from the provided context. If the answer is not in the context, say you do not know.")
top_k = int(inp.get("top_k", 4))
max_tokens = int(inp.get("max_tokens", 1024))

# Compute
from typing import List
import torch
from transformers import AutoTokenizer, AutoModel
from pymilvus import MilvusClient
from openai import OpenAI

class RagKbAnswer:
    """
    Single-node RAG: embed question -> Milvus search -> grounded OpenAI answer.
    """

    def __init__(self, model_name: str, hugging_face_token: str, milvus_uri: str, openai_api_key: str):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=hugging_face_token)
        self.model = AutoModel.from_pretrained(model_name, token=hugging_face_token)
        self.milvus = MilvusClient(uri=milvus_uri)
        self.openai = OpenAI(api_key=openai_api_key)

    def embed(self, text: str) -> List[float]:
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        with torch.no_grad():
            outputs = self.model(**inputs)
        emb = outputs.last_hidden_state[:, 0, :]
        return torch.nn.functional.normalize(emb, p=2, dim=1).squeeze().tolist()

    def retrieve(self, collection_name: str, vector: List[float], top_k: int) -> str:
        hits = self.milvus.search(
            collection_name=collection_name,
            data=[vector],
            limit=top_k,
            output_fields=["text"],
        )
        passages = []
        for group in hits:
            for item in group:
                entity = item.get("entity", {}) if isinstance(item, dict) else {}
                text = entity.get("text")
                if text:
                    passages.append(text)
        return "\n\n".join(passages)

    def generate(self, system_message: str, context: str, question: str, model: str, max_tokens: int) -> str:
        user_content = f"Context:\n{context}\n\nQuestion: {question}"
        completion = self.openai.chat.completions.create(
            model=model,
            max_tokens=max_tokens,
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_content},
            ],
        )
        return completion.choices[0].message.content or ""

    def run(self, question: str, collection_name: str, system_message: str, model: str, top_k: int, max_tokens: int):
        vector = self.embed(question)
        context = self.retrieve(collection_name, vector, top_k)
        answer = self.generate(system_message, context, question, model, max_tokens)
        return answer, context

rag = RagKbAnswer(model_name=model_name, hugging_face_token=hugging_face_token, milvus_uri=milvus_uri, openai_api_key=openai_api_key)
answer, context = rag.run(question, collection_name, system_message, openai_model, top_k, max_tokens)

# Output collection (runner reads __ivy_ctx__)
out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["answer"] = answer
out["context"] = context

Tests

Requires: python:3.11

• basic-rag-answer

  Answer a question against a knowledge base (network; tolerate success or environment error).

  ragkbllmbasic
• missing-question

  Missing required question should fail.

  ragerror