PyStreamMCP 0.3.0__py3-none-any.whl

pystreammcp/core.py

@@ -0,0 +1,567 @@
+"""Core metadata indexing and discovery engine for StreamMCP."""
+
+import json
+import uuid
+from datetime import datetime, timezone
+from typing import Dict, List, Optional, Set, Tuple
+from enum import Enum
+from dataclasses import dataclass, field, asdict
+
+
+class CapabilityClassification(Enum):
+    """Classification of tool/capability types."""
+
+    Search = "search"
+    Read = "read"
+    Write = "write"
+    Compute = "compute"
+    Transform = "transform"
+    Aggregate = "aggregate"
+    Route = "route"
+    Filter = "filter"
+    Sort = "sort"
+    Join = "join"
+    Unknown = "unknown"
+
+    @staticmethod
+    def from_name(name: str) -> "CapabilityClassification":
+        """Classify from tool name/description."""
+        lower = name.lower()
+        classifications = {
+            ("search", "find", "query", "lookup"): CapabilityClassification.Search,
+            ("read", "get", "fetch", "retrieve"): CapabilityClassification.Read,
+            ("write", "create", "insert", "post"): CapabilityClassification.Write,
+            ("compute", "calculate", "process"): CapabilityClassification.Compute,
+            ("transform", "convert", "map"): CapabilityClassification.Transform,
+            ("aggregate", "summarize", "group"): CapabilityClassification.Aggregate,
+            ("route", "dispatch", "forward"): CapabilityClassification.Route,
+            ("filter", "select", "where"): CapabilityClassification.Filter,
+            ("sort", "order", "rank"): CapabilityClassification.Sort,
+            ("join", "merge", "combine"): CapabilityClassification.Join,
+        }
+        for keywords, classification in classifications.items():
+            if any(kw in lower for kw in keywords):
+                return classification
+        return CapabilityClassification.Unknown
+
+
+class RelationType(Enum):
+    """Types of relationships between tools/systems."""
+
+    Produces = "produces"
+    Consumes = "consumes"
+    Transforms = "transforms"
+    Depends = "depends"
+    Triggers = "triggers"
+    Related = "related"
+    Compatible = "compatible"
+    Incompatible = "incompatible"
+
+
+class JsonType(Enum):
+    """JSON schema types."""
+
+    String = "string"
+    Integer = "integer"
+    Number = "number"
+    Boolean = "boolean"
+    Array = "array"
+    Object = "object"
+    Null = "null"
+    Unknown = "unknown"
+
+    @staticmethod
+    def from_string(s: str) -> "JsonType":
+        """Convert string to JsonType."""
+        mapping = {
+            "string": JsonType.String,
+            "integer": JsonType.Integer,
+            "number": JsonType.Number,
+            "boolean": JsonType.Boolean,
+            "array": JsonType.Array,
+            "object": JsonType.Object,
+            "null": JsonType.Null,
+        }
+        return mapping.get(s.lower(), JsonType.Unknown)
+
+
+@dataclass
+class Metadata:
+    """Tool metadata."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    name: str = ""
+    description: str = ""
+    server_id: str = ""
+    tool_id: str = ""
+    input_schema: dict = field(default_factory=dict)
+    output_schema: dict = field(default_factory=dict)
+    tags: List[str] = field(default_factory=list)
+    created_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+    updated_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
+
+
+@dataclass
+class Capability:
+    """Tool capability."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    metadata_id: str = ""
+    classification: CapabilityClassification = CapabilityClassification.Unknown
+    inputs: List[str] = field(default_factory=list)
+    outputs: List[str] = field(default_factory=list)
+    cost_estimate: float = 0.0
+    latency_estimate_ms: int = 100
+    reliability_score: float = 0.95
+
+
+@dataclass
+class SchemaProperty:
+    """Schema property definition."""
+
+    name: str
+    prop_type: JsonType
+    description: str = ""
+    required: bool = False
+
+
+@dataclass
+class Schema:
+    """Tool schema."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    metadata_id: str = ""
+    base_type: JsonType = JsonType.Object
+    properties: List[SchemaProperty] = field(default_factory=list)
+    required_fields: List[str] = field(default_factory=list)
+    is_array: bool = False
+    extraction_confidence: float = 1.0
+
+
+@dataclass
+class TypeHierarchy:
+    """Type hierarchy definition."""
+
+    parent_type: Optional[str] = None
+    child_types: List[str] = field(default_factory=list)
+    compatible_types: List[str] = field(default_factory=list)
+
+
+@dataclass
+class Relationship:
+    """Relationship between tools/systems."""
+
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    from_id: str = ""
+    to_id: str = ""
+    relation_type: RelationType = RelationType.Related
+    confidence: float = 0.8
+    description: str = ""
+
+
+class MetadataStore:
+    """In-memory metadata storage with search."""
+
+    def __init__(self):
+        self._data: Dict[str, Metadata] = {}
+
+    def insert(self, metadata: Metadata) -> None:
+        """Insert metadata."""
+        self._data[metadata.id] = metadata
+
+    def get(self, id: str) -> Optional[Metadata]:
+        """Get metadata by ID."""
+        return self._data.get(id)
+
+    def get_by_server(self, server_id: str) -> List[Metadata]:
+        """Get all metadata for a server."""
+        return [m for m in self._data.values() if m.server_id == server_id]
+
+    def get_by_tool(self, tool_id: str) -> List[Metadata]:
+        """Get metadata for a tool."""
+        return [m for m in self._data.values() if m.tool_id == tool_id]
+
+    def search_by_name(self, query: str) -> List[Metadata]:
+        """Search metadata by name/description."""
+        query_lower = query.lower()
+        return [
+            m
+            for m in self._data.values()
+            if query_lower in m.name.lower() or query_lower in m.description.lower()
+        ]
+
+    def all(self) -> List[Metadata]:
+        """Get all metadata."""
+        return list(self._data.values())
+
+    def count(self) -> int:
+        """Get count of metadata items."""
+        return len(self._data)
+
+    def clear(self) -> None:
+        """Clear all metadata."""
+        self._data.clear()
+
+
+class CapabilityDiscovery:
+    """Discover and classify tool capabilities."""
+
+    def __init__(self):
+        self._capabilities: Dict[str, Capability] = {}
+
+    def discover_from_metadata(self, metadata: Metadata) -> Capability:
+        """Discover capability from metadata."""
+        classification = self._classify_by_schema(metadata.name, metadata.description)
+        inputs = self._extract_parameters(metadata.input_schema)
+        outputs = self._extract_parameters(metadata.output_schema)
+
+        capability = Capability(
+            metadata_id=metadata.id,
+            classification=classification,
+            inputs=inputs,
+            outputs=outputs,
+        )
+        self._capabilities[capability.id] = capability
+        return capability
+
+    def get_capability(self, id: str) -> Optional[Capability]:
+        """Get capability by ID."""
+        return self._capabilities.get(id)
+
+    def get_by_classification(
+        self, classification: CapabilityClassification
+    ) -> List[Capability]:
+        """Get capabilities by classification."""
+        return [
+            c for c in self._capabilities.values() if c.classification == classification
+        ]
+
+    def get_by_metadata(self, metadata_id: str) -> List[Capability]:
+        """Get capabilities for metadata."""
+        return [
+            c for c in self._capabilities.values() if c.metadata_id == metadata_id
+        ]
+
+    def all(self) -> List[Capability]:
+        """Get all capabilities."""
+        return list(self._capabilities.values())
+
+    def count(self) -> int:
+        """Get count of capabilities."""
+        return len(self._capabilities)
+
+    def _classify_by_schema(self, name: str, description: str) -> CapabilityClassification:
+        """Classify by tool name and description."""
+        combined = f"{name} {description}".lower()
+
+        if "search" in combined or "find" in combined or "query" in combined:
+            return CapabilityClassification.Search
+        elif "read" in combined or "get" in combined or "fetch" in combined:
+            return CapabilityClassification.Read
+        elif "write" in combined or "create" in combined or "insert" in combined:
+            return CapabilityClassification.Write
+        elif "compute" in combined or "calculate" in combined:
+            return CapabilityClassification.Compute
+        elif "transform" in combined or "convert" in combined:
+            return CapabilityClassification.Transform
+        elif "aggregate" in combined or "summarize" in combined:
+            return CapabilityClassification.Aggregate
+        elif "filter" in combined or "select" in combined:
+            return CapabilityClassification.Filter
+        elif "sort" in combined or "order" in combined:
+            return CapabilityClassification.Sort
+        elif "join" in combined or "merge" in combined:
+            return CapabilityClassification.Join
+        return CapabilityClassification.Unknown
+
+    def _extract_parameters(self, schema: dict) -> List[str]:
+        """Extract parameters from schema."""
+        if not isinstance(schema, dict):
+            return []
+        props = schema.get("properties", {})
+        if isinstance(props, dict):
+            return list(props.keys())
+        return []
+
+
+class SchemaExtractor:
+    """Extract and manage schemas from tool definitions."""
+
+    def __init__(self):
+        self._schemas: Dict[str, Schema] = {}
+        self._hierarchies: Dict[str, TypeHierarchy] = {}
+
+    def extract(self, metadata_id: str, schema_json: dict) -> Schema:
+        """Extract schema from JSON."""
+        base_type = self._infer_type(schema_json)
+        properties = self._extract_properties(schema_json)
+        required = self._extract_required_fields(schema_json)
+
+        schema = Schema(
+            metadata_id=metadata_id,
+            base_type=base_type,
+            properties=properties,
+            required_fields=required,
+        )
+        self._schemas[schema.id] = schema
+        return schema
+
+    def extract_input_output(
+        self, metadata_id: str, input_json: dict, output_json: dict
+    ) -> Tuple[Schema, Schema]:
+        """Extract input and output schemas."""
+        input_schema = self.extract(f"{metadata_id}_input", input_json)
+        output_schema = self.extract(f"{metadata_id}_output", output_json)
+        return input_schema, output_schema
+
+    def get_schema(self, id: str) -> Optional[Schema]:
+        """Get schema by ID."""
+        return self._schemas.get(id)
+
+    def get_by_metadata(self, metadata_id: str) -> List[Schema]:
+        """Get schemas for metadata."""
+        return [s for s in self._schemas.values() if metadata_id in s.metadata_id]
+
+    def build_hierarchy(self, type_name: str, parent: Optional[str] = None) -> None:
+        """Build type hierarchy."""
+        hierarchy = TypeHierarchy(parent_type=parent)
+        self._hierarchies[type_name] = hierarchy
+
+    def get_hierarchy(self, type_name: str) -> Optional[TypeHierarchy]:
+        """Get type hierarchy."""
+        return self._hierarchies.get(type_name)
+
+    def all_schemas(self) -> List[Schema]:
+        """Get all schemas."""
+        return list(self._schemas.values())
+
+    def count(self) -> int:
+        """Get count of schemas."""
+        return len(self._schemas)
+
+    def _infer_type(self, schema: dict) -> JsonType:
+        """Infer JSON type from schema."""
+        if isinstance(schema, dict):
+            if "type" in schema:
+                return JsonType.from_string(schema["type"])
+            if "properties" in schema:
+                return JsonType.Object
+            if "items" in schema:
+                return JsonType.Array
+        return JsonType.Unknown
+
+    def _extract_properties(self, schema: dict) -> List[SchemaProperty]:
+        """Extract properties from schema."""
+        if not isinstance(schema, dict):
+            return []
+
+        props_dict = schema.get("properties", {})
+        if not isinstance(props_dict, dict):
+            return []
+
+        properties = []
+        for name, prop in props_dict.items():
+            if isinstance(prop, dict):
+                prop_type = JsonType.from_string(prop.get("type", "unknown"))
+                description = prop.get("description", "")
+                properties.append(
+                    SchemaProperty(name=name, prop_type=prop_type, description=description)
+                )
+        return properties
+
+    def _extract_required_fields(self, schema: dict) -> List[str]:
+        """Extract required fields from schema."""
+        if not isinstance(schema, dict):
+            return []
+        required = schema.get("required", [])
+        return required if isinstance(required, list) else []
+
+
+class RelationshipGraph:
+    """Build and query relationship graphs between tools/systems."""
+
+    def __init__(self):
+        self._relationships: Dict[str, Relationship] = {}
+        self._edges: Dict[str, List[str]] = {}
+
+    def add_relationship(self, relationship: Relationship) -> None:
+        """Add relationship to graph."""
+        self._relationships[relationship.id] = relationship
+
+        if relationship.from_id not in self._edges:
+            self._edges[relationship.from_id] = []
+        self._edges[relationship.from_id].append(relationship.to_id)
+
+        in_key = f"_{relationship.to_id}"
+        if in_key not in self._edges:
+            self._edges[in_key] = []
+        self._edges[in_key].append(relationship.from_id)
+
+    def get_relationship(self, id: str) -> Optional[Relationship]:
+        """Get relationship by ID."""
+        return self._relationships.get(id)
+
+    def get_outgoing(self, from_id: str) -> List[Relationship]:
+        """Get outgoing relationships."""
+        if from_id not in self._edges:
+            return []
+
+        to_ids = self._edges[from_id]
+        result = []
+        for rel in self._relationships.values():
+            if rel.from_id == from_id and rel.to_id in to_ids:
+                result.append(rel)
+        return result
+
+    def get_incoming(self, to_id: str) -> List[Relationship]:
+        """Get incoming relationships."""
+        key = f"_{to_id}"
+        if key not in self._edges:
+            return []
+
+        from_ids = self._edges[key]
+        result = []
+        for rel in self._relationships.values():
+            if rel.to_id == to_id and rel.from_id in from_ids:
+                result.append(rel)
+        return result
+
+    def get_by_type(self, relation_type: RelationType) -> List[Relationship]:
+        """Get relationships by type."""
+        return [r for r in self._relationships.values() if r.relation_type == relation_type]
+
+    def find_path(self, from_id: str, to_id: str, max_depth: int = 5) -> Optional[List[str]]:
+        """Find path between two nodes."""
+        visited: Set[str] = set()
+        path: List[str] = []
+        result = self._dfs(from_id, to_id, visited, path, max_depth)
+        return result if result else None
+
+    def all_relationships(self) -> List[Relationship]:
+        """Get all relationships."""
+        return list(self._relationships.values())
+
+    def count(self) -> int:
+        """Get count of relationships."""
+        return len(self._relationships)
+
+    def _dfs(
+        self,
+        current: str,
+        target: str,
+        visited: Set[str],
+        path: List[str],
+        depth: int,
+    ) -> Optional[List[str]]:
+        """Depth-first search for path."""
+        if depth == 0:
+            return None
+
+        if current == target:
+            return path + [target]
+
+        if current in visited:
+            return None
+
+        visited.add(current)
+        path = path + [current]
+
+        for relationship in self.get_outgoing(current):
+            result = self._dfs(relationship.to_id, target, visited, path, depth - 1)
+            if result:
+                return result
+
+        return None
+
+
+class MetadataIndex:
+    """High-level metadata indexing and discovery engine."""
+
+    def __init__(self):
+        self.metadata_store = MetadataStore()
+        self.capability_discovery = CapabilityDiscovery()
+        self.schema_extractor = SchemaExtractor()
+        self.relationship_graph = RelationshipGraph()
+
+    def index_tool(
+        self,
+        name: str,
+        description: str,
+        server_id: str,
+        tool_id: str,
+        input_schema: dict,
+        output_schema: dict,
+    ) -> None:
+        """Index a tool."""
+        metadata = Metadata(
+            name=name,
+            description=description,
+            server_id=server_id,
+            tool_id=tool_id,
+            input_schema=input_schema,
+            output_schema=output_schema,
+        )
+        self.metadata_store.insert(metadata)
+        self.capability_discovery.discover_from_metadata(metadata)
+        self.schema_extractor.extract_input_output(metadata.id, input_schema, output_schema)
+
+    def search_by_name(self, query: str) -> List[Metadata]:
+        """Search tools by name."""
+        return self.metadata_store.search_by_name(query)
+
+    def get_tools_by_server(self, server_id: str) -> List[Metadata]:
+        """Get tools for a server."""
+        return self.metadata_store.get_by_server(server_id)
+
+    def get_capabilities_by_classification(
+        self, classification: CapabilityClassification
+    ) -> List[Capability]:
+        """Get capabilities by classification."""
+        return self.capability_discovery.get_by_classification(classification)
+
+    def get_tool_schema(self, metadata_id: str) -> List[Schema]:
+        """Get schema for a tool."""
+        return self.schema_extractor.get_by_metadata(metadata_id)
+
+    def get_all_metadata(self) -> List[Metadata]:
+        """Get all metadata."""
+        return self.metadata_store.all()
+
+    def stats(self) -> dict:
+        """Get index statistics."""
+        return {
+            "total_metadata": self.metadata_store.count(),
+            "total_capabilities": self.capability_discovery.count(),
+            "total_schemas": self.schema_extractor.count(),
+            "total_relationships": self.relationship_graph.count(),
+        }
+
+    def bulk_index_tools(self, tools: List[dict]) -> dict:
+        """Bulk index tools."""
+        import time
+
+        start = time.time()
+        success = 0
+        errors = 0
+
+        for tool in tools:
+            try:
+                self.index_tool(
+                    name=tool["name"],
+                    description=tool["description"],
+                    server_id=tool["server_id"],
+                    tool_id=tool["tool_id"],
+                    input_schema=tool.get("input_schema", {}),
+                    output_schema=tool.get("output_schema", {}),
+                )
+                success += 1
+            except Exception:
+                errors += 1
+
+        duration_ms = int((time.time() - start) * 1000)
+        return {
+            "total": success + errors,
+            "success": success,
+            "errors": errors,
+            "duration_ms": duration_ms,
+        }