mdb-engine 0.1.6__py3-none-any.whl → 0.4.12__py3-none-any.whl

mdb_engine/embeddings/dependencies.py

@@ -1,97 +1,60 @@
 """
-Embedding Service Dependency Injection for FastAPI
+Embedding Service Utilities
 
-This module provides FastAPI dependency functions to inject embedding services
-into route handlers. The embedding service is automatically initialized from
-the app's manifest.json configuration.
-"""
+This module provides utility functions for creating embedding services.
+For FastAPI dependency injection, use the request-scoped dependencies
+from `mdb_engine.dependencies` instead.
 
-from typing import Any, Optional
+Usage:
+    # For FastAPI routes (RECOMMENDED):
+    from mdb_engine.dependencies import get_embedding_service
 
-# Optional FastAPI import (only needed if FastAPI is available)
-try:
-    from fastapi import Depends, HTTPException
+    @app.post("/embed")
+    async def embed(embedding_service=Depends(get_embedding_service)):
+        ...
 
-    FASTAPI_AVAILABLE = True
-except ImportError:
-    FASTAPI_AVAILABLE = False
+    # For standalone/utility usage:
+    from mdb_engine.embeddings.dependencies import get_embedding_service_for_app
 
-    # Stub for when FastAPI is not available
-    def Depends(*args, **kwargs):
-        return None
+    service = get_embedding_service_for_app("my_app", engine)
+"""
 
-    class HTTPException(Exception):
-        pass
+from typing import TYPE_CHECKING
 
+if TYPE_CHECKING:
+    from ..core.engine import MongoDBEngine
 
 from .service import EmbeddingService, get_embedding_service
 
-# Global engine registry (for apps that don't pass engine explicitly)
-_global_engine: Optional[Any] = None
-_global_app_slug: Optional[str] = None
-
-
-def set_global_engine(engine: Any, app_slug: Optional[str] = None) -> None:
-    """
-    Set global MongoDBEngine instance for embedding dependency injection.
-
-    This is useful when you have a single engine instance that you want
-    to use across all apps. Call this during application startup.
-
-    Args:
-        engine: MongoDBEngine instance
-        app_slug: Optional app slug
-    """
-    global _global_engine, _global_app_slug
-    _global_engine = engine
-    _global_app_slug = app_slug
-
-
-def get_global_engine() -> Optional[Any]:
-    """
-    Get global MongoDBEngine instance.
-
-    Returns:
-        MongoDBEngine instance if set, None otherwise
-    """
-    return _global_engine
-
 
 def get_embedding_service_for_app(
-    app_slug: str, engine: Optional[Any] = None
-) -> Optional[EmbeddingService]:
+    app_slug: str, engine: "MongoDBEngine"
+) -> EmbeddingService | None:
     """
-    Get embedding service for a specific app.
+    Get embedding service for a specific app using the engine instance.
 
-    This is a helper function that can be used with FastAPI's Depends()
-    to inject the embedding service into route handlers.
+    This is a utility function for cases where you need to create an
+    embedding service outside of a FastAPI request context (e.g., in
+    background tasks, CLI tools, or tests).
+
+    For FastAPI routes, use `mdb_engine.dependencies.get_embedding_service` instead.
 
     Args:
-        app_slug: App slug (typically extracted from route context)
-        engine: MongoDBEngine instance (optional, will try to get from context)
+        app_slug: App slug to get embedding config from
+        engine: MongoDBEngine instance
 
     Returns:
-        EmbeddingService instance if embedding is enabled for this app, None otherwise
+        EmbeddingService instance if embedding is enabled, None otherwise
 
     Example:
-        ```python
-        from fastapi import Depends
-        from mdb_engine.embeddings.dependencies import get_embedding_service_for_app
-
-        @app.post("/embed")
-        async def embed_endpoint(
-            embedding_service = Depends(lambda: get_embedding_service_for_app("my_app"))
-        ):
-            if not embedding_service:
-                raise HTTPException(503, "Embedding service not available")
-            embeddings = await embedding_service.embed_chunks(["Hello world"])
-            return {"embeddings": embeddings}
-        ```
-    """
-    # Try to get engine from context if not provided
-    if engine is None:
-        engine = _global_engine
+        # In a background task or CLI
+        engine = MongoDBEngine(...)
+        await engine.initialize()
 
+        service = get_embedding_service_for_app("my_app", engine)
+        if service:
+            embeddings = await service.embed_chunks(["Hello world"])
+    """
     if engine is None:
         return None
 
@@ -108,86 +71,6 @@ def get_embedding_service_for_app(
     return get_embedding_service(config=embedding_config)
 
 
-def create_embedding_dependency(app_slug: str, engine: Optional[Any] = None):
-    """
-    Create a FastAPI dependency function for embedding service.
-
-    This creates a dependency function that can be used with Depends()
-    to inject the embedding service into route handlers.
-
-    Args:
-        app_slug: App slug
-        engine: MongoDBEngine instance (optional)
-
-    Returns:
-        Dependency function that returns EmbeddingService or raises HTTPException
-
-    Example:
-        ```python
-        from fastapi import Depends
-        from mdb_engine.embeddings.dependencies import create_embedding_dependency
-
-        embedding_dep = create_embedding_dependency("my_app", engine)
-
-        @app.post("/embed")
-        async def embed_endpoint(embedding_service = Depends(embedding_dep)):
-            embeddings = await embedding_service.embed_chunks(["Hello world"])
-            return {"embeddings": embeddings}
-        ```
-    """
-
-    def _get_embedding_service() -> EmbeddingService:
-        embedding_service = get_embedding_service_for_app(app_slug, engine)
-        if embedding_service is None:
-            if FASTAPI_AVAILABLE:
-                raise HTTPException(
-                    status_code=503,
-                    detail=f"Embedding service not available for app '{app_slug}'. "
-                    "Ensure 'embedding_config.enabled' is true in manifest.json and "
-                    "embedding dependencies are installed.",
-                )
-            else:
-                raise RuntimeError(
-                    f"Embedding service not available for app '{app_slug}'"
-                )
-        return embedding_service
-
-    return _get_embedding_service
-
-
-def get_embedding_service_dependency(app_slug: str):
-    """
-    Get embedding service dependency using global engine.
-
-    This is a convenience function that uses the global engine registry.
-    Set the engine with set_global_engine() during app startup.
-
-    Args:
-        app_slug: App slug
-
-    Returns:
-        Dependency function for FastAPI Depends()
-
-    Example:
-        ```python
-        from fastapi import FastAPI, Depends
-        from mdb_engine.embeddings.dependencies import (
-            set_global_engine, get_embedding_service_dependency
-        )
-
-        app = FastAPI()
-
-        # During startup
-        set_global_engine(engine, app_slug="my_app")
-
-        # In routes
-        @app.post("/embed")
-        async def embed(embedding_service = Depends(get_embedding_service_dependency("my_app"))):
-            return await embedding_service.embed_chunks(["Hello world"])
-        ```
-    """
-    return create_embedding_dependency(app_slug, _global_engine)
-
-
-# Alias for backward compatibility
-get_embedding_service_dep = get_embedding_service_dependency
+__all__ = [
+    "get_embedding_service_for_app",
+]

mdb_engine/embeddings/service.py

@@ -23,7 +23,7 @@ import os
 import time
 from abc import ABC, abstractmethod
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 # Optional OpenAI SDK import
 try:
@@ -59,9 +59,7 @@ class BaseEmbeddingProvider(ABC):
     """
 
     @abstractmethod
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
         Generate embeddings for text.
 
@@ -84,7 +82,7 @@ class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
+        api_key: str | None = None,
         default_model: str = "text-embedding-3-small",
     ):
         """
@@ -108,9 +106,7 @@ class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
         self.client = AsyncOpenAI(api_key=api_key)
         self.default_model = default_model
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """Generate embeddings using OpenAI."""
         model = model or self.default_model
 
@@ -134,7 +130,7 @@ class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
             ConnectionError,
             OSError,
         ) as e:
-            logger.error(f"OpenAI embedding failed: {e}")
+            logger.exception(f"OpenAI embedding failed: {e}")
             raise EmbeddingServiceError(f"OpenAI embedding failed: {str(e)}") from e
 
 
@@ -149,9 +145,9 @@ class AzureOpenAIEmbeddingProvider(BaseEmbeddingProvider):
 
     def __init__(
         self,
-        api_key: Optional[str] = None,
-        endpoint: Optional[str] = None,
-        api_version: Optional[str] = None,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        api_version: str | None = None,
         default_model: str = "text-embedding-3-small",
     ):
         """
@@ -191,9 +187,7 @@ class AzureOpenAIEmbeddingProvider(BaseEmbeddingProvider):
         )
         self.default_model = default_model
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """Generate embeddings using Azure OpenAI."""
         model = model or self.default_model
 
@@ -217,10 +211,8 @@ class AzureOpenAIEmbeddingProvider(BaseEmbeddingProvider):
             ConnectionError,
             OSError,
         ) as e:
-            logger.error(f"Azure OpenAI embedding failed: {e}")
-            raise EmbeddingServiceError(
-                f"Azure OpenAI embedding failed: {str(e)}"
-            ) from e
+            logger.exception(f"Azure OpenAI embedding failed: {e}")
+            raise EmbeddingServiceError(f"Azure OpenAI embedding failed: {str(e)}") from e
 
 
 def _detect_provider_from_env() -> str:
@@ -257,8 +249,8 @@ class EmbeddingProvider:
 
     def __init__(
         self,
-        embedding_provider: Optional[BaseEmbeddingProvider] = None,
-        config: Optional[Dict[str, Any]] = None,
+        embedding_provider: BaseEmbeddingProvider | None = None,
+        config: dict[str, Any] | None = None,
     ):
         """
         Initialize Embedding Provider.
@@ -281,31 +273,21 @@ class EmbeddingProvider:
         else:
             # Auto-detect provider from environment variables
             provider_type = _detect_provider_from_env()
-            default_model = (config or {}).get(
-                "default_embedding_model", "text-embedding-3-small"
-            )
+            default_model = (config or {}).get("default_embedding_model", "text-embedding-3-small")
 
             if provider_type == "azure":
-                self.embedding_provider = AzureOpenAIEmbeddingProvider(
-                    default_model=default_model
-                )
+                self.embedding_provider = AzureOpenAIEmbeddingProvider(default_model=default_model)
                 logger.info(
                     f"Auto-detected Azure OpenAI embedding provider (model: {default_model})"
                 )
             else:
-                self.embedding_provider = OpenAIEmbeddingProvider(
-                    default_model=default_model
-                )
-                logger.info(
-                    f"Auto-detected OpenAI embedding provider (model: {default_model})"
-                )
+                self.embedding_provider = OpenAIEmbeddingProvider(default_model=default_model)
+                logger.info(f"Auto-detected OpenAI embedding provider (model: {default_model})")
 
         # Store config for potential future use
         self.config = config or {}
 
-    async def embed(
-        self, text: Union[str, List[str]], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
         Generates vector embeddings for a string or list of strings.
 
@@ -341,7 +323,7 @@ class EmbeddingProvider:
             return vectors
 
         except (AttributeError, TypeError, ValueError, RuntimeError, KeyError) as e:
-            logger.error(f"EMBED_FAILED: {str(e)}")
+            logger.exception(f"EMBED_FAILED: {str(e)}")
             raise EmbeddingServiceError(f"Embedding failed: {str(e)}") from e
 
 
@@ -371,10 +353,10 @@ class EmbeddingService:
 
     def __init__(
         self,
-        embedding_provider: Optional[EmbeddingProvider] = None,
+        embedding_provider: EmbeddingProvider | None = None,
         default_max_tokens: int = 1000,
         default_tokenizer_model: str = "gpt-3.5-turbo",
-        config: Optional[Dict[str, Any]] = None,
+        config: dict[str, Any] | None = None,
     ):
         """
         Initialize Embedding Service.
@@ -407,9 +389,7 @@ class EmbeddingService:
         self.default_max_tokens = default_max_tokens
         self.default_tokenizer_model = default_tokenizer_model
 
-    def _create_splitter(
-        self, max_tokens: int, tokenizer_model: Optional[str] = None
-    ) -> TextSplitter:
+    def _create_splitter(self, max_tokens: int, tokenizer_model: str | None = None) -> TextSplitter:
         """
         Create a TextSplitter instance.
 
@@ -429,9 +409,9 @@ class EmbeddingService:
     async def chunk_text(
         self,
         text_content: str,
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-    ) -> List[str]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+    ) -> list[str]:
         """
         Split text into semantic chunks.
 
@@ -465,32 +445,39 @@ class EmbeddingService:
             logger.error(f"Error chunking text: {e}", exc_info=True)
             raise EmbeddingServiceError(f"Chunking failed: {str(e)}") from e
 
-    async def embed_chunks(
-        self, chunks: List[str], model: Optional[str] = None
-    ) -> List[List[float]]:
+    async def embed(self, text: str | list[str], model: str | None = None) -> list[list[float]]:
         """
-        Generate embeddings for text chunks.
+        Generate embeddings for text or a list of texts.
 
-        Uses the user-provided embedding provider/function.
+        Natural API that works with both single strings and lists.
 
         Args:
-            chunks: List of text chunks to embed
+            text: A single string or list of strings to embed
             model: Optional model identifier (passed to embedding provider)
 
         Returns:
-            List of embedding vectors (each is a list of floats)
+            List of embedding vectors (each is a list of floats).
+            If input was a single string, returns a list containing one vector.
 
         Example:
-            chunks = ["chunk 1", "chunk 2"]
-            vectors = await service.embed_chunks(chunks, model="text-embedding-3-small")
+            # Single string
+            vectors = await service.embed("Hello world", model="text-embedding-3-small")
+            # vectors is [[0.1, 0.2, ...]]
+
+            # List of strings (batch - more efficient)
+            vectors = await service.embed(["chunk 1", "chunk 2"], model="text-embedding-3-small")
+            # vectors is [[0.1, ...], [0.2, ...]]
         """
+        # Normalize to list
+        chunks = [text] if isinstance(text, str) else text
+
         if not chunks:
             return []
 
         try:
             # Use EmbeddingProvider's embed method (handles retries, logging, etc.)
             vectors = await self.embedding_provider.embed(chunks, model=model)
-            logger.info(f"Generated {len(vectors)} embeddings")
+            logger.info(f"Generated {len(vectors)} embedding(s)")
             return vectors
         except (
             AttributeError,
@@ -503,16 +490,36 @@ class EmbeddingService:
             logger.error(f"Error generating embeddings: {e}", exc_info=True)
             raise EmbeddingServiceError(f"Embedding generation failed: {str(e)}") from e
 
+    async def embed_chunks(self, chunks: list[str], model: str | None = None) -> list[list[float]]:
+        """
+        Generate embeddings for text chunks (list only).
+
+        DEPRECATED: Use embed() instead, which accepts both strings and lists.
+        This method is kept for backward compatibility.
+
+        Args:
+            chunks: List of text chunks to embed
+            model: Optional model identifier (passed to embedding provider)
+
+        Returns:
+            List of embedding vectors (each is a list of floats)
+
+        Example:
+            chunks = ["chunk 1", "chunk 2"]
+            vectors = await service.embed_chunks(chunks, model="text-embedding-3-small")
+        """
+        return await self.embed(chunks, model=model)
+
     async def process_and_store(
         self,
         text_content: str,
         source_id: str,
         collection: Any,  # MongoDB collection (AppDB Collection or Motor collection)
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-        embedding_model: Optional[str] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Any]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+        embedding_model: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
         """
         Process text and store chunks with embeddings in MongoDB.
 
@@ -573,7 +580,7 @@ class EmbeddingService:
             ConnectionError,
             OSError,
         ) as e:
-            logger.error(f"Failed to generate embeddings for {source_id}: {e}")
+            logger.exception(f"Failed to generate embeddings for {source_id}: {e}")
             raise EmbeddingServiceError(f"Embedding generation failed: {str(e)}") from e
 
         if len(vectors) != len(chunks):
@@ -583,7 +590,7 @@ class EmbeddingService:
 
         # Step 3: Prepare documents for insertion
         documents_to_insert = []
-        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors)):
+        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors, strict=False)):
             doc = {
                 "source_id": source_id,
                 "chunk_index": i,
@@ -614,9 +621,7 @@ class EmbeddingService:
                 result = await collection.insert_many(documents_to_insert)
                 inserted_count = len(result.inserted_ids)
 
-            logger.info(
-                f"Successfully inserted {inserted_count} documents for source: {source_id}"
-            )
+            logger.info(f"Successfully inserted {inserted_count} documents for source: {source_id}")
 
             return {
                 "chunks_created": len(chunks),
@@ -632,18 +637,16 @@ class EmbeddingService:
             KeyError,
             ConnectionError,
         ) as e:
-            logger.error(
-                f"Failed to store documents for {source_id}: {e}", exc_info=True
-            )
+            logger.error(f"Failed to store documents for {source_id}: {e}", exc_info=True)
             raise EmbeddingServiceError(f"Storage failed: {str(e)}") from e
 
     async def process_text(
         self,
         text_content: str,
-        max_tokens: Optional[int] = None,
-        tokenizer_model: Optional[str] = None,
-        embedding_model: Optional[str] = None,
-    ) -> List[Dict[str, Any]]:
+        max_tokens: int | None = None,
+        tokenizer_model: str | None = None,
+        embedding_model: str | None = None,
+    ) -> list[dict[str, Any]]:
         """
         Process text and return chunks with embeddings (without storing).
 
@@ -687,7 +690,7 @@ class EmbeddingService:
 
         # Prepare results
         results = []
-        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors)):
+        for i, (chunk_text, vector) in enumerate(zip(chunks, vectors, strict=False)):
             results.append(
                 {
                     "chunk_index": i,
@@ -706,8 +709,8 @@ class EmbeddingService:
 
 # Dependency injection helper
 def get_embedding_service(
-    embedding_provider: Optional[BaseEmbeddingProvider] = None,
-    config: Optional[Dict[str, Any]] = None,
+    embedding_provider: BaseEmbeddingProvider | None = None,
+    config: dict[str, Any] | None = None,
 ) -> EmbeddingService:
     """
     Create EmbeddingService instance with auto-detected or provided embedding provider.