PyStreamMCP 3.2.0__py3-none-any.whl → 3.3.0__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.0"
 
 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/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

@@ -21,6 +21,7 @@ from pystreammcp import (
 )
 
 from pystreammcp.adapters import QueryResult as AdapterQueryResult
+from pystreammcp.discovery import SourceRegistry
 
 
 class LangchainAdapter(AgentFrameworkAdapter):
@@ -30,11 +31,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 +48,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 +98,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 +110,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/mcp_server.py

@@ -4,17 +4,86 @@ Implements the Model Context Protocol standard to expose
 query planning, discovery, and optimization as MCP tools.
 """
 
-import json
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Literal, Optional, Type
 from dataclasses import dataclass
 
+from pydantic import BaseModel, Field, ValidationError
+
 from .agent import Agent
+from .discovery import SourceRegistry
 
 
 logger = logging.getLogger(__name__)
 
 
+# ── MCP tool argument schemas ────────────────────────────────────────────────
+#
+# Previously `call_tool`/`_tool_*` only did manual `if not text: return error`
+# checks per field -- no type checking (a `max_tokens` of "ten" would pass
+# straight through to `int()` arithmetic downstream), no enum validation (an
+# `intent`/`strategy` of anything at all was accepted despite the JSON schema
+# below advertising a fixed enum), and no consistent required-field handling
+# across the four tools. These models are now both the single source of
+# truth for `input_schema` (via `model_json_schema()`) and the actual runtime
+# validation, so the two can no longer silently drift apart the way a
+# hand-duplicated dict schema could.
+
+
+class QueryToolArgs(BaseModel):
+    text: str = Field(..., min_length=1, description="The query text")
+    intent: Literal["retrieve", "discover", "aggregate", "synthesize", "analyze"] = Field(
+        "retrieve", description="Query intent - what the agent is trying to do"
+    )
+    max_tokens: Optional[int] = Field(
+        None, gt=0, description="Maximum tokens for context (overrides default)"
+    )
+
+
+class DiscoverToolArgs(BaseModel):
+    context: str = Field(..., min_length=1, description="Context for discovery")
+    limit: int = Field(10, gt=0, le=1000, description="Maximum sources to return")
+
+
+class RegisterSourceToolArgs(BaseModel):
+    name: str = Field(..., min_length=1, description="Unique source name")
+    description: str = Field(
+        ..., min_length=1, 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 OptimizeToolArgs(BaseModel):
+    text: str = Field(..., min_length=1, description="Query to optimize")
+    strategy: Literal["balanced", "token_efficient", "quality_first"] = Field(
+        "balanced", description="Optimization strategy"
+    )
+
+
+_TOOL_ARG_MODELS: Dict[str, Type[BaseModel]] = {
+    "pystreammcp_query": QueryToolArgs,
+    "pystreammcp_discover": DiscoverToolArgs,
+    "pystreammcp_register_source": RegisterSourceToolArgs,
+    "pystreammcp_optimize": OptimizeToolArgs,
+}
+
+
+def _validation_error_response(name: str, error: ValidationError) -> Dict[str, Any]:
+    return {
+        "status": "error",
+        "message": f"Invalid arguments for tool '{name}'",
+        "validation_errors": [
+            {
+                "field": ".".join(str(p) for p in e["loc"]),
+                "issue": e["msg"],
+                "type": e["type"],
+            }
+            for e in error.errors()
+        ],
+    }
+
+
 @dataclass
 class MCPTool:
     """MCP tool definition."""
@@ -27,89 +96,51 @@ class MCPTool:
 class PyStreamMCPServer:
     """MCP server implementing PyStreamMCP protocol."""
 
-    def __init__(self, agent_id: str = "mcp_agent"):
+    def __init__(self, agent_id: str = "mcp_agent", agent: Optional[Agent] = None):
         """Initialize MCP server.
 
         Args:
-            agent_id: Agent ID for this server
+            agent_id: Agent ID for this server (ignored if `agent` is given)
+            agent: Optional pre-built `Agent` (or any object exposing a
+                compatible `.query(text)` method) to use instead of
+                constructing a default one. Mainly for tests that need to
+                inject controlled failure behavior (rate limits, malformed
+                responses, etc.) into the tool-execution path.
         """
         self.agent_id = agent_id
-        self.agent = Agent(
+        self.agent = agent or Agent(
             agent_id=agent_id,
             name=f"MCP Agent: {agent_id}",
             optimization_strategy="balanced",
             max_tokens=2000,
         )
+        self.registry = SourceRegistry()
         self.tools = self._define_tools()
 
     def _define_tools(self) -> Dict[str, MCPTool]:
-        """Define MCP tools."""
+        """Define MCP tools. `input_schema` is generated from the Pydantic
+        arg models above, so the advertised schema and the schema actually
+        enforced at call time can't drift apart."""
         return {
             "query": MCPTool(
                 name="pystreammcp_query",
                 description="Execute a query with PyStreamMCP optimization and discover relevant context",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "text": {
-                            "type": "string",
-                            "description": "The query text",
-                        },
-                        "intent": {
-                            "type": "string",
-                            "enum": [
-                                "retrieve",
-                                "discover",
-                                "aggregate",
-                                "synthesize",
-                                "analyze",
-                            ],
-                            "description": "Query intent - what the agent is trying to do",
-                        },
-                        "max_tokens": {
-                            "type": "integer",
-                            "description": "Maximum tokens for context (overrides default)",
-                        },
-                    },
-                    "required": ["text"],
-                },
+                input_schema=QueryToolArgs.model_json_schema(),
             ),
             "discover": MCPTool(
                 name="pystreammcp_discover",
                 description="Discover relevant data sources and context for a query",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "context": {
-                            "type": "string",
-                            "description": "Context for discovery",
-                        },
-                        "limit": {
-                            "type": "integer",
-                            "description": "Maximum sources to return",
-                        },
-                    },
-                    "required": ["context"],
-                },
+                input_schema=DiscoverToolArgs.model_json_schema(),
+            ),
+            "register_source": MCPTool(
+                name="pystreammcp_register_source",
+                description="Register a data source (name, description, tags) so pystreammcp_discover can find it",
+                input_schema=RegisterSourceToolArgs.model_json_schema(),
             ),
             "optimize": MCPTool(
                 name="pystreammcp_optimize",
                 description="Optimize a query for cost reduction while maintaining quality",
-                input_schema={
-                    "type": "object",
-                    "properties": {
-                        "text": {
-                            "type": "string",
-                            "description": "Query to optimize",
-                        },
-                        "strategy": {
-                            "type": "string",
-                            "enum": ["balanced", "token_efficient", "quality_first"],
-                            "description": "Optimization strategy",
-                        },
-                    },
-                    "required": ["text"],
-                },
+                input_schema=OptimizeToolArgs.model_json_schema(),
             ),
         }
 
@@ -131,6 +162,16 @@ class PyStreamMCPServer:
     def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
         """Call a tool by name.
 
+        Validates `arguments` against the tool's Pydantic arg model before
+        dispatch (returning a structured `validation_errors` response on
+        failure, rather than letting a wrong type or invalid enum value
+        flow through to the handler unchecked), and isolates the handler
+        call in a try/except so a failure during execution -- a bug, or a
+        downstream dependency (e.g. a real LLM/framework adapter wired in
+        via `agent=`) hitting a rate limit, context-window overrun, or
+        returning malformed data -- produces a structured error response
+        instead of propagating out of `call_tool` uncaught.
+
         Args:
             name: Tool name
             arguments: Tool arguments
@@ -141,64 +182,59 @@ class PyStreamMCPServer:
         Raises:
             ValueError: If tool not found
         """
-        if name == "pystreammcp_query":
-            return self._tool_query(arguments)
-        elif name == "pystreammcp_discover":
-            return self._tool_discover(arguments)
-        elif name == "pystreammcp_optimize":
-            return self._tool_optimize(arguments)
-        else:
+        model = _TOOL_ARG_MODELS.get(name)
+        if model is None:
             raise ValueError(f"Unknown tool: {name}")
 
-    def _tool_query(self, args: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute query tool."""
-        text = args.get("text")
-        intent = args.get("intent", "retrieve")
-        max_tokens = args.get("max_tokens")
-
-        if not text:
+        try:
+            validated = model(**arguments)
+        except ValidationError as e:
+            return _validation_error_response(name, e)
+
+        handler = {
+            "pystreammcp_query": self._tool_query,
+            "pystreammcp_discover": self._tool_discover,
+            "pystreammcp_register_source": self._tool_register_source,
+            "pystreammcp_optimize": self._tool_optimize,
+        }[name]
+
+        try:
+            return handler(validated)
+        except Exception as e:
+            logger.exception(f"Tool '{name}' failed during execution")
             return {
                 "status": "error",
-                "message": "text parameter required",
+                "message": f"Tool '{name}' failed during execution: {e}",
+                "error_type": type(e).__name__,
             }
 
-        result = self.agent.query(text)
+    def _tool_query(self, args: QueryToolArgs) -> Dict[str, Any]:
+        """Execute query tool."""
+        result = self.agent.query(args.text)
 
         return {
             "status": "success",
             "query_id": result.query_id,
-            "text": text,
-            "intent": intent,
+            "text": args.text,
+            "intent": args.intent,
             "baseline_tokens": result.baseline_tokens,
             "optimized_tokens": result.optimized_tokens,
             "cost_reduction_percent": result.cost_reduction_percent,
             "execution_time_ms": result.execution_time_ms,
             "context": {
                 "description": "Context optimized for the query",
-                "recommendation": f"Use this context for {intent} queries",
+                "recommendation": f"Use this context for {args.intent} queries",
             },
         }
 
-    def _tool_discover(self, args: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute discovery tool."""
-        context = args.get("context")
-        limit = args.get("limit", 10)
+    def _tool_discover(self, args: DiscoverToolArgs) -> Dict[str, Any]:
+        """Execute discovery tool.
 
-        if not context:
-            return {
-                "status": "error",
-                "message": "context parameter required",
-            }
-
-        # TODO: Implement actual discovery logic
-        sources = [
-            {
-                "name": f"source_{i}",
-                "relevance": 0.95 - (i * 0.05),
-                "type": "database",
-            }
-            for i in range(min(limit, 5))
-        ]
+        Ranks *registered* sources (see pystreammcp_register_source) by
+        real token overlap against `context`. Returns an empty list if
+        nothing is registered or nothing overlaps — never fabricated data.
+        """
+        sources = self.registry.discover(args.context, limit=args.limit)
 
         return {
             "status": "success",
@@ -206,18 +242,20 @@ class PyStreamMCPServer:
             "total_sources": len(sources),
         }
 
-    def _tool_optimize(self, args: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute optimization tool."""
-        text = args.get("text")
-        strategy = args.get("strategy", "balanced")
+    def _tool_register_source(self, args: RegisterSourceToolArgs) -> Dict[str, Any]:
+        """Register a data source for discovery."""
+        self.registry.register(
+            name=args.name,
+            description=args.description,
+            type=args.type,
+            tags=args.tags,
+        )
 
-        if not text:
-            return {
-                "status": "error",
-                "message": "text parameter required",
-            }
+        return {"status": "success", "name": args.name}
 
-        result = self.agent.query(text)
+    def _tool_optimize(self, args: OptimizeToolArgs) -> Dict[str, Any]:
+        """Execute optimization tool."""
+        result = self.agent.query(args.text)
 
         return {
             "status": "success",
@@ -225,7 +263,7 @@ class PyStreamMCPServer:
             "baseline_tokens": result.baseline_tokens,
             "optimized_tokens": result.optimized_tokens,
             "cost_reduction_percent": result.cost_reduction_percent,
-            "strategy_used": strategy,
+            "strategy_used": args.strategy,
             "techniques": [
                 "pruning",
                 "summarization",
@@ -248,7 +286,16 @@ class PyStreamMCPServer:
         if message_type == "call_tool":
             tool_name = message.get("name")
             tool_args = message.get("arguments", {})
-            return self.call_tool(tool_name, tool_args)
+            try:
+                return self.call_tool(tool_name, tool_args)
+            except ValueError as e:
+                # Unknown tool name -- call_tool's own validation/execution
+                # error paths already return structured dicts; this is the
+                # one case it still raises (see its docstring), and
+                # process_message is the actual message-transport entry
+                # point, so an unrecognized tool name shouldn't crash the
+                # whole call.
+                return {"status": "error", "message": str(e)}
 
         elif message_type == "list_tools":
             return {

pystreammcp/testing.py

@@ -0,0 +1,123 @@
+"""Reusable test doubles for PyStreamMCP, exported for downstream use.
+
+Previously the only mock (`MockAdapter` in `tests/test_sprint1_foundation.py`)
+was test-local -- not importable by a downstream project's own test suite,
+and not usable for simulating real LLM/framework failure states (rate
+limits, context-window overruns, malformed model responses) at all, since it
+only ever returned a fixed happy-path result.
+
+This module provides:
+  - `MockAdapter`: the same happy-path mock, promoted here so it's a real,
+    importable, reusable fixture instead of copy-pasted per test file.
+  - `RateLimitError`, `ContextWindowExceededError`,
+    `MalformedModelResponseError`: the three LLM-failure classes named in
+    the "no end-to-end tests for LLM failure states" gap. PyStreamMCP's own
+    code never calls a model provider directly (see README "Known Issues"),
+    so these aren't wrapping any provider-specific SDK exception -- they
+    exist to let tests exercise how PyStreamMCP's *own* code (the MCP tool
+    dispatcher, adapter query paths) behaves when something it's
+    orchestrating (a real `Agent`/`AgentFrameworkAdapter` a caller plugs in)
+    fails in one of these ways, without needing network access or a real
+    API key.
+  - `FailingAgent` / `FailingAdapter`: drop-in stand-ins for `Agent` /
+    `AgentFrameworkAdapter` that raise a configured failure on `query()`
+    instead of computing a result.
+"""
+
+from typing import Any, Dict, Optional
+
+from .adapters import AdapterConfig, AgentFrameworkAdapter, QueryResult
+
+
+class RateLimitError(Exception):
+    """Raised by a test double to simulate a model provider rate limit."""
+
+
+class ContextWindowExceededError(Exception):
+    """Raised by a test double to simulate exceeding the model's context window."""
+
+
+class MalformedModelResponseError(Exception):
+    """Raised by a test double to simulate a model returning invalid/unparseable JSON."""
+
+
+class MockAdapter(AgentFrameworkAdapter):
+    """Happy-path mock adapter: always returns a fixed, valid result.
+
+    Useful for tests that need *an* adapter wired in but don't care about
+    its output -- for failure-injection tests, use `FailingAdapter` instead.
+    """
+
+    def query(self, text: str, intent: str = "retrieve", **kwargs) -> QueryResult:
+        return QueryResult(
+            query_id="mock-query-1",
+            text=text,
+            intent=intent,
+            baseline_tokens=1000,
+            optimized_tokens=400,
+            cost_reduction_percent=60.0,
+            execution_time_ms=150.5,
+            context={"result": "mock"},
+        )
+
+    async def query_async(self, text: str, intent: str = "retrieve", **kwargs) -> QueryResult:
+        return self.query(text, intent, **kwargs)
+
+    def discover(self, context: str, **kwargs) -> Dict[str, Any]:
+        return {
+            "sources": [{"name": "source_1", "relevance": 0.95}],
+            "total": 1,
+        }
+
+    async def discover_async(self, context: str, **kwargs) -> Dict[str, Any]:
+        return self.discover(context, **kwargs)
+
+    def optimize(self, query_text: str, strategy: Optional[str] = None, **kwargs) -> QueryResult:
+        return self.query(query_text)
+
+    async def optimize_async(
+        self, query_text: str, strategy: Optional[str] = None, **kwargs
+    ) -> QueryResult:
+        return self.optimize(query_text, strategy, **kwargs)
+
+
+class FailingAgent:
+    """Drop-in stand-in for `pystreammcp.agent.Agent`: raises `failure` from
+    `query()` instead of computing a result. Pass to
+    `PyStreamMCPServer(agent=FailingAgent(RateLimitError(...)))` to test how
+    the MCP tool-call path handles a failure during execution."""
+
+    def __init__(self, failure: Exception):
+        self.failure = failure
+
+    def query(self, text: str, *args, **kwargs):
+        raise self.failure
+
+
+class FailingAdapter(AgentFrameworkAdapter):
+    """Drop-in `AgentFrameworkAdapter` that raises `failure` from every
+    operation, for framework-integration failure-path tests."""
+
+    def __init__(self, config: AdapterConfig, failure: Exception):
+        super().__init__(config)
+        self.failure = failure
+
+    def query(self, text: str, intent: str = "retrieve", **kwargs) -> QueryResult:
+        raise self.failure
+
+    async def query_async(self, text: str, intent: str = "retrieve", **kwargs) -> QueryResult:
+        raise self.failure
+
+    def discover(self, context: str, **kwargs) -> Dict[str, Any]:
+        raise self.failure
+
+    async def discover_async(self, context: str, **kwargs) -> Dict[str, Any]:
+        raise self.failure
+
+    def optimize(self, query_text: str, strategy: Optional[str] = None, **kwargs) -> QueryResult:
+        raise self.failure
+
+    async def optimize_async(
+        self, query_text: str, strategy: Optional[str] = None, **kwargs
+    ) -> QueryResult:
+        raise self.failure

{pystreammcp-3.2.0.dist-info → pystreammcp-3.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyStreamMCP
-Version: 3.2.0
+Version: 3.3.0
 Summary: Intelligent MCP orchestration hub - Intent understanding, capability matching, tool ranking with cohesive foundation
 Author-email: Georgi Mammen Mullassery <mullassery@gmail.com>
 License: Proprietary License — Free to use with explicit attribution
@@ -91,7 +91,11 @@ pip install "PyStreamMCP[semantic-kernel]"  # Semantic Kernel adapter
 pip install "PyStreamMCP[all-integrations]" # every framework adapter
 ```
 
-Requires Python 3.9+.
+Requires Python 3.9+ for the core SDK (`Agent`, `SourceRegistry`,
+`Orchestrator`). The `mcp` and `semantic-kernel` extras require Python
+3.10+ — their upstream packages don't publish Python 3.9-compatible
+releases, so `pip install` for those extras will fail on 3.9. CI tests
+3.10, 3.11, and 3.12 only; 3.9 is not covered by CI.
 
 ## Quick start
 
@@ -120,6 +124,33 @@ for a full runnable walkthrough, including framework adapters for
 LangChain, LlamaIndex, CrewAI, Semantic Kernel, PydanticAI, and Haystack
 under [`pystreammcp.integrations`](python/pystreammcp/integrations/).
 
+## Data source discovery
+
+`SourceRegistry` (`pystreammcp.SourceRegistry`) ranks *registered* data
+sources by real word-overlap relevance to a query — not a fabricated or
+hardcoded result:
+
+```python
+from pystreammcp import SourceRegistry
+
+registry = SourceRegistry()
+registry.register(
+    "orders_db",
+    description="customer order history and revenue by quarter",
+    type="database",
+    tags=["orders", "revenue"],
+)
+
+results = registry.discover("customer revenue by quarter")
+# [{"name": "orders_db", "type": "database", "relevance": 0.4444, "matched_terms": [...]}]
+```
+
+The same registry backs `POST /discover` (REST), the `pystreammcp_discover`
+MCP tool, and `LangchainAdapter.discover()` — register sources once via
+`POST /sources` / `pystreammcp_register_source` / `registry.register()`,
+then discover against them from any of the three surfaces. This is a
+simple Jaccard token-overlap heuristic, not semantic/embedding search.
+
 ## MCP orchestration & federation
 
 `Orchestrator` (`pystreammcp.Orchestrator`) discovers and routes across
@@ -217,6 +248,49 @@ pip install -e ".[dev,api,mcp,langchain,llamaindex,semantic-kernel]"
 pytest tests/ -v
 ```
 
+Note: this exact command needs Python 3.10+ (see [Install](#install)); on
+3.9 the `mcp` and `semantic-kernel` extras fail to resolve, so drop them
+from the extras list if you're on 3.9.
+
+## Known Issues
+
+- ~~No end-to-end tests for LLM failure states~~ **Fixed.** PyStreamMCP's
+  own code never calls a model provider directly (no code path here to add
+  `except RateLimitError`/`APIError` handling *to* — see `AgentFrameworkAdapter`
+  below), so what's now tested is the part this repo actually owns: that the
+  MCP tool-call path doesn't crash when whatever it's orchestrating fails.
+  `pystreammcp.testing` exports `RateLimitError`, `ContextWindowExceededError`,
+  `MalformedModelResponseError`, and `FailingAgent`/`FailingAdapter` stand-ins
+  that raise them; `mcp_server.py`'s `call_tool`/`process_message` now catch
+  exceptions raised during tool execution and return a structured
+  `{"status": "error", "error_type": ..., "message": ...}` response instead
+  of propagating uncaught. See `tests/test_mcp_server_resilience.py`.
+- ~~No Pydantic validation for the MCP tool protocol itself~~ **Fixed.**
+  Each of the four MCP tools now has a Pydantic arg model
+  (`QueryToolArgs`/`DiscoverToolArgs`/`RegisterSourceToolArgs`/`OptimizeToolArgs`
+  in `mcp_server.py`) validated in `call_tool` before dispatch (wrong types,
+  invalid enum values, and missing required fields all now return a
+  structured `validation_errors` response instead of reaching the handler
+  unchecked), and `MCPTool.input_schema` is generated from those same models
+  via `model_json_schema()` — the advertised schema and the schema actually
+  enforced can no longer drift apart the way a hand-duplicated dict could.
+- ~~No shipped mock MCP server / mock LLM fixture~~ **Fixed.** Added
+  `tests/conftest.py` (centralizing the `sys.path` setup 5 test files were
+  duplicating, plus `mock_adapter`/`failing_agent`/`failing_adapter`/
+  `llm_failure` pytest fixtures) and `pystreammcp/testing.py`, exporting
+  `MockAdapter` (promoted from the old test-local copy) and the
+  failure-injection stand-ins above for downstream projects' own test
+  suites to import directly.
+- **Per-query context discovery** now has a real implementation
+  (`pystreammcp.discovery.SourceRegistry`): register data sources (name,
+  description, type, tags), and `/discover` (REST), `pystreammcp_discover`
+  (MCP), and `LangchainAdapter.discover()` all rank *registered* sources
+  by real Jaccard token-overlap against the query — an empty registry or
+  a non-matching context returns an empty list, not fabricated data. This
+  is a simple word-overlap heuristic, not semantic/embedding search — if
+  you need that, wrap `SourceRegistry` with your own similarity scoring.
+- No open GitHub issues as of this pass.
+
 ## Rust workspace (not shipped)
 
 `core/` and `python/src/lib.rs` are an in-progress Rust workspace intended

{pystreammcp-3.2.0.dist-info → pystreammcp-3.3.0.dist-info}/RECORD

@@ -1,16 +1,16 @@
-pystreammcp/__init__.py,sha256=XBZOos263gyAHVhxNXsaeXIDt9hcS77iVdAGoFZ02to,2813
+pystreammcp/__init__.py,sha256=paZxqfxmLptS1tkZkL5LBU_MDKMTxCirHp8vGiOguS4,2881
 pystreammcp/_mcp_connector.py,sha256=CTPdQHT64aVxmolwG0l0irCj2uQshwJvhbdbGgyBEu0,31208
 pystreammcp/_mcp_tools.py,sha256=7s3DtzIbJxHVFMmISdd9EcFbmtKrSMIDXH27j0onhMY,32540
 pystreammcp/adapters.py,sha256=uB0zGKJU6jcTyRsUDIiiOgGg6K-8dbJw1YLt0ZiAiao,6801
 pystreammcp/agent.py,sha256=LJUPxnvcrT01rG8l3pqVXQ58OZyeqBn41GYM-hyHIWo,6385
-pystreammcp/api.py,sha256=Qk90v7NwxsKp21bfwuzFv_tPWeRevda-nOhpmSYSe0I,11004
+pystreammcp/api.py,sha256=4_9i86lQFKxzHdPB1TXzYVBJS6jna9Q2_s7Ea3-v6vg,12860
 pystreammcp/cli.py,sha256=Tk_vlwZT5Rv6OB4QXOaWUhv_JxlqijiNK5fL2EU1O2s,8569
 pystreammcp/cli_daemon.py,sha256=_ZSNtifMGOIkcR2mZaycdDDYYes4jtv0cdeMdMKU0Jw,1999
 pystreammcp/cli_dashboard.py,sha256=J5pnUUpkvQGjltUOFN6SPbNF1yyXsjHQR0VCI14ty6M,10562
 pystreammcp/context.py,sha256=WIgwbb0MAZTnPfQiZW0PLunEtxJatsHmwTnCRsvSqLU,2213
 pystreammcp/decomposition.py,sha256=peRR1MZbH29DCWMSvtL4Btzp1QhpXsFuVo5LOXPnoIs,6532
-pystreammcp/discovery.py,sha256=dEXhMA6AEHy2Ipo5LyL_dJMvHwqvcqTj4nNOyEMwbI4,2472
-pystreammcp/mcp_server.py,sha256=pv_l20kjZTW9KMSGU4izADGlsfF0ggrU7Papta0fGro,8808
+pystreammcp/discovery.py,sha256=FWfuYraXjBBkSReU-kxcKg0MJvnfmV3TH1wf5mvedH8,7312
+pystreammcp/mcp_server.py,sha256=hL5z9op-X62kdtFXLCM213pJP_9LTr1GsEYSt5d4Zrs,11972
 pystreammcp/ml.py,sha256=SZMX0DGak4SHLoXbVw2n85OrbrHvnqqSzHIkvgEjS80,4005
 pystreammcp/multi_agent.py,sha256=oUoyjix6WqiE7yX2limQAwWEQkMbwcrFV-vSKK-iDYc,4641
 pystreammcp/okf_core.py,sha256=JtyZ-KoMwoM7fSzOt9lEA-7uSvUT-fN7qggE7KtGNgQ,10039
@@ -21,12 +21,13 @@ pystreammcp/quality.py,sha256=l9_iW5YTWLtYlA7UbX915Z_pJSQf7dJUXk_MtFDUfpc,12832
 pystreammcp/query.py,sha256=rtWYmoOM6hbZ1W2MJpiBOoQ6Y-x0vyr3fZqSHagZzuI,3441
 pystreammcp/server.py,sha256=3BplCTjmguETK3fk1ObGayM-wvET1e-Uujs5wwwfJxY,20473
 pystreammcp/streaming.py,sha256=yKurMNPgPpF4D_jZIXk_TZaIt4nHFcV1ODrHXylna_g,3283
+pystreammcp/testing.py,sha256=b4E1JXfBYNqGPEVB5275lLPbfNDMqwutnnSoXLLnjis,4841
 pystreammcp/webhook_handlers.py,sha256=iawVSJJBj-DzDz7bd4ik4TnykdF2tU0tyYkCQhUwYq4,28084
 pystreammcp/webhook_router.py,sha256=a1lFKSluQA-y46O2wDAxwtZdg27GL7YgUYWBgySl1sY,20583
 pystreammcp/integrations/__init__.py,sha256=QOEpm_EcA7bLr7w0NXAWO-NLN5i37b_Q3vjA6-WDXxQ,101
 pystreammcp/integrations/crewai.py,sha256=v7qYQ4xo6FAaIUgOuNJYainF_7XkWGY9vX1kTO7EtrM,5724
 pystreammcp/integrations/haystack.py,sha256=bFVvzNeKmR7NVc4__8Z_aQBnWqMpHP8dzZms5so-NS4,7025
-pystreammcp/integrations/langchain.py,sha256=ngx9V5jDhigRzMzvOnqFs3vTkWC0OdP_OIiruPkXMAI,13966
+pystreammcp/integrations/langchain.py,sha256=v5obvQdM06KHOpy7-npABvTb0hdAMxYKAgDM9av9RVM,14542
 pystreammcp/integrations/langsmith.py,sha256=ePQpQG_Zs3xCG-PbOeDY2CYDclMYFLEPNbQ-He-Ll5A,8642
 pystreammcp/integrations/llamaindex.py,sha256=Ra6p6ptABPiIJQ2E88Ta7FY1ALig-XXIuJxaijqwyuw,12329
 pystreammcp/integrations/pydantic_ai.py,sha256=NY0qFw6aL4ygUnt5bLfkJgFzeIPc5Yycp0SVhwARr2U,4675
@@ -47,8 +48,8 @@ pystreammcp/qa/__init__.py,sha256=NvUUjyWNIyMfckW4J34Cn1xKYX5CnDlM22mXZF3pfnY,17
 pystreammcp/qa/validator.py,sha256=qTpdkFE7Z7l16fE1emwqfbr8W2_zh1m4Zvzz3hjnGt8,10694
 pystreammcp/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pystreammcp/scripts/post_install.py,sha256=lnxiXH8Owq_F4M4KFHNL-GJ5JMYVevUcVo3M5JXOpRo,2719
-pystreammcp-3.2.0.dist-info/licenses/LICENSE,sha256=0SAYRjt9bcQPGaEkH6VFFLo9JWVIZ-cNcOZM89kJNhQ,970
-pystreammcp-3.2.0.dist-info/METADATA,sha256=_sVkccAmfYu165c0MRwSHgYKcZhhYzk5vUTPIHLotnc,8643
-pystreammcp-3.2.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
-pystreammcp-3.2.0.dist-info/top_level.txt,sha256=SHNryMw9AAzKhuqpS2tCkxIfTpx4z623tlW8BazwTNM,12
-pystreammcp-3.2.0.dist-info/RECORD,,
+pystreammcp-3.3.0.dist-info/licenses/LICENSE,sha256=0SAYRjt9bcQPGaEkH6VFFLo9JWVIZ-cNcOZM89kJNhQ,970
+pystreammcp-3.3.0.dist-info/METADATA,sha256=JtwjENWHTPTcsBRud2CJTeHaSJIVhGMKZKXsnxA1xR4,12669
+pystreammcp-3.3.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+pystreammcp-3.3.0.dist-info/top_level.txt,sha256=SHNryMw9AAzKhuqpS2tCkxIfTpx4z623tlW8BazwTNM,12
+pystreammcp-3.3.0.dist-info/RECORD,,