qwen

Dockerfile

@@ -1,24 +1,20 @@
-FROM python:3.9
-
-# Create user
-RUN useradd -m -u 1000 user
-
-# Create /data/chroma as root and give ownership to user BEFORE switching
-RUN mkdir -p /data/chroma && chown -R user:user /data
-
-USER user
-ENV PATH="/home/user/.local/bin:$PATH"
+FROM chromadb/chroma:latest
 
 WORKDIR /app
 
-COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+# Install FastAPI and Uvicorn for the health check app
+RUN pip install fastapi uvicorn requests
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
 
-COPY --chown=user . /app
+COPY . .
 
-# Make start.sh executable
-RUN chmod +x /app/start.sh
+RUN chmod +x start.sh
 
-VOLUME ["/data/chroma"]
+EXPOSE 8000
+EXPOSE 7860
 
-CMD ["/app/start.sh"]
+# Start both ChromaDB and the Health Check App
+# We use a simple script to run both in background
+CMD ["sh", "-c", "./start.sh & python3 app.py"]

app.py

@@ -1,94 +1,38 @@
-import json
 import os
-import chromadb
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import Optional
-
-app = FastAPI()
-
-# ── Config ────────────────────────────────────────────────────────────────────
-CONFIG_PATH = "/data/config.json"
-DEFAULT_CONFIG = {"text": "Hello, World!"}
-
-CHROMA_PATH = "/data/chroma"  # same /data persistent volume, survives restarts
-
-# ── ChromaDB client ───────────────────────────────────────────────────────────
-os.makedirs(CHROMA_PATH, exist_ok=True)
-chroma_client = chromadb.PersistentClient(path=CHROMA_PATH)
-collection = chroma_client.get_or_create_collection(name="my_collection")
-
-# ── Pydantic models ───────────────────────────────────────────────────────────
-class Document(BaseModel):
-    id: str
-    text: str
-    metadata: Optional[dict] = {}
-
-class QueryRequest(BaseModel):
-    query_text: str
-    n_results: int = 3
-
-# ── Config helpers ────────────────────────────────────────────────────────────
-def create_default_config():
-    os.makedirs(os.path.dirname(CONFIG_PATH), exist_ok=True)
-    with open(CONFIG_PATH, "w") as f:
-        json.dump(DEFAULT_CONFIG, f, indent=2)
-
-def load_config():
-    if not os.path.exists(CONFIG_PATH):
-        create_default_config()
-    with open(CONFIG_PATH, "r") as f:
-        return json.load(f)
-
-# ── Routes ────────────────────────────────────────────────────────────────────
-@app.get("/")
-def greet_json():
-    config = load_config()
-    return {"message": config.get("text")}
-
-@app.post("/documents/add")
-def add_document(doc: Document):
-    existing_ids = collection.get(ids=[doc.id])["ids"]
-    if existing_ids:
-        raise HTTPException(status_code=409, detail=f"Document '{doc.id}' already exists.")
-    collection.add(
-        documents=[doc.text],
-        ids=[doc.id],
-        metadatas=[doc.metadata],
-    )
-    return {"message": f"Document '{doc.id}' added successfully."}
-
-@app.post("/documents/query")
-def query_documents(req: QueryRequest):
-    count = collection.count()
-    if count == 0:
-        return {"results": [], "message": "Collection is empty."}
-    n = min(req.n_results, count)
-    results = collection.query(query_texts=[req.query_text], n_results=n)
-    hits = [
-        {
-            "id": results["ids"][0][i],
-            "text": results["documents"][0][i],
-            "metadata": results["metadatas"][0][i],
-            "distance": results["distances"][0][i],
-        }
-        for i in range(len(results["ids"][0]))
-    ]
-    return {"results": hits}
-
-@app.get("/documents/list")
-def list_documents():
-    data = collection.get()
-    docs = [
-        {"id": data["ids"][i], "text": data["documents"][i], "metadata": data["metadatas"][i]}
-        for i in range(len(data["ids"]))
-    ]
-    return {"count": len(docs), "documents": docs}
-
-@app.delete("/documents/{doc_id}")
-def delete_document(doc_id: str):
-    existing = collection.get(ids=[doc_id])["ids"]
-    if not existing:
-        raise HTTPException(status_code=404, detail=f"Document '{doc_id}' not found.")
-    collection.delete(ids=[doc_id])
-    return {"message": f"Document '{doc_id}' deleted."}
+import requests
+from fastapi import FastAPI
+from fastapi.responses import HTMLResponse
+
+app = FastAPI(title="ChromaDB on HF Space")
+
+CHROMA_URL = "http://localhost:8000"
+
+@app.get("/", response_class=HTMLResponse)
+def read_root():
+    return """
+    <html>
+        <head><title>ChromaDB HF Space</title></head>
+        <body>
+            <h1>✅ ChromaDB is Running!</h1>
+            <p>Persistent Storage: /data/chroma_db</p>
+            <p>API Endpoint: <a href="/docs">Swagger UI</a> (if enabled)</p>
+            <p>Direct ChromaDB API: <code>{}/api/v1/heartbeat</code></p>
+        </body>
+    </html>
+    """.format(CHROMA_URL)
+
+@app.get("/health")
+def health_check():
+    try:
+        response = requests.get(f"{CHROMA_URL}/api/v1/heartbeat")
+        if response.status_code == 200:
+            return {"status": "healthy", "chromadb": "up"}
+        else:
+            return {"status": "unhealthy", "chromadb": "error"}
+    except Exception as e:
+        return {"status": "error", "detail": str(e)}
+
+if __name__ == "__main__":
+    import uvicorn
+    # Run on port 7860 for HF Spaces UI
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))

requirements.txt

@@ -1,4 +1,4 @@
+chromadb==0.5.0
 fastapi
-uvicorn[standard]
-chromadb
-pydantic
+uvicorn
+requests

start.sh

@@ -1,13 +1,24 @@
 #!/bin/bash
-set -e
 
-echo "===== Application Startup at $(date '+%Y-%m-%d %H:%M:%S') ====="
+echo "=== [HF Space] Starting ChromaDB with Persistent Storage ==="
+echo "Time: $(date)"
 
-mkdir -p /data/chroma
-mkdir -p /data
+# Define the persistent directory inside the mounted bucket
+PERSIST_DIR="/data/chroma_db"
 
-echo "[startup] Directories ready."
-echo "[startup] HF_DATASET_REPO=$HF_DATASET_REPO"
-echo "[startup] Starting uvicorn..."
+# Create the directory if it doesn't exist
+mkdir -p $PERSIST_DIR
 
-exec uvicorn app:app --host 0.0.0.0 --port 7860 --log-level info
+echo "📂 Persistence directory: $PERSIST_DIR"
+echo "✅ Checking permissions..."
+ls -ld $PERSIST_DIR
+
+# Set environment variables for ChromaDB
+export CHROMA_SERVER_HOST=0.0.0.0
+export CHROMA_SERVER_HTTP_PORT=8000
+export CHROMA_DB_IMPL=chromadb.db.impl.sqlite
+export PERSIST_DIRECTORY=$PERSIST_DIR
+
+# If you are using the newer ChromaDB version that uses the 'chroma run' command:
+echo "🚀 Starting ChromaDB Server..."
+exec chroma run --path $PERSIST_DIR --host 0.0.0.0 --port 8000