ok

Dockerfile

@@ -1,41 +1,43 @@
-# Use Python 3.11 Slim
+# Use the official Python 3.11 slim image
 FROM python:3.11-slim
 
 # Set environment variables
 ENV PYTHONDONTWRITEBYTECODE=1 \
     PYTHONUNBUFFERED=1 \
     PYTHONIOENCODING=UTF-8 \
-    HF_HOME=/app/cache 
+    HF_HOME=/app/cache
 
-# Install basic tools
-RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends curl \
+    && rm -rf /var/lib/apt/lists/* \
+    && useradd -m -u 1000 user
 
-# Create user to avoid permission issues
-RUN useradd -m -u 1000 user
 WORKDIR /app
 
-# --- LAYER 1: Install Dependencies ---
+# --- LAYER 1: Dependencies ---
 COPY --chown=user:user requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # --- LAYER 2: Download Models (Cached) ---
-# This ensures models are baked into the image and load instantly
+# We download models for 384, 768, and 1024 dimensions.
+# 384 dim: BAAI/bge-small-en-v1.5
+# 768 dim: BAAI/bge-base-en-v1.5
+# 1024 dim: BAAI/bge-large-en-v1.5
 RUN python3 -c "from huggingface_hub import snapshot_download; \
     snapshot_download(repo_id='BAAI/bge-small-en-v1.5', local_dir='./models/bge-384'); \
     snapshot_download(repo_id='BAAI/bge-base-en-v1.5', local_dir='./models/bge-768'); \
     snapshot_download(repo_id='BAAI/bge-large-en-v1.5', local_dir='./models/bge-1024')"
 
-# --- LAYER 3: App Code ---
+# --- LAYER 3: Application Code ---
 COPY --chown=user:user . .
 
-# Create storage directory for the database and set permissions
-RUN mkdir -p /app/storage && chown -R user:user /app/storage && chmod 777 /app/storage
-RUN mkdir -p $HF_HOME && chown -R user:user $HF_HOME
+# Ensure permissions
+RUN mkdir -p $HF_HOME && chown -R user:user /app/cache && chown -R user:user /app/models
 
-# Switch to non-root user
+# Switch user
 USER user
 
-# Expose Port
+# Expose port
 EXPOSE 7860
 
 # Start script

main.py

@@ -11,9 +11,8 @@ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 
-# Import Custom Modules
+# Import the new MultiEmbeddingService
 from model_service import MultiEmbeddingService
-from vector_store import SmartVectorStore  # <--- MUST HAVE THIS FILE
 
 # ============================================================================
 # LOGGING
@@ -33,57 +32,34 @@ ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*').split(',')
 # Global context container
 ml_context = {
     "service": None,
-    "executor": None,
-    "store": None
+    "executor": None
 }
 
-# ============================================================================
-# BACKGROUND TASKS
-# ============================================================================
-async def background_cleanup_task():
-    """Runs continuously to clean up data older than 24 hours."""
-    while True:
-        await asyncio.sleep(3600)  # Sleep for 1 hour
-        if ml_context["store"]:
-            logger.info("⏰ Running scheduled storage cleanup...")
-            ml_context["store"].prune_expired()
-
 # ============================================================================
 # LIFESPAN MANAGER
 # ============================================================================
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Lifecycle manager: Loads models, DB, and thread pool."""
-    logger.info("🚀 Initializing Multi-Dimensional Embedding Service...")
+    """Lifecycle manager: Loads models and thread pool."""
+    # --- Startup ---
+    logger.info("Initializing Multi-Dimensional Embedding Service...")
 
     # 1. Thread Pool
     cpu_count = multiprocessing.cpu_count()
     max_workers = cpu_count * 2
     executor = ThreadPoolExecutor(max_workers=max_workers)
     ml_context["executor"] = executor
+    logger.info(f"Thread pool ready: {max_workers} workers")
 
     # 2. Load Models
     try:
         service = MultiEmbeddingService()
-        service.load_all_models() 
+        service.load_all_models() # Loads 384, 768, 1024 models
         ml_context["service"] = service
     except Exception as e:
-        logger.critical(f"❌ Critical error loading models: {e}", exc_info=True)
+        logger.critical(f"Critical error loading models: {e}", exc_info=True)
         raise e
 
-    # 3. Load Vector Store (Database)
-    try:
-        # ttl_hours=24 ensures data is deleted after 24 hours
-        store = SmartVectorStore(storage_path="./storage", ttl_hours=24)
-        ml_context["store"] = store
-        logger.info("✅ Vector Store loaded with 24h retention policy.")
-    except Exception as e:
-        logger.critical(f"❌ Critical error loading Vector Store: {e}")
-        raise e
-
-    # 4. Start Cleanup Task
-    cleanup_task = asyncio.create_task(background_cleanup_task())
-
     if AUTH_TOKEN:
         logger.info("🔒 Auth enabled.")
     
@@ -91,7 +67,6 @@ async def lifespan(app: FastAPI):
     
     # --- Shutdown ---
     logger.info("Shutting down...")
-    cleanup_task.cancel()
     if ml_context["executor"]:
         ml_context["executor"].shutdown(wait=True)
     ml_context.clear()
@@ -100,8 +75,8 @@ async def lifespan(app: FastAPI):
 # APP SETUP
 # ============================================================================
 app = FastAPI(
-    title="Multi-Dim Embedding & Retrieval API",
-    version="3.1.0",
+    title="Multi-Dim Embedding API",
+    version="3.0.0",
     lifespan=lifespan
 )
 
@@ -123,7 +98,7 @@ async def verify_token(credentials: Optional[HTTPAuthorizationCredentials] = Sec
     return True
 
 # ============================================================================
-# Pydantic MODELS
+# MODELS
 # ============================================================================
 class EmbedRequest(BaseModel):
     data: Union[str, List[str]] = Field(..., description="Text string or list of strings")
@@ -139,14 +114,12 @@ class EmbedRequest(BaseModel):
     }
 
 class EmbedResponse(BaseModel):
-    id: Union[int, List[int]] = Field(..., description="Unique ID(s) for retrieval")
     embeddings: Union[List[float], List[List[float]]] = Field(...)
     dimension: int
     count: int
 
 class DeEmbedRequest(BaseModel):
     vector: List[float] = Field(..., description="The embedding vector to decode")
-    dimension: int = Field(768, description="The dimension of the vector")
 
 # ============================================================================
 # ENDPOINTS
@@ -165,91 +138,62 @@ async def health_check():
 @app.post("/embed", response_model=EmbedResponse, dependencies=[Depends(verify_token)])
 async def create_embeddings(request: EmbedRequest):
     """
-    Generate embeddings AND store them for later retrieval.
-    Accepts Single String OR Array of Strings.
+    Generate embeddings for specific dimensions.
+    Supported dimensions: 384, 768, 1024.
     """
     service = ml_context.get("service")
     executor = ml_context.get("executor")
-    store = ml_context.get("store")
 
     if not service or not executor:
         raise HTTPException(status_code=503, detail="Service unavailable")
 
-    # Validate Dimension
     if request.dimension not in service.models:
-        raise HTTPException(status_code=400, detail=f"Dimension {request.dimension} not supported.")
+        raise HTTPException(
+            status_code=400, 
+            detail=f"Dimension {request.dimension} not supported. Use 384, 768, or 1024."
+        )
 
     try:
-        # 1. Normalize Input
         is_single = isinstance(request.data, str)
-        inputs = [request.data] if is_single else request.data
-        count = len(inputs)
+        count = 1 if is_single else len(request.data)
 
-        # 2. Generate Embeddings (CPU Thread Pool)
         loop = asyncio.get_running_loop()
         embeddings = await loop.run_in_executor(
             executor,
             service.generate_embedding,
-            inputs, # Pass list for batch processing
+            request.data,
             request.dimension
         )
 
-        # 3. Store in Vector DB (Get Unique IDs)
-        stored_ids = []
-        
-        # If batch processing, embeddings is a list of lists. 
-        # If single, it might be a list of floats, so we wrap it to iterate consistently.
-        vectors_to_process = [embeddings] if is_single else embeddings
-
-        for text, vec in zip(inputs, vectors_to_process):
-            # store.add generates a UNIQUE ID (does not overwrite old data)
-            new_id = store.add(text, vec, request.dimension)
-            stored_ids.append(new_id)
-
-        # 4. Return Response
         return EmbedResponse(
-            id=stored_ids[0] if is_single else stored_ids,
             embeddings=embeddings,
             dimension=request.dimension,
             count=count
         )
 
     except Exception as e:
-        logger.error(f"Inference error: {e}", exc_info=True)
+        logger.error(f"Inference error: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/deembed", dependencies=[Depends(verify_token)])
 async def de_embed_vector(request: DeEmbedRequest):
     """
-    Retrieve original text using the vector.
-    Lookups are done via FAISS (Exact/Nearest Neighbor).
+    Experimental: Reverse vector to text.
+    
+    NOTE: Mathematically, standard embedding models (BERT, BGE) are NOT reversible 
+    because they are lossy compression algorithms. 
+    
+    To retrieve text from a vector, you must use a Vector Database (retrieval), 
+    not a direct model inversion.
     """
-    store = ml_context.get("store")
+    # In a real scenario, this would look like:
+    # result = vector_db.search(vector=request.vector, top_k=1)
+    # return {"text": result.text}
     
-    if not store:
-        raise HTTPException(status_code=503, detail="Store unavailable")
-
-    # Search in the store
-    result = store.search(request.vector, request.dimension)
-
-    if result:
-        return {
-            "found": True,
-            "text": result["text"],
-            "created_at_timestamp": result["created_at"],
-            "note": "Data expires 24h after creation."
-        }
-    else:
-        raise HTTPException(
-            status_code=404, 
-            detail="Vector not found. It may have expired (24h limit) or was never stored."
+    raise HTTPException(
+        status_code=501, 
+        detail=(
+            "De-embedding (Vector-to-Text) is not possible with standalone embedding models. "
+            "This endpoint requires a connected Vector Database to perform a similarity search."
         )
-
-@app.get("/check_id/{dimension}/{uid}")
-async def check_by_id(dimension: int, uid: int):
-    """Debug endpoint: Check if an ID exists without the vector."""
-    store = ml_context.get("store")
-    data = store.get_by_id(uid, dimension)
-    if data:
-        return data
-    raise HTTPException(status_code=404, detail="ID not found")
+    )

requirements.txt

@@ -10,7 +10,4 @@ numpy==1.26.4
 
 # Production dependencies
 python-multipart==0.0.20
-aiofiles==24.1.0
-
-# Vector database dependencies
-faiss-cpu==1.9.0.post1
+aiofiles==24.1.0

vector_store.py

@@ -1,116 +0,0 @@
-import faiss
-import numpy as np
-import pickle
-import os
-import time
-import logging
-import random
-
-logger = logging.getLogger("SmartStore")
-
-class SmartVectorStore:
-    def __init__(self, storage_path="./storage", ttl_hours=24):
-        self.storage_path = storage_path
-        self.ttl_seconds = ttl_hours * 3600
-        os.makedirs(storage_path, exist_ok=True)
-        
-        self.indices = {} 
-        self.metadata = {} # Maps ID -> { "text": str, "created_at": float }
-        
-        # Initialize indexes for all dimensions
-        for dim in [384, 768, 1024]:
-            self._load_or_create_index(dim)
-
-    def _load_or_create_index(self, dim):
-        index_file = os.path.join(self.storage_path, f"index_{dim}.faiss")
-        meta_file = os.path.join(self.storage_path, f"meta_{dim}.pkl")
-
-        if os.path.exists(index_file) and os.path.exists(meta_file):
-            try:
-                self.indices[dim] = faiss.read_index(index_file)
-                with open(meta_file, "rb") as f:
-                    self.metadata[dim] = pickle.load(f)
-                logger.info(f"📂 Loaded DB for dim {dim} with {self.indices[dim].ntotal} items.")
-            except Exception:
-                logger.warning(f"⚠️ Corrupt DB for {dim}, creating new.")
-                self._create_new_index(dim)
-        else:
-            self._create_new_index(dim)
-
-    def _create_new_index(self, dim):
-        # IndexIDMap lets us assign our own IDs
-        self.indices[dim] = faiss.IndexIDMap(faiss.IndexFlatL2(dim))
-        self.metadata[dim] = {}
-
-    def add(self, text: str, vector: list[float], dim: int):
-        """Adds text, assigns a unique ID, and saves timestamp."""
-        
-        # Generate Unique ID (Time based + Random)
-        unique_id = int(time.time() * 1000) + random.randint(0, 999)
-        
-        vector_np = np.array([vector], dtype=np.float32)
-        id_np = np.array([unique_id], dtype=np.int64)
-
-        # Add to FAISS
-        self.indices[dim].add_with_ids(vector_np, id_np)
-        
-        # Add to Metadata
-        self.metadata[dim][unique_id] = {
-            "text": text,
-            "created_at": time.time()
-        }
-        
-        # Save to disk
-        self._save(dim)
-        return unique_id
-
-    def search(self, vector: list[float], dim: int):
-        """Finds closest text by vector."""
-        if self.indices[dim].ntotal == 0:
-            return None
-
-        vector_np = np.array([vector], dtype=np.float32)
-        D, I = self.indices[dim].search(vector_np, 1)
-        
-        found_id = I[0][0]
-        distance = D[0][0] # 0.0 is exact match
-
-        if found_id != -1 and distance < 1e-4:
-            if found_id in self.metadata[dim]:
-                return self.metadata[dim][found_id]
-        
-        return None
-
-    def get_by_id(self, unique_id: int, dim: int):
-        """Direct lookup by ID."""
-        return self.metadata[dim].get(unique_id)
-
-    def prune_expired(self):
-        """Deletes items older than 24 hours."""
-        current_time = time.time()
-        
-        for dim in self.indices:
-            ids_to_remove = []
-            
-            for uid, data in list(self.metadata[dim].items()):
-                age = current_time - data["created_at"]
-                if age > self.ttl_seconds:
-                    ids_to_remove.append(uid)
-            
-            if ids_to_remove:
-                logger.info(f"🧹 Purging {len(ids_to_remove)} expired items from Dim {dim}...")
-                
-                # Remove from Metadata
-                for uid in ids_to_remove:
-                    del self.metadata[dim][uid]
-                
-                # Remove from FAISS
-                ids_np = np.array(ids_to_remove, dtype=np.int64)
-                self.indices[dim].remove_ids(ids_np)
-                
-                self._save(dim)
-
-    def _save(self, dim):
-        faiss.write_index(self.indices[dim], os.path.join(self.storage_path, f"index_{dim}.faiss"))
-        with open(os.path.join(self.storage_path, f"meta_{dim}.pkl"), "wb") as f:
-            pickle.dump(self.metadata[dim], f)