rag files

rag.py

@@ -0,0 +1,60 @@
+import os
+import faiss
+import json
+from sentence_transformers import SentenceTransformer
+
+KNOWLEDGE_BASE_DIR = "./knowledge_base"
+FAISS_INDEX_PATH = "./faiss_index.bin"
+METADATA_PATH = "./faiss_metadata.json"
+EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"
+
+print(f"--- Loading Embedding Model: {EMBEDDING_MODEL} ---")
+embedder = SentenceTransformer(EMBEDDING_MODEL)
+text_chunks = []
+
+def chunk_document(text):
+    paragraphs = text.split("\n\n")
+    return [p.strip() for p in paragraphs if len(p.strip()) > 20]
+
+def initialize_vector_store():
+    global text_chunks
+    embedding_dim = embedder.get_sentence_embedding_dimension()
+    
+    if os.path.exists(FAISS_INDEX_PATH) and os.path.exists(METADATA_PATH):
+        print("--- Loading existing FAISS index ---")
+        index = faiss.read_index(FAISS_INDEX_PATH)
+        with open(METADATA_PATH, "r") as f:
+            text_chunks = json.load(f)
+        return index
+
+    print("--- Building FAISS index from scratch ---")
+    os.makedirs(KNOWLEDGE_BASE_DIR, exist_ok=True)
+    
+    all_text = ""
+    for filename in os.listdir(KNOWLEDGE_BASE_DIR):
+        if filename.endswith(".txt") or filename.endswith(".md"):
+            with open(os.path.join(KNOWLEDGE_BASE_DIR, filename), "r", encoding="utf-8") as f:
+                all_text += f.read() + "\n\n"
+                
+    if not all_text.strip():
+        all_text = "Flow matching is a generative modeling framework."
+
+    text_chunks = chunk_document(all_text)
+    
+    print(f"--- Embedding {len(text_chunks)} chunks ---")
+    embeddings = embedder.encode(text_chunks, convert_to_numpy=True)
+    
+    index = faiss.IndexFlatL2(embedding_dim)
+    index.add(embeddings)
+    
+    faiss.write_index(index, FAISS_INDEX_PATH)
+    with open(METADATA_PATH, "w") as f:
+        json.dump(text_chunks, f)
+        
+    return index
+
+def retrieve_context(index, query, top_k=3):
+    query_vector = embedder.encode([query], convert_to_numpy=True)
+    distances, indices = index.search(query_vector, k=top_k)
+    retrieved_chunks = [text_chunks[i] for i in indices[0] if i != -1]
+    return "\n\n---\n\n".join(retrieved_chunks)

schemas.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+class BotRequest(BaseModel):
+    page_text: str
+    user_query: str