Spaces:
Running
Running
File size: 8,932 Bytes
e70c280 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 | 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()
@classmethod
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,
)
@classmethod
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,
)
@classmethod
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"
)
@wraps(fn)
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
|