hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/redis_utils.py

@@ -0,0 +1,82 @@
+from typing import Dict, Optional
+
+from redis.asyncio.cluster import RedisCluster
+
+
+class ManualStreamKeysRedisCluster(RedisCluster):
+    """When connecting to a cluster, we need to know the key of each command to route it to the correct node. Stream functions (XREADGROUP, XREAD, etc) have a complicated syntax and the client does not extract the key, rather it sends the query to the server for parsing. If the server is down, this fails and the library does not have retries/reconnects for this code path. This can cause the client to get stuck and never reconnect. Since we know the keys, we can do this calculation ourselves to workaround this problem."""
+
+    def get_node_id_from_key(self, key: str):
+        keyslot = self.keyslot(key)
+        return self.nodes_manager.get_node_from_slot(keyslot)
+
+    @classmethod
+    def from_url(cls, url, **kwargs):
+        return super().from_url(
+            url=url,
+            **kwargs,
+        )
+
+    async def xreadgroup(
+        self,
+        groupname: str,
+        consumername: str,
+        streams: Dict[str, str],
+        block: int,
+        count: int,
+        noack: bool,
+    ):
+        # Get the node from the first stream key
+        if len(streams) != 1:
+            raise ValueError("Only one stream key is supported")
+        node = self.get_node_from_key(list(streams.keys())[0])
+        cmd_args = [
+            "XREADGROUP",
+            "GROUP",
+            groupname,
+            consumername,
+        ]
+        if noack:
+            cmd_args.append("NOACK")
+        cmd_args.extend(
+            [
+                "BLOCK",
+                str(block),
+                "COUNT",
+                str(count),
+                "STREAMS",
+                *streams.keys(),
+                *streams.values(),
+            ]
+        )
+
+        return await self.execute_command(*cmd_args, target_nodes=[node])
+
+    async def xread(
+        self,
+        streams: Dict[str, str],
+        block: int,
+        count: Optional[int] = None,
+    ):
+        # Get the node from the first stream key
+        # Get the node from the first stream key
+        if len(streams) != 1:
+            raise ValueError("Only one stream key is supported")
+        node = self.get_node_from_key(list(streams.keys())[0])
+
+        cmd_args = [
+            "XREAD",
+            "BLOCK",
+            str(block),
+        ]
+        if count is not None:
+            cmd_args.extend(["COUNT", str(count)])
+
+        cmd_args.extend(
+            [
+                "STREAMS",
+                *streams.keys(),
+                *streams.values(),
+            ]
+        )
+        return await self.execute_command(*cmd_args, target_nodes=[node])

hyperforge/retrieval/agent.py

@@ -0,0 +1,169 @@
+import asyncio
+from functools import wraps
+from typing import Callable, List, Literal, Optional
+from uuid import uuid4
+
+from sentry_sdk import capture_exception
+
+from hyperforge import logger
+from hyperforge.agent import Agent
+from hyperforge.configure import get_agent_klass
+from hyperforge.context.agent import ContextAgent
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import QuestionMemory
+from hyperforge.retrieval.config import RetrievalAgentConfig
+
+
+def handle_stage_error(stage_name: str):
+    """Decorator to handle errors in retrieval agent stages"""
+
+    def decorator(func: Callable):
+        @wraps(func)
+        async def wrapper(self, memory: QuestionMemory, *args, **kwargs):
+            try:
+                return await func(self, memory, *args, **kwargs)
+            except Exception as e:
+                capture_exception(e)
+                logger.exception(f"Error in {stage_name}")
+                memory.final_answer = f"Error in {stage_name}"
+                raise  # Re-raise to stop execution
+
+        return wrapper
+
+    return decorator
+
+
+class RetrievalAgent(Agent):
+    module: Literal["retrieval"] = "retrieval"
+    debug: bool = False
+    preprocess: Optional[list[Agent]] = None
+    context: Optional[list[ContextAgent]] = None
+    generation: Optional[list[Agent]] = None
+    postprocess: Optional[list[Agent]] = None
+
+    def __init__(
+        self,
+        debug: bool = False,
+        preprocess: Optional[list[Agent]] = None,
+        context: Optional[list[ContextAgent]] = None,
+        generation: Optional[list[Agent]] = None,
+        postprocess: Optional[list[Agent]] = None,
+    ):
+        self.debug = debug
+        self.preprocess = preprocess or []
+        self.context = context or []
+        self.generation = generation or []
+        self.postprocess = postprocess or []
+
+    async def inner_from_config(self, config: object, agent_id: object = None) -> None:  # type: ignore[override]
+        """No-op: RetrievalAgent is constructed via from_config_class, not from_config."""
+        pass
+
+    @classmethod
+    async def from_config_class(cls, config: RetrievalAgentConfig):
+        preprocess = []
+
+        for agent_obj in config.preprocess:
+            agent_class = get_agent_klass(agent_obj.module)
+            preprocess.append(await agent_class.from_config(agent_obj))
+
+        context: list[ContextAgent] = []
+        for context_agent_obj in config.context:
+            agent_class = get_agent_klass(context_agent_obj.module)
+            context.append(await agent_class.from_config(context_agent_obj))  # type: ignore[arg-type]
+
+        generation = []
+        for generation_agent_obj in config.generation:
+            agent_class = get_agent_klass(generation_agent_obj.module)
+            generation.append(await agent_class.from_config(generation_agent_obj))
+
+        postprocess = []
+        for post_agent_obj in config.postprocess:
+            agent_class = get_agent_klass(post_agent_obj.module)
+            postprocess.append(await agent_class.from_config(post_agent_obj))
+
+        return cls(
+            preprocess=preprocess,
+            context=context,
+            postprocess=postprocess,
+            generation=generation,
+        )
+
+    @handle_stage_error("preprocess")
+    async def _run_preprocess(self, memory: QuestionMemory, manager: Manager):
+        if self.preprocess:
+            await asyncio.gather(
+                *[
+                    preprocess(memory=memory, manager=manager)
+                    for preprocess in self.preprocess
+                ]
+            )
+
+    @handle_stage_error("context")
+    async def _run_context(self, memory: QuestionMemory, manager: Manager):
+        if self.context and self.debug is False:
+            questions: List[tuple[str, str]] = memory.get_questions()
+            # Launch all context agents for all questions in parallel
+            tasks = [
+                generation.get_question_context(
+                    memory,
+                    manager,
+                    question_uuid=question_uuid,
+                    question=question,
+                    flow_id=str(uuid4()),
+                )
+                for question_uuid, question in questions
+                for generation in self.context
+            ]
+            await asyncio.gather(*tasks)
+        elif self.context and self.debug is True:
+            questions = memory.get_questions()
+            for generation in self.context:
+                for question_uuid, question in questions:
+                    await generation.get_question_context(
+                        memory,
+                        manager,
+                        question_uuid=question_uuid,
+                        question=question,
+                        flow_id=str(uuid4()),
+                    )
+
+    @handle_stage_error("generation")
+    async def _run_generation(self, memory: QuestionMemory, manager: Manager):
+        if self.generation:
+            await asyncio.gather(
+                *[agent(memory, manager) for agent in self.generation],
+            )
+
+    @handle_stage_error("postprocess")
+    async def _run_postprocess(self, memory: QuestionMemory, manager: Manager):
+        if self.postprocess:
+            await asyncio.gather(
+                *[postprocess(memory, manager) for postprocess in self.postprocess],
+            )
+
+    async def __call__(
+        self,
+        memory: QuestionMemory,
+        manager: Manager,
+    ):
+        while memory.restart:
+            memory.restart = False
+
+            await self._run_preprocess(memory=memory, manager=manager)
+            if memory.secure is False:
+                memory.final_answer = "Insecure query"
+                break
+
+            await self._run_context(memory, manager)
+
+            if memory.final_answer is None:
+                await self._run_generation(memory, manager)
+
+            if memory.restart is False:
+                await memory.add_final_answer()
+                await self._run_postprocess(memory, manager)
+
+            if memory.secure is False:
+                memory.final_answer = "Insecure context retrieved"
+                break

hyperforge/retrieval/config.py

@@ -0,0 +1,94 @@
+import datetime
+from typing import Annotated, Any, Dict, Literal, Union
+
+from pydantic import BaseModel, Field, TypeAdapter, field_serializer, field_validator
+
+from hyperforge.agent import AgentConfig
+from hyperforge.configure import get_agent_config_klass, get_driver_config_klass
+from hyperforge.driver import DriverConfig
+from hyperforge.models import MemoryConfig, Rules
+from hyperforge.prompts import PromptConfig
+from hyperforge.workflows import WorkflowData
+
+
+class RetrievalAgentConfig(BaseModel):
+    drivers: list[DriverConfig]
+    rules: Rules
+    memory: MemoryConfig
+    workflow: WorkflowData
+
+    preprocess: list[AgentConfig]
+    context: list[AgentConfig]
+    generation: list[AgentConfig]
+    postprocess: list[AgentConfig]
+
+    @field_serializer("preprocess", "context", "generation", "postprocess", "drivers")
+    def serialize_agents(self, agents: list[AgentConfig]) -> list[Dict[str, Any]]:
+        return [agent.model_dump() for agent in agents]
+
+    @field_validator("drivers", mode="before")
+    def validate_drivers(cls, value: list[Dict[str, Any]], field):
+        if len(value) == 0:
+            return []
+        if all([isinstance(agent, DriverConfig) for agent in value]):
+            return value
+        result = []
+        for agent_config in value:
+            module = agent_config["provider"]
+            agent_klass = get_driver_config_klass(module)
+            result.append(agent_klass.model_validate(agent_config))
+
+        return result
+
+    @field_validator(
+        "preprocess", "context", "generation", "postprocess", mode="before"
+    )
+    def validate_agents(cls, value: list[Dict[str, Any]], field):
+        if len(value) == 0:
+            return []
+        if all([isinstance(agent, AgentConfig) for agent in value]):
+            return value
+        result = []
+        for agent_config in value:
+            module = agent_config["module"]
+            agent_klass = get_agent_config_klass(module)
+            result.append(agent_klass.model_validate(agent_config))
+
+        return result
+
+    def is_empty(self) -> bool:
+        return (
+            len(self.drivers) == 0
+            and len(self.preprocess) == 0
+            and len(self.context) == 0
+            and len(self.generation) == 0
+            and len(self.postprocess) == 0
+        )
+
+
+class RetrievalAgentExportV1(BaseModel):
+    version: Literal["1"] = Field(default="1")
+    agent_config: RetrievalAgentConfig | None
+    agent_config_workflows: dict[str, RetrievalAgentConfig] = Field(
+        default_factory=dict
+    )
+    prompts: list[PromptConfig]
+    timestamp: str = Field(
+        default_factory=lambda: datetime.datetime.now(datetime.UTC).isoformat(
+            timespec="seconds"
+        ),
+        description="Timestamp of when the export was created in UTC",
+    )
+
+
+RetrievalAgentExport = Annotated[
+    Union[RetrievalAgentExportV1], Field(discriminator="version")
+]
+retrievalAgentAdapter = TypeAdapter(RetrievalAgentExport)
+
+
+class RetrievalAgentExportRequest(BaseModel):
+    passphrase: str = Field(
+        description="Passphrase to encrypt the exported configuration. Will be required for import.",
+        min_length=16,
+    )

hyperforge/server/__init__.py

@@ -0,0 +1,5 @@
+import logging
+
+logger = logging.getLogger("hyperforge.server")
+
+SERVICE_NAME = "hyperforge.server"

hyperforge/server/cache.py

@@ -0,0 +1,131 @@
+from lru import LRU
+from redis.asyncio import Redis
+
+from hyperforge.models import HistoryQuestionAnswer, Source
+
+
+class Cache:
+    async def get(self, key: str) -> str | None:
+        raise NotImplementedError()
+
+    async def set(self, key: str, value: str, expire: int):
+        raise NotImplementedError()
+
+    async def get_list(self, key: str) -> list[str] | None:
+        raise NotImplementedError()
+
+    async def append(self, key: str, values: list[str], limit: int, expire: int):
+        raise NotImplementedError()
+
+
+class NoCache(Cache):
+    async def get(self, key: str) -> str | None:
+        return None
+
+    async def get_list(self, key: str) -> list[str] | None:
+        return None
+
+    async def set(self, key: str, value: str, expire: int):
+        pass
+
+    async def append(self, key: str, values: list[str], limit: int, expire: int):
+        pass
+
+
+class InMemoryCache(Cache):
+    _store: LRU
+    _list_store: LRU
+
+    def __init__(self, size: int = 3000):
+        self._store = LRU(size)
+        self._list_store = LRU(size)
+
+    async def get(self, key: str) -> str | None:
+        return self._store.get(key)
+
+    async def get_list(self, key: str) -> list[str] | None:
+        return self._list_store.get(key)
+
+    async def set(self, key: str, value: str, expire: int):
+        self._store[key] = value
+
+    async def append(self, key: str, values: list[str], limit: int, expire: int):
+        if key not in self._list_store:
+            self._list_store[key] = []
+        self._list_store[key].extend(values)
+        # Enforce limit
+        if len(self._list_store[key]) > limit:
+            self._list_store[key] = self._list_store[key][-limit:]
+
+
+class ValkeyCache(Cache):
+    def __init__(self, client: Redis):
+        self.client = client
+
+    async def get(self, key: str) -> str | None:
+        return await self.client.get(key)
+
+    async def get_list(self, key: str) -> list[str] | None:
+        list = await self.client.lrange(key, 0, -1)
+        if len(list) == 0:
+            # This can mean empty list or non-cached, be specific in return
+            if await self.client.exists(key):
+                return []
+            else:
+                return None
+        else:
+            return list
+
+    async def set(self, key: str, value: str, expire: int):
+        await self.client.set(key, value, ex=expire)
+
+    async def append(self, key: str, values: list[str], limit: int, expire: int):
+        async with self.client.pipeline() as multi:
+            await (
+                multi.rpush(key, *values)
+                .ltrim(key, -limit, -1)
+                .expire(key, expire)
+                .execute()
+            )
+
+
+# The following are small wrapper classes for each cache item, to ensure they are
+# always accessed in a consistent way.
+#
+# Key design notes:
+# - Components separated by dot (.)
+# - From greater to lesser specificity (so it's possible to invalidate by prefix)
+# - Keys include descriptive components before IDs so it's easy to see what a random hex number means
+class CachedNucliaDBSource:
+    def __init__(self, cache: Cache, agent_id: str, source: str):
+        self._cache = cache
+        self._key = f"cache.agent.{agent_id}.source.{source}"
+
+    async def get(self) -> Source | None:
+        value = await self._cache.get(self._key)
+        if value is None:
+            return None
+        return Source.model_validate_json(value)
+
+    async def set(self, source: Source):
+        await self._cache.set(self._key, source.model_dump_json(), expire=900)
+
+
+class CachedSessionQA:
+    def __init__(self, cache: Cache, agent_id: str, session: str):
+        self._cache = cache
+        self._key = f"cache.agent.{agent_id}.session.{session}.qa_history"
+
+    async def get(self) -> list[HistoryQuestionAnswer] | None:
+        value = await self._cache.get_list(self._key)
+        if value is None:
+            return None
+        return [HistoryQuestionAnswer.model_validate_json(v) for v in value]
+
+    async def append(self, qa: HistoryQuestionAnswer):
+        await self.append_all([qa])
+
+    async def append_all(self, qas: list[HistoryQuestionAnswer]):
+        await self._cache.append(
+            self._key, [qa.model_dump_json() for qa in qas], limit=20, expire=3600
+        )

hyperforge/server/run.py

@@ -0,0 +1,109 @@
+import asyncio
+from importlib.metadata import version
+from typing import Optional
+
+import sentry_sdk
+from hyperforge_database.agents import AgentManager
+from hyperforge_database.settings import DataManagerSettings
+from hyperforge_server import SERVICE_NAME
+from hyperforge_server.cache import ValkeyCache
+from hyperforge_server.session import SessionManager
+from hyperforge_server.settings import Settings
+from nucliadb_telemetry.fastapi import application_metrics
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogLevel, LogSettings
+from nucliadb_telemetry.tracerprovider import AsyncTracerProvider
+from nucliadb_telemetry.utils import get_telemetry, setup_telemetry
+from sentry_sdk.integrations.excepthook import ExcepthookIntegration
+
+from hyperforge.broker.redis import RedisBroker
+
+
+def set_sentry(zone: str, environment: str, sentry_url: str):
+    sentry_exception = ExcepthookIntegration(always_run=True)
+    sentry_sdk.init(
+        release=version("hyperforge"),
+        environment=environment,
+        dsn=sentry_url,
+        integrations=[sentry_exception],
+    )
+    sentry_sdk.set_tag("zone", zone)
+
+
+async def run_metrics_server(port: int):
+    import uvicorn
+
+    config = uvicorn.Config(
+        application_metrics, host="0.0.0.0", port=port, log_level="info"
+    )
+    server = uvicorn.Server(config)
+    await server.serve()
+
+
+async def run_server(
+    settings: Settings,
+    tracer: Optional[AsyncTracerProvider],
+    data_manager_settings: DataManagerSettings,
+) -> SessionManager:
+    if tracer:
+        await setup_telemetry(SERVICE_NAME)
+    # Connect to Valkey
+    broker = RedisBroker.from_url(
+        url=settings.valkey_url,
+        activate_subject=settings.activate_subject,
+        keepalive_ms=int(settings.pubsub_keepalive_seconds * 1000),
+        cluster_mode=settings.valkey_cluster_mode,
+    )
+
+    agent_manager = await AgentManager.from_settings(
+        settings=data_manager_settings,
+    )
+    await agent_manager.initialize()
+
+    session = SessionManager(
+        settings=settings,
+        broker=broker,
+        agent_manager=agent_manager,
+        cache=ValkeyCache(broker._client),
+    )
+
+    return session
+
+
+def run():  # pragma: no cover
+    settings = Settings()
+    setup_logging(
+        settings=LogSettings(
+            debug=settings.debug,
+            log_level=LogLevel(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,
+            },
+        )
+    )
+    data_manager_settings = DataManagerSettings()
+    tracer = get_telemetry("nuclia-arag-server")
+    if settings.sentry_url is not None:
+        set_sentry(
+            settings.zone,
+            settings.running_environment,
+            settings.sentry_url,
+        )
+    loop = asyncio.get_event_loop()
+
+    loop.create_task(run_metrics_server(settings.metrics_port))
+
+    session = loop.run_until_complete(
+        run_server(settings, tracer, data_manager_settings)
+    )
+    loop.run_until_complete(session.initialize())
+    try:
+        loop.run_forever()
+    finally:
+        loop.run_until_complete(session.finalize())
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.close()

hyperforge/server/sandbox.py

@@ -0,0 +1,60 @@
+import asyncio
+from importlib.metadata import version
+
+import sentry_sdk
+from nucliadb_telemetry.fastapi import application_metrics
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogLevel, LogSettings
+from sentry_sdk.integrations.excepthook import ExcepthookIntegration
+
+from agents.restricted.src.hyperforge_restricted import sandbox
+from hyperforge.server.settings import Settings
+
+
+def set_sentry(zone: str, environment: str, sentry_url: str):
+    sentry_exception = ExcepthookIntegration(always_run=True)
+    sentry_sdk.init(
+        release=version("hyperforge"),
+        environment=environment,
+        dsn=sentry_url,
+        integrations=[sentry_exception],
+    )
+    sentry_sdk.set_tag("zone", zone)
+
+
+async def run_metrics_server(port: int):
+    import uvicorn
+
+    config = uvicorn.Config(
+        application_metrics, host="0.0.0.0", port=port, log_level="info"
+    )
+    server = uvicorn.Server(config)
+    await server.serve()
+
+
+def run():  # pragma: no cover
+    settings = Settings()
+    sandbox_settings = sandbox.SandboxSettings()
+    setup_logging(
+        settings=LogSettings(
+            debug=settings.debug,
+            log_level=LogLevel(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,
+            },
+        )
+    )
+
+    loop = asyncio.get_event_loop()
+
+    loop.create_task(run_metrics_server(sandbox_settings.sandbox_metrics_port))
+    loop.create_task(sandbox.run_sandbox_server())
+    try:
+        loop.run_forever()
+    finally:
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.close()