Build a RAG knowledge-base agent

export AIWAVE_API_KEY="sk-your-api-key"
python -m pip install openai
python rag_knowledge_base.py "How long are password reset links valid?"

import math
import os
import sys

from openai import OpenAI

BASE_URL = os.environ.get("AIWAVE_BASE_URL", "https://aiwave.live/v1")
CHAT_MODEL = os.environ.get("AIWAVE_CHAT_MODEL", "deepseek-v4-flash")
EMBEDDING_MODEL = "qwen3.7-text-embedding"

KNOWLEDGE_BASE = [
    {"title": "Password reset", "text": "Password reset links expire after 30 minutes. Request a new link after expiry."},
    {"title": "API key rotation", "text": "Create a replacement key, update the client, verify one bounded request, then revoke the old key."},
    {"title": "Invoice timing", "text": "Usage appears in the Console request record after the request completes. Review the dated rate before forecasting."},
    {"title": "Support handoff", "text": "For an unresolved account issue, include the request ID and error class, never the API key or request body."},
]

def cosine_similarity(left, right):
    numerator = sum(a * b for a, b in zip(left, right))
    left_norm = math.sqrt(sum(value * value for value in left))
    right_norm = math.sqrt(sum(value * value for value in right))
    if not left_norm or not right_norm:
        return 0.0
    return numerator / (left_norm * right_norm)

def retrieve(client, question, top_k=2):
    texts = [item["text"] for item in KNOWLEDGE_BASE] + [question]
    response = client.embeddings.create(model=EMBEDDING_MODEL, input=texts)
    vectors = sorted(response.data, key=lambda item: item.index)
    question_vector = vectors[-1].embedding
    scored = []
    for item, vector in zip(KNOWLEDGE_BASE, vectors[:-1]):
        scored.append((cosine_similarity(vector.embedding, question_vector), item))
    scored.sort(key=lambda pair: pair[0], reverse=True)
    return [item for _score, item in scored[:top_k]]

def answer(question):
    client = OpenAI(api_key=os.environ["AIWAVE_API_KEY"], base_url=BASE_URL)
    passages = retrieve(client, question)
    context = "\n\n".join(
        f"[{index}] {item['title']}: {item['text']}"
        for index, item in enumerate(passages, start=1)
    )
    response = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=[
            {"role": "system", "content": "Answer only from the supplied context. Cite supporting passages as [1], [2]. If the context is insufficient, say so."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
        ],
        max_tokens=180,
    )
    return response.choices[0].message.content or "No answer returned."

if __name__ == "__main__":
    if len(sys.argv) != 2:
        raise SystemExit('Usage: python rag_knowledge_base.py "your question"')
    print(answer(sys.argv[1]))