Add app\db\milvus.py

app//db//milvus.py

@@ -0,0 +1,117 @@
+import logging
+import asyncio
+import re
+from typing import List, Dict, Any, Optional
+
+from pymilvus import (
+    connections, 
+    utility, 
+    FieldSchema, 
+    CollectionSchema, 
+    DataType, 
+    Collection
+)
+from app.core.config import settings
+
+logger = logging.getLogger("rm_research.db.milvus")
+
+class MilvusVectorDB:
+    """
+    Institutional Scale Vector Intelligence Layer.
+    Optimized for high-recall academic searches with non-blocking I/O 
+    and strict input sanitization to prevent expression injection.
+    """
+
+    def __init__(self):
+        self.collection_name = "academic_knowledge_corpus"
+        self.dim = 768  # Tuned for scholarly transformer embeddings
+        self.alias = "default"
+        # Regex to ensure IDs are alphanumeric or standard UUID/Slug formats
+        self._sanitizer = re.compile(r"^[a-zA-Z0-9_\-]+$")
+
+    async def connect(self):
+        """Establishes thread-safe connection to Milvus cluster."""
+        loop = asyncio.get_running_loop()
+        try:
+            if not connections.has_connection(self.alias):
+                await loop.run_in_executor(
+                    None,
+                    lambda: connections.connect(
+                        alias=self.alias,
+                        host=settings.MILVUS_HOST,
+                        port=settings.MILVUS_PORT,
+                        user=settings.MILVUS_USER,
+                        password=settings.MILVUS_PASSWORD,
+                        secure=True,
+                        timeout=30
+                    )
+                )
+                logger.info(f"Connected to Milvus: {settings.MILVUS_HOST}")
+        except Exception as e:
+            logger.critical(f"Milvus Auth Failure: {str(e)}")
+            raise
+
+    async def search_ann(
+        self, 
+        query_vector: List[float], 
+        limit: int = 10, 
+        institution_id: Optional[str] = None,
+        disciplines: Optional[List[str]] = None
+    ) -> List[Dict[str, Any]]:
+        """
+        Executes Secure Approximate Nearest Neighbor (ANN) search.
+        Includes a whitelist-based filter builder to prevent injection attacks.
+        """
+        await self.connect()
+        collection = Collection(self.collection_name)
+        loop = asyncio.get_running_loop()
+
+        # 1. Build & Sanitize Expression (Security Fix)
+        filters = []
+        
+        if institution_id:
+            if self._sanitizer.match(institution_id):
+                filters.append(f"attributes['institution_id'] == '{institution_id}'")
+            else:
+                logger.warning(f"Sanitization block: Invalid institution_id '{institution_id}'")
+
+        if disciplines:
+            valid_dis = [d for d in disciplines if self._sanitizer.match(d)]
+            if valid_dis:
+                filters.append(f"attributes['discipline'] in {valid_dis}")
+
+        expr = " and ".join(filters) if filters else None
+
+        # 2. Execute Search in Executor
+        results = await loop.run_in_executor(
+            None,
+            lambda: collection.search(
+                data=[query_vector],
+                anns_field="embedding",
+                param={"metric_type": "COSINE", "params": {"ef": 128}},
+                limit=limit,
+                expr=expr,
+                output_fields=["paper_id", "attributes"]
+            )
+        )
+
+        return [
+            {
+                "paper_id": hit.entity.get("paper_id"),
+                "score": round(1.0 - hit.distance, 4), # Normalized similarity
+                "metadata": hit.entity.get("attributes")
+            } for hit in results[0]
+        ]
+
+    async def insert_batch(self, vectors: List[List[float]], ids: List[str], metadata: List[Dict]):
+        """Ingest batch into Milvus and flush to disk for persistence."""
+        await self.connect()
+        collection = Collection(self.collection_name)
+        loop = asyncio.get_running_loop()
+        
+        await loop.run_in_executor(None, lambda: collection.insert([ids, vectors, metadata]))
+        await loop.run_in_executor(None, collection.flush)
+        logger.info(f"Ingested {len(ids)} artifacts.")
+
+# Singleton instance
+milvus_db = MilvusVectorDB()