PyStreamMCP 0.3.0__py3-none-any.whl

pystreammcp/persistence.py

@@ -0,0 +1,276 @@
+"""SQLite persistence layer for PyStreamMCP metadata."""
+
+import sqlite3
+import json
+from pathlib import Path
+from typing import List, Optional
+from datetime import datetime, timezone
+from pystreammcp.core import Metadata, Capability, CapabilityClassification
+
+
+class SqlitePersistence:
+    """SQLite-based persistence for metadata indexing."""
+
+    def __init__(self, path: str = ":memory:"):
+        """Initialize SQLite store."""
+        self.path = path
+        self.conn = sqlite3.connect(path, check_same_thread=False, timeout=30.0)
+        self.conn.row_factory = sqlite3.Row
+        self._init_schema()
+
+    def _init_schema(self) -> None:
+        """Initialize database schema."""
+        self.conn.executescript(
+            """
+            PRAGMA journal_mode = WAL;
+            PRAGMA synchronous = NORMAL;
+            PRAGMA temp_store = MEMORY;
+            PRAGMA foreign_keys = ON;
+
+            CREATE TABLE IF NOT EXISTS metadata (
+                id TEXT PRIMARY KEY,
+                name TEXT NOT NULL,
+                description TEXT,
+                server_id TEXT NOT NULL,
+                tool_id TEXT NOT NULL,
+                input_schema TEXT,
+                output_schema TEXT,
+                tags TEXT,
+                created_at TEXT,
+                updated_at TEXT
+            );
+
+            CREATE TABLE IF NOT EXISTS capabilities (
+                id TEXT PRIMARY KEY,
+                metadata_id TEXT NOT NULL,
+                classification TEXT NOT NULL,
+                inputs TEXT,
+                outputs TEXT,
+                cost_estimate REAL,
+                latency_estimate_ms INTEGER,
+                reliability_score REAL,
+                FOREIGN KEY(metadata_id) REFERENCES metadata(id)
+            );
+
+            CREATE TABLE IF NOT EXISTS schemas (
+                id TEXT PRIMARY KEY,
+                metadata_id TEXT NOT NULL,
+                base_type TEXT NOT NULL,
+                properties TEXT,
+                required_fields TEXT,
+                is_array BOOLEAN,
+                extraction_confidence REAL,
+                FOREIGN KEY(metadata_id) REFERENCES metadata(id)
+            );
+
+            CREATE TABLE IF NOT EXISTS search_index (
+                metadata_id TEXT PRIMARY KEY,
+                name_lower TEXT,
+                description_lower TEXT,
+                all_text TEXT,
+                FOREIGN KEY(metadata_id) REFERENCES metadata(id)
+            );
+
+            CREATE INDEX IF NOT EXISTS idx_server_id ON metadata(server_id);
+            CREATE INDEX IF NOT EXISTS idx_tool_id ON metadata(tool_id);
+            CREATE INDEX IF NOT EXISTS idx_capability_classification ON capabilities(classification);
+            CREATE INDEX IF NOT EXISTS idx_capability_metadata ON capabilities(metadata_id);
+            CREATE INDEX IF NOT EXISTS idx_schema_metadata ON schemas(metadata_id);
+        """
+        )
+        self.conn.commit()
+
+    def insert_metadata(self, metadata: Metadata) -> None:
+        """Insert or update metadata."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """INSERT OR REPLACE INTO metadata
+               (id, name, description, server_id, tool_id, input_schema, output_schema, tags, created_at, updated_at)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+            (
+                metadata.id,
+                metadata.name,
+                metadata.description,
+                metadata.server_id,
+                metadata.tool_id,
+                json.dumps(metadata.input_schema),
+                json.dumps(metadata.output_schema),
+                ",".join(metadata.tags),
+                metadata.created_at,
+                metadata.updated_at,
+            ),
+        )
+
+        # Update search index
+        cursor.execute(
+            """INSERT OR REPLACE INTO search_index
+               (metadata_id, name_lower, description_lower, all_text)
+               VALUES (?, ?, ?, ?)""",
+            (
+                metadata.id,
+                metadata.name.lower(),
+                metadata.description.lower(),
+                f"{metadata.name} {metadata.description} {' '.join(metadata.tags)}",
+            ),
+        )
+        self.conn.commit()
+
+    def get_metadata(self, id: str) -> Optional[Metadata]:
+        """Get metadata by ID."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """SELECT id, name, description, server_id, tool_id, input_schema, output_schema, tags, created_at, updated_at
+               FROM metadata WHERE id = ?""",
+            (id,),
+        )
+        row = cursor.fetchone()
+        if row:
+            return Metadata(
+                id=row["id"],
+                name=row["name"],
+                description=row["description"],
+                server_id=row["server_id"],
+                tool_id=row["tool_id"],
+                input_schema=json.loads(row["input_schema"]),
+                output_schema=json.loads(row["output_schema"]),
+                tags=row["tags"].split(",") if row["tags"] else [],
+                created_at=row["created_at"],
+                updated_at=row["updated_at"],
+            )
+        return None
+
+    def search_metadata(self, query: str, limit: int = 100) -> List[Metadata]:
+        """Search metadata by name/description."""
+        cursor = self.conn.cursor()
+        pattern = f"%{query.lower()}%"
+        cursor.execute(
+            """SELECT m.id, m.name, m.description, m.server_id, m.tool_id,
+                      m.input_schema, m.output_schema, m.tags, m.created_at, m.updated_at
+               FROM metadata m
+               INNER JOIN search_index s ON m.id = s.metadata_id
+               WHERE s.name_lower LIKE ? OR s.description_lower LIKE ?
+               LIMIT ?""",
+            (pattern, pattern, limit),
+        )
+
+        results = []
+        for row in cursor.fetchall():
+            results.append(
+                Metadata(
+                    id=row["id"],
+                    name=row["name"],
+                    description=row["description"],
+                    server_id=row["server_id"],
+                    tool_id=row["tool_id"],
+                    input_schema=json.loads(row["input_schema"]),
+                    output_schema=json.loads(row["output_schema"]),
+                    tags=row["tags"].split(",") if row["tags"] else [],
+                    created_at=row["created_at"],
+                    updated_at=row["updated_at"],
+                )
+            )
+        return results
+
+    def get_by_server(self, server_id: str) -> List[Metadata]:
+        """Get all metadata for a server."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """SELECT id, name, description, server_id, tool_id, input_schema, output_schema, tags, created_at, updated_at
+               FROM metadata WHERE server_id = ?""",
+            (server_id,),
+        )
+
+        results = []
+        for row in cursor.fetchall():
+            results.append(
+                Metadata(
+                    id=row["id"],
+                    name=row["name"],
+                    description=row["description"],
+                    server_id=row["server_id"],
+                    tool_id=row["tool_id"],
+                    input_schema=json.loads(row["input_schema"]),
+                    output_schema=json.loads(row["output_schema"]),
+                    tags=row["tags"].split(",") if row["tags"] else [],
+                    created_at=row["created_at"],
+                    updated_at=row["updated_at"],
+                )
+            )
+        return results
+
+    def insert_capability(self, capability: Capability) -> None:
+        """Insert or update capability."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """INSERT OR REPLACE INTO capabilities
+               (id, metadata_id, classification, inputs, outputs, cost_estimate, latency_estimate_ms, reliability_score)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+            (
+                capability.id,
+                capability.metadata_id,
+                capability.classification.value,
+                ",".join(capability.inputs),
+                ",".join(capability.outputs),
+                capability.cost_estimate,
+                capability.latency_estimate_ms,
+                capability.reliability_score,
+            ),
+        )
+        self.conn.commit()
+
+    def get_capabilities_by_classification(
+        self, classification: CapabilityClassification
+    ) -> List[Capability]:
+        """Get capabilities by classification."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """SELECT id, metadata_id, classification, inputs, outputs, cost_estimate, latency_estimate_ms, reliability_score
+               FROM capabilities WHERE classification = ?""",
+            (classification.value,),
+        )
+
+        results = []
+        for row in cursor.fetchall():
+            results.append(
+                Capability(
+                    id=row["id"],
+                    metadata_id=row["metadata_id"],
+                    classification=CapabilityClassification(row["classification"]),
+                    inputs=[s for s in row["inputs"].split(",") if s],
+                    outputs=[s for s in row["outputs"].split(",") if s],
+                    cost_estimate=row["cost_estimate"],
+                    latency_estimate_ms=row["latency_estimate_ms"],
+                    reliability_score=row["reliability_score"],
+                )
+            )
+        return results
+
+    def count_metadata(self) -> int:
+        """Get count of metadata entries."""
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT COUNT(*) as count FROM metadata")
+        return cursor.fetchone()["count"]
+
+    def clear(self) -> None:
+        """Clear all data."""
+        self.conn.executescript(
+            """
+            DELETE FROM capabilities;
+            DELETE FROM schemas;
+            DELETE FROM search_index;
+            DELETE FROM metadata;
+        """
+        )
+        self.conn.commit()
+
+    def close(self) -> None:
+        """Close database connection."""
+        self.conn.close()
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit."""
+        self.close()

pystreammcp/routing.py

@@ -0,0 +1,239 @@
+"""Semantic routing engine for intelligent tool selection."""
+
+from dataclasses import dataclass
+from typing import List
+from pystreammcp.core import Metadata, Capability, CapabilityClassification
+
+
+@dataclass
+class RoutingScore:
+    """Score for routing a query to a tool."""
+
+    metadata_id: str
+    tool_name: str
+    server_id: str
+    relevance_score: float
+    capability_match: float
+    schema_match: float
+    cost_efficiency: float
+    overall_score: float
+
+    def __lt__(self, other):
+        """Sort by overall_score descending."""
+        return self.overall_score > other.overall_score
+
+
+class SemanticRouter:
+    """Route queries to the most relevant tools."""
+
+    @staticmethod
+    def score_tool(
+        query: str, metadata: Metadata, capabilities: List[Capability]
+    ) -> RoutingScore:
+        """Score a tool for a given query."""
+        relevance = SemanticRouter._calculate_relevance(query, metadata)
+        capability_match = SemanticRouter._calculate_capability_match(query, capabilities)
+        schema_match = SemanticRouter._calculate_schema_match(query, metadata)
+        cost_efficiency = SemanticRouter._calculate_cost_efficiency(capabilities)
+
+        overall_score = (
+            (relevance * 0.4)
+            + (capability_match * 0.35)
+            + (schema_match * 0.15)
+            + (cost_efficiency * 0.1)
+        )
+
+        return RoutingScore(
+            metadata_id=metadata.id,
+            tool_name=metadata.name,
+            server_id=metadata.server_id,
+            relevance_score=relevance,
+            capability_match=capability_match,
+            schema_match=schema_match,
+            cost_efficiency=cost_efficiency,
+            overall_score=min(max(overall_score, 0.0), 1.0),
+        )
+
+    @staticmethod
+    def rank_tools(
+        query: str, metadata_list: List[Metadata], capabilities_map: dict
+    ) -> List[RoutingScore]:
+        """Rank all tools for a query."""
+        scores = []
+        for metadata in metadata_list:
+            caps = capabilities_map.get(metadata.id, [])
+            score = SemanticRouter.score_tool(query, metadata, caps)
+            scores.append(score)
+
+        scores.sort()
+        return scores
+
+    @staticmethod
+    def select_best_tools(
+        scores: List[RoutingScore], threshold: float = 0.7, max_count: int = 5
+    ) -> List[RoutingScore]:
+        """Select best tools above threshold."""
+        return [s for s in scores if s.overall_score >= threshold][:max_count]
+
+    @staticmethod
+    def _calculate_relevance(query: str, metadata: Metadata) -> float:
+        """Calculate relevance score."""
+        query_lower = query.lower()
+        name_lower = metadata.name.lower()
+        desc_lower = metadata.description.lower()
+
+        relevance = 0.0
+
+        # Exact name match
+        if query_lower in name_lower:
+            relevance += 0.8
+
+        # Word-level matches in name
+        for word in query_lower.split():
+            if len(word) > 2 and word in name_lower:
+                relevance += 0.3
+
+        # Description matches
+        desc_matches = sum(
+            1 for word in query_lower.split() if len(word) > 2 and word in desc_lower
+        )
+        relevance += desc_matches * 0.15
+
+        # Tag matches
+        tag_matches = sum(
+            1 for tag in metadata.tags if query_lower in tag.lower()
+        )
+        relevance += tag_matches * 0.2
+
+        return min(relevance, 1.0)
+
+    @staticmethod
+    def _calculate_capability_match(
+        query: str, capabilities: List[Capability]
+    ) -> float:
+        """Calculate capability match score."""
+        query_lower = query.lower()
+        match_score = 0.0
+
+        classifications = [
+            (CapabilityClassification.Search, "search"),
+            (CapabilityClassification.Read, "read"),
+            (CapabilityClassification.Write, "write"),
+            (CapabilityClassification.Compute, "compute"),
+            (CapabilityClassification.Transform, "transform"),
+            (CapabilityClassification.Aggregate, "aggregate"),
+        ]
+
+        for cap in capabilities:
+            for class_type, keyword in classifications:
+                if cap.classification == class_type and keyword in query_lower:
+                    match_score += 0.6
+
+            # Check for input/output parameter matches
+            for input_param in cap.inputs:
+                if input_param.lower() in query_lower:
+                    match_score += 0.2
+
+        return min(match_score, 1.0)
+
+    @staticmethod
+    def _calculate_schema_match(query: str, metadata: Metadata) -> float:
+        """Calculate schema match score."""
+        query_lower = query.lower()
+        schema_score = 0.0
+
+        # Check input schema
+        if "properties" in metadata.input_schema:
+            props = metadata.input_schema["properties"]
+            for key in props.keys():
+                if key.lower() in query_lower:
+                    schema_score += 0.25
+
+        # Check output schema
+        if "properties" in metadata.output_schema:
+            props = metadata.output_schema["properties"]
+            for key in props.keys():
+                if key.lower() in query_lower:
+                    schema_score += 0.15
+
+        return min(schema_score, 1.0)
+
+    @staticmethod
+    def _calculate_cost_efficiency(capabilities: List[Capability]) -> float:
+        """Calculate cost efficiency score."""
+        if not capabilities:
+            return 0.5
+
+        avg_cost = sum(c.cost_estimate for c in capabilities) / len(capabilities)
+        avg_reliability = sum(c.reliability_score for c in capabilities) / len(
+            capabilities
+        )
+
+        cost_score = 1.0 - min(avg_cost, 1.0)
+        reliability_score = avg_reliability
+
+        return min(
+            max((cost_score * 0.5 + reliability_score * 0.5), 0.0), 1.0
+        )
+
+
+class RoutingPlan:
+    """Execution plan for a query."""
+
+    def __init__(
+        self,
+        query: str,
+        primary_tools: List[RoutingScore],
+        secondary_tools: List[RoutingScore],
+    ):
+        self.query = query
+        self.primary_tools = primary_tools
+        self.secondary_tools = secondary_tools
+        self.estimated_cost = len(primary_tools) * 0.1
+        self.estimated_latency_ms = min(len(primary_tools) * 50, 1000)
+        self.confidence = (
+            sum(t.overall_score for t in primary_tools) / len(primary_tools)
+            if primary_tools
+            else 0.0
+        )
+
+    def to_dict(self):
+        """Convert to dict."""
+        return {
+            "query": self.query,
+            "primary_tools": [
+                {
+                    "tool_name": t.tool_name,
+                    "server_id": t.server_id,
+                    "overall_score": t.overall_score,
+                    "relevance": t.relevance_score,
+                    "capability_match": t.capability_match,
+                }
+                for t in self.primary_tools
+            ],
+            "secondary_tools": [
+                {
+                    "tool_name": t.tool_name,
+                    "server_id": t.server_id,
+                    "overall_score": t.overall_score,
+                }
+                for t in self.secondary_tools[:3]
+            ],
+            "estimated_cost": self.estimated_cost,
+            "estimated_latency_ms": self.estimated_latency_ms,
+            "confidence": self.confidence,
+        }
+
+
+class RoutingPlanner:
+    """Plan routing for queries."""
+
+    @staticmethod
+    def plan(
+        query: str, all_scores: List[RoutingScore], max_primary: int = 3
+    ) -> RoutingPlan:
+        """Create routing plan."""
+        primary = SemanticRouter.select_best_tools(all_scores, 0.7, max_primary)
+        secondary = SemanticRouter.select_best_tools(all_scores, 0.4, 5)
+
+        return RoutingPlan(query, primary, secondary)