PyStreamMCP 3.2.0__py3-none-any.whl → 3.3.1__py3-none-any.whl

pystreammcp/__init__.py

@@ -7,12 +7,12 @@ Reduces token usage by 60-75% while maintaining quality.
 Includes StatGuardian integration for data quality validation.
 """
 
-__version__ = "3.2.0"
+__version__ = "3.3.1"
 
 from .agent import Agent, QueryResult
 from .query import Query, QueryIntent, QueryConstraints
 from .context import Context, ContextType, ContextWindow
-from .discovery import Discovery, DiscoveredSource, SourceType
+from .discovery import Discovery, DiscoveredSource, SourceType, SourceRegistry, DataSource
 
 # Optimization imports
 try:
@@ -79,6 +79,8 @@ __all__ = [
     "Discovery",
     "DiscoveredSource",
     "SourceType",
+    "SourceRegistry",
+    "DataSource",
     # Legacy optimization (optional, phase 1-2)
     *((["OptimizationStrategy"]) if OptimizationStrategy else []),
     *((["StrategyType"]) if StrategyType else []),

pystreammcp/api.py

@@ -13,6 +13,7 @@ import logging
 
 from .adapters import AdapterRegistry, AdapterConfig, FrameworkType, QueryResult
 from .agent import Agent
+from .discovery import SourceRegistry
 
 
 logger = logging.getLogger(__name__)
@@ -59,6 +60,15 @@ class DiscoveryResponse(BaseModel):
     total_sources: int
 
 
+class SourceRegistration(BaseModel):
+    """Register a data source for discovery."""
+
+    name: str = Field(..., description="Unique source name")
+    description: str = Field(..., description="What this source contains, in plain text")
+    type: str = Field("unknown", description="Source type, e.g. database, api, docs")
+    tags: List[str] = Field(default_factory=list, description="Optional keyword tags")
+
+
 class OptimizationRequest(BaseModel):
     """Query optimization request."""
 
@@ -117,6 +127,7 @@ class PyStreamMCPAPI:
         self.title = title
         self.version = version
         self.agents: Dict[str, Agent] = {}
+        self.registry = SourceRegistry()
         self.app = self._create_app()
 
     def _create_app(self) -> FastAPI:
@@ -260,13 +271,45 @@ class PyStreamMCPAPI:
             }
 
         # Discovery Endpoints
+        @app.post("/sources", response_model=Dict[str, Any])
+        async def register_source(source: SourceRegistration):
+            """Register a data source so /discover can find it."""
+            self.registry.register(
+                name=source.name,
+                description=source.description,
+                type=source.type,
+                tags=source.tags,
+            )
+            return {"status": "success", "name": source.name}
+
+        @app.delete("/sources/{name}")
+        async def unregister_source(name: str):
+            """Remove a registered data source."""
+            removed = self.registry.unregister(name)
+            if not removed:
+                raise HTTPException(status_code=404, detail=f"Source {name} not found")
+            return {"status": "success", "name": name}
+
+        @app.get("/sources")
+        async def list_sources():
+            """List all registered data sources."""
+            return {
+                "status": "success",
+                "sources": [
+                    {"name": s.name, "type": s.type, "tags": s.tags}
+                    for s in self.registry.list_sources()
+                ],
+            }
+
         @app.post("/discover", response_model=DiscoveryResponse)
         async def discover(request: DiscoveryRequest):
-            """Discover relevant data sources."""
-            agent = self._get_agent(request.agent_id)
+            """Discover data sources relevant to `context`.
 
-            # TODO: Implement discovery logic
-            sources = []
+            Ranks *registered* sources (see POST /sources) by real token
+            overlap against `context`. Returns an empty list if nothing is
+            registered or nothing overlaps — never a fabricated result.
+            """
+            sources = self.registry.discover(request.context)
 
             return {
                 "status": "success",

pystreammcp/cli.py

@@ -2,11 +2,9 @@
 
 import json
 import sys
-from typing import Optional
 import click
 
 from pystreammcp import Agent, __version__
-from pystreammcp.api import PyStreamMCPAPI
 from pystreammcp.cli_dashboard import PyStreamMCPDashboard
 
 
@@ -83,6 +81,16 @@ def query(
 @click.option("--reload", is_flag=True)
 def server(host: str, port: int, reload: bool):
     """Start PyStreamMCP API server."""
+    try:
+        from pystreammcp.api import PyStreamMCPAPI
+    except ImportError:
+        click.echo(
+            "Error: the 'server' command requires the 'api' extra. "
+            "Install with: pip install 'pystreammcp[api]'",
+            err=True,
+        )
+        raise SystemExit(1)
+
     click.echo(f"Starting server on {host}:{port}...")
     api = PyStreamMCPAPI()
     api.run(host=host, port=port, reload=reload)
@@ -124,162 +132,5 @@ def dashboard(
         sys.exit(1)
 
 
-# Legacy CLI interface for compatibility
-class CLIInterface:
-    """Command-line interface for PyStreamMCP (legacy)."""
-
-    def __init__(self):
-        self.agent = None
-
-    def create_agent(
-        self, agent_id: str, strategy: str = "balanced", max_tokens: int = 2000
-    ):
-        """Create an agent."""
-        self.agent = Agent(
-            agent_id=agent_id,
-            optimization_strategy=strategy,
-            max_tokens=max_tokens,
-        )
-        return {"status": "success", "agent_id": agent_id}
-
-    def query(
-        self, text: str, intent: str = "retrieve", agent_id: Optional[str] = None
-    ) -> dict:
-        """Execute a query."""
-        if agent_id:
-            self.create_agent(agent_id)
-
-        if not self.agent:
-            return {"status": "error", "message": "No agent created"}
-
-        result = self.agent.query(text)
-
-        return {
-            "status": "success",
-            "query_id": result.query_id,
-            "query_text": text,
-            "baseline_tokens": result.baseline_tokens,
-            "optimized_tokens": result.optimized_tokens,
-            "cost_reduction_percent": result.cost_reduction_percent,
-            "cost_saved": (result.baseline_tokens - result.optimized_tokens) * 0.00001,
-            "execution_time_ms": result.execution_time_ms,
-            "meets_target": 60 <= result.cost_reduction_percent <= 75,
-        }
-
-    def get_metrics(self) -> dict:
-        """Get agent metrics."""
-        if not self.agent:
-            return {"status": "error", "message": "No agent created"}
-
-        metrics = self.agent.get_metrics()
-        return {
-            "status": "success",
-            "metrics": metrics,
-        }
-
-
-def main():
-    """Main CLI entry point."""
-    cli = CLIInterface()
-
-    if len(sys.argv) < 2:
-        print_help()
-        sys.exit(1)
-
-    command = sys.argv[1]
-
-    try:
-        if command == "query":
-            if len(sys.argv) < 3:
-                print(json.dumps({"error": "Missing query text"}))
-                sys.exit(1)
-
-            text = sys.argv[2]
-            intent = sys.argv[3] if len(sys.argv) > 3 else "retrieve"
-            agent_id = sys.argv[4] if len(sys.argv) > 4 else "cli_agent"
-
-            result = cli.query(text, intent=intent, agent_id=agent_id)
-            print(json.dumps(result))
-
-        elif command == "create-agent":
-            if len(sys.argv) < 3:
-                print(json.dumps({"error": "Missing agent_id"}))
-                sys.exit(1)
-
-            agent_id = sys.argv[2]
-            strategy = sys.argv[3] if len(sys.argv) > 3 else "balanced"
-            max_tokens = int(sys.argv[4]) if len(sys.argv) > 4 else 2000
-
-            result = cli.create_agent(agent_id, strategy, max_tokens)
-            print(json.dumps(result))
-
-        elif command == "metrics":
-            result = cli.get_metrics()
-            print(json.dumps(result))
-
-        elif command == "help":
-            print_help()
-
-        else:
-            print(json.dumps({"error": f"Unknown command: {command}"}))
-            sys.exit(1)
-
-    except Exception as e:
-        print(json.dumps({"error": str(e), "status": "error"}))
-        sys.exit(1)
-
-
-def print_help():
-    """Print help message."""
-    help_text = """
-PyStreamMCP CLI - Intelligence Layer for Workflows
-
-USAGE:
-    pystreammcp <command> [options]
-
-COMMANDS:
-    query <text> [intent] [agent_id]
-        Execute a query with optimization
-        - text: The query text (required)
-        - intent: Query type (retrieve|discover|aggregate) (default: retrieve)
-        - agent_id: Agent identifier (default: cli_agent)
-
-        Example:
-            pystreammcp query "Top 10 customers by LTV" retrieve my_agent
-
-    create-agent <agent_id> [strategy] [max_tokens]
-        Create an agent with specific configuration
-        - agent_id: Agent identifier (required)
-        - strategy: Optimization strategy (balanced|token_efficient|quality_first) (default: balanced)
-        - max_tokens: Token budget (default: 2000)
-
-        Example:
-            pystreammcp create-agent workflow_agent token_efficient 1500
-
-    metrics
-        Get current agent metrics
-
-        Example:
-            pystreammcp metrics
-
-    help
-        Show this help message
-
-OUTPUT FORMAT:
-    All commands return JSON output for easy parsing in workflows
-
-EXAMPLES:
-    # Query with cost tracking
-    pystreammcp query "Find churn risk customers"
-
-    # Efficient query for token-constrained environments
-    pystreammcp query "List segments" discover automation_agent
-
-    # Get metrics for reporting
-    pystreammcp metrics | jq '.metrics.total_cost_saved'
-"""
-    print(help_text)
-
-
 if __name__ == "__main__":
-    main()
+    cli()

pystreammcp/discovery.py

@@ -1,9 +1,33 @@
-"""Discovery types for PyStreamMCP."""
-
-from typing import Optional, List
-from enum import Enum
+"""Discovery types and logic for PyStreamMCP.
+
+`SourceRegistry` is real discovery logic: register data sources (name,
+description, type, tags), then find the ones most relevant to a query
+context using word-overlap relevance scoring. This is deliberately a
+simple, honest heuristic — Jaccard token overlap between the query context
+and each source's registered description/tags/name — not a semantic or
+embedding-based search. Sources with zero token overlap are excluded from
+results rather than padded in with a fabricated score, so an empty
+registry or a context that matches nothing returns no sources, not
+placeholder data.
+
+`Discovery` / `DiscoveredSource` / `SourceType` are the result types
+`SourceRegistry.discover()` returns results as, so callers get real,
+structured discovery results rather than a bare dict.
+"""
+
+from __future__ import annotations
+
+import re
 from dataclasses import dataclass, field
 from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Optional, Set
+
+_WORD_RE = re.compile(r"[a-z0-9]+")
+
+
+def _tokenize(text: str) -> Set[str]:
+    return set(_WORD_RE.findall(text.lower()))
 
 
 class SourceType(str, Enum):
@@ -14,11 +38,12 @@ class SourceType(str, Enum):
     CACHE = "cache"
     EXTERNAL = "external"
     COMPUTED = "computed"
+    UNKNOWN = "unknown"
 
 
 @dataclass
 class DiscoveredSource:
-    """A discovered data source."""
+    """A discovered data source, returned by SourceRegistry.discover()."""
 
     name: str
     source_type: SourceType
@@ -80,3 +105,112 @@ class Discovery:
     def high_relevance_sources(self, threshold: float = 0.8) -> List[DiscoveredSource]:
         """Get high-relevance sources."""
         return [s for s in self.discovered_sources if s.relevance_score >= threshold]
+
+
+@dataclass
+class DataSource:
+    """A data source registered with a SourceRegistry."""
+
+    name: str
+    description: str
+    type: str = "unknown"
+    tags: List[str] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    estimated_tokens: int = 0
+
+    def _text(self) -> str:
+        return " ".join([self.name, self.description, *self.tags])
+
+
+class SourceRegistry:
+    """In-memory registry of data sources, with real relevance scoring."""
+
+    def __init__(self) -> None:
+        self._sources: Dict[str, DataSource] = {}
+
+    def register(
+        self,
+        name: str,
+        description: str,
+        type: str = "unknown",
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        estimated_tokens: int = 0,
+    ) -> DataSource:
+        """Register (or replace) a data source by name."""
+        source = DataSource(
+            name=name,
+            description=description,
+            type=type,
+            tags=list(tags) if tags else [],
+            metadata=dict(metadata) if metadata else {},
+            estimated_tokens=estimated_tokens,
+        )
+        self._sources[name] = source
+        return source
+
+    def unregister(self, name: str) -> bool:
+        """Remove a registered source. Returns True if it existed."""
+        return self._sources.pop(name, None) is not None
+
+    def list_sources(self) -> List[DataSource]:
+        return list(self._sources.values())
+
+    def discover(self, context: str, limit: int = 10) -> List[Dict[str, Any]]:
+        """Return registered sources ranked by real relevance to `context`.
+
+        Relevance is Jaccard token overlap: |context tokens ∩ source
+        tokens| / |context tokens ∪ source tokens|. Sources that share no
+        tokens with `context` are excluded, not returned with a made-up
+        low score. Returns plain dicts (the wire-format shape used by the
+        REST/MCP surfaces) — use `discover_typed` for `DiscoveredSource`
+        objects.
+        """
+        return [
+            {
+                "name": s.name,
+                "type": s.type,
+                "relevance": round(score, 4),
+                "matched_terms": matched,
+            }
+            for score, matched, s in self._score(context, limit)
+        ]
+
+    def discover_typed(self, query_id: str, context: str, limit: int = 10) -> Discovery:
+        """Same ranking as `discover`, returned as a `Discovery` result of
+        real `DiscoveredSource` objects instead of plain dicts."""
+        result = Discovery(query_id=query_id)
+        for score, _matched, s in self._score(context, limit):
+            try:
+                source_type = SourceType(s.type)
+            except ValueError:
+                source_type = SourceType.UNKNOWN
+            result.add_source(
+                DiscoveredSource(
+                    name=s.name,
+                    source_type=source_type,
+                    relevance_score=score,
+                    estimated_tokens=s.estimated_tokens,
+                )
+            )
+        return result
+
+    def _score(self, context: str, limit: int):
+        context_tokens = _tokenize(context)
+        if not context_tokens or not self._sources:
+            return []
+
+        scored = []
+        for source in self._sources.values():
+            source_tokens = _tokenize(source._text())
+            if not source_tokens:
+                continue
+            overlap = context_tokens & source_tokens
+            if not overlap:
+                continue
+            union = context_tokens | source_tokens
+            relevance = len(overlap) / len(union)
+            scored.append((relevance, sorted(overlap), source))
+
+        scored.sort(key=lambda item: item[0], reverse=True)
+        return scored[:limit]

pystreammcp/integrations/langchain.py

@@ -8,7 +8,7 @@ Provides both:
 2. Legacy PyStreamMCPTool & PyStreamMCPRetriever (backward compatible)
 """
 
-from typing import Optional, Dict, Any, List, Type
+from typing import Optional, Dict, Any, List, Type, TYPE_CHECKING
 import asyncio
 from pystreammcp import (
     Agent,
@@ -21,6 +21,11 @@ from pystreammcp import (
 )
 
 from pystreammcp.adapters import QueryResult as AdapterQueryResult
+from pystreammcp.discovery import SourceRegistry
+
+if TYPE_CHECKING:
+    from langchain_core.tools import Tool as LangchainTool
+    from langchain_core.retrievers import BaseRetriever
 
 
 class LangchainAdapter(AgentFrameworkAdapter):
@@ -30,11 +35,15 @@ class LangchainAdapter(AgentFrameworkAdapter):
     Supports both sync and async execution.
     """
 
-    def __init__(self, config: AdapterConfig):
+    def __init__(self, config: AdapterConfig, registry: Optional[SourceRegistry] = None):
         """Initialize Langchain adapter.
 
         Args:
             config: Adapter configuration
+            registry: Optional shared SourceRegistry (e.g. the same one used
+                by a PyStreamMCPAPI or PyStreamMCPServer instance, so sources
+                registered on one surface are discoverable from another). A
+                new, empty registry is created if not given.
         """
         super().__init__(config)
         self.agent = Agent(
@@ -43,6 +52,7 @@ class LangchainAdapter(AgentFrameworkAdapter):
             optimization_strategy=config.optimization_strategy,
             max_tokens=config.max_tokens,
         )
+        self.registry = registry if registry is not None else SourceRegistry()
 
     def query(
         self, text: str, intent: str = "retrieve", **kwargs
@@ -92,6 +102,11 @@ class LangchainAdapter(AgentFrameworkAdapter):
     def discover(self, context: str, **kwargs) -> Dict[str, Any]:
         """Discover relevant data sources for context.
 
+        Ranks sources registered on `self.registry` (see
+        `SourceRegistry.register` / pass a shared registry to `__init__`)
+        by real token overlap against `context`. Returns an empty list if
+        nothing is registered or nothing overlaps — never fabricated data.
+
         Args:
             context: Context for discovery
             **kwargs: Additional arguments
@@ -99,17 +114,10 @@ class LangchainAdapter(AgentFrameworkAdapter):
         Returns:
             Dictionary with discovered sources
         """
-        # TODO: Integrate with discovery module
+        sources = self.registry.discover(context, limit=kwargs.get("limit", 10))
         return {
-            "sources": [
-                {
-                    "name": "data_warehouse",
-                    "relevance": 0.95,
-                    "type": "database",
-                    "estimated_tokens": 500,
-                }
-            ],
-            "total_sources": 1,
+            "sources": sources,
+            "total_sources": len(sources),
             "search_context": context,
         }
 

pystreammcp/integrations/llamaindex.py

@@ -8,7 +8,7 @@ Provides both:
 2. Legacy StreamMCPRetriever (backward compatible)
 """
 
-from typing import Any, List, Optional, Dict
+from typing import Any, List, Optional, Dict, TYPE_CHECKING
 import asyncio
 from pystreammcp import (
     Agent,
@@ -22,6 +22,9 @@ from pystreammcp import (
 
 from pystreammcp.adapters import QueryResult as AdapterQueryResult
 
+if TYPE_CHECKING:
+    from llama_index.core.retrievers import BaseRetriever
+
 
 class LlamaIndexAdapter(AgentFrameworkAdapter):
     """LlamaIndex adapter using new AgentFrameworkAdapter base class.

pystreammcp/integrations/semantic_kernel.py

@@ -8,7 +8,7 @@ Provides both:
 2. Legacy PyStreamMCPPlugin (backward compatible)
 """
 
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, TYPE_CHECKING
 import asyncio
 from pystreammcp import (
     Agent,
@@ -22,6 +22,9 @@ from pystreammcp import (
 
 from pystreammcp.adapters import QueryResult as AdapterQueryResult
 
+if TYPE_CHECKING:
+    from semantic_kernel.functions import KernelPlugin
+
 
 class SemanticKernelAdapter(AgentFrameworkAdapter):
     """Semantic Kernel adapter using new AdapterFrameworkAdapter base class.