49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
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)
|