mcp-hangar 0.2.0__py3-none-any.whl

mcp_hangar/infrastructure/knowledge_base/contracts.py

@@ -0,0 +1,202 @@
+"""Knowledge Base abstractions and contracts.
+
+Defines interfaces for knowledge base operations that can be implemented
+by different backends (PostgreSQL, SQLite, etc.).
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Any, Optional
+
+
+class KnowledgeBaseDriver(Enum):
+    """Supported knowledge base drivers."""
+
+    POSTGRES = "postgres"
+    SQLITE = "sqlite"
+    MEMORY = "memory"  # For testing
+
+
+@dataclass
+class KnowledgeBaseConfig:
+    """Knowledge base configuration."""
+
+    enabled: bool = False
+    driver: KnowledgeBaseDriver = KnowledgeBaseDriver.SQLITE
+    dsn: str = ""
+    pool_size: int = 5
+    cache_ttl_s: int = 300
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "KnowledgeBaseConfig":
+        """Create config from dictionary."""
+        if not data.get("enabled", False):
+            return cls(enabled=False)
+
+        dsn = data.get("dsn", "")
+
+        # Auto-detect driver from DSN
+        if dsn.startswith("postgresql://") or dsn.startswith("postgres://"):
+            driver = KnowledgeBaseDriver.POSTGRES
+        elif dsn.startswith("sqlite://") or dsn.endswith(".db"):
+            driver = KnowledgeBaseDriver.SQLITE
+        else:
+            # Explicit driver override
+            driver_str = data.get("driver", "sqlite")
+            driver = KnowledgeBaseDriver(driver_str)
+
+        return cls(
+            enabled=True,
+            driver=driver,
+            dsn=dsn,
+            pool_size=data.get("pool_size", 5),
+            cache_ttl_s=data.get("cache_ttl_s", 300),
+        )
+
+
+@dataclass
+class AuditEntry:
+    """Audit log entry."""
+
+    event_type: str
+    provider: Optional[str] = None
+    tool: Optional[str] = None
+    arguments: Optional[dict] = None
+    result_summary: Optional[str] = None
+    duration_ms: Optional[int] = None
+    success: bool = True
+    error_message: Optional[str] = None
+    correlation_id: Optional[str] = None
+    timestamp: Optional[datetime] = None
+
+
+@dataclass
+class ProviderStateEntry:
+    """Provider state history entry."""
+
+    provider_id: str
+    old_state: Optional[str]
+    new_state: str
+    reason: Optional[str] = None
+    timestamp: Optional[datetime] = None
+
+
+@dataclass
+class MetricEntry:
+    """Provider metric entry."""
+
+    provider_id: str
+    metric_name: str
+    metric_value: float
+    labels: Optional[dict] = None
+    timestamp: Optional[datetime] = None
+
+
+class IKnowledgeBase(ABC):
+    """Abstract interface for knowledge base operations.
+
+    Implementations must be thread-safe and handle their own connection management.
+    """
+
+    @abstractmethod
+    async def initialize(self) -> bool:
+        """Initialize the knowledge base (create tables, run migrations).
+
+        Returns True if successful.
+        """
+        pass
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close connections and cleanup resources."""
+        pass
+
+    @abstractmethod
+    async def is_healthy(self) -> bool:
+        """Check if knowledge base is operational."""
+        pass
+
+    # === Cache Operations ===
+
+    @abstractmethod
+    async def cache_get(self, provider: str, tool: str, arguments: dict) -> Optional[dict]:
+        """Get cached tool result. Returns None if not found or expired."""
+        pass
+
+    @abstractmethod
+    async def cache_set(
+        self,
+        provider: str,
+        tool: str,
+        arguments: dict,
+        result: Any,
+        ttl_s: Optional[int] = None,
+    ) -> bool:
+        """Cache tool result. Returns True if successful."""
+        pass
+
+    @abstractmethod
+    async def cache_invalidate(self, provider: Optional[str] = None, tool: Optional[str] = None) -> int:
+        """Invalidate cache entries. Returns count of invalidated entries."""
+        pass
+
+    @abstractmethod
+    async def cache_cleanup(self) -> int:
+        """Remove expired cache entries. Returns count of removed entries."""
+        pass
+
+    # === Audit Operations ===
+
+    @abstractmethod
+    async def audit_log(self, entry: AuditEntry) -> bool:
+        """Log audit entry. Returns True if successful."""
+        pass
+
+    @abstractmethod
+    async def audit_query(
+        self,
+        provider: Optional[str] = None,
+        tool: Optional[str] = None,
+        success: Optional[bool] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[AuditEntry]:
+        """Query audit log entries."""
+        pass
+
+    @abstractmethod
+    async def audit_stats(self, hours: int = 24) -> dict:
+        """Get audit statistics for last N hours."""
+        pass
+
+    # === Provider State Operations ===
+
+    @abstractmethod
+    async def record_state_change(self, entry: ProviderStateEntry) -> bool:
+        """Record provider state change. Returns True if successful."""
+        pass
+
+    @abstractmethod
+    async def get_state_history(self, provider_id: str, limit: int = 100) -> list[ProviderStateEntry]:
+        """Get provider state history."""
+        pass
+
+    # === Metrics Operations ===
+
+    @abstractmethod
+    async def record_metric(self, entry: MetricEntry) -> bool:
+        """Record provider metric. Returns True if successful."""
+        pass
+
+    @abstractmethod
+    async def get_metrics(
+        self,
+        provider_id: str,
+        metric_name: Optional[str] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[MetricEntry]:
+        """Get provider metrics."""
+        pass

mcp_hangar/infrastructure/knowledge_base/memory.py

@@ -0,0 +1,177 @@
+"""In-memory implementation of IKnowledgeBase for testing."""
+
+from collections import defaultdict
+from datetime import datetime, timedelta, timezone
+import hashlib
+import json
+from typing import Any, Optional
+
+from ...logging_config import get_logger
+from .contracts import AuditEntry, IKnowledgeBase, KnowledgeBaseConfig, MetricEntry, ProviderStateEntry
+
+logger = get_logger(__name__)
+
+
+class MemoryKnowledgeBase(IKnowledgeBase):
+    """In-memory implementation for testing."""
+
+    def __init__(self, config: KnowledgeBaseConfig):
+        self._config = config
+        self._cache: dict[str, tuple[Any, datetime]] = {}  # key -> (result, expires_at)
+        self._audit_log: list[AuditEntry] = []
+        self._state_history: dict[str, list[ProviderStateEntry]] = defaultdict(list)
+        self._metrics: dict[str, list[MetricEntry]] = defaultdict(list)
+        self._initialized = False
+
+    async def initialize(self) -> bool:
+        self._initialized = True
+        logger.info("memory_kb_initialized")
+        return True
+
+    async def close(self) -> None:
+        self._cache.clear()
+        self._audit_log.clear()
+        self._state_history.clear()
+        self._metrics.clear()
+        self._initialized = False
+        logger.info("memory_kb_closed")
+
+    async def is_healthy(self) -> bool:
+        return self._initialized
+
+    def _hash_arguments(self, arguments: dict) -> str:
+        serialized = json.dumps(arguments, sort_keys=True, default=str)
+        return hashlib.sha256(serialized.encode()).hexdigest()[:32]
+
+    def _cache_key(self, provider: str, tool: str, arguments: dict) -> str:
+        return f"{provider}:{tool}:{self._hash_arguments(arguments)}"
+
+    # === Cache Operations ===
+
+    async def cache_get(self, provider: str, tool: str, arguments: dict) -> Optional[dict]:
+        key = self._cache_key(provider, tool, arguments)
+        if key in self._cache:
+            result, expires_at = self._cache[key]
+            if expires_at > datetime.now(timezone.utc):
+                logger.debug("cache_hit", provider=provider, tool=tool)
+                return result
+            else:
+                del self._cache[key]
+        return None
+
+    async def cache_set(
+        self,
+        provider: str,
+        tool: str,
+        arguments: dict,
+        result: Any,
+        ttl_s: Optional[int] = None,
+    ) -> bool:
+        key = self._cache_key(provider, tool, arguments)
+        ttl = ttl_s or self._config.cache_ttl_s
+        expires_at = datetime.now(timezone.utc) + timedelta(seconds=ttl)
+        self._cache[key] = (result, expires_at)
+        return True
+
+    async def cache_invalidate(self, provider: Optional[str] = None, tool: Optional[str] = None) -> int:
+        if not provider and not tool:
+            count = len(self._cache)
+            self._cache.clear()
+            return count
+
+        prefix = f"{provider}:" if provider else ""
+        keys_to_delete = [k for k in self._cache if k.startswith(prefix) and (not tool or f":{tool}:" in k)]
+        for k in keys_to_delete:
+            del self._cache[k]
+        return len(keys_to_delete)
+
+    async def cache_cleanup(self) -> int:
+        now = datetime.now(timezone.utc)
+        expired = [k for k, (_, exp) in self._cache.items() if exp < now]
+        for k in expired:
+            del self._cache[k]
+        logger.info("cache_cleanup", deleted=len(expired))
+        return len(expired)
+
+    # === Audit Operations ===
+
+    async def audit_log(self, entry: AuditEntry) -> bool:
+        entry.timestamp = entry.timestamp or datetime.now(timezone.utc)
+        self._audit_log.append(entry)
+        return True
+
+    async def audit_query(
+        self,
+        provider: Optional[str] = None,
+        tool: Optional[str] = None,
+        success: Optional[bool] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[AuditEntry]:
+        results = []
+        for entry in reversed(self._audit_log):
+            if provider and entry.provider != provider:
+                continue
+            if tool and entry.tool != tool:
+                continue
+            if success is not None and entry.success != success:
+                continue
+            if since and entry.timestamp and entry.timestamp < since:
+                continue
+            results.append(entry)
+            if len(results) >= limit:
+                break
+        return results
+
+    async def audit_stats(self, hours: int = 24) -> dict:
+        since = datetime.now(timezone.utc) - timedelta(hours=hours)
+        recent = [e for e in self._audit_log if e.timestamp and e.timestamp >= since]
+
+        durations = [e.duration_ms for e in recent if e.duration_ms is not None]
+
+        return {
+            "total": len(recent),
+            "success_count": sum(1 for e in recent if e.success),
+            "error_count": sum(1 for e in recent if not e.success),
+            "providers": len(set(e.provider for e in recent if e.provider)),
+            "tools": len(set(e.tool for e in recent if e.tool)),
+            "avg_duration_ms": sum(durations) / len(durations) if durations else None,
+        }
+
+    # === Provider State Operations ===
+
+    async def record_state_change(self, entry: ProviderStateEntry) -> bool:
+        entry.timestamp = entry.timestamp or datetime.now(timezone.utc)
+        self._state_history[entry.provider_id].append(entry)
+        return True
+
+    async def get_state_history(self, provider_id: str, limit: int = 100) -> list[ProviderStateEntry]:
+        history = self._state_history.get(provider_id, [])
+        return list(reversed(history[-limit:]))
+
+    # === Metrics Operations ===
+
+    async def record_metric(self, entry: MetricEntry) -> bool:
+        entry.timestamp = entry.timestamp or datetime.now(timezone.utc)
+        self._metrics[entry.provider_id].append(entry)
+        return True
+
+    async def get_metrics(
+        self,
+        provider_id: str,
+        metric_name: Optional[str] = None,
+        since: Optional[datetime] = None,
+        limit: int = 100,
+    ) -> list[MetricEntry]:
+        metrics = self._metrics.get(provider_id, [])
+
+        results = []
+        for m in reversed(metrics):
+            if metric_name and m.metric_name != metric_name:
+                continue
+            if since and m.timestamp and m.timestamp < since:
+                continue
+            results.append(m)
+            if len(results) >= limit:
+                break
+        return results