pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/storage/connectors/base.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+
+from pulse_engine.storage.schemas import (
+    ConnectorHealth,
+    Document,
+    SearchQuery,
+    SearchResult,
+    StoreResult,
+)
+
+
+class BaseStorageConnector(ABC):
+    @abstractmethod
+    async def initialize(self) -> None: ...
+
+    @abstractmethod
+    async def store(self, tenant_id: str, documents: list[Document]) -> StoreResult: ...
+
+    @abstractmethod
+    async def retrieve(self, tenant_id: str, doc_id: str) -> Document | None: ...
+
+    @abstractmethod
+    async def search(self, tenant_id: str, query: SearchQuery) -> SearchResult: ...
+
+    @abstractmethod
+    async def delete(self, tenant_id: str, doc_id: str) -> bool: ...
+
+    @abstractmethod
+    async def health_check(self) -> ConnectorHealth: ...
+
+    @abstractmethod
+    async def teardown(self) -> None: ...

pulse_engine/storage/connectors/opensearch.py

@@ -0,0 +1,344 @@
+import time
+from typing import Any
+
+import structlog
+from opensearchpy import ConnectionError as OSConnectionError
+from opensearchpy import NotFoundError as OSNotFoundError
+
+from pulse_engine.core.exceptions import ServiceUnavailableError
+from pulse_engine.services.opensearch import OpenSearchService
+from pulse_engine.storage.connectors.base import BaseStorageConnector
+from pulse_engine.storage.schemas import (
+    ConnectorHealth,
+    Document,
+    DocumentSource,
+    SearchHit,
+    SearchQuery,
+    SearchQueryType,
+    SearchResult,
+    StoreResult,
+)
+
+logger = structlog.get_logger()
+
+
+class OpenSearchConnector(BaseStorageConnector):
+    def __init__(
+        self,
+        opensearch: OpenSearchService,
+        embedding_dimension: int = 768,
+        index_prefix: str = "pulse_kb",
+    ) -> None:
+        self._opensearch = opensearch
+        self._embedding_dimension = embedding_dimension
+        self._index_prefix = index_prefix
+        self._existing_indices: set[str] = set()
+
+    def _index_name(self, tenant_id: str) -> str:
+        return f"{self._index_prefix}_{tenant_id}"
+
+    def _get_mapping(self) -> dict[str, Any]:
+        return {
+            "settings": {
+                "index": {"knn": True},
+            },
+            "mappings": {
+                "properties": {
+                    "doc_id": {"type": "keyword"},
+                    "source": {"type": "keyword"},
+                    "chunk_id": {"type": "keyword"},
+                    "content": {"type": "text"},
+                    "metadata": {"type": "object", "enabled": True},
+                    "metadata_chunk": {"type": "object", "enabled": True},
+                    "content_summary": {"type": "text"},
+                    "content_title": {"type": "text"},
+                    "embedding": {
+                        "type": "knn_vector",
+                        "dimension": self._embedding_dimension,
+                        "method": {
+                            "name": "hnsw",
+                            "space_type": "l2",
+                            "engine": "faiss",
+                        },
+                    },
+                    "content_embeddings": {
+                        "type": "knn_vector",
+                        "dimension": self._embedding_dimension,
+                        "method": {
+                            "name": "hnsw",
+                            "space_type": "l2",
+                            "engine": "faiss",
+                        },
+                    },
+                    "content_summary_embeddings": {
+                        "type": "knn_vector",
+                        "dimension": self._embedding_dimension,
+                        "method": {
+                            "name": "hnsw",
+                            "space_type": "l2",
+                            "engine": "faiss",
+                        },
+                    },
+                    "content_title_embeddings": {
+                        "type": "knn_vector",
+                        "dimension": self._embedding_dimension,
+                        "method": {
+                            "name": "hnsw",
+                            "space_type": "l2",
+                            "engine": "faiss",
+                        },
+                    },
+                }
+            },
+        }
+
+    async def _ensure_index(self, tenant_id: str) -> None:
+        index = self._index_name(tenant_id)
+        if index in self._existing_indices:
+            return
+        try:
+            exists = await self._opensearch.index_exists(index)
+            if not exists:
+                await self._opensearch.create_index(index, self._get_mapping())
+            self._existing_indices.add(index)
+        except OSConnectionError as e:
+            raise ServiceUnavailableError("OpenSearch unavailable") from e
+
+    def _doc_to_body(self, doc: Document) -> dict[str, Any]:
+        body: dict[str, Any] = {
+            "doc_id": doc.doc_id,
+            "source": doc.source.value,
+            "content": doc.content,
+            "metadata": doc.metadata,
+            "metadata_chunk": doc.metadata_chunk,
+        }
+        if doc.chunk_id is not None:
+            body["chunk_id"] = doc.chunk_id
+        if doc.content_summary is not None:
+            body["content_summary"] = doc.content_summary
+        if doc.content_title is not None:
+            body["content_title"] = doc.content_title
+        if doc.embedding is not None:
+            body["embedding"] = doc.embedding
+        if doc.content_embeddings is not None:
+            body["content_embeddings"] = doc.content_embeddings
+        if doc.content_summary_embeddings is not None:
+            body["content_summary_embeddings"] = doc.content_summary_embeddings
+        if doc.content_title_embeddings is not None:
+            body["content_title_embeddings"] = doc.content_title_embeddings
+        return body
+
+    def _body_to_doc(self, source: dict[str, Any]) -> Document:
+        return Document(
+            doc_id=source["doc_id"],
+            source=DocumentSource(source["source"]),
+            chunk_id=source.get("chunk_id"),
+            content=source["content"],
+            metadata=source.get("metadata", {}),
+            metadata_chunk=source.get("metadata_chunk", {}),
+            content_summary=source.get("content_summary"),
+            content_title=source.get("content_title"),
+            embedding=source.get("embedding"),
+            content_embeddings=source.get("content_embeddings"),
+            content_summary_embeddings=source.get("content_summary_embeddings"),
+            content_title_embeddings=source.get("content_title_embeddings"),
+        )
+
+    async def initialize(self) -> None:
+        logger.info("opensearch_connector_initialized")
+
+    async def store(self, tenant_id: str, documents: list[Document]) -> StoreResult:
+        await self._ensure_index(tenant_id)
+        index = self._index_name(tenant_id)
+        stored = 0
+        failed = 0
+        errors: list[str] = []
+
+        if len(documents) == 1:
+            doc = documents[0]
+            try:
+                await self._opensearch.index_document(
+                    index=index, doc_id=doc.doc_id, body=self._doc_to_body(doc)
+                )
+                stored = 1
+            except Exception as e:
+                failed = 1
+                errors.append(str(e))
+        else:
+            bulk_body: list[dict[str, Any]] = []
+            for doc in documents:
+                bulk_body.append({"index": {"_index": index, "_id": doc.doc_id}})
+                bulk_body.append(self._doc_to_body(doc))
+
+            try:
+                result = await self._opensearch.bulk(bulk_body)
+                for item in result.get("items", []):
+                    action = item.get("index", {})
+                    if action.get("status", 500) < 300:
+                        stored += 1
+                    else:
+                        failed += 1
+                        errors.append(
+                            action.get("error", {}).get("reason", "Unknown error")
+                        )
+            except OSConnectionError as e:
+                raise ServiceUnavailableError("OpenSearch unavailable") from e
+
+        return StoreResult(stored_count=stored, failed_count=failed, errors=errors)
+
+    async def retrieve(self, tenant_id: str, doc_id: str) -> Document | None:
+        index = self._index_name(tenant_id)
+        try:
+            result = await self._opensearch.get_document(index=index, doc_id=doc_id)
+            return self._body_to_doc(result["_source"])
+        except OSNotFoundError:
+            return None
+        except OSConnectionError as e:
+            raise ServiceUnavailableError("OpenSearch unavailable") from e
+
+    def _build_text_query(self, query: SearchQuery) -> dict[str, Any]:
+        must: dict[str, Any] = {
+            "multi_match": {
+                "query": query.query_text,
+                "fields": ["content", "content_title", "content_summary"],
+            }
+        }
+        body: dict[str, Any] = {
+            "query": {"bool": {"must": [must]}},
+            "highlight": {"fields": {"content": {}}},
+            "size": query.size,
+            "from": query.from_,
+        }
+        self._apply_filters(body, query.filters)
+        return body
+
+    def _build_semantic_query(self, query: SearchQuery) -> dict[str, Any]:
+        body: dict[str, Any] = {
+            "size": query.size,
+            "from": query.from_,
+            "query": {
+                "knn": {
+                    "embedding": {
+                        "vector": query.vector,
+                        "k": query.size,
+                    }
+                }
+            },
+        }
+        self._apply_filters(body, query.filters)
+        return body
+
+    def _build_hybrid_query(self, query: SearchQuery) -> dict[str, Any]:
+        text_clause: dict[str, Any] = {
+            "multi_match": {
+                "query": query.query_text,
+                "fields": ["content", "content_title", "content_summary"],
+            }
+        }
+        knn_clause: dict[str, Any] = {
+            "knn": {
+                "embedding": {
+                    "vector": query.vector,
+                    "k": query.size,
+                }
+            }
+        }
+        body: dict[str, Any] = {
+            "size": query.size,
+            "from": query.from_,
+            "query": {
+                "bool": {
+                    "should": [text_clause, knn_clause],
+                }
+            },
+            "highlight": {"fields": {"content": {}}},
+        }
+        self._apply_filters(body, query.filters)
+        return body
+
+    @staticmethod
+    def _apply_filters(body: dict[str, Any], filters: dict[str, Any]) -> None:
+        if not filters:
+            return
+        filter_clauses = [{"term": {k: v}} for k, v in filters.items()]
+        query = body.get("query", {})
+        if "bool" not in query:
+            body["query"] = {"bool": {"must": [query], "filter": filter_clauses}}
+        else:
+            query["bool"]["filter"] = filter_clauses
+
+    async def search(self, tenant_id: str, query: SearchQuery) -> SearchResult:
+        index = self._index_name(tenant_id)
+
+        if query.query_type == SearchQueryType.text:
+            body = self._build_text_query(query)
+        elif query.query_type == SearchQueryType.semantic:
+            body = self._build_semantic_query(query)
+        else:
+            body = self._build_hybrid_query(query)
+
+        try:
+            start = time.monotonic()
+            raw = await self._opensearch.search(index=index, query=body)
+            took_ms = int((time.monotonic() - start) * 1000)
+        except OSConnectionError as e:
+            raise ServiceUnavailableError("OpenSearch unavailable") from e
+
+        hits: list[SearchHit] = []
+        for h in raw.get("hits", {}).get("hits", []):
+            src = h["_source"]
+            highlights = []
+            if "highlight" in h:
+                for field_highlights in h["highlight"].values():
+                    highlights.extend(field_highlights)
+
+            hits.append(
+                SearchHit(
+                    doc_id=src["doc_id"],
+                    score=h.get("_score", 0.0),
+                    source=DocumentSource(src["source"]),
+                    content=src["content"],
+                    highlights=highlights,
+                    metadata=src.get("metadata", {}),
+                )
+            )
+
+        total_value = raw.get("hits", {}).get("total", {})
+        if isinstance(total_value, dict):
+            total = total_value.get("value", 0)
+        else:
+            total = total_value
+
+        return SearchResult(hits=hits, total=total, took_ms=took_ms)
+
+    async def delete(self, tenant_id: str, doc_id: str) -> bool:
+        index = self._index_name(tenant_id)
+        try:
+            await self._opensearch.delete_document(index=index, doc_id=doc_id)
+            return True
+        except OSNotFoundError:
+            return False
+        except OSConnectionError as e:
+            raise ServiceUnavailableError("OpenSearch unavailable") from e
+
+    async def health_check(self) -> ConnectorHealth:
+        start = time.monotonic()
+        try:
+            ok = await self._opensearch.ping()
+            latency = (time.monotonic() - start) * 1000
+            return ConnectorHealth(
+                connector="opensearch",
+                status="up" if ok else "down",
+                latency_ms=latency,
+            )
+        except Exception:
+            latency = (time.monotonic() - start) * 1000
+            return ConnectorHealth(
+                connector="opensearch",
+                status="down",
+                latency_ms=latency,
+            )
+
+    async def teardown(self) -> None:
+        self._existing_indices.clear()
+        logger.info("opensearch_connector_teardown")

pulse_engine/storage/knowledge_base.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+import structlog
+
+from pulse_engine.core.exceptions import NotFoundError, ServiceUnavailableError
+from pulse_engine.storage.connectors.athena import AthenaConnector
+from pulse_engine.storage.connectors.opensearch import OpenSearchConnector
+from pulse_engine.storage.schemas import (
+    Document,
+    KBStatsResponse,
+    QueryResult,
+    SearchQuery,
+    SearchResult,
+    StoreResult,
+)
+
+logger = structlog.get_logger()
+
+
+class KnowledgeBaseService:
+    def __init__(
+        self,
+        opensearch_connector: OpenSearchConnector,
+        athena_connector: AthenaConnector | None = None,
+    ) -> None:
+        self._opensearch = opensearch_connector
+        self._athena = athena_connector
+
+    async def store_documents(
+        self, tenant_id: str, documents: list[Document]
+    ) -> StoreResult:
+        return await self._opensearch.store(tenant_id, documents)
+
+    async def retrieve_document(self, tenant_id: str, doc_id: str) -> Document:
+        result = await self._opensearch.retrieve(tenant_id, doc_id)
+        if result is None:
+            raise NotFoundError(f"Document {doc_id} not found")
+        return result
+
+    async def search(self, tenant_id: str, query: SearchQuery) -> SearchResult:
+        return await self._opensearch.search(tenant_id, query)
+
+    async def delete_document(self, tenant_id: str, doc_id: str) -> None:
+        deleted = await self._opensearch.delete(tenant_id, doc_id)
+        if not deleted:
+            raise NotFoundError(f"Document {doc_id} not found")
+
+    async def get_stats(self, tenant_id: str) -> KBStatsResponse:
+        index = self._opensearch._index_name(tenant_id)
+        try:
+            stats = await self._opensearch._opensearch.index_stats(index)
+            total = stats.get("_all", {}).get("primaries", {})
+            doc_count = total.get("docs", {}).get("count", 0)
+            size_bytes = total.get("store", {}).get("size_in_bytes", 0)
+            return KBStatsResponse(
+                document_count=doc_count,
+                index_size_bytes=size_bytes,
+            )
+        except Exception as e:
+            logger.warning("stats_fetch_failed", error=str(e))
+            return KBStatsResponse(document_count=0, index_size_bytes=0)
+
+    async def execute_query(
+        self, tenant_id: str, sql: str, parameters: dict[str, Any] | None = None
+    ) -> QueryResult:
+        if self._athena is None:
+            raise ServiceUnavailableError("Athena connector is not configured")
+        return await self._athena.execute_query(tenant_id, sql, parameters)

pulse_engine/storage/router.py

@@ -0,0 +1,78 @@
+from fastapi import APIRouter, Depends
+from starlette.responses import JSONResponse
+
+from pulse_engine.core.scope import require_scope
+from pulse_engine.dependencies import get_knowledge_base
+from pulse_engine.middleware.tenant import get_tenant_id
+from pulse_engine.storage.knowledge_base import KnowledgeBaseService
+from pulse_engine.storage.schemas import (
+    Document,
+    KBStatsResponse,
+    QueryRequest,
+    QueryResult,
+    SearchQuery,
+    SearchResult,
+    StoreRequest,
+    StoreResult,
+)
+
+router = APIRouter(prefix="/kb", tags=["Knowledge Base"])
+
+
+@router.post(
+    "/documents",
+    response_model=StoreResult,
+    status_code=201,
+    dependencies=[require_scope("kb:write")],
+)
+async def store_documents(
+    body: StoreRequest,
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> StoreResult:
+    return await kb.store_documents(tenant_id, body.documents)
+
+
+@router.get("/documents/{document_id}", response_model=Document)
+async def retrieve_document(
+    document_id: str,
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> Document:
+    return await kb.retrieve_document(tenant_id, document_id)
+
+
+@router.post("/search", response_model=SearchResult)
+async def search_documents(
+    body: SearchQuery,
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> SearchResult:
+    return await kb.search(tenant_id, body)
+
+
+@router.delete("/documents/{document_id}", status_code=204)
+async def delete_document(
+    document_id: str,
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> JSONResponse:
+    await kb.delete_document(tenant_id, document_id)
+    return JSONResponse(status_code=204, content=None)
+
+
+@router.get("/stats", response_model=KBStatsResponse)
+async def get_stats(
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> KBStatsResponse:
+    return await kb.get_stats(tenant_id)
+
+
+@router.post("/query", response_model=QueryResult)
+async def run_query(
+    body: QueryRequest,
+    tenant_id: str = Depends(get_tenant_id),
+    kb: KnowledgeBaseService = Depends(get_knowledge_base),
+) -> QueryResult:
+    return await kb.execute_query(tenant_id, body.sql, body.parameters)

pulse_engine/storage/schemas.py

@@ -0,0 +1,93 @@
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class DocumentSource(str, Enum):
+    speech = "speech"
+    tweet = "tweet"
+    youtube = "youtube"
+    processed = "processed"
+
+
+class Document(BaseModel):
+    doc_id: str
+    source: DocumentSource
+    chunk_id: str | None = None
+    content: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    metadata_chunk: dict[str, Any] = Field(default_factory=dict)
+    content_summary: str | None = None
+    content_title: str | None = None
+    embedding: list[float] | None = None
+    content_embeddings: list[float] | None = None
+    content_summary_embeddings: list[float] | None = None
+    content_title_embeddings: list[float] | None = None
+
+
+class SearchQueryType(str, Enum):
+    text = "text"
+    semantic = "semantic"
+    hybrid = "hybrid"
+
+
+class SearchQuery(BaseModel):
+    query_type: SearchQueryType = SearchQueryType.text
+    query_text: str | None = None
+    vector: list[float] | None = None
+    filters: dict[str, Any] = Field(default_factory=dict)
+    size: int = Field(default=10, ge=1, le=100)
+    from_: int = Field(default=0, ge=0, alias="from")
+
+    model_config = {"populate_by_name": True}
+
+
+class StoreRequest(BaseModel):
+    documents: list[Document]
+
+
+class StoreResult(BaseModel):
+    stored_count: int
+    failed_count: int
+    errors: list[str] = Field(default_factory=list)
+
+
+class SearchHit(BaseModel):
+    doc_id: str
+    score: float
+    source: DocumentSource
+    content: str
+    highlights: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class SearchResult(BaseModel):
+    hits: list[SearchHit]
+    total: int
+    took_ms: int
+
+
+class ConnectorHealth(BaseModel):
+    connector: str
+    status: str = Field(pattern=r"^(up|down|degraded)$")
+    latency_ms: float | None = None
+    details: dict[str, Any] = Field(default_factory=dict)
+
+
+class QueryRequest(BaseModel):
+    sql: str
+    parameters: dict[str, Any] = Field(default_factory=dict)
+
+
+class QueryResult(BaseModel):
+    columns: list[str]
+    rows: list[list[Any]]
+    row_count: int
+    execution_time_ms: int
+
+
+class KBStatsResponse(BaseModel):
+    document_count: int
+    index_size_bytes: int
+    sources: dict[str, int] = Field(default_factory=dict)

pulse_engine/testing/__init__.py

@@ -0,0 +1,13 @@
+"""Pulse Engine testing utilities — import fixtures and mocks for product tests."""
+
+from pulse_engine.testing.mocks import (
+    MockExtractor,
+    MockOrchestratorAdapter,
+    MockStorageConnector,
+)
+
+__all__ = [
+    "MockExtractor",
+    "MockOrchestratorAdapter",
+    "MockStorageConnector",
+]

pulse_engine/testing/fixtures.py

@@ -0,0 +1,50 @@
+"""Reusable pytest fixtures for products that depend on pulse-engine.
+
+Products import these in their ``conftest.py``::
+
+    from pulse_engine.testing.fixtures import *  # noqa: F401, F403
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from pulse_engine.processor.pipeline import ProcessingPipeline
+from pulse_engine.storage.knowledge_base import KnowledgeBaseService
+from pulse_engine.testing.mocks import (
+    MockExtractor,
+    MockOrchestratorAdapter,
+    MockStorageConnector,
+)
+
+
+@pytest.fixture()
+def mock_storage_connector() -> MockStorageConnector:
+    return MockStorageConnector()
+
+
+@pytest.fixture()
+def mock_orchestrator() -> MockOrchestratorAdapter:
+    return MockOrchestratorAdapter()
+
+
+@pytest.fixture()
+def mock_extractor() -> MockExtractor:
+    return MockExtractor()
+
+
+@pytest.fixture()
+def kb_service(mock_storage_connector: MockStorageConnector) -> KnowledgeBaseService:
+    """KnowledgeBaseService backed by an in-memory mock connector."""
+    return KnowledgeBaseService(
+        opensearch_connector=mock_storage_connector,  # type: ignore[arg-type]
+        athena_connector=None,
+    )
+
+
+@pytest.fixture()
+def processing_pipeline(
+    kb_service: KnowledgeBaseService,
+) -> ProcessingPipeline:
+    """Pipeline using all default stages and the mock KB service."""
+    return ProcessingPipeline(kb_service=kb_service)