rebuild chroma on startup

Dockerfile

@@ -5,10 +5,10 @@ WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY app.py fiscal.py auth.py plaid_client.py ./
-COPY chroma_db/ ./chroma_db/
+COPY app.py fiscal.py auth.py plaid_client.py rebuild_index.py ./
 COPY docs/ ./docs/
 
 EXPOSE 7860
 
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+# Rebuild ChromaDB then start the server
+CMD ["sh", "-c", "python rebuild_index.py && uvicorn app:app --host 0.0.0.0 --port 7860"]

rebuild_index.py

@@ -0,0 +1,39 @@
+from langchain_community.document_loaders import TextLoader
+from langchain.text_splitter import MarkdownHeaderTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+import shutil
+import os
+
+# Wipe old index
+if os.path.exists("./chroma_db"):
+    shutil.rmtree("./chroma_db")
+    print("Deleted old chroma_db/")
+
+# Load markdown
+loader = TextLoader("./docs/fiscal_knowledge.md", encoding="utf-8")
+docs = loader.load()
+text = docs[0].page_content
+print(f"Loaded {len(text):,} characters")
+
+splitter = MarkdownHeaderTextSplitter(
+    headers_to_split_on=[("##", "question")],
+    strip_headers=False,
+)
+chunks = splitter.split_text(text)
+
+print(f"Split into {len(chunks)} chunks")
+print("Loading embedding model (first run downloads ~270MB)...")
+embeddings = HuggingFaceEmbeddings(
+    model_name="nomic-ai/nomic-embed-text-v1",
+    model_kwargs={"trust_remote_code": True},
+)
+
+vectorstore = Chroma.from_documents(
+    chunks,
+    embedding=embeddings,
+    persist_directory="./chroma_db",
+)
+
+print(f"Indexed {len(chunks)} chunks into chroma_db/")
+print("Done. Ready to deploy.")