aiagents4pharma 0.0.0__py3-none-any.whl

aiagents4pharma/talk2scholars/tools/pdf/utils/singleton_manager.py

@@ -0,0 +1,140 @@
+"""
+Singleton manager for Milvus connections and vector stores.
+Handles connection reuse, event loops, and GPU detection caching.
+"""
+
+import asyncio
+import logging
+import threading
+from typing import Any
+
+from langchain_core.embeddings import Embeddings
+from langchain_milvus import Milvus
+from pymilvus import connections, db, utility
+from pymilvus.exceptions import MilvusException
+
+from .gpu_detection import detect_nvidia_gpu
+
+logger = logging.getLogger(__name__)
+
+
+class VectorstoreSingleton:
+    """Singleton manager for Milvus connections and vector stores."""
+
+    _instance = None
+    _lock = threading.Lock()
+    _connections = {}  # Store connections by connection string
+    _vector_stores = {}  # Store vector stores by collection name
+    _event_loops = {}  # Store event loops by thread ID
+    _gpu_detected = None  # Cache GPU detection result
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def get_event_loop(self) -> asyncio.AbstractEventLoop:
+        """Get or create event loop for current thread."""
+        thread_id = threading.get_ident()
+
+        if thread_id not in self._event_loops:
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_closed():
+                    raise RuntimeError("Event loop is closed")
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            self._event_loops[thread_id] = loop
+            logger.info("Created new event loop for thread %s", thread_id)
+
+        return self._event_loops[thread_id]
+
+    def detect_gpu_once(self) -> bool:
+        """Detect GPU availability once and cache the result."""
+        if self._gpu_detected is None:
+            self._gpu_detected = detect_nvidia_gpu()
+            gpu_status = "available" if self._gpu_detected else "not available"
+            logger.info("GPU detection completed: NVIDIA GPU %s", gpu_status)
+        return self._gpu_detected
+
+    def get_connection(self, host: str, port: int, db_name: str) -> str:
+        """Get or create a Milvus connection."""
+        conn_key = f"{host}:{port}/{db_name}"
+
+        if conn_key not in self._connections:
+            try:
+                # Check if already connected
+                if connections.has_connection("default"):
+                    connections.remove_connection("default")
+
+                # Connect to Milvus
+                connections.connect(
+                    alias="default",
+                    host=host,
+                    port=port,
+                )
+                logger.info("Connected to Milvus at %s:%s", host, port)
+
+                # Check if database exists, create if not
+                existing_dbs = db.list_database()
+                if db_name not in existing_dbs:
+                    db.create_database(db_name)
+                    logger.info("Created database: %s", db_name)
+
+                # Use the database
+                db.using_database(db_name)
+                logger.info("Using database: %s", db_name)
+                logger.debug(
+                    "Milvus DB switched to: %s, available collections: %s",
+                    db_name,
+                    utility.list_collections(),
+                )
+
+                self._connections[conn_key] = "default"
+
+            except MilvusException as e:
+                logger.error("Failed to connect to Milvus: %s", e)
+                raise
+
+        return self._connections[conn_key]
+
+    def get_vector_store(
+        self,
+        collection_name: str,
+        embedding_model: Embeddings,
+        connection_args: dict[str, Any],
+    ) -> Milvus:
+        """Get or create a vector store for a collection."""
+        if collection_name not in self._vector_stores:
+            # Ensure event loop exists for this thread
+            self.get_event_loop()
+
+            # Create LangChain Milvus instance with explicit URI format
+            # This ensures LangChain uses the correct host
+            milvus_uri = f"http://{connection_args['host']}:{connection_args['port']}"
+
+            vector_store = Milvus(
+                embedding_function=embedding_model,
+                collection_name=collection_name,
+                connection_args={
+                    "uri": milvus_uri,  # Use URI format instead of host/port
+                    "host": connection_args["host"],
+                    "port": connection_args["port"],
+                },
+                text_field="text",
+                auto_id=False,
+                drop_old=False,
+                consistency_level="Strong",
+            )
+
+            self._vector_stores[collection_name] = vector_store
+            logger.info(
+                "Created new vector store for collection: %s with URI: %s",
+                collection_name,
+                milvus_uri,
+            )
+
+        return self._vector_stores[collection_name]

aiagents4pharma/talk2scholars/tools/pdf/utils/tool_helper.py

@@ -0,0 +1,86 @@
+"""
+Helper class for question and answer tool in PDF processing.
+"""
+
+import logging
+from typing import Any
+
+from .get_vectorstore import get_vectorstore
+
+logger = logging.getLogger(__name__)
+
+
+class QAToolHelper:
+    """
+    Encapsulates helper routines for the PDF Question & Answer tool.
+    Enhanced with automatic GPU/CPU detection and optimization.
+    """
+
+    def __init__(self) -> None:
+        self.config: Any = None
+        self.call_id: str = ""
+        self.has_gpu: bool = False  # Track GPU availability
+        logger.debug("Initialized QAToolHelper")
+
+    def start_call(self, config: Any, call_id: str) -> None:
+        """Initialize helper with current config and call identifier."""
+        self.config = config
+        self.call_id = call_id
+        logger.debug("QAToolHelper started call %s", call_id)
+
+    def get_state_models_and_data(self, state: dict) -> tuple[Any, Any, dict[str, Any]]:
+        """Retrieve embedding model, LLM, and article data from agent state."""
+        text_emb = state.get("text_embedding_model")
+        if not text_emb:
+            msg = "No text embedding model found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        llm = state.get("llm_model")
+        if not llm:
+            msg = "No LLM model found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        articles = state.get("article_data", {})
+        if not articles:
+            msg = "No article_data found in state."
+            logger.error("%s: %s", self.call_id, msg)
+            raise ValueError(msg)
+        return text_emb, llm, articles
+
+    def init_vector_store(self, emb_model: Any) -> Any:
+        """Get the singleton Milvus vector store instance with GPU/CPU optimization."""
+        logger.info(
+            "%s: Getting singleton vector store instance with hardware optimization",
+            self.call_id,
+        )
+        vs = get_vectorstore(embedding_model=emb_model, config=self.config)
+
+        # Track GPU availability from vector store
+        self.has_gpu = getattr(vs, "has_gpu", False)
+        hardware_type = "GPU-accelerated" if self.has_gpu else "CPU-only"
+
+        logger.info(
+            "%s: Vector store initialized (%s mode)",
+            self.call_id,
+            hardware_type,
+        )
+
+        # Log hardware-specific configuration
+        if hasattr(vs, "index_params"):
+            index_type = vs.index_params.get("index_type", "Unknown")
+            logger.info(
+                "%s: Using %s index type for %s processing",
+                self.call_id,
+                index_type,
+                hardware_type,
+            )
+
+        return vs
+
+    def get_hardware_stats(self) -> dict[str, Any]:
+        """Get current hardware configuration stats for monitoring."""
+        return {
+            "gpu_available": self.has_gpu,
+            "hardware_mode": "GPU-accelerated" if self.has_gpu else "CPU-only",
+            "call_id": self.call_id,
+        }

aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py

@@ -0,0 +1,150 @@
+"""
+Vector normalization utilities for GPU COSINE similarity support.
+Since GPU indexes don't support COSINE distance, we normalize vectors
+and use IP (Inner Product) distance instead.
+"""
+
+import logging
+
+import numpy as np
+from langchain_core.embeddings import Embeddings
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_vector(vector: list[float] | np.ndarray) -> list[float]:
+    """
+    Normalize a single vector to unit length.
+
+    Args:
+        vector: Input vector as list or numpy array
+
+    Returns:
+        Normalized vector as list
+    """
+    vector = np.asarray(vector, dtype=np.float32)
+    norm = np.linalg.norm(vector)
+
+    if norm == 0:
+        logger.warning("Zero vector encountered during normalization")
+        return vector.tolist()
+
+    normalized = vector / norm
+    return normalized.tolist()
+
+
+def normalize_vectors_batch(vectors: list[list[float]]) -> list[list[float]]:
+    """
+    Normalize a batch of vectors to unit length.
+
+    Args:
+        vectors: List of vectors
+
+    Returns:
+        List of normalized vectors
+    """
+    if not vectors:
+        return vectors
+
+    # Convert to numpy array for efficient computation
+    vectors_array = np.asarray(vectors, dtype=np.float32)
+
+    # Calculate norms for each vector
+    norms = np.linalg.norm(vectors_array, axis=1, keepdims=True)
+
+    # Handle zero vectors
+    zero_mask = norms.flatten() == 0
+    if np.any(zero_mask):
+        logger.warning("Found %d zero vectors during batch normalization", np.sum(zero_mask))
+        norms[zero_mask] = 1.0  # Avoid division by zero
+
+    # Normalize
+    normalized = vectors_array / norms
+
+    return normalized.tolist()
+
+
+class NormalizingEmbeddings(Embeddings):
+    """
+    Wrapper around an embedding model that automatically normalizes outputs.
+    This is needed for GPU indexes when using COSINE similarity.
+    """
+
+    def __init__(self, embedding_model: Embeddings, normalize_for_gpu: bool = True):
+        """
+        Initialize the normalizing wrapper.
+
+        Args:
+            embedding_model: The underlying embedding model
+            normalize_for_gpu: Whether to normalize embeddings (for GPU compatibility)
+        """
+        self.embedding_model = embedding_model
+        self.normalize_for_gpu = normalize_for_gpu
+
+        if normalize_for_gpu:
+            logger.info("Embedding model wrapped with normalization for GPU compatibility")
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed documents and optionally normalize."""
+        embeddings = self.embedding_model.embed_documents(texts)
+
+        if self.normalize_for_gpu:
+            embeddings = normalize_vectors_batch(embeddings)
+            logger.debug("Normalized %d document embeddings for GPU", len(embeddings))
+
+        return embeddings
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query and optionally normalize."""
+        embedding = self.embedding_model.embed_query(text)
+
+        if self.normalize_for_gpu:
+            embedding = normalize_vector(embedding)
+            logger.debug("Normalized query embedding for GPU")
+
+        return embedding
+
+    def __getattr__(self, name):
+        """Delegate other attributes to the underlying model."""
+        return getattr(self.embedding_model, name)
+
+
+def should_normalize_vectors(has_gpu: bool, use_cosine: bool) -> bool:
+    """
+    Determine if vectors should be normalized based on hardware and similarity metric.
+
+    Args:
+        has_gpu: Whether GPU is being used
+        use_cosine: Whether COSINE similarity is desired
+
+    Returns:
+        True if vectors should be normalized
+    """
+    needs_normalization = has_gpu and use_cosine
+
+    if needs_normalization:
+        logger.info("Vector normalization ENABLED: GPU detected with COSINE similarity request")
+    else:
+        logger.info("Vector normalization DISABLED: GPU=%s, COSINE=%s", has_gpu, use_cosine)
+
+    return needs_normalization
+
+
+def wrap_embedding_model_if_needed(
+    embedding_model: Embeddings, has_gpu: bool, use_cosine: bool = True
+) -> Embeddings:
+    """
+    Wrap embedding model with normalization if needed for GPU compatibility.
+
+    Args:
+        embedding_model: Original embedding model
+        has_gpu: Whether GPU is being used
+        use_cosine: Whether COSINE similarity is desired
+
+    Returns:
+        Original or wrapped embedding model
+    """
+    if should_normalize_vectors(has_gpu, use_cosine):
+        return NormalizingEmbeddings(embedding_model, normalize_for_gpu=True)
+
+    return embedding_model