from pathlib import Path
from openai import OpenAI
import faiss
import numpy as np
import os

client = OpenAI()

def embed(text):
    res = client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return np.array(res.data[0].embedding, dtype=np.float32)

class RAGAgent:
    def __init__(self, index_path="vector.index"):
        if Path(index_path).exists():
            self.index = faiss.read_index(index_path)
        else:
            self.index = faiss.IndexFlatL2(1536)
        self.docs = []

    def add_document(self, text):
        vec = embed(text)
        self.index.add(np.array([vec]))
        self.docs.append(text)

    def query(self, q, top_k=3):
        vec = embed(q)
        distances, indices = self.index.search(np.array([vec]), top_k)
        return [self.docs[i] for i in indices[0] if i < len(self.docs)]

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--add", help="Add document text")
    parser.add_argument("--query", help="Query text")
    args = parser.parse_args()
    agent = RAGAgent()
    if args.add:
        agent.add_document(args.add)
        print("Document added.")
    if args.query:
        results = agent.query(args.query)
        print("Results:")
        for r in results:
            print("-", r)