gnosisllm-knowledge 0.2.0__py3-none-any.whl

gnosisllm_knowledge/backends/opensearch/config.py

@@ -0,0 +1,195 @@
+"""OpenSearch backend configuration."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class OpenSearchConfig:
+    """OpenSearch backend configuration.
+
+    Example:
+        ```python
+        # From environment variables
+        config = OpenSearchConfig.from_env()
+
+        # Explicit configuration
+        config = OpenSearchConfig(
+            host="localhost",
+            port=9200,
+            embedding_model="text-embedding-3-small",
+        )
+
+        # AWS OpenSearch Service
+        config = OpenSearchConfig(
+            host="search-domain.us-east-1.es.amazonaws.com",
+            port=443,
+            use_ssl=True,
+            use_aws_sigv4=True,
+            aws_region="us-east-1",
+        )
+        ```
+    """
+
+    # === Connection ===
+    host: str = "localhost"
+    port: int = 9200
+    use_ssl: bool = False
+    verify_certs: bool = True
+    ca_certs: str | None = None
+
+    # Authentication
+    username: str | None = None
+    password: str | None = None
+
+    # AWS OpenSearch Service
+    use_aws_sigv4: bool = False
+    aws_region: str | None = None
+    aws_service: str = "es"  # "es" for OpenSearch, "aoss" for Serverless
+
+    # === Cluster (High Availability) ===
+    nodes: tuple[str, ...] | None = None  # Multiple nodes for HA
+    sniff_on_start: bool = False
+    sniff_on_node_failure: bool = True
+    sniff_timeout: float = 10.0
+    sniffer_timeout: float = 60.0
+
+    # === Embedding ===
+    embedding_model: str = "text-embedding-3-small"
+    embedding_dimension: int = 1536
+    openai_api_key: str | None = None
+    embedding_batch_size: int = 100
+
+    # === Index Settings ===
+    index_prefix: str = "gnosisllm"
+    number_of_shards: int = 5
+    number_of_replicas: int = 1
+    refresh_interval: str = "1s"
+
+    # Pipeline names (auto-generated if not set)
+    ingest_pipeline_name: str | None = None
+    search_pipeline_name: str | None = None
+
+    # === k-NN Settings ===
+    knn_engine: str = "lucene"  # lucene (recommended for OpenSearch 2.9+), faiss
+    knn_space_type: str = "l2"  # l2, cosinesimil, innerproduct
+    knn_algo_param_ef_search: int = 512
+    knn_algo_param_ef_construction: int = 512
+    knn_algo_param_m: int = 16
+
+    # === Neural Search (OpenSearch 2.9+) ===
+    # model_id is the deployed ML model in OpenSearch for embeddings
+    model_id: str | None = None
+    model_group_id: str | None = None
+    embedding_field: str = "content_embedding"  # Field name for embeddings
+
+    # === Agentic Search ===
+    # Agent IDs from 'gnosisllm-knowledge agentic setup'
+    flow_agent_id: str | None = None
+    conversational_agent_id: str | None = None
+    # LLM for agent reasoning (OpenAI model ID)
+    agentic_llm_model: str = "gpt-4o"
+    # Agent execution limits
+    agentic_max_iterations: int = 5
+    agentic_timeout_seconds: int = 60
+    # Conversation memory settings
+    memory_window_size: int = 10  # Messages to keep in context
+
+    # === Timeouts ===
+    connect_timeout: float = 5.0
+    read_timeout: float = 30.0
+    bulk_timeout: float = 120.0
+
+    # === Bulk Indexing ===
+    bulk_batch_size: int = 500
+    bulk_max_concurrent: int = 3
+
+    @property
+    def url(self) -> str:
+        """Get the full OpenSearch URL."""
+        scheme = "https" if self.use_ssl else "http"
+        return f"{scheme}://{self.host}:{self.port}"
+
+    @property
+    def knowledge_index_name(self) -> str:
+        """Get the default knowledge index name."""
+        return f"{self.index_prefix}-knowledge"
+
+    @property
+    def agentic_memory_index_name(self) -> str:
+        """Get the agentic memory index name."""
+        return f"{self.index_prefix}-memory"
+
+    def get_index_name(self, collection_id: str | None = None) -> str:
+        """Get index name for a collection.
+
+        Args:
+            collection_id: Optional collection ID for multi-tenant indexing.
+
+        Returns:
+            Index name pattern.
+        """
+        if collection_id:
+            return f"{self.index_prefix}-{collection_id}"
+        return self.knowledge_index_name
+
+    @classmethod
+    def from_env(cls) -> OpenSearchConfig:
+        """Create config from environment variables.
+
+        Environment variables:
+            OPENSEARCH_HOST: Host (default: localhost)
+            OPENSEARCH_PORT: Port (default: 9200)
+            OPENSEARCH_USE_SSL: Use SSL (default: false)
+            OPENSEARCH_VERIFY_CERTS: Verify certificates (default: true)
+            OPENSEARCH_USERNAME: Username
+            OPENSEARCH_PASSWORD: Password
+            OPENSEARCH_USE_AWS_SIGV4: Use AWS Sig v4 auth (default: false)
+            AWS_REGION: AWS region for Sig v4
+            OPENSEARCH_NODES: Comma-separated list of nodes
+            EMBEDDING_MODEL: OpenAI embedding model
+            EMBEDDING_DIMENSION: Embedding vector dimension
+            OPENAI_API_KEY: OpenAI API key
+            OPENSEARCH_INDEX_PREFIX: Index name prefix
+            OPENSEARCH_SHARDS: Number of shards
+            OPENSEARCH_REPLICAS: Number of replicas
+            OPENSEARCH_FLOW_AGENT_ID: Flow agent ID for agentic search
+            OPENSEARCH_CONVERSATIONAL_AGENT_ID: Conversational agent ID
+            AGENTIC_LLM_MODEL: LLM model for agent reasoning (default: gpt-4o)
+            AGENTIC_MAX_ITERATIONS: Maximum agent iterations (default: 5)
+            AGENTIC_TIMEOUT_SECONDS: Agent execution timeout (default: 60)
+
+        Returns:
+            Configuration from environment.
+        """
+        nodes_str = os.getenv("OPENSEARCH_NODES", "")
+        nodes = tuple(nodes_str.split(",")) if nodes_str else None
+
+        return cls(
+            host=os.getenv("OPENSEARCH_HOST", "localhost"),
+            port=int(os.getenv("OPENSEARCH_PORT", "9200")),
+            use_ssl=os.getenv("OPENSEARCH_USE_SSL", "").lower() == "true",
+            verify_certs=os.getenv("OPENSEARCH_VERIFY_CERTS", "true").lower() == "true",
+            username=os.getenv("OPENSEARCH_USERNAME"),
+            password=os.getenv("OPENSEARCH_PASSWORD"),
+            use_aws_sigv4=os.getenv("OPENSEARCH_USE_AWS_SIGV4", "").lower() == "true",
+            aws_region=os.getenv("AWS_REGION"),
+            nodes=nodes,
+            embedding_model=os.getenv("EMBEDDING_MODEL", "text-embedding-3-small"),
+            embedding_dimension=int(os.getenv("EMBEDDING_DIMENSION", "1536")),
+            openai_api_key=os.getenv("OPENAI_API_KEY"),
+            index_prefix=os.getenv("OPENSEARCH_INDEX_PREFIX", "gnosisllm"),
+            number_of_shards=int(os.getenv("OPENSEARCH_SHARDS", "5")),
+            number_of_replicas=int(os.getenv("OPENSEARCH_REPLICAS", "1")),
+            model_id=os.getenv("OPENSEARCH_MODEL_ID"),
+            ingest_pipeline_name=os.getenv("OPENSEARCH_INGEST_PIPELINE"),
+            search_pipeline_name=os.getenv("OPENSEARCH_SEARCH_PIPELINE"),
+            # Agentic search configuration
+            flow_agent_id=os.getenv("OPENSEARCH_FLOW_AGENT_ID"),
+            conversational_agent_id=os.getenv("OPENSEARCH_CONVERSATIONAL_AGENT_ID"),
+            agentic_llm_model=os.getenv("AGENTIC_LLM_MODEL", "gpt-4o"),
+            agentic_max_iterations=int(os.getenv("AGENTIC_MAX_ITERATIONS", "5")),
+            agentic_timeout_seconds=int(os.getenv("AGENTIC_TIMEOUT_SECONDS", "60")),
+        )

gnosisllm_knowledge/backends/opensearch/indexer.py

@@ -0,0 +1,499 @@
+"""OpenSearch document indexer implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Sequence
+
+# Note: Callable is still needed for on_batch_complete callback type
+
+from gnosisllm_knowledge.backends.opensearch.config import OpenSearchConfig
+from gnosisllm_knowledge.backends.opensearch.mappings import (
+    get_knowledge_index_mappings,
+    get_knowledge_index_settings,
+)
+from gnosisllm_knowledge.core.domain.document import Document
+from gnosisllm_knowledge.core.domain.result import BatchResult, IndexResult
+from gnosisllm_knowledge.core.exceptions import IndexError
+
+if TYPE_CHECKING:
+    from opensearchpy import AsyncOpenSearch
+
+logger = logging.getLogger(__name__)
+
+
+class OpenSearchIndexer:
+    """OpenSearch document indexer.
+
+    Implements the IDocumentIndexer protocol for indexing documents
+    into OpenSearch with k-NN vector embeddings.
+
+    Example:
+        ```python
+        config = OpenSearchConfig.from_env()
+        client = AsyncOpenSearch(hosts=[config.url])
+        indexer = OpenSearchIndexer(client, config)
+
+        # Index a single document
+        result = await indexer.index(document, "my-index")
+
+        # Bulk index documents
+        result = await indexer.bulk_index(documents, "my-index")
+        ```
+    """
+
+    def __init__(
+        self,
+        client: AsyncOpenSearch,
+        config: OpenSearchConfig,
+    ) -> None:
+        """Initialize the indexer.
+
+        Args:
+            client: OpenSearch async client.
+            config: OpenSearch configuration.
+
+        Note:
+            Embeddings are generated automatically by OpenSearch ingest pipeline.
+            Run 'gnosisllm-knowledge setup' to configure the ML model and pipeline.
+        """
+        self._client = client
+        self._config = config
+
+    async def index(
+        self,
+        document: Document,
+        index_name: str,
+        **options: Any,
+    ) -> IndexResult:
+        """Index a single document.
+
+        Args:
+            document: Document to index.
+            index_name: Target index name.
+            **options: Additional options (e.g., refresh).
+
+        Returns:
+            Index result with success/failure info.
+
+        Raises:
+            IndexError: If indexing fails.
+        """
+        try:
+            # Prepare document for indexing
+            # Embeddings are generated by OpenSearch ingest pipeline
+            doc_body = self._prepare_document(document)
+
+            # Index the document
+            refresh = options.get("refresh", False)
+            await self._client.index(
+                index=index_name,
+                id=document.doc_id,
+                body=doc_body,
+                refresh=refresh,
+            )
+
+            return IndexResult(
+                success=True,
+                document_id=document.doc_id,
+                index_name=index_name,
+                indexed_count=1,
+                failed_count=0,
+            )
+        except Exception as e:
+            logger.error(f"Failed to index document {document.doc_id}: {e}")
+            raise IndexError(
+                message=f"Failed to index document: {e}",
+                details={"document_id": document.doc_id},
+                cause=e,
+            ) from e
+
+    async def bulk_index(
+        self,
+        documents: Sequence[Document],
+        index_name: str,
+        **options: Any,
+    ) -> IndexResult:
+        """Index multiple documents efficiently.
+
+        Args:
+            documents: Documents to index.
+            index_name: Target index name.
+            **options: Additional options (batch_size, refresh).
+
+        Returns:
+            Index result with counts.
+        """
+        if not documents:
+            return IndexResult(success=True, index_name=index_name, indexed_count=0, failed_count=0)
+
+        batch_size = options.get("batch_size", self._config.bulk_batch_size)
+        refresh = options.get("refresh", False)
+
+        total_indexed = 0
+        total_failed = 0
+        errors: list[str] = []
+
+        # Process in batches
+        for i in range(0, len(documents), batch_size):
+            batch = documents[i : i + batch_size]
+            batch_result = await self._index_batch(batch, index_name)
+
+            total_indexed += batch_result.indexed_count
+            total_failed += batch_result.failed_count
+            if batch_result.errors:
+                errors.extend(batch_result.errors)
+
+        # Final refresh if requested
+        if refresh:
+            await self._client.indices.refresh(index=index_name)
+
+        return IndexResult(
+            success=total_failed == 0,
+            index_name=index_name,
+            indexed_count=total_indexed,
+            failed_count=total_failed,
+            errors=errors if errors else None,
+        )
+
+    async def bulk_index_streaming(
+        self,
+        documents: AsyncIterator[Document],
+        index_name: str,
+        *,
+        batch_size: int = 500,
+        max_concurrent_batches: int = 3,
+        on_batch_complete: Callable[[BatchResult], None] | None = None,
+    ) -> IndexResult:
+        """Stream-index documents with backpressure handling.
+
+        Args:
+            documents: Async iterator of documents.
+            index_name: Target index name.
+            batch_size: Documents per batch.
+            max_concurrent_batches: Max concurrent batch operations.
+            on_batch_complete: Callback for batch completion.
+
+        Returns:
+            Index result with totals.
+        """
+        semaphore = asyncio.Semaphore(max_concurrent_batches)
+        total_indexed = 0
+        total_failed = 0
+        batch_count = 0
+        errors: list[str] = []
+
+        current_batch: list[Document] = []
+
+        async def process_batch(batch: list[Document], batch_num: int) -> IndexResult:
+            async with semaphore:
+                result = await self._index_batch(batch, index_name)
+
+                if on_batch_complete:
+                    batch_result = BatchResult(
+                        total=len(batch),
+                        succeeded=result.indexed_count,
+                        failed=result.failed_count,
+                        errors=result.errors or [],
+                    )
+                    on_batch_complete(batch_result)
+
+                return result
+
+        tasks: list[asyncio.Task[IndexResult]] = []
+
+        async for doc in documents:
+            current_batch.append(doc)
+
+            if len(current_batch) >= batch_size:
+                batch_count += 1
+                task = asyncio.create_task(
+                    process_batch(current_batch.copy(), batch_count)
+                )
+                tasks.append(task)
+                current_batch = []
+
+        # Process remaining documents
+        if current_batch:
+            batch_count += 1
+            task = asyncio.create_task(
+                process_batch(current_batch.copy(), batch_count)
+            )
+            tasks.append(task)
+
+        # Wait for all tasks
+        results = await asyncio.gather(*tasks)
+        for result in results:
+            total_indexed += result.indexed_count
+            total_failed += result.failed_count
+            if result.errors:
+                errors.extend(result.errors)
+
+        return IndexResult(
+            success=total_failed == 0,
+            index_name=index_name,
+            indexed_count=total_indexed,
+            failed_count=total_failed,
+            errors=errors if errors else None,
+        )
+
+    async def upsert(
+        self,
+        document: Document,
+        index_name: str,
+    ) -> IndexResult:
+        """Upsert (update or insert) a document.
+
+        Args:
+            document: Document to upsert.
+            index_name: Target index name.
+
+        Returns:
+            Index result.
+        """
+        # Embeddings are generated by OpenSearch ingest pipeline
+        doc_body = self._prepare_document(document)
+
+        await self._client.index(
+            index=index_name,
+            id=document.doc_id,
+            body=doc_body,
+            refresh=False,
+        )
+
+        return IndexResult(
+            success=True,
+            document_id=document.doc_id,
+            index_name=index_name,
+            indexed_count=1,
+            failed_count=0,
+        )
+
+    async def delete(
+        self,
+        doc_id: str,
+        index_name: str,
+    ) -> bool:
+        """Delete a document by ID.
+
+        Args:
+            doc_id: Document ID to delete.
+            index_name: Index name.
+
+        Returns:
+            True if deleted, False if not found.
+        """
+        try:
+            await self._client.delete(index=index_name, id=doc_id)
+            return True
+        except Exception as e:
+            if "not_found" in str(e).lower():
+                return False
+            raise IndexError(
+                message=f"Failed to delete document: {e}",
+                details={"document_id": doc_id},
+                cause=e,
+            ) from e
+
+    async def bulk_delete(
+        self,
+        doc_ids: Sequence[str],
+        index_name: str,
+    ) -> int:
+        """Delete multiple documents.
+
+        Args:
+            doc_ids: Document IDs to delete.
+            index_name: Index name.
+
+        Returns:
+            Count of deleted documents.
+        """
+        if not doc_ids:
+            return 0
+
+        actions = [
+            {"delete": {"_index": index_name, "_id": doc_id}} for doc_id in doc_ids
+        ]
+
+        response = await self._client.bulk(body=actions)
+        deleted = sum(
+            1
+            for item in response.get("items", [])
+            if item.get("delete", {}).get("result") == "deleted"
+        )
+
+        return deleted
+
+    async def delete_by_query(
+        self,
+        query: dict[str, Any],
+        index_name: str,
+    ) -> int:
+        """Delete documents matching query.
+
+        Args:
+            query: OpenSearch query dictionary.
+            index_name: Index name.
+
+        Returns:
+            Count of deleted documents.
+        """
+        response = await self._client.delete_by_query(
+            index=index_name,
+            body=query,
+            refresh=True,
+        )
+        return response.get("deleted", 0)
+
+    async def ensure_index(
+        self,
+        index_name: str,
+        **options: Any,
+    ) -> bool:
+        """Ensure index exists with correct mapping.
+
+        Args:
+            index_name: Index name to create.
+            **options: Override settings/mappings.
+
+        Returns:
+            True if created, False if already exists.
+        """
+        exists = await self._client.indices.exists(index=index_name)
+        if exists:
+            return False
+
+        settings = options.get("settings", get_knowledge_index_settings(self._config))
+        mappings = options.get("mappings", get_knowledge_index_mappings(self._config))
+
+        await self._client.indices.create(
+            index=index_name,
+            body={
+                "settings": settings,
+                "mappings": mappings,
+            },
+        )
+
+        logger.info(f"Created index: {index_name}")
+        return True
+
+    async def delete_index(self, index_name: str) -> bool:
+        """Delete an index.
+
+        Args:
+            index_name: Index to delete.
+
+        Returns:
+            True if deleted, False if not found.
+        """
+        try:
+            await self._client.indices.delete(index=index_name)
+            return True
+        except Exception as e:
+            if "index_not_found" in str(e).lower():
+                return False
+            raise
+
+    async def refresh_index(self, index_name: str) -> bool:
+        """Refresh index to make documents searchable.
+
+        Args:
+            index_name: Index to refresh.
+
+        Returns:
+            True on success.
+        """
+        await self._client.indices.refresh(index=index_name)
+        return True
+
+    async def _index_batch(
+        self,
+        documents: Sequence[Document],
+        index_name: str,
+    ) -> IndexResult:
+        """Index a batch of documents.
+
+        Args:
+            documents: Documents to index.
+            index_name: Target index.
+
+        Returns:
+            Batch index result.
+        """
+        # Embeddings are generated by OpenSearch ingest pipeline
+        actions = []
+        for doc in documents:
+            doc_body = self._prepare_document(doc)
+            actions.append({"index": {"_index": index_name, "_id": doc.doc_id}})
+            actions.append(doc_body)
+
+        if not actions:
+            return IndexResult(success=True, index_name=index_name, indexed_count=0, failed_count=0)
+
+        response = await self._client.bulk(body=actions)
+
+        indexed = 0
+        failed = 0
+        errors: list[str] = []
+
+        for item in response.get("items", []):
+            index_result = item.get("index", {})
+            if index_result.get("status") in (200, 201):
+                indexed += 1
+            else:
+                failed += 1
+                error = index_result.get("error", {})
+                errors.append(f"{index_result.get('_id')}: {error.get('reason', 'Unknown error')}")
+
+        return IndexResult(
+            success=failed == 0,
+            index_name=index_name,
+            indexed_count=indexed,
+            failed_count=failed,
+            errors=errors if errors else None,
+        )
+
+    def _prepare_document(self, document: Document) -> dict[str, Any]:
+        """Prepare document for indexing.
+
+        Args:
+            document: Document to prepare.
+
+        Returns:
+            Dictionary for indexing.
+        """
+        # Calculate content hash if not present
+        content_hash = document.content_hash
+        if not content_hash:
+            content_hash = hashlib.sha256(document.content.encode()).hexdigest()
+
+        now = datetime.now(timezone.utc)
+
+        return {
+            "id": document.doc_id,
+            "content": document.content,
+            "url": document.url,
+            "title": document.title,
+            "source": document.source,
+            "account_id": document.account_id,
+            "collection_id": document.collection_id,
+            "source_id": document.source_id,
+            "chunk_index": document.chunk_index,
+            "total_chunks": document.total_chunks,
+            "parent_doc_id": document.parent_doc_id,
+            "quality_score": document.quality_score,
+            "language": document.language,
+            "content_hash": content_hash,
+            "word_count": document.word_count or len(document.content.split()),
+            "status": document.status.value,
+            "pii_detected": document.pii_detected,
+            "pii_redacted": document.pii_redacted,
+            "metadata": document.metadata,
+            "created_at": document.created_at.isoformat() if document.created_at else now.isoformat(),
+            "updated_at": document.updated_at.isoformat() if document.updated_at else None,
+            "indexed_at": now.isoformat(),
+        }