gnosisllm-knowledge 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

gnosisllm_knowledge/backends/opensearch/memory/config.py

@@ -0,0 +1,127 @@
+"""Memory-specific configuration."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from gnosisllm_knowledge.core.domain.memory import MemoryStrategy
+
+
+@dataclass(frozen=True)
+class MemoryConfig:
+    """Configuration for Agentic Memory.
+
+    Example:
+        ```python
+        # From environment
+        config = MemoryConfig.from_env()
+
+        # Explicit configuration
+        config = MemoryConfig(
+            host="localhost",
+            port=9200,
+            llm_model_id="model-123",
+            embedding_model_id="model-456",
+        )
+        ```
+    """
+
+    # === OpenSearch Connection ===
+    host: str = "localhost"
+    port: int = 9200
+    username: str | None = None
+    password: str | None = None
+    use_ssl: bool = False
+    verify_certs: bool = True
+
+    # === Model IDs (Required for inference) ===
+    llm_model_id: str | None = None
+    embedding_model_id: str | None = None
+
+    # === LLM Response Parsing ===
+    # OpenAI: $.choices[0].message.content
+    # Bedrock Claude: $.output.message.content[0].text
+    llm_result_path: str = "$.choices[0].message.content"
+
+    # === Connector Configuration ===
+    # For setup: OpenAI API key
+    openai_api_key: str | None = None
+    llm_model: str = "gpt-4o"
+    embedding_model: str = "text-embedding-3-small"
+    embedding_dimension: int = 1536
+
+    # === Timeouts ===
+    connect_timeout: float = 5.0
+    inference_timeout: float = 60.0
+
+    # === Default Strategies ===
+    default_strategies: tuple[MemoryStrategy, ...] = (
+        MemoryStrategy.SEMANTIC,
+        MemoryStrategy.USER_PREFERENCE,
+    )
+
+    @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 auth(self) -> tuple[str, str] | None:
+        """Get auth tuple if credentials are configured."""
+        if self.username and self.password:
+            return (self.username, self.password)
+        return None
+
+    @property
+    def is_configured(self) -> bool:
+        """Check if memory is properly configured for inference."""
+        return bool(self.llm_model_id and self.embedding_model_id)
+
+    @classmethod
+    def from_env(cls) -> MemoryConfig:
+        """Create config from environment variables.
+
+        Environment Variables:
+            OPENSEARCH_HOST: OpenSearch host (default: localhost)
+            OPENSEARCH_PORT: OpenSearch port (default: 9200)
+            OPENSEARCH_USERNAME: Username
+            OPENSEARCH_PASSWORD: Password
+            OPENSEARCH_USE_SSL: Use SSL (default: false)
+            OPENSEARCH_VERIFY_CERTS: Verify certs (default: true)
+            OPENSEARCH_LLM_MODEL_ID: LLM model ID for inference
+            OPENSEARCH_EMBEDDING_MODEL_ID: Embedding model ID
+            OPENSEARCH_LLM_RESULT_PATH: JSONPath for LLM response
+            OPENAI_API_KEY: OpenAI API key (for setup)
+            MEMORY_LLM_MODEL: LLM model name (default: gpt-4o)
+            MEMORY_EMBEDDING_MODEL: Embedding model (default: text-embedding-3-small)
+            MEMORY_EMBEDDING_DIMENSION: Embedding dimension (default: 1536)
+            MEMORY_INFERENCE_TIMEOUT: Inference timeout (default: 60)
+            OPENSEARCH_CONNECT_TIMEOUT: Connect timeout (default: 5)
+        """
+        return cls(
+            # Connection
+            host=os.getenv("OPENSEARCH_HOST", "localhost"),
+            port=int(os.getenv("OPENSEARCH_PORT", "9200")),
+            username=os.getenv("OPENSEARCH_USERNAME"),
+            password=os.getenv("OPENSEARCH_PASSWORD"),
+            use_ssl=os.getenv("OPENSEARCH_USE_SSL", "false").lower() == "true",
+            verify_certs=os.getenv("OPENSEARCH_VERIFY_CERTS", "true").lower() == "true",
+            # Model IDs
+            llm_model_id=os.getenv("OPENSEARCH_LLM_MODEL_ID"),
+            embedding_model_id=os.getenv("OPENSEARCH_EMBEDDING_MODEL_ID"),
+            # LLM parsing
+            llm_result_path=os.getenv(
+                "OPENSEARCH_LLM_RESULT_PATH",
+                "$.choices[0].message.content",
+            ),
+            # Connector setup
+            openai_api_key=os.getenv("OPENAI_API_KEY"),
+            llm_model=os.getenv("MEMORY_LLM_MODEL", "gpt-4o"),
+            embedding_model=os.getenv("MEMORY_EMBEDDING_MODEL", "text-embedding-3-small"),
+            embedding_dimension=int(os.getenv("MEMORY_EMBEDDING_DIMENSION", "1536")),
+            # Timeouts
+            connect_timeout=float(os.getenv("OPENSEARCH_CONNECT_TIMEOUT", "5.0")),
+            inference_timeout=float(os.getenv("MEMORY_INFERENCE_TIMEOUT", "60.0")),
+        )

gnosisllm_knowledge/backends/opensearch/memory/setup.py

@@ -0,0 +1,322 @@
+"""Memory setup operations - Connector and Model creation.
+
+CRITICAL: The LLM connector MUST use both system_prompt AND user_prompt.
+If only system_prompt is used, zero facts will be extracted.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import httpx
+
+from gnosisllm_knowledge.core.exceptions import MemoryConfigurationError
+
+if TYPE_CHECKING:
+    from gnosisllm_knowledge.backends.opensearch.memory.config import MemoryConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class SetupStatus:
+    """Result of setup verification."""
+
+    is_ready: bool
+    checks: dict[str, bool]
+
+
+class MemorySetup:
+    """Setup operations for Agentic Memory.
+
+    Creates the required OpenSearch connectors and models for memory to work.
+
+    Example:
+        ```python
+        setup = MemorySetup(config)
+
+        # Create connectors and models
+        llm_model_id = await setup.setup_llm_model()
+        embedding_model_id = await setup.setup_embedding_model()
+
+        # Verify setup
+        status = await setup.verify_setup()
+        if status.is_ready:
+            print("Memory is ready!")
+        ```
+    """
+
+    def __init__(self, config: MemoryConfig) -> None:
+        """Initialize setup.
+
+        Args:
+            config: Memory configuration.
+        """
+        self._config = config
+        self._base_url = config.url
+        self._auth = config.auth
+
+    async def setup_llm_model(self) -> str:
+        """Create OpenAI LLM connector and model for fact extraction.
+
+        CRITICAL: The connector uses BOTH system_prompt AND user_prompt.
+
+        Returns:
+            The deployed LLM model ID.
+
+        Raises:
+            MemoryConfigurationError: If OpenAI API key is not configured.
+        """
+        if not self._config.openai_api_key:
+            raise MemoryConfigurationError(
+                "OpenAI API key required for LLM setup",
+                missing_config=["openai_api_key"],
+            )
+
+        connector_id = await self._create_llm_connector()
+        model_id = await self._register_model(
+            name="OpenAI LLM for Agentic Memory",
+            connector_id=connector_id,
+            function_name="remote",
+        )
+        await self._deploy_model(model_id)
+
+        logger.info(f"LLM model deployed: {model_id}")
+        return model_id
+
+    async def _create_llm_connector(self) -> str:
+        """Create OpenAI chat connector.
+
+        CRITICAL: Uses BOTH system_prompt AND user_prompt parameters.
+        This is required for Agentic Memory fact extraction to work.
+
+        Returns:
+            The connector ID.
+        """
+        # CRITICAL: Both system_prompt AND user_prompt are required
+        request_body = (
+            '{"model": "${parameters.model}", '
+            '"messages": ['
+            '{"role": "system", "content": "${parameters.system_prompt}"}, '
+            '{"role": "user", "content": "${parameters.user_prompt}"}'
+            "]}"
+        )
+
+        connector_body: dict[str, Any] = {
+            "name": "OpenAI Chat Connector for Agentic Memory",
+            "description": "Connector for OpenAI with system_prompt AND user_prompt support",
+            "version": "1",
+            "protocol": "http",
+            "parameters": {
+                "model": self._config.llm_model,
+            },
+            "credential": {
+                "openAI_key": self._config.openai_api_key,
+            },
+            "actions": [
+                {
+                    "action_type": "predict",
+                    "method": "POST",
+                    "url": "https://api.openai.com/v1/chat/completions",
+                    "headers": {
+                        "Authorization": "Bearer ${credential.openAI_key}",
+                        "Content-Type": "application/json",
+                    },
+                    "request_body": request_body,
+                }
+            ],
+        }
+
+        async with httpx.AsyncClient(
+            verify=self._config.verify_certs,
+            timeout=self._config.connect_timeout,
+        ) as client:
+            response = await client.post(
+                f"{self._base_url}/_plugins/_ml/connectors/_create",
+                json=connector_body,
+                auth=self._auth,
+            )
+            response.raise_for_status()
+            result = response.json()
+
+        connector_id = result.get("connector_id")
+        logger.info(f"LLM connector created: {connector_id}")
+        return connector_id
+
+    async def setup_embedding_model(self) -> str:
+        """Create OpenAI embedding connector and model.
+
+        Returns:
+            The deployed embedding model ID.
+
+        Raises:
+            MemoryConfigurationError: If OpenAI API key is not configured.
+        """
+        if not self._config.openai_api_key:
+            raise MemoryConfigurationError(
+                "OpenAI API key required for embedding setup",
+                missing_config=["openai_api_key"],
+            )
+
+        connector_id = await self._create_embedding_connector()
+        model_id = await self._register_model(
+            name="OpenAI Embedding for Agentic Memory",
+            connector_id=connector_id,
+            function_name="remote",
+        )
+        await self._deploy_model(model_id)
+
+        logger.info(f"Embedding model deployed: {model_id}")
+        return model_id
+
+    async def _create_embedding_connector(self) -> str:
+        """Create OpenAI embedding connector.
+
+        Returns:
+            The connector ID.
+        """
+        connector_body: dict[str, Any] = {
+            "name": "OpenAI Embedding Connector",
+            "description": "Connector for OpenAI text-embedding models",
+            "version": "1",
+            "protocol": "http",
+            "parameters": {
+                "model": self._config.embedding_model,
+            },
+            "credential": {
+                "openAI_key": self._config.openai_api_key,
+            },
+            "actions": [
+                {
+                    "action_type": "predict",
+                    "method": "POST",
+                    "url": "https://api.openai.com/v1/embeddings",
+                    "headers": {
+                        "Authorization": "Bearer ${credential.openAI_key}",
+                        "Content-Type": "application/json",
+                    },
+                    "request_body": '{"model": "${parameters.model}", "input": ${parameters.input}}',
+                    "post_process_function": "connector.post_process.openai.embedding",
+                }
+            ],
+        }
+
+        async with httpx.AsyncClient(
+            verify=self._config.verify_certs,
+            timeout=self._config.connect_timeout,
+        ) as client:
+            response = await client.post(
+                f"{self._base_url}/_plugins/_ml/connectors/_create",
+                json=connector_body,
+                auth=self._auth,
+            )
+            response.raise_for_status()
+            result = response.json()
+
+        connector_id = result.get("connector_id")
+        logger.info(f"Embedding connector created: {connector_id}")
+        return connector_id
+
+    async def _register_model(
+        self,
+        name: str,
+        connector_id: str,
+        function_name: str = "remote",
+    ) -> str:
+        """Register a model with OpenSearch.
+
+        Args:
+            name: Model name.
+            connector_id: Connector ID to use.
+            function_name: Model function name (default: remote).
+
+        Returns:
+            The registered model ID.
+        """
+        model_body: dict[str, Any] = {
+            "name": name,
+            "function_name": function_name,
+            "connector_id": connector_id,
+        }
+
+        async with httpx.AsyncClient(
+            verify=self._config.verify_certs,
+            timeout=self._config.connect_timeout,
+        ) as client:
+            response = await client.post(
+                f"{self._base_url}/_plugins/_ml/models/_register",
+                json=model_body,
+                auth=self._auth,
+            )
+            response.raise_for_status()
+            result = response.json()
+
+        return result.get("model_id")
+
+    async def _deploy_model(self, model_id: str) -> None:
+        """Deploy a model.
+
+        Args:
+            model_id: Model ID to deploy.
+        """
+        async with httpx.AsyncClient(
+            verify=self._config.verify_certs,
+            timeout=60.0,  # Deployment can be slow
+        ) as client:
+            response = await client.post(
+                f"{self._base_url}/_plugins/_ml/models/{model_id}/_deploy",
+                auth=self._auth,
+            )
+            response.raise_for_status()
+
+    async def verify_setup(self) -> SetupStatus:
+        """Verify that memory is properly configured.
+
+        Returns:
+            SetupStatus with verification results.
+        """
+        checks: dict[str, bool] = {}
+
+        # Check LLM model
+        if self._config.llm_model_id:
+            llm_ok = await self._check_model(self._config.llm_model_id)
+            checks["llm_model"] = llm_ok
+        else:
+            checks["llm_model"] = False
+
+        # Check embedding model
+        if self._config.embedding_model_id:
+            embed_ok = await self._check_model(self._config.embedding_model_id)
+            checks["embedding_model"] = embed_ok
+        else:
+            checks["embedding_model"] = False
+
+        is_ready = all(checks.values())
+        return SetupStatus(is_ready=is_ready, checks=checks)
+
+    async def _check_model(self, model_id: str) -> bool:
+        """Check if a model is deployed and responding.
+
+        Args:
+            model_id: Model ID to check.
+
+        Returns:
+            True if model is deployed and ready.
+        """
+        try:
+            async with httpx.AsyncClient(
+                verify=self._config.verify_certs,
+                timeout=self._config.connect_timeout,
+            ) as client:
+                response = await client.get(
+                    f"{self._base_url}/_plugins/_ml/models/{model_id}",
+                    auth=self._auth,
+                )
+                if response.status_code == 200:
+                    data = response.json()
+                    return data.get("model_state") == "DEPLOYED"
+        except Exception:
+            pass
+        return False

gnosisllm_knowledge/backends/opensearch/searcher.py

@@ -296,12 +296,25 @@ class OpenSearchKnowledgeSearcher:
             Parsed search results.
         """
         try:
+            logger.debug(f"OpenSearch query: {os_query}")
+            logger.debug(f"Query include_highlights: {query.include_highlights}")
+
             response = await self._client.search(
                 index=index_name,
                 body=os_query,
                 **params,
             )
 
+            # Debug: Log first hit to see if highlights are present
+            hits = response.get("hits", {}).get("hits", [])
+            if hits:
+                first_hit = hits[0]
+                logger.debug(f"First hit keys: {first_hit.keys()}")
+                if "highlight" in first_hit:
+                    logger.debug(f"Highlight data: {first_hit['highlight']}")
+                else:
+                    logger.debug("No 'highlight' key in response hit")
+
             duration_ms = (time.perf_counter() - start_time) * 1000
             return self._parse_response(query, response, duration_ms)
 
@@ -381,3 +394,225 @@ class OpenSearchKnowledgeSearcher:
             search_after_token=search_after_token,
             has_more=len(items) == query.limit and total_hits > query.offset + len(items),
         )
+
+    async def get_collections(self, index_name: str) -> list[dict[str, Any]]:
+        """Get unique collections with document counts via aggregation.
+
+        Args:
+            index_name: Index to query.
+
+        Returns:
+            List of collections with id, name, and document_count.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                logger.debug(f"Index {index_name} does not exist")
+                return []
+
+            # Aggregation query for unique collection_ids with counts
+            # Also aggregate collection_name for display
+            query = {
+                "size": 0,
+                "aggs": {
+                    "collections": {
+                        "terms": {
+                            "field": "collection_id",
+                            "size": 1000,  # Max collections to return
+                        },
+                        "aggs": {
+                            "collection_name": {
+                                "terms": {
+                                    "field": "collection_name",
+                                    "size": 1,
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            response = await self._client.search(index=index_name, body=query)
+
+            collections = []
+            buckets = response.get("aggregations", {}).get("collections", {}).get("buckets", [])
+
+            for bucket in buckets:
+                collection_id = bucket.get("key")
+                if not collection_id:
+                    continue
+
+                doc_count = bucket.get("doc_count", 0)
+
+                # Get collection name from nested agg or use ID as fallback
+                name_buckets = bucket.get("collection_name", {}).get("buckets", [])
+                collection_name = name_buckets[0].get("key") if name_buckets else collection_id
+
+                collections.append({
+                    "id": collection_id,
+                    "name": collection_name or collection_id,
+                    "document_count": doc_count,
+                })
+
+            logger.debug(f"Found {len(collections)} collections in {index_name}")
+            return collections
+
+        except Exception as e:
+            logger.error(f"Failed to get collections from {index_name}: {e}")
+            return []
+
+    async def get_stats(self, index_name: str) -> dict[str, Any]:
+        """Get index statistics.
+
+        Args:
+            index_name: Index to query.
+
+        Returns:
+            Dictionary with document_count and index info.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                return {
+                    "document_count": 0,
+                    "index_name": index_name,
+                    "exists": False,
+                }
+
+            # Get index stats
+            stats = await self._client.indices.stats(index=index_name)
+            index_stats = stats.get("indices", {}).get(index_name, {})
+            primaries = index_stats.get("primaries", {})
+            docs = primaries.get("docs", {})
+
+            return {
+                "document_count": docs.get("count", 0),
+                "index_name": index_name,
+                "exists": True,
+                "size_bytes": primaries.get("store", {}).get("size_in_bytes", 0),
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to get stats for {index_name}: {e}")
+            return {
+                "document_count": 0,
+                "index_name": index_name,
+                "error": str(e),
+            }
+
+    async def list_documents(
+        self,
+        index_name: str,
+        *,
+        account_id: str | None = None,
+        source_id: str | None = None,
+        collection_id: str | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> dict[str, Any]:
+        """List documents with optional filters.
+
+        Args:
+            index_name: Index to query.
+            account_id: Optional account ID filter.
+            source_id: Optional source ID filter.
+            collection_id: Optional collection ID filter.
+            limit: Maximum documents to return.
+            offset: Number of documents to skip.
+
+        Returns:
+            Dictionary with documents, total, limit, offset.
+        """
+        try:
+            # Check if index exists
+            exists = await self._client.indices.exists(index=index_name)
+            if not exists:
+                logger.debug(f"Index {index_name} does not exist")
+                return {
+                    "documents": [],
+                    "total": 0,
+                    "limit": limit,
+                    "offset": offset,
+                }
+
+            # Build filter clauses
+            filters: list[dict[str, Any]] = []
+
+            if account_id:
+                filters.append({"term": {"account_id": account_id}})
+
+            if source_id:
+                filters.append({"term": {"source_id": source_id}})
+
+            if collection_id:
+                filters.append({"term": {"collection_id": collection_id}})
+
+            # Build query with match_all and filters
+            query: dict[str, Any] = {
+                "size": limit,
+                "from": offset,
+                "sort": [
+                    {"created_at": {"order": "desc", "unmapped_type": "date"}},
+                    {"_id": {"order": "asc"}},
+                ],
+                "_source": {
+                    "excludes": ["content_embedding"],
+                },
+            }
+
+            if filters:
+                query["query"] = {
+                    "bool": {
+                        "must": [{"match_all": {}}],
+                        "filter": filters,
+                    }
+                }
+            else:
+                query["query"] = {"match_all": {}}
+
+            response = await self._client.search(index=index_name, body=query)
+
+            hits = response.get("hits", {})
+            total = hits.get("total", {})
+            total_hits = total.get("value", 0) if isinstance(total, dict) else total
+
+            documents = []
+            for hit in hits.get("hits", []):
+                source = hit.get("_source", {})
+                doc = {
+                    "id": hit.get("_id", ""),
+                    "title": source.get("title"),
+                    "url": source.get("url"),
+                    "content_preview": (source.get("content", ""))[:200],
+                    "content": source.get("content"),
+                    "chunk_index": source.get("chunk_index"),
+                    "total_chunks": source.get("total_chunks"),
+                    "source_id": source.get("source_id"),
+                    "collection_id": source.get("collection_id"),
+                    "created_at": source.get("created_at"),
+                    "metadata": source.get("metadata"),
+                }
+                documents.append(doc)
+
+            logger.debug(
+                f"Listed {len(documents)} documents from {index_name} "
+                f"(total: {total_hits}, source_id: {source_id})"
+            )
+
+            return {
+                "documents": documents,
+                "total": total_hits,
+                "limit": limit,
+                "offset": offset,
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to list documents from {index_name}: {e}")
+            return {
+                "documents": [],
+                "total": 0,
+                "limit": limit,
+                "offset": offset,
+            }