hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/__init__.py

@@ -0,0 +1,16 @@
+import logging
+
+import fire  # type: ignore
+import jinja2
+
+# Jinja vulnerability is not a concern since this is used to format a string prompt, not to render HTML
+PROMPT_ENVIRONMENT = jinja2.Environment()  # nosemgrep
+
+
+logger = logging.getLogger("hyperforge")
+
+
+def cli():
+    from hyperforge.arag import ARAG
+
+    fire.Fire(ARAG)

hyperforge/agent.py

@@ -0,0 +1,81 @@
+import abc
+import uuid
+from typing import Any, Generic, List, Optional, Self, TypeVar
+
+from pydantic import BaseModel, Field
+
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import QuestionMemory
+from hyperforge.utils import WidgetType
+
+
+class AgentConfig(BaseModel):
+    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
+    title: str = "agent"
+    rules: Optional[List[str]] = Field(
+        default=None,
+        title="Agent rules",
+        description="List of rules to follow when executing this agent",
+        json_schema_extra={
+            "widget": WidgetType.NOT_SHOWN,
+        },
+    )
+    max_retries: int = 1
+    module: Any = Field(
+        ..., title="Agent module", description="Module/type of the agent"
+    )
+
+    def subagent_configure(self, config: dict):
+        pass
+
+
+T_Config = TypeVar("T_Config", bound=AgentConfig)
+
+
+class Agent(Generic[T_Config]):
+    __root_agent__: bool = False
+    agent_id: str
+    config: T_Config
+
+    def __init__(self, config: T_Config, agent_id: Optional[str] = None):
+        self.config: T_Config = config
+        self.agent_id: str = (
+            agent_id
+            if agent_id is not None
+            else config.id
+            if hasattr(config, "id") and config.id is not None
+            else str(uuid.uuid4())
+        )
+
+    @abc.abstractmethod
+    async def inner_from_config(self, config: T_Config, agent_id: Optional[str] = None):
+        """Initialize the agent from the config. This is where you should put any
+        async initialization code that needs to run when the agent is created.
+        """
+        pass
+
+    @classmethod
+    async def from_config(
+        cls, config: T_Config, agent_id: Optional[str] = None
+    ) -> Self:
+        instance = cls(config=config, agent_id=agent_id)
+        await instance.inner_from_config(config, agent_id)
+        return instance
+
+    def step_title(self, description: str) -> str:
+        """Format a step title as '<agent title>: <description>'.
+
+        Uses the user-configured instance title if set, otherwise falls back to
+        the Pydantic model_config title (e.g. 'MCP', 'SQL query') which describes
+        the agent type.
+        """
+        config = self.config  # type: ignore[attr-defined]
+        title = type(config).model_config.get("title") or config.title  # type: ignore[attr-defined]
+        return f"{title}: {description}"
+
+    async def __call__(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+    ):
+        raise NotImplementedError()

hyperforge/api/__init__.py

@@ -0,0 +1,20 @@
+import logging
+
+logger = logging.getLogger("hyperforge.api")
+
+SERVICE_NAME = "hyperforge_api"
+
+
+# Define the filter
+class EndpointFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        return (
+            record.args is not None
+            and len(record.args) >= 3
+            and record.args[2]  # type: ignore
+            not in ("/", "/metrics", "/health/alive", "/health/ready")
+        )
+
+
+# Add filter to the logger
+logging.getLogger("uvicorn.access").addFilter(EndpointFilter())

hyperforge/api/app.py

@@ -0,0 +1,155 @@
+from typing import Any, Optional, Tuple
+
+import prometheus_client  # type: ignore
+from fastapi import APIRouter, FastAPI
+from lru import LRU
+from mcp.server.lowlevel.server import Server as MCPServer
+from mcp.server.streamable_http import (
+    StreamableHTTPServerTransport,
+)
+from nucliadb_sdk.v2.sdk import NucliaDBAsync
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogLevel, LogSettings
+from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
+from prometheus_client import CONTENT_TYPE_LATEST  # type: ignore
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.responses import PlainTextResponse
+
+from hyperforge.api import SERVICE_NAME, internal, logger, v1
+from hyperforge.api.authentication import RaoAuthenticationBackend
+from hyperforge.api.logging import set_sentry
+from hyperforge.api.settings import Settings
+from hyperforge.broker import Broker
+from hyperforge.broker.redis import RedisBroker
+from hyperforge.configure import GLOBAL_REGISTRY, load_all_configurations, scan
+from hyperforge.db.agents import AgentManager
+from hyperforge.db.settings import DataManagerSettings
+from hyperforge.feature_flag import get_flag_service
+
+router = APIRouter()
+
+
+@router.get("/metrics")
+async def serve_metrics():  # pragma: no cover
+    output = prometheus_client.exposition.generate_latest()
+    return PlainTextResponse(
+        output.decode("utf8"), headers={"Content-Type": CONTENT_TYPE_LATEST}
+    )
+
+
+@router.get("/health/ready")
+async def health_ready():
+    return {"status": "ok"}
+
+
+@router.get("/health/alive")
+async def health_alive():
+    return {"status": "ok"}
+
+
+class HTTPApplication(FastAPI):
+    agent_manager: AgentManager
+    arag_search: NucliaDBAsync
+    arag_writer: NucliaDBAsync
+    arag_reader: NucliaDBAsync
+    broker: Broker
+    extra_middlewares: Optional[list[Any]] = None
+
+    def __init__(
+        self,
+        settings: Settings,
+        data_manager_settings: DataManagerSettings,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.settings = settings
+        self.data_manager_settings = data_manager_settings
+        self.include_router(internal.router)
+        self.include_router(v1.router)
+        self.include_router(router)
+        self.add_middleware(
+            AuthenticationMiddleware,
+            backend=RaoAuthenticationBackend(),
+        )
+        if self.extra_middlewares is not None:
+            for extra_middleware in self.extra_middlewares:
+                self.add_middleware(extra_middleware)
+        self.add_event_handler("startup", self.startup)
+        self.add_event_handler("shutdown", self.shutdown)
+
+    async def startup(self) -> None:
+        GLOBAL_REGISTRY.clear()
+        await setup_telemetry(SERVICE_NAME)
+        setup_logging(
+            settings=LogSettings(
+                debug=self.settings.debug,
+                log_level=LogLevel(self.settings.log_level),
+                logger_levels={
+                    "uvicorn.error": LogLevel.ERROR,
+                    "nucliadb_telemetry": LogLevel.ERROR,
+                    "mcp.client.streamable_http": LogLevel.WARNING,
+                    "mcp.server.lowlevel.server": LogLevel.WARNING,
+                    "hyperforge.configure": LogLevel.WARNING,
+                },
+            )
+        )
+        if self.settings.sentry_url is not None:
+            set_sentry(
+                self.settings.zone,
+                self.settings.running_environment,
+                self.settings.sentry_url,
+            )
+
+        get_flag_service()  # precache the flag service
+
+        if self.settings.memory_apikey_nucliadb is None:
+            api_key = None
+            headers = {"X-NUCLIADB-ROLES": "WRITER;READER"}
+        else:
+            api_key = self.settings.memory_apikey_nucliadb
+            headers = None
+
+        self.arag_writer = NucliaDBAsync(
+            url=self.settings.memory_writer_nucliadb,
+            api_key=api_key,
+            headers=headers,
+        )
+        self.arag_reader = NucliaDBAsync(
+            url=self.settings.memory_reader_nucliadb,
+            api_key=api_key,
+            headers=headers,
+        )
+        self.arag_search = NucliaDBAsync(
+            url=self.settings.memory_search_nucliadb,
+            api_key=api_key,
+            headers=headers,
+        )
+
+        self.broker = RedisBroker.from_url(
+            url=self.settings.valkey_url,
+            activate_subject=self.settings.activate_subject,
+            keepalive_ms=int(self.settings.pubsub_keepalive_seconds * 1000),
+            cluster_mode=self.settings.valkey_cluster_mode,
+        )
+
+        self.sses: LRU[Tuple[str, str], StreamableHTTPServerTransport] = LRU(size=100)
+        self.mcp_servers: LRU[str, MCPServer] = LRU(size=100)
+
+        self.agent_manager = await AgentManager.from_settings(
+            settings=self.data_manager_settings
+        )
+        await self.agent_manager.initialize()
+
+        for load_module in self.settings.load_modules:
+            try:
+                scan(load_module)
+                load_all_configurations(load_module)
+            except ImportError:
+                logger.error(f"Module {load_module} could not be loaded")
+
+    async def shutdown(self) -> None:
+        await self.agent_manager.finalize()
+        await self.broker.finalize()
+        await clean_telemetry(SERVICE_NAME)
+        GLOBAL_REGISTRY.clear()

hyperforge/api/authentication.py

@@ -0,0 +1,271 @@
+import asyncio
+import functools
+import inspect
+import typing
+from enum import Enum
+from typing import Optional
+
+from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
+from starlette.exceptions import HTTPException
+from starlette.requests import HTTPConnection, Request
+from starlette.responses import RedirectResponse, Response
+from starlette.websockets import WebSocket
+
+
+class User(BaseUser):
+    def __init__(self, username: str, security_groups: list[str] | None = None) -> None:
+        self.username = username
+        self._security_groups = security_groups
+
+    @property
+    def is_authenticated(self) -> bool:
+        return True
+
+    @property
+    def display_name(self) -> str:
+        return self.username
+
+    @property
+    def security_groups(self) -> list[str] | None:
+        return self._security_groups
+
+
+class RaoAuthenticationBackend(AuthenticationBackend):
+    """Authentication backend with a mixture of RAO and NucliaDB auth.
+
+    This mixture is required while migrating /ask endpoint from NucliaDB to RAO,
+    as roles are injected by authorizer and it resolves by path instead of path
+    and service. Thus, we handle NucliaDB auth headers (X-NUCLIADB-*) as well as
+    the regular learning headers (X-STF-*)
+
+    """
+
+    def __init__(self) -> None:
+        self.roles_headers = [
+            "X-STF-ROLES",
+            "X-NUCLIADB-ROLES",
+        ]
+        self.user_headers = [
+            "X-STF-USER",
+            "X-NUCLIADB-USER",
+        ]
+        self.security_groups_headers = ["X-NUCLIADB-SECURITY-GROUPS"]
+
+    async def authenticate(self, request) -> tuple[AuthCredentials, BaseUser] | None:
+        # There are two groups of headers to authenticate: X-STF-* and
+        # X-NUCLIADB-*. As authorizer should only resolve to one set of headers,
+        # we scan and try to find any of both. While endpoint roles are properly
+        # synchronized with authorizer rules, we don't really care which one
+        # there is, nor we will mix them
+
+        auth_creds = None
+        for roles_header in self.roles_headers:
+            if roles_header in request.headers:
+                header_roles = request.headers[roles_header]
+                roles = header_roles.split(";")
+                auth_creds = AuthCredentials(roles)
+                break
+
+        if auth_creds is None:
+            return None
+
+        user = None
+        for user_header in self.user_headers:
+            if user_header in request.headers:
+                user = request.headers[user_header]
+
+                raw_security_groups: str | None = None
+                for security_group_header in self.security_groups_headers:
+                    if security_group_header in request.headers:
+                        raw_security_groups = request.headers[security_group_header]
+                        break
+
+                security_groups: list[str] | None = None
+                if raw_security_groups is not None:
+                    security_groups = raw_security_groups.split(";")
+
+                user = User(username=user, security_groups=security_groups)
+                break
+
+        if user is None:
+            user = User(username="Anonymous")
+
+        return auth_creds, user
+
+
+def has_required_scope(conn: HTTPConnection, scopes: typing.Sequence[str]) -> bool:
+    if conn.auth is None or conn.auth.scopes is None:
+        raise HTTPException(status_code=403, detail="Missing authorizer headers.")
+
+    for scope in scopes:
+        if scope in conn.auth.scopes:
+            return True
+    return False
+
+
+def requires(
+    scopes: typing.Union[str, typing.Sequence[str]],
+    status_code: int = 403,
+    redirect: Optional[str] = None,
+) -> typing.Callable:
+    # As a fastapi requirement, custom Enum classes have to inherit also from
+    # string, so we MUST check for Enum before str
+    if isinstance(scopes, Enum):
+        scopes_list = [scopes.value]
+    elif isinstance(scopes, str):
+        scopes_list = [scopes]
+    elif isinstance(scopes, list):
+        scopes_list = [
+            scope.value if isinstance(scope, Enum) else scope for scope in scopes
+        ]
+
+    def decorator(func: typing.Callable) -> typing.Callable:
+        func.__required_scopes__ = scopes_list  # type: ignore
+        type = None
+        sig = inspect.signature(func)
+        for idx, parameter in enumerate(sig.parameters.values()):
+            if parameter.name == "request" or parameter.name == "websocket":
+                type = parameter.name
+                break
+        else:
+            raise Exception(
+                f'No "request" or "websocket" argument on function "{func}"'
+            )
+
+        if type == "websocket":
+            # Handle websocket functions. (Always async)
+            @functools.wraps(func)
+            async def websocket_wrapper(
+                *args: typing.Any, **kwargs: typing.Any
+            ) -> None:
+                websocket = kwargs.get("websocket", None)
+                assert isinstance(websocket, WebSocket)
+
+                if not has_required_scope(websocket, scopes_list):
+                    await websocket.close()
+                else:
+                    await func(*args, **kwargs)
+
+            return websocket_wrapper
+
+        elif asyncio.iscoroutinefunction(func):
+            # Handle async request/response functions.
+            @functools.wraps(func)
+            async def async_wrapper(
+                *args: typing.Any, **kwargs: typing.Any
+            ) -> Response:
+                request = kwargs.get("request", None)
+                assert isinstance(request, Request)
+
+                if not has_required_scope(request, scopes_list):
+                    if redirect is not None:
+                        return RedirectResponse(
+                            url=request.url_for(redirect), status_code=303
+                        )
+                    raise HTTPException(status_code=status_code)
+                return await func(*args, **kwargs)
+
+            return async_wrapper
+
+        else:
+            # Handle sync request/response functions.
+            @functools.wraps(func)
+            def sync_wrapper(*args: typing.Any, **kwargs: typing.Any) -> Response:
+                request = kwargs.get("request", args[idx])
+                assert isinstance(request, Request)
+
+                if not has_required_scope(request, scopes_list):
+                    if redirect is not None:
+                        return RedirectResponse(
+                            url=request.url_for(redirect), status_code=303
+                        )
+                    raise HTTPException(status_code=status_code)
+                return func(*args, **kwargs)
+
+            return sync_wrapper
+
+    return decorator
+
+
+def requires_one(
+    scopes: typing.Union[str, typing.Sequence[str]],
+    status_code: int = 403,
+    redirect: Optional[str] = None,
+) -> typing.Callable:
+    # As a fastapi requirement, custom Enum classes have to inherit also from
+    # string, so we MUST check for Enum before str
+    if isinstance(scopes, Enum):
+        scopes_list = [scopes.value]
+    elif isinstance(scopes, str):
+        scopes_list = [scopes]
+    elif isinstance(scopes, list):
+        scopes_list = [
+            scope.value if isinstance(scope, Enum) else scope for scope in scopes
+        ]
+
+    def decorator(func: typing.Callable) -> typing.Callable:
+        func.__required_scopes__ = scopes_list  # type: ignore
+        type = None
+        sig = inspect.signature(func)
+        for idx, parameter in enumerate(sig.parameters.values()):
+            if parameter.name == "request" or parameter.name == "websocket":
+                type = parameter.name
+                break
+        else:
+            raise Exception(
+                f'No "request" or "websocket" argument on function "{func}"'
+            )
+
+        if type == "websocket":
+            # Handle websocket functions. (Always async)
+            @functools.wraps(func)
+            async def websocket_wrapper(
+                *args: typing.Any, **kwargs: typing.Any
+            ) -> None:
+                websocket = kwargs.get("websocket", None)
+                assert isinstance(websocket, WebSocket)
+
+                if not has_required_scope(websocket, scopes_list):
+                    await websocket.close()
+                else:
+                    await func(*args, **kwargs)
+
+            return websocket_wrapper
+
+        elif asyncio.iscoroutinefunction(func):
+            # Handle async request/response functions.
+            @functools.wraps(func)
+            async def async_wrapper(
+                *args: typing.Any, **kwargs: typing.Any
+            ) -> Response:
+                request = kwargs.get("request", None)
+                assert isinstance(request, Request)
+
+                if not has_required_scope(request, scopes_list):
+                    if redirect is not None:
+                        return RedirectResponse(
+                            url=request.url_for(redirect), status_code=303
+                        )
+                    raise HTTPException(status_code=status_code)
+                return await func(*args, **kwargs)
+
+            return async_wrapper
+
+        else:
+            # Handle sync request/response functions.
+            @functools.wraps(func)
+            def sync_wrapper(*args: typing.Any, **kwargs: typing.Any) -> Response:
+                request = kwargs.get("request", args[idx])
+                assert isinstance(request, Request)
+
+                if not has_required_scope(request, scopes_list):
+                    if redirect is not None:
+                        return RedirectResponse(
+                            url=request.url_for(redirect), status_code=303
+                        )
+                    raise HTTPException(status_code=status_code)
+                return func(*args, **kwargs)
+
+            return sync_wrapper
+
+    return decorator

hyperforge/api/commands.py

@@ -0,0 +1,33 @@
+import uvicorn
+from nucliadb_telemetry.fastapi import instrument_app
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.utils import get_telemetry
+
+from hyperforge import openapi
+from hyperforge.api import SERVICE_NAME
+from hyperforge.api.app import HTTPApplication
+from hyperforge.api.settings import Settings
+from hyperforge.api.v1.router import router
+from hyperforge.db.settings import DataManagerSettings
+
+
+def run():  # pragma: no cover
+    setup_logging()
+    settings = Settings()
+    data_manager_settings = DataManagerSettings()
+    app = HTTPApplication(
+        settings,
+        data_manager_settings=data_manager_settings,
+    )
+    instrument_app(
+        app,
+        tracer_provider=get_telemetry(SERVICE_NAME),
+        excluded_urls=["/", "/metrics", "/health/ready", "/health/alive"],
+        metrics=True,
+        trace_id_on_responses=True,
+    )
+    uvicorn.run(app, host=settings.http_host, port=settings.http_port)
+
+
+def extract_openapi():
+    openapi.extract_openapi_command("arag", "ARAG API", router)

hyperforge/api/internal/__init__.py

@@ -0,0 +1,4 @@
+from . import inspect
+from .router import router
+
+__all__ = ["inspect", "router"]

hyperforge/api/internal/inspect.py

@@ -0,0 +1,30 @@
+from typing import TYPE_CHECKING
+
+from starlette.requests import Request
+
+from hyperforge.api.internal.router import router
+from hyperforge.api.models import InspectData
+from hyperforge.db.agents import AgentManager
+
+if TYPE_CHECKING:
+    from hyperforge.api.app import HTTPApplication
+
+
+@router.get(
+    "/api/internal/v1/agent/{kbid}",
+    status_code=200,
+    description="Report task is done",
+    tags=["Task"],
+    include_in_schema=False,
+)
+async def inspect_agent_info(request: Request, kbid: str, account: str) -> InspectData:
+    app: HTTPApplication = request.app
+    agent_manager: AgentManager = app.agent_manager
+
+    return InspectData(
+        contexts=await agent_manager.get_context(account, kbid),
+        driver=await agent_manager.get_drivers(account, kbid),
+        postprocess=await agent_manager.get_postprocess(account, kbid),
+        preprocess=await agent_manager.get_preprocess(account, kbid),
+        rules=await agent_manager.get_rules(account, kbid),
+    )

hyperforge/api/internal/router.py

@@ -0,0 +1,3 @@
+from fastapi.routing import APIRouter
+
+router = APIRouter()

hyperforge/api/logging.py

@@ -0,0 +1,18 @@
+from importlib.metadata import version
+from typing import Optional
+
+import sentry_sdk
+from sentry_sdk.integrations.excepthook import ExcepthookIntegration
+
+
+def set_sentry(zone: str, environment: str, sentry_url: Optional[str] = None):
+    if sentry_url:
+        sentry_exception = ExcepthookIntegration(always_run=True)
+        version_num = version("hyperforge")
+        sentry_sdk.init(
+            release=version_num,
+            environment=environment,
+            dsn=sentry_url,
+            integrations=[sentry_exception],
+        )
+        sentry_sdk.set_tag("zone", zone)