diff --git a/rag_agent.py b/rag_agent.py new file mode 100644 index 0000000..3423b9f --- /dev/null +++ b/rag_agent.py @@ -0,0 +1,48 @@ +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)