Spaces:
Running
Running
| from __future__ import annotations | |
| import inspect | |
| from functools import wraps | |
| from typing import TYPE_CHECKING, Any, Callable, Optional | |
| from src.adapters.cache import CacheAdapter | |
| from src.adapters.data import DataAdapter | |
| from src.adapters.embeddings import EmbeddingsAdapter, VectorStoreAdapter | |
| from src.adapters.graph import GraphAdapter | |
| from src.adapters.llm import LLMAdapter | |
| from src.core.plugins.prompts import PromptRegistry, prompt_registry | |
| if TYPE_CHECKING: | |
| from fastapi import APIRouter, FastAPI | |
| from mcp.server import FastMCP | |
| from src.config import Config | |
| class PluginAdapters: | |
| def __init__( | |
| self, | |
| cache: CacheAdapter, | |
| graph: GraphAdapter, | |
| data: DataAdapter, | |
| vector_store: VectorStoreAdapter, | |
| llm_small: LLMAdapter, | |
| llm_large: LLMAdapter, | |
| embeddings: EmbeddingsAdapter, | |
| embeddings_small: EmbeddingsAdapter, | |
| ): | |
| self.cache = cache | |
| self.graph = graph | |
| self.data = data | |
| self.vector_store = vector_store | |
| self.llm_small = llm_small | |
| self.llm_large = llm_large | |
| self.embeddings = embeddings | |
| self.embeddings_small = embeddings_small | |
| class PluginContext: | |
| def __init__( | |
| self, | |
| adapters: PluginAdapters, | |
| prompts: PromptRegistry, | |
| config: "Config", | |
| app: Optional["FastAPI"] = None, | |
| mcp: Optional["FastMCP"] = None, | |
| workspace_mcp: Optional["FastMCP"] = None, | |
| ): | |
| self._app = app | |
| self.mcp = mcp | |
| self.workspace_mcp = workspace_mcp or mcp | |
| self.adapters = adapters | |
| self.prompts = prompts | |
| self.config = config | |
| self._routers: list[tuple["APIRouter", dict[str, Any]]] = [] | |
| self._event_handlers: dict[str, list[Callable]] = {} | |
| self._search_retrievers: dict[str, Callable] = {} | |
| self._search_rerankers: dict[str, Callable] = {} | |
| self._workspace_route_paths: set[tuple[str, str]] = set() | |
| self._workspace_mcp_tools: set[str] = set() | |
| def _build_adapters(cls) -> PluginAdapters: | |
| from src.core.instances import ( | |
| cache_adapter, | |
| data_adapter, | |
| embeddings_adapter, | |
| embeddings_small_adapter, | |
| graph_adapter, | |
| llm_large_adapter, | |
| llm_small_adapter, | |
| vector_store_adapter, | |
| ) | |
| return PluginAdapters( | |
| cache=cache_adapter, | |
| graph=graph_adapter, | |
| data=data_adapter, | |
| vector_store=vector_store_adapter, | |
| llm_small=llm_small_adapter, | |
| llm_large=llm_large_adapter, | |
| embeddings=embeddings_adapter, | |
| embeddings_small=embeddings_small_adapter, | |
| ) | |
| def from_app(cls, app: "FastAPI") -> "PluginContext": | |
| from src.config import config | |
| return cls( | |
| app=app, | |
| adapters=cls._build_adapters(), | |
| prompts=prompt_registry, | |
| config=config, | |
| ) | |
| def from_mcp( | |
| cls, mcp: "FastMCP", *, workspace_mcp: Optional["FastMCP"] = None | |
| ) -> "PluginContext": | |
| from src.config import config | |
| return cls( | |
| mcp=mcp, | |
| workspace_mcp=workspace_mcp, | |
| adapters=cls._build_adapters(), | |
| prompts=prompt_registry, | |
| config=config, | |
| ) | |
| def include_router( | |
| self, | |
| router: "APIRouter", | |
| skip_pat: bool = False, | |
| skip_brain: bool = False, | |
| **kwargs: Any, | |
| ) -> None: | |
| if self._app is None: | |
| return | |
| self._app.include_router(router, **kwargs) | |
| self._routers.append((router, kwargs)) | |
| prefix = kwargs.get("prefix", "") or getattr(router, "prefix", "") | |
| if prefix: | |
| if skip_pat: | |
| from src.services.api.middlewares.auth import BrainPATMiddleware | |
| BrainPATMiddleware.excluded_prefixes.add(prefix) | |
| if skip_brain: | |
| from src.services.api.middlewares.brains import BrainMiddleware | |
| BrainMiddleware.excluded_prefixes.add(prefix) | |
| def add_middleware(self, middleware_cls: type, **kwargs: Any) -> None: | |
| if self._app is None: | |
| return | |
| self._app.add_middleware(middleware_cls, **kwargs) | |
| def register_mcp_tool(self, fn: Callable, **kwargs: Any) -> None: | |
| if self.mcp is None: | |
| return | |
| self.mcp.tool(**kwargs)(fn) | |
| def register_workspace_route( | |
| self, | |
| router: "APIRouter", | |
| *, | |
| capability: str, | |
| ) -> None: | |
| """Explicitly expose a plugin router below the hosted workspace API.""" | |
| if self._app is None: | |
| return | |
| from fastapi.routing import APIRoute | |
| from src.services.api.dependencies import get_brain_id | |
| if not capability or not capability.replace("-", "").isalnum(): | |
| raise ValueError("Workspace plugin capability must be URL-safe") | |
| if "/system" in (getattr(router, "prefix", "") or ""): | |
| raise ValueError("System routes cannot be workspace-hosted") | |
| routes = [route for route in router.routes if isinstance(route, APIRoute)] | |
| if not routes: | |
| raise ValueError("Workspace plugin router has no API routes") | |
| for route in routes: | |
| dependencies = { | |
| dependency.call for dependency in route.dependant.dependencies | |
| } | |
| if get_brain_id not in dependencies: | |
| raise ValueError( | |
| f"Workspace plugin route {route.path} must declare " | |
| "brain_id: str = Depends(get_brain_id)" | |
| ) | |
| for method in route.methods: | |
| full_path = ( | |
| f"/brains/{{workspace_slug}}/api/plugins/{capability}{route.path}" | |
| ) | |
| key = (method, full_path) | |
| existing = { | |
| (candidate_method, candidate.path) | |
| for candidate in self._app.routes | |
| if isinstance(candidate, APIRoute) | |
| for candidate_method in candidate.methods | |
| } | |
| if key in existing or key in self._workspace_route_paths: | |
| raise ValueError(f"Workspace plugin route collision: {method} {full_path}") | |
| self._workspace_route_paths.add(key) | |
| self._app.include_router( | |
| router, | |
| prefix=f"/brains/{{workspace_slug}}/api/plugins/{capability}", | |
| tags=[f"workspace-plugin:{capability}"], | |
| ) | |
| def register_workspace_mcp_tool( | |
| self, | |
| fn: Callable, | |
| *, | |
| name: str | None = None, | |
| **kwargs: Any, | |
| ) -> None: | |
| """Register an opt-in MCP tool with immutable workspace context injection.""" | |
| if self.workspace_mcp is None: | |
| return | |
| from src.services.brain_context import brain_context_var | |
| tool_name = name or fn.__name__ | |
| existing_tools = getattr( | |
| getattr(self.workspace_mcp, "_tool_manager", None), "_tools", {} | |
| ) | |
| if tool_name in self._workspace_mcp_tools or tool_name in existing_tools: | |
| raise ValueError(f"Workspace MCP tool collision: {tool_name}") | |
| if "brain_context" not in inspect.signature(fn).parameters: | |
| raise ValueError( | |
| "Workspace MCP tools must declare a brain_context parameter" | |
| ) | |
| async def scoped(*args: Any, **call_kwargs: Any): | |
| context = brain_context_var.get() | |
| if context is None: | |
| raise PermissionError("Workspace MCP context is required") | |
| call_kwargs["brain_context"] = context | |
| result = fn(*args, **call_kwargs) | |
| return await result if inspect.isawaitable(result) else result | |
| signature = inspect.signature(fn) | |
| scoped.__signature__ = signature.replace( # type: ignore[attr-defined] | |
| parameters=[ | |
| parameter | |
| for parameter in signature.parameters.values() | |
| if parameter.name != "brain_context" | |
| ] | |
| ) | |
| self._workspace_mcp_tools.add(tool_name) | |
| self.workspace_mcp.tool(name=tool_name, **kwargs)(scoped) | |
| def add_event_handler(self, event: str, handler: Callable) -> None: | |
| self._event_handlers.setdefault(event, []).append(handler) | |
| def register_search_retriever(self, name: str, fn: Callable) -> None: | |
| from src.core.search.hooks import register_search_retriever | |
| register_search_retriever(name, fn) | |
| self._search_retrievers[name] = fn | |
| def register_search_reranker(self, name: str, fn: Callable) -> None: | |
| from src.core.search.hooks import register_search_reranker | |
| register_search_reranker(name, fn) | |
| self._search_rerankers[name] = fn | |