ok

Dockerfile

@@ -1,43 +1,41 @@
-# Use the official Python 3.11 slim image
+# Use Python 3.11 Slim
 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 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
+# Install basic tools
+RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
 
+# Create user to avoid permission issues
+RUN useradd -m -u 1000 user
 WORKDIR /app
 
-# --- LAYER 1: Dependencies ---
+# --- LAYER 1: Install Dependencies ---
 COPY --chown=user:user requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 # --- LAYER 2: Download Models (Cached) ---
-# 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
+# This ensures models are baked into the image and load instantly
 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: Application Code ---
+# --- LAYER 3: App Code ---
 COPY --chown=user:user . .
 
-# Ensure permissions
-RUN mkdir -p $HF_HOME && chown -R user:user /app/cache && chown -R user:user /app/models
+# 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
 
-# Switch user
+# Switch to non-root user
 USER user
 
-# Expose port
+# Expose Port
 EXPOSE 7860
 
 # Start script

main.py

@@ -11,8 +11,9 @@ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 
-# Import the new MultiEmbeddingService
+# Import Custom Modules
 from model_service import MultiEmbeddingService
+from vector_store import SmartVectorStore  # <--- MUST HAVE THIS FILE
 
 # ============================================================================
 # LOGGING
@@ -32,34 +33,57 @@ ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*').split(',')
 # Global context container
 ml_context = {
     "service": None,
-    "executor": None
+    "executor": None,
+    "store": 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 and thread pool."""
-    # --- Startup ---
-    logger.info("Initializing Multi-Dimensional Embedding Service...")
+    """Lifecycle manager: Loads models, DB, and thread pool."""
+    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() # Loads 384, 768, 1024 models
+        service.load_all_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.")
     
@@ -67,6 +91,7 @@ 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()
@@ -75,8 +100,8 @@ async def lifespan(app: FastAPI):
 # APP SETUP
 # ============================================================================
 app = FastAPI(
-    title="Multi-Dim Embedding API",
-    version="3.0.0",
+    title="Multi-Dim Embedding & Retrieval API",
+    version="3.1.0",
     lifespan=lifespan
 )
 
@@ -98,7 +123,7 @@ async def verify_token(credentials: Optional[HTTPAuthorizationCredentials] = Sec
     return True
 
 # ============================================================================
-# MODELS
+# Pydantic MODELS
 # ============================================================================
 class EmbedRequest(BaseModel):
     data: Union[str, List[str]] = Field(..., description="Text string or list of strings")
@@ -114,12 +139,14 @@ 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
@@ -138,62 +165,91 @@ async def health_check():
 @app.post("/embed", response_model=EmbedResponse, dependencies=[Depends(verify_token)])
 async def create_embeddings(request: EmbedRequest):
     """
-    Generate embeddings for specific dimensions.
-    Supported dimensions: 384, 768, 1024.
+    Generate embeddings AND store them for later retrieval.
+    Accepts Single String OR Array of Strings.
     """
     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. Use 384, 768, or 1024."
-        )
+        raise HTTPException(status_code=400, detail=f"Dimension {request.dimension} not supported.")
 
     try:
+        # 1. Normalize Input
         is_single = isinstance(request.data, str)
-        count = 1 if is_single else len(request.data)
+        inputs = [request.data] if is_single else request.data
+        count = len(inputs)
 
+        # 2. Generate Embeddings (CPU Thread Pool)
         loop = asyncio.get_running_loop()
         embeddings = await loop.run_in_executor(
             executor,
             service.generate_embedding,
-            request.data,
+            inputs, # Pass list for batch processing
             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}")
+        logger.error(f"Inference error: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/deembed", dependencies=[Depends(verify_token)])
 async def de_embed_vector(request: DeEmbedRequest):
     """
-    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.
+    Retrieve original text using the vector.
+    Lookups are done via FAISS (Exact/Nearest Neighbor).
     """
-    # In a real scenario, this would look like:
-    # result = vector_db.search(vector=request.vector, top_k=1)
-    # return {"text": result.text}
+    store = ml_context.get("store")
     
-    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."
+    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."
         )
-    )
+
+@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")

model_service.py

@@ -1,16 +1,14 @@
 import logging
-from sentence_transformers import SentenceTransformer
 import torch
+from sentence_transformers import SentenceTransformer
 
 logger = logging.getLogger("EmbedService")
 
 class MultiEmbeddingService:
     def __init__(self):
         self.models = {}
-        # Auto-detect GPU, otherwise use CPU
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
-        # Map dimensions to local folders (downloaded in Dockerfile)
         self.model_map = {
             384: "./models/bge-384",
             768: "./models/bge-768",
@@ -18,31 +16,27 @@ class MultiEmbeddingService:
         }
 
     def load_all_models(self):
-        """Loads all defined models into memory ONCE at startup."""
+        """Loads all defined models into memory."""
         logger.info(f"🚀 Acceleration Device: {self.device.upper()}")
         
         for dim, path in self.model_map.items():
             try:
-                logger.info(f"Loading {dim}-dimension model from {path}...")
+                logger.info(f"Loading {dim}-dimension model...")
                 model = SentenceTransformer(path, device=self.device)
-                model.eval() # Set to evaluation mode (faster inference)
+                model.eval()
                 self.models[dim] = model
-                logger.info(f"✅ Loaded model for dimension {dim}")
             except Exception as e:
                 logger.error(f"❌ Failed to load {dim}-dim model: {e}")
 
-    def generate_embedding(self, text: str | list[str], dimension: int):
-        """Generates embeddings using the specific model for the requested dimension."""
+    def generate_embedding(self, text, dimension):
         if dimension not in self.models:
-            raise ValueError(f"Dimension {dimension} not supported. Available: {list(self.models.keys())}")
+            raise ValueError(f"Dimension {dimension} not supported.")
         
-        # --- OPTIMIZATION FIX ---
-        # show_progress_bar=False prevents the logs you saw
-        # batch_size=32 ensures efficient processing for lists
+        # show_progress_bar=False stops the spam
         return self.models[dimension].encode(
             text, 
             normalize_embeddings=True,
             convert_to_numpy=True,
-            show_progress_bar=False,  # <--- THIS STOPS THE LOG SPAM
+            show_progress_bar=False, 
             batch_size=32
         ).tolist()

requirements.txt

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

vector_store.py

@@ -0,0 +1,116 @@
+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)