ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/services/o365_remote_auth.py

@@ -0,0 +1,235 @@
+"""Utilities to manage remote Office 365 interactive login sessions."""
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from typing import Any, Dict, Optional, List
+from uuid import uuid4
+
+from navconfig.logging import logging
+
+from ..interfaces.o365 import O365Client
+
+
+@dataclass
+class RemoteAuthSession:
+    """Metadata for a remote interactive login session."""
+
+    session_id: str
+    status: str = "pending"
+    login_url: Optional[str] = None
+    device_flow: Optional[Dict[str, Any]] = None
+    created_at: datetime = field(default_factory=datetime.utcnow)
+    updated_at: datetime = field(default_factory=datetime.utcnow)
+    expires_at: Optional[datetime] = None
+    result_summary: Optional[Dict[str, Any]] = None
+    error: Optional[str] = None
+    task: Optional[asyncio.Task] = None
+    client: Optional[O365Client] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        data: Dict[str, Any] = {
+            "session_id": self.session_id,
+            "status": self.status,
+            "created_at": self.created_at.isoformat() + "Z",
+            "updated_at": self.updated_at.isoformat() + "Z",
+        }
+        if self.login_url:
+            data["auth_url"] = self.login_url
+        if self.device_flow:
+            data["device_flow"] = self.device_flow
+        if self.expires_at:
+            data["expires_at"] = self.expires_at.isoformat() + "Z"
+        if self.result_summary:
+            data["result"] = self.result_summary
+        if self.error:
+            data["error"] = self.error
+        return data
+
+
+class RemoteAuthManager:
+    """Manage remote Office 365 interactive login sessions."""
+
+    def __init__(self):
+        self._sessions: Dict[str, RemoteAuthSession] = {}
+        self._lock = asyncio.Lock()
+        self.logger = logging.getLogger("Parrot.O365.RemoteAuth")
+
+    async def start_session(
+        self,
+        *,
+        credentials: Optional[Dict[str, Any]] = None,
+        scopes: Optional[List[str]] = None,
+        redirect_uri: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Start a new remote interactive login session."""
+
+        session_id = str(uuid4())
+        session = RemoteAuthSession(session_id=session_id)
+        async with self._lock:
+            self._sessions[session_id] = session
+
+        client = O365Client(credentials=credentials or {})
+        session.client = client
+        client.processing_credentials()
+        client.set_auth_mode("delegated")
+
+        loop = asyncio.get_running_loop()
+        login_url_future: asyncio.Future = loop.create_future()
+        device_flow_future: asyncio.Future = loop.create_future()
+
+        def _login_callback(url: str) -> bool:
+            def _set_login_url() -> None:
+                session.login_url = url
+                session.updated_at = datetime.utcnow()
+                if not login_url_future.done():
+                    login_url_future.set_result(url)
+
+            loop.call_soon_threadsafe(_set_login_url)
+            return False
+
+        def _device_flow_callback(flow: Dict[str, Any]) -> None:
+            def _set_device_flow() -> None:
+                session.device_flow = {
+                    "user_code": flow.get("user_code"),
+                    "verification_uri": flow.get("verification_uri"),
+                    "verification_uri_complete": flow.get("verification_uri_complete"),
+                    "message": flow.get("message"),
+                    "expires_in": flow.get("expires_in"),
+                    "interval": flow.get("interval"),
+                }
+                if flow.get("expires_in"):
+                    session.expires_at = datetime.utcnow() + timedelta(seconds=int(flow["expires_in"]))
+                session.updated_at = datetime.utcnow()
+                if not device_flow_future.done():
+                    device_flow_future.set_result(session.device_flow)
+
+            loop.call_soon_threadsafe(_set_device_flow)
+
+        async def _run_login() -> None:
+            try:
+                result = await client.interactive_login(
+                    scopes=scopes,
+                    redirect_uri=redirect_uri or "http://localhost",
+                    open_browser=False,
+                    login_callback=_login_callback,
+                    device_flow_callback=_device_flow_callback,
+                )
+                summary = self._summarize_result(result)
+                session.result_summary = summary if summary else None
+                session.status = "authorized"
+                session.updated_at = datetime.utcnow()
+            except Exception as exc:  # pragma: no cover - network/auth errors
+                session.status = "failed"
+                session.error = str(exc)
+                session.updated_at = datetime.utcnow()
+                self.logger.error("Remote login session %s failed: %s", session_id, exc)
+            finally:
+                with contextlib.suppress(Exception):
+                    await client.close()
+                session.client = None
+
+        session.task = asyncio.create_task(_run_login())
+
+        waiters = []
+        if not login_url_future.done():
+            waiters.append(login_url_future)
+        if not device_flow_future.done():
+            waiters.append(device_flow_future)
+
+        if waiters:
+            try:
+                await asyncio.wait(waiters, timeout=5, return_when=asyncio.FIRST_COMPLETED)
+            except Exception:  # pragma: no cover - defensive
+                pass
+
+        return session.to_dict()
+
+    async def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+        if not session:
+            return None
+
+        if session.status == "pending" and session.expires_at and datetime.utcnow() > session.expires_at:
+            await self._expire_session(session)
+
+        return session.to_dict()
+
+    async def cancel_session(self, session_id: str) -> bool:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+        if not session:
+            return False
+
+        await self._finalize_session(session, status="cancelled", error="Session cancelled by user")
+        return True
+
+    async def shutdown(self) -> None:
+        async with self._lock:
+            sessions = list(self._sessions.values())
+
+        for session in sessions:
+            await self._finalize_session(session, status=session.status)
+
+    async def _expire_session(self, session: RemoteAuthSession) -> None:
+        await self._finalize_session(
+            session,
+            status="expired",
+            error="Authorization session expired before completion",
+        )
+
+    async def _finalize_session(
+        self,
+        session: RemoteAuthSession,
+        *,
+        status: str,
+        error: Optional[str] = None,
+    ) -> None:
+        if session.task and not session.task.done():
+            session.task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await session.task
+        if session.client:
+            with contextlib.suppress(Exception):
+                await session.client.close()
+            session.client = None
+        session.status = status
+        session.error = error
+        session.updated_at = datetime.utcnow()
+
+    def _summarize_result(self, result: Dict[str, Any]) -> Dict[str, Any]:
+        summary: Dict[str, Any] = {}
+        if not result:
+            return summary
+
+        if result.get("expires_on"):
+            summary["expires_on"] = result["expires_on"]
+        if result.get("scope"):
+            summary["scope"] = result["scope"]
+        if result.get("token_source"):
+            summary["token_source"] = result["token_source"]
+
+        account = result.get("account")
+        if isinstance(account, dict):
+            summary["account"] = {
+                "home_account_id": account.get("home_account_id"),
+                "environment": account.get("environment"),
+                "username": account.get("username"),
+            }
+
+        claims = result.get("id_token_claims")
+        if isinstance(claims, dict):
+            summary["user"] = {
+                "name": claims.get("name"),
+                "preferred_username": claims.get("preferred_username"),
+                "oid": claims.get("oid"),
+                "tid": claims.get("tid"),
+            }
+
+        return summary
+
+
+__all__ = ["RemoteAuthManager", "RemoteAuthSession"]

parrot/stores/__init__.py

@@ -0,0 +1,7 @@
+from .abstract import AbstractStore
+from .postgres import PgVectorStore
+
+supported_stores = {
+    'postgres': 'PgVectorStore',
+    'kb': 'KnowledgeBaseStore',
+}

parrot/stores/abstract.py

@@ -0,0 +1,352 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union
+import importlib
+from collections.abc import Callable
+from datamodel.parsers.json import JSONContent  # pylint: disable=E0611 # noqa
+from navconfig.logging import logging
+from ..conf import (
+    EMBEDDING_DEFAULT_MODEL
+)
+from ..exceptions import ConfigError  # pylint: disable=E0611
+from ..embeddings import supported_embeddings
+
+
+logging.getLogger(name='datasets').setLevel(logging.WARNING)
+
+class AbstractStore(ABC):
+    """AbstractStore class.
+
+        Base class for all Database Vector Stores.
+    Args:
+        embeddings (str): Embedding name.
+
+    Supported Vector Stores:
+        - Qdrant
+        - Milvus
+        - Faiss
+        - Chroma
+        - PgVector
+    """
+
+    def __init__(
+        self,
+        embedding_model: Union[dict, str] = None,
+        embedding: Union[dict, Callable] = None,
+        **kwargs
+    ):
+        self.client: Callable = None
+        self.vector: Callable = None
+        self._embed_: Callable = None
+        self._connected: bool = False
+        if embedding_model is not None:
+            if isinstance(embedding_model, str):
+                self.embedding_model = {
+                    'model_name': embedding_model,
+                    'model_type': 'huggingface'
+                }
+            elif isinstance(embedding_model, dict):
+                self.embedding_model = embedding_model
+            else:
+                raise ValueError(
+                    "Embedding Model must be a string or a dictionary."
+                )
+        # Use or not connection to a vector database:
+        self._use_database: bool = kwargs.get('use_database', True)
+        # Database Information:
+        self.collection_name: str = kwargs.get('collection_name', 'my_collection')
+        self.dimension: int = kwargs.get("dimension", 384)
+        self._metric_type: str = kwargs.get("metric_type", 'COSINE')
+        self._index_type: str = kwargs.get("index_type", 'IVF_FLAT')
+        self.database: str = kwargs.get('database', '')
+        self.index_name = kwargs.get("index_name", "my_index")
+        if embedding is not None:
+            if isinstance(embedding, str):
+                self.embedding_model = {
+                    'model_name': embedding,
+                    'model_type': 'huggingface'
+                }
+            elif isinstance(embedding, dict):
+                self.embedding_model = embedding
+            else:
+                # is a callable:
+                self.embedding_model = {
+                    'model_name': EMBEDDING_DEFAULT_MODEL,
+                    'model_type': 'huggingface'
+                }
+                self._embed_ = embedding
+        self.logger = logging.getLogger(
+            f"Store.{__name__}"
+        )
+        # Client Connection (if required):
+        self._connection = None
+        # Create the Embedding Model:
+        self._embed_ = self.create_embedding(
+            embedding_model=self.embedding_model
+        )
+        # Track context depth
+        self._context_depth = 0
+        # JSON parser (based on orjson):
+        self._json = JSONContent()
+
+    @property
+    def connected(self) -> bool:
+        return self._connected
+
+    def is_connected(self):
+        return self._connected
+
+    @abstractmethod
+    async def connection(self) -> tuple:
+        pass
+
+    def get_connection(self) -> Any:
+        return self._connection
+
+    def engine(self):
+        return self._connection
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        pass
+
+    # Async Context Manager
+    async def __aenter__(self):
+        if self._use_database:
+            if not self._connection:
+                await self.connection()
+        self._context_depth += 1
+        return self
+
+    async def _free_resources(self):
+        if self._embed_:
+            self._embed_.free()
+        self._embed_ = None
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        # closing Embedding
+        if self._embed_:
+            await self._free_resources()
+        try:
+            # Only disconnect if we're exiting the outermost context
+            if self._context_depth <= 0:
+                await self.disconnect()
+                self._context_depth = 0
+        except RuntimeError:
+            pass
+
+    @abstractmethod
+    def get_vector(self, metric_type: str = None, **kwargs):
+        pass
+
+    def get_vectorstore(self):
+        return self.get_vector()
+
+    @abstractmethod
+    async def similarity_search(
+        self,
+        query: str,
+        collection: Union[str, None] = None,
+        limit: int = 2,
+        similarity_threshold: float = 0.0,
+        search_strategy: str = "auto",
+        metadata_filters: Union[dict, None] = None,
+        **kwargs
+    ) -> list:  # noqa
+        pass
+
+    @abstractmethod
+    async def from_documents(
+        self,
+        documents: List[Any],
+        collection: Union[str, None] = None,
+        **kwargs
+    ) -> Callable:
+        """
+        Create Vector Store from Documents.
+
+        Args:
+            documents (List[Any]): List of Documents.
+            collection (str): Collection Name.
+            kwargs: Additional Arguments.
+
+        Returns:
+            Callable VectorStore.
+        """
+
+    @abstractmethod
+    async def create_collection(self, collection: str) -> None:
+        """
+        Create Collection in Vector Store.
+
+        Args:
+            collection (str): Collection Name.
+
+        Returns:
+            None.
+        """
+        pass
+
+    @abstractmethod
+    async def add_documents(
+        self,
+        documents: List[Any],
+        collection: Union[str, None] = None,
+        **kwargs
+    ) -> None:
+        """
+        Add Documents to Vector Store.
+
+        Args:
+            documents (List[Any]): List of Documents.
+            collection (str): Collection Name.
+            kwargs: Additional Arguments.
+
+        Returns:
+            None.
+        """
+
+    def create_embedding(
+        self,
+
+        embedding_model: dict,
+        **kwargs
+    ):
+        """
+        Create Embedding Model.
+
+        Args:
+            embedding_model (dict): Embedding Model Configuration.
+            kwargs: Additional Arguments.
+
+        Returns:
+            Callable: Embedding Model.
+
+        """
+        model_type = embedding_model.get('model_type', 'huggingface')
+        model_name = embedding_model.get('model', EMBEDDING_DEFAULT_MODEL)
+        if model_type not in supported_embeddings:
+            raise ConfigError(
+                f"Embedding Model Type: {model_type} not supported."
+            )
+        embed_cls = supported_embeddings[model_type]
+        cls_path = f"..embeddings.{model_type}"  # Relative module path
+        try:
+            embed_module = importlib.import_module(
+                cls_path,
+                package=__package__
+            )
+            embed_obj = getattr(embed_module, embed_cls)
+            return embed_obj(
+                model_name=model_name,
+                **kwargs
+            )
+        except ImportError as e:
+            raise ConfigError(
+                f"Error Importing Embedding Model: {model_type}"
+            ) from e
+
+    def get_default_embedding(self):
+        embed_model = {
+            'model_name': EMBEDDING_DEFAULT_MODEL,
+            'model_type': 'huggingface'
+        }
+        return self.create_embedding(
+            embedding_model=embed_model
+        )
+
+    def generate_embedding(self, documents: List[Any]) -> List[Any]:
+        if not self._embed_:
+            self._embed_ = self.get_default_embedding()
+
+        # Using the Embed Model to Generate Embeddings:
+        embeddings = self._embed_.embed_documents(documents)
+        return embeddings
+
+    @abstractmethod
+    async def prepare_embedding_table(
+        self,
+        tablename: str,
+        conn: Any = None,
+        embedding_column: str = 'embedding',
+        document_column: str = 'document',
+        metadata_column: str = 'cmetadata',
+        dimension: int = None,
+        id_column: str = 'id',
+        use_jsonb: bool = True,
+        drop_columns: bool = False,
+        create_all_indexes: bool = True,
+        **kwargs
+    ):
+        """
+        Prepare a Table as an embedding table with advanced features.
+        This method prepares a table with the following columns:
+        - id: unique identifier (String)
+        - embedding: the vector column (Vector(dimension) or JSONB)
+        - document: text column containing the document
+        - collection_id: UUID column for collection identification.
+        - metadata: JSONB column for metadata
+        - Additional columns based on the provided `columns` list
+        - Enhanced indexing strategies for efficient querying
+        - Support for multiple distance strategies (COSINE, L2, IP, etc.)
+        Args:
+        - tablename (str): Name of the table to create.
+        - embedding_column (str): Name of the column for storing embeddings.
+        - document_column (str): Name of the column for storing document text.
+        - metadata_column (str): Name of the column for storing metadata.
+        - dimension (int): Dimension of the embedding vector.
+        - id_column (str): Name of the column for storing unique identifiers.
+        - use_jsonb (bool): Whether to use JSONB for metadata storage.
+        - drop_columns (bool): Whether to drop existing columns.
+        - create_all_indexes (bool): Whether to create all distance strategies.
+    """
+        pass
+
+
+    @abstractmethod
+    async def delete_documents(
+        self,
+        documents: Optional[Any] = None,
+        pk: str = 'source_type',
+        values: Optional[Union[str, List[str]]] = None,
+        table: Optional[str] = None,
+        schema: Optional[str] = None,
+        collection: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Delete Documents from the Vector Store.
+        Args:
+            documents (Optional[Any]): Documents to delete.
+            pk (str): Primary key field.
+            values (Optional[Union[str, List[str]]]): Values to match for deletion.
+            table (Optional[str]): Table name.
+            schema (Optional[str]): Schema name.
+            collection (Optional[str]): Collection name.
+            kwargs: Additional arguments.
+        Returns:
+            int: Number of deleted documents.
+        """
+        pass
+
+
+    @abstractmethod
+    async def delete_documents_by_filter(
+        self,
+        search_filter: Dict[str, Union[str, List[str]]],
+        table: Optional[str] = None,
+        schema: Optional[str] = None,
+        collection: Optional[str] = None,
+        **kwargs
+    ) -> int:
+        """
+        Delete Documents by filter.
+        Args:
+            search_filter (Dict[str, Union[str, List[str]]]): Filter criteria.
+            table (Optional[str]): Table name.
+            schema (Optional[str]): Schema name.
+            collection (Optional[str]): Collection name.
+            kwargs: Additional arguments.
+        Returns:
+            int: Number of deleted documents.
+        """
+        pass