Add rag_agent.py

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)