refactor: 调整 vectordb 位 faiss 并更新依赖

.gitignore

@@ -2,3 +2,4 @@
 .env
 **/__pycache__
 chroma_db
+faiss_index

agent.py

@@ -11,7 +11,7 @@ from langchain_community.document_loaders import ArxivLoader
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_core.tools import tool
 from langchain.tools.retriever import create_retriever_tool
-from langchain_community.vectorstores import Chroma
+from langchain_community.vectorstores import FAISS
 from langchain_core.documents import Document
 import shutil
 import pandas as pd # Ny import för pandas
@@ -118,24 +118,22 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
 
 # Retrieval
-CHROMA_DIR = "./chroma_db"
+INDEX_DIR = "./faiss_index"
 CSV_PATH = "./supabase_docs.csv"
-EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
+# Use a lightweight 384-dim model to avoid excessive RAM and potential onnxruntime crashes
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 _SIMILARITY_THRESHOLD = 0.2  # lower distance means more similar
 
-embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
+embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL, model_kwargs={"device": "cpu"})
 
-if os.path.exists(CHROMA_DIR):
-    print(f"Loading existing ChromaDB from {CHROMA_DIR}")
-    vector_store = Chroma(
-        persist_directory=CHROMA_DIR,
-        embedding_function=embeddings,
-    )
+if os.path.exists(INDEX_DIR):
+    print(f"Loading existing FAISS index from {INDEX_DIR}")
+    vector_store = FAISS.load_local(INDEX_DIR, embeddings, allow_dangerous_deserialization=True)
 else:
-    print(f"Creating new ChromaDB at {CHROMA_DIR}, and loading documents from {CSV_PATH}")
-    if os.path.exists(CHROMA_DIR):
-        shutil.rmtree(CHROMA_DIR)
-    os.makedirs(CHROMA_DIR)
+    print(f"Creating new FAISS index at {INDEX_DIR}, and loading documents from {CSV_PATH}")
+    if os.path.exists(INDEX_DIR):
+        shutil.rmtree(INDEX_DIR)
+    os.makedirs(INDEX_DIR)
 
     if not os.path.exists(CSV_PATH):
         raise FileNotFoundError(f"CSV file {CSV_PATH} does not exist")
@@ -157,28 +155,27 @@ else:
 
         documents.append(Document(page_content=question_part, metadata=metadata))
 
+        # Simple progress indicator every 200 docs
+        if (i + 1) % 200 == 0:
+            print(f"Prepared {i + 1}/{len(df)} documents for embedding…")
+
     if not documents:
-        print("No documents loaded from CSV. ChromaDB will be empty.")
+        print("No documents loaded from CSV. FAISS index will be empty.")
 
-        vector_store = Chroma(
-            persist_directory=CHROMA_DIR,
-            embedding_function=embeddings
-        )
+        vector_store = FAISS.from_documents(documents=[], embedding=embeddings)
+        vector_store.save_local(INDEX_DIR)
     else:
-        vector_store = Chroma.from_documents(
-            documents=documents,
-            embedding=embeddings,
-            persist_directory=CHROMA_DIR,
-        )
-        vector_store.persist()
-        print(f"ChromaDB initialized and persisted with {len(documents)} documents from CSV.")
+        print("Embedding documents and building FAISS index — this may take a few minutes…")
+        vector_store = FAISS.from_documents(documents=documents, embedding=embeddings)
+        vector_store.save_local(INDEX_DIR)
+        print(f"FAISS index built and saved with {len(documents)} documents from CSV.")
         
 
 # Retriever tool
 retriever_tool = create_retriever_tool(
-    retriever = vector_store.as_retriever(),
-    name = "Question_Search",
-    description = "A tool to retrieve similar questions from a vector store. The retrieved document's metadata contains the 'final_answer' to the question."
+    retriever=vector_store.as_retriever(),
+    name="Question_Search",
+    description="Retrieve similar questions from FAISS index; metadata includes 'final_answer'."
 )
 
 # Agent

pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "chromadb>=1.0.15",
     "ddgs>=9.0.0",
     "duckduckgo-search>=8.1.1",
+    "faiss-cpu>=1.11.0",
     "gradio[oauth]>=5.36.2",
     "langchain>=0.3.26",
     "langchain-chroma>=0.2.4",

requirements.txt

@@ -3,6 +3,7 @@ requests
 chromadb
 ddgs
 duckduckgo-search
+faiss-cpu
 langchain
 langchain-chroma
 langchain-community

uv.lock

@@ -538,6 +538,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/0d/9feae160378a3553fa9a339b0e9c1a048e147a4127210e286ef18b730f03/durationpy-0.10-py3-none-any.whl", hash = "sha256:3b41e1b601234296b4fb368338fdcd3e13e0b4fb5b67345948f4f2bf9868b286", size = 3922 },
 ]
 
+[[package]]
+name = "faiss-cpu"
+version = "1.11.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+    { name = "packaging" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e7/9a/e33fc563f007924dd4ec3c5101fe5320298d6c13c158a24a9ed849058569/faiss_cpu-1.11.0.tar.gz", hash = "sha256:44877b896a2b30a61e35ea4970d008e8822545cb340eca4eff223ac7f40a1db9", size = 70218 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3b/d3/7178fa07047fd770964a83543329bb5e3fc1447004cfd85186ccf65ec3ee/faiss_cpu-1.11.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:356437b9a46f98c25831cdae70ca484bd6c05065af6256d87f6505005e9135b9", size = 3313807 },
+    { url = "https://files.pythonhosted.org/packages/9e/71/25f5f7b70a9f22a3efe19e7288278da460b043a3b60ad98e4e47401ed5aa/faiss_cpu-1.11.0-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:c4a3d35993e614847f3221c6931529c0bac637a00eff0d55293e1db5cb98c85f", size = 7913537 },
+    { url = "https://files.pythonhosted.org/packages/b0/c8/a5cb8466c981ad47750e1d5fda3d4223c82f9da947538749a582b3a2d35c/faiss_cpu-1.11.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:8f9af33e0b8324e8199b93eb70ac4a951df02802a9dcff88e9afc183b11666f0", size = 3785180 },
+    { url = "https://files.pythonhosted.org/packages/7f/37/eaf15a7d80e1aad74f56cf737b31b4547a1a664ad3c6e4cfaf90e82454a8/faiss_cpu-1.11.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:48b7e7876829e6bdf7333041800fa3c1753bb0c47e07662e3ef55aca86981430", size = 31287630 },
+    { url = "https://files.pythonhosted.org/packages/ff/5c/902a78347e9c47baaf133e47863134e564c39f9afe105795b16ee986b0df/faiss_cpu-1.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:bdc199311266d2be9d299da52361cad981393327b2b8aa55af31a1b75eaaf522", size = 15005398 },
+    { url = "https://files.pythonhosted.org/packages/92/90/d2329ce56423cc61f4c20ae6b4db001c6f88f28bf5a7ef7f8bbc246fd485/faiss_cpu-1.11.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:0c98e5feff83b87348e44eac4d578d6f201780dae6f27f08a11d55536a20b3a8", size = 3313807 },
+    { url = "https://files.pythonhosted.org/packages/24/14/8af8f996d54e6097a86e6048b1a2c958c52dc985eb4f935027615079939e/faiss_cpu-1.11.0-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:796e90389427b1c1fb06abdb0427bb343b6350f80112a2e6090ac8f176ff7416", size = 7913539 },
+    { url = "https://files.pythonhosted.org/packages/b2/2b/437c2f36c3aa3cffe041479fced1c76420d3e92e1f434f1da3be3e6f32b1/faiss_cpu-1.11.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:2b6e355dda72b3050991bc32031b558b8f83a2b3537a2b9e905a84f28585b47e", size = 3785181 },
+    { url = "https://files.pythonhosted.org/packages/66/75/955527414371843f558234df66fa0b62c6e86e71e4022b1be9333ac6004c/faiss_cpu-1.11.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6c482d07194638c169b4422774366e7472877d09181ea86835e782e6304d4185", size = 31287635 },
+    { url = "https://files.pythonhosted.org/packages/50/51/35b7a3f47f7859363a367c344ae5d415ea9eda65db0a7d497c7ea2c0b576/faiss_cpu-1.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:13eac45299532b10e911bff1abbb19d1bf5211aa9e72afeade653c3f1e50e042", size = 15005455 },
+]
+
 [[package]]
 name = "fastapi"
 version = "0.116.0"
@@ -1098,6 +1120,7 @@ dependencies = [
     { name = "chromadb" },
     { name = "ddgs" },
     { name = "duckduckgo-search" },
+    { name = "faiss-cpu" },
     { name = "gradio", extra = ["oauth"] },
     { name = "langchain" },
     { name = "langchain-chroma" },
@@ -1115,6 +1138,7 @@ requires-dist = [
     { name = "chromadb", specifier = ">=1.0.15" },
     { name = "ddgs", specifier = ">=9.0.0" },
     { name = "duckduckgo-search", specifier = ">=8.1.1" },
+    { name = "faiss-cpu", specifier = ">=1.11.0" },
     { name = "gradio", extras = ["oauth"], specifier = ">=5.36.2" },
     { name = "langchain", specifier = ">=0.3.26" },
     { name = "langchain-chroma", specifier = ">=0.2.4" },