hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/db/encryption.py

@@ -0,0 +1,91 @@
+import base64
+import os
+from functools import cache
+from typing import Any
+
+from cryptography.fernet import Fernet, InvalidToken
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+from hyperforge.db.settings import EncryptionSettings
+from hyperforge.driver import DriverConfig, EncryptedPayload
+
+
+@cache
+def get_fernet() -> Fernet:
+    settings = EncryptionSettings()
+    return Fernet(settings.encryption_secret_key)
+
+
+def encrypt_data(data: str) -> str:
+    f = get_fernet()
+    return f.encrypt(data.encode()).decode()
+
+
+def decrypt_data(data: str) -> str:
+    f = get_fernet()
+    try:
+        return f.decrypt(token=data, ttl=None).decode()
+    except InvalidToken:
+        raise ValueError("Invalid encryption token.")
+
+
+def fernet_key_from_passphrase(
+    passphrase: str, salt: bytes | None
+) -> tuple[bytes, bytes]:
+    """Generate a Fernet key from a passphrase and salt (if provided).
+    From https://cryptography.io/en/latest/fernet/#using-passwords-with-fernet
+    """
+    if salt is None:
+        salt = os.urandom(16)
+    kdf = PBKDF2HMAC(
+        algorithm=hashes.SHA256(),
+        length=32,
+        salt=salt,
+        iterations=1_200_000,
+    )
+    key = base64.urlsafe_b64encode(kdf.derive(passphrase.encode()))
+    return key, salt
+
+
+def dump_without_encrypted_fields(
+    model: DriverConfig[EncryptedPayload],
+) -> dict[str, Any]:
+    data = model.model_dump()  # type: ignore
+    for field in model.config.encrypted_fields:  # type: ignore
+        if field in data["config"]:
+            del data["config"][field]
+    return {"config": data}
+
+
+def encrypt_fields(model: EncryptedPayload) -> dict[str, Any]:
+    data = model.model_dump()  # type: ignore
+    for field in model.encrypted_fields:  # type: ignore
+        if field in data:
+            if isinstance(data[field], str):
+                data[field] = encrypt_data(data=data[field])
+            elif isinstance(data[field], dict):
+                for k, v in data[field].items():
+                    if isinstance(v, str):
+                        data[field][k] = encrypt_data(data=v)
+    return data
+
+
+def decrypt_fields(model: EncryptedPayload) -> None:
+    for field in model.encrypted_fields:
+        if not hasattr(model, field):
+            raise AttributeError(f"Field '{field}' not found in {type(model).__name__}")
+        value = getattr(model, field)
+        if value is not None:
+            try:
+                if isinstance(value, str):
+                    value = decrypt_data(data=value)
+                elif isinstance(value, dict):
+                    for k, v in value.items():
+                        if isinstance(v, str):
+                            value[k] = decrypt_data(data=v)
+            except ValueError:
+                # We ignore the error to support current unencrypted data
+                continue
+
+            setattr(model, field, value)

hyperforge/db/exceptions.py

@@ -0,0 +1,26 @@
+class RetrievalAgentError(Exception):
+    """Base class for exceptions in Retrieval Agents."""
+
+
+class NotFoundError(RetrievalAgentError):
+    """Exception raised when a requested resource is not found."""
+
+
+class DriverNotFoundError(NotFoundError):
+    """Exception raised when the specified driver is not found."""
+
+
+class ProtectedWorkflowError(RetrievalAgentError):
+    """Exception raised when a protected workflow is modified."""
+
+
+class ParseExportError(RetrievalAgentError):
+    """Exception raised when there is an error parsing the export file."""
+
+
+class ExportEncryptionError(RetrievalAgentError):
+    """Exception raised when there is an error with encryption."""
+
+
+class InvalidTargetAgentError(RetrievalAgentError):
+    """Exception raised when the target agent is invalid for an operation."""

hyperforge/db/settings.py

@@ -0,0 +1,16 @@
+from pydantic_settings import BaseSettings
+
+
+class DataManagerSettings(BaseSettings):
+    postgresql_dsn: str
+    export_read_chunk_size: int = 1024 * 1024  # 1 MB
+    export_read_max_size: int = 10 * 1024 * 1024  # 10 MB
+
+
+class EncryptionSettings(BaseSettings):
+    encryption_secret_key: str
+
+
+class IDPSettings(BaseSettings):
+    dummy_idp: bool = False
+    idp_regional_grpc: str = "idp-grpc.idp-regional.svc.cluster.local:9090"

hyperforge/db/workflow_cleanup.py

@@ -0,0 +1,69 @@
+import asyncio
+import datetime
+from importlib.metadata import version
+
+from nucliadb_telemetry import errors
+from nucliadb_telemetry.errors import setup_error_handling
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogLevel, LogSettings
+
+from hyperforge.db import logger
+from hyperforge.db.agents import WORKFLOW_PURGE_RETENTION, AgentManager
+from hyperforge.db.settings import DataManagerSettings
+
+
+async def cleanup_deleted_workflows(
+    manager: AgentManager,
+    older_than: datetime.timedelta = WORKFLOW_PURGE_RETENTION,
+):
+    logger.info("Cleaning up deleted workflows")
+    workflows = await manager.get_expired_deleted_workflows(older_than=older_than)
+    logger.info("Found deleted workflows to clean up", extra={"count": len(workflows)})
+
+    for workflow in workflows:
+        try:
+            logger.info(
+                "Purging deleted workflow",
+                extra={
+                    "account": workflow["account"],
+                    "agent_id": workflow["agent_id"],
+                    "workflow_id": workflow["workflow_id"],
+                },
+            )
+            await manager.purge_deleted_workflow(
+                account=workflow["account"],
+                agent_id=workflow["agent_id"],
+                workflow_id=workflow["workflow_id"],
+            )
+        except Exception as exc:
+            errors.capture_exception(exc)
+            logger.error(
+                "Failed to purge deleted workflow",
+                exc_info=exc,
+                extra={
+                    "account": workflow["account"],
+                    "agent_id": workflow["agent_id"],
+                    "workflow_id": workflow["workflow_id"],
+                },
+            )
+
+
+async def cronjob(manager: AgentManager):
+    await cleanup_deleted_workflows(manager)
+
+
+def run():  # pragma: no cover
+    asyncio.run(_main())
+
+
+async def _main():
+    log_settings = LogSettings(logger_levels={"hyperforge.server": LogLevel.INFO})
+    setup_logging(settings=log_settings)
+    setup_error_handling(version("hyperforge"))
+    data_manager_settings = DataManagerSettings()
+    manager = await AgentManager.from_settings(settings=data_manager_settings)
+    await manager.initialize()
+    try:
+        await cronjob(manager)
+    finally:
+        await manager.finalize()

hyperforge/definition.py

@@ -0,0 +1,13 @@
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+
+class NUAConfig(BaseModel):
+    key: str
+
+
+class FunctionDefinition(BaseModel):
+    name: str
+    description: str
+    parameters: Dict[str, Any]

hyperforge/driver.py

@@ -0,0 +1,31 @@
+from typing import Any, ClassVar, Generic, Self, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EncryptedPayload(BaseModel):
+    encrypted_fields: ClassVar[list[str]] = []
+
+
+T = TypeVar("T", bound="EncryptedPayload")
+
+
+class DriverConfig(BaseModel, Generic[T]):
+    id: str | None = None
+    identifier: str
+    name: str
+    provider: Any = Field(
+        ..., description="The type of driver, e.g., 'google', 'marklogic', etc."
+    )
+    config: T = Field(..., description="The configuration specific to the driver.")
+
+
+class Driver(BaseModel):
+    name: str
+    provider: str
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @classmethod
+    async def init(cls, driver: Any) -> Self:
+        raise NotImplementedError()

hyperforge/dummy.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+from stashify_protos.protos import idp_pb2
+
+
+class DummyIDPRegionalGRPCUtility:
+    requests: list[Any]
+
+    def __init__(self) -> None:
+        self.requests = []
+
+    async def SendRAOExportEmail(
+        self, payload: idp_pb2.SendRAOExportEmailRequest
+    ) -> idp_pb2.SendRAOExportEmailResponse:
+        self.requests.append(payload)
+        return idp_pb2.SendRAOExportEmailResponse(
+            status=idp_pb2.SendRAOExportEmailResponse.Status.OK,
+            message="Dummy IDP Regional GRPC Utility: Email sent successfully",
+        )
+
+    async def SendEmail(
+        self, payload: idp_pb2.SendEmailRequest
+    ) -> idp_pb2.SendEmailResponse:
+        self.requests.append(payload)
+        return idp_pb2.SendEmailResponse(
+            status=idp_pb2.SendEmailResponse.Status.OK,
+            message="Dummy IDP Regional GRPC Utility: Email sent successfully",
+        )

hyperforge/engine.py

@@ -0,0 +1,189 @@
+import logging
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, cast
+
+from nuclia.lib.nua import AsyncNuaClient
+
+from hyperforge.configure import GLOBAL_REGISTRY, load_all_configurations
+from hyperforge.interaction import AragAnswer
+from hyperforge.llm import NoopNuaClient, NuaBaseModel, NUAConnection
+from hyperforge.manager import Manager
+from hyperforge.memory.memory import BaseSessionMemory, QuestionMemory, SessionMemory
+from hyperforge.retrieval.agent import RetrievalAgent
+from hyperforge.retrieval.config import RetrievalAgentConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class State:
+    agent: RetrievalAgent
+    manager: Manager
+
+
+async def init(
+    config: Optional[Dict[str, Any]] = None,
+    agent_id: str = "default",
+    internal_nua: bool = False,
+    internal_nua_api: str = "http://predict.learning.svc.cluster.local:8080",
+    local_openai: Optional[str] = None,
+    local_openai_model: Optional[str] = None,
+    external_nua_api_key: Optional[str] = None,
+    loaded_modules: list[str] = [],
+    retrieval_config: Optional[RetrievalAgentConfig] = None,
+    session_id: str = "default_session",
+    memory_klass: type[BaseSessionMemory] = SessionMemory,
+) -> Tuple[State, SessionMemory]:
+    from hyperforge.configure import scan
+
+    for load_module in loaded_modules:
+        try:
+            scan(load_module)
+            load_all_configurations(load_module)
+        except ImportError:
+            logger.error(f"Module {load_module} could not be loaded")
+
+    if retrieval_config is None:
+        if config is None:
+            raise ValueError("Either config or retrieval_config must be provided")
+        retrieval_config = RetrievalAgentConfig.model_validate(config)
+
+    state = await get_state(
+        agent_id=agent_id,
+        config=retrieval_config,
+        internal_nua=internal_nua,
+        internal_nua_api=internal_nua_api,
+        local_openai=local_openai,
+        local_openai_model=local_openai_model,
+        external_nua_api_key=external_nua_api_key,
+    )
+    session_memory = memory_klass.from_config(
+        retrieval_config.memory,
+        agent_id=agent_id,
+        workflow_id="default",
+        rules=retrieval_config.rules,
+    )
+    session_memory.init(session_id)
+    return state, session_memory
+
+
+async def main(
+    config: Optional[Dict[str, Any]] = None,
+    agent_id: str = "default",
+    internal_nua: bool = False,
+    internal_nua_api: str = "http://predict.learning.svc.cluster.local:8080",
+    local_openai: Optional[str] = None,
+    local_openai_model: Optional[str] = None,
+    external_nua_api_key: Optional[str] = None,
+    question: str = "",
+    loaded_modules: list[str] = [],
+    retrieval_config: Optional[RetrievalAgentConfig] = None,
+    callback: Optional[Callable[[AragAnswer], Awaitable[None]]] = None,
+    session_id: str = "default_session",
+    user_metadata: Optional[Dict[str, str]] = None,
+    headers: Optional[Dict[str, str]] = None,
+    memory_klass: type[BaseSessionMemory] = SessionMemory,
+    streaming: bool = False,
+) -> QuestionMemory:
+    try:
+        state, session_memory = await init(
+            config=config,
+            agent_id=agent_id,
+            internal_nua=internal_nua,
+            internal_nua_api=internal_nua_api,
+            local_openai=local_openai,
+            local_openai_model=local_openai_model,
+            external_nua_api_key=external_nua_api_key,
+            loaded_modules=loaded_modules,
+            retrieval_config=retrieval_config,
+            session_id=session_id,
+            memory_klass=memory_klass,
+        )
+        question_memory = session_memory.start_question(question, streaming=streaming)
+        if callback is not None:
+            question_memory.set_callback_fn(callback)
+
+        if user_metadata:
+            question_memory.session.user_info.update(user_metadata)
+
+        if headers:
+            question_memory.headers.update(headers)
+
+        if state.agent is None:
+            raise ValueError("Agent could not be initialized")
+
+        await state.agent(
+            question_memory,
+            state.manager,
+        )
+    except Exception as e:
+        raise e
+    finally:
+        GLOBAL_REGISTRY.clear()
+    return question_memory
+
+
+async def engine(
+    manager: Manager,
+    agent: RetrievalAgent,
+    question_memory: QuestionMemory,
+    user_metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    if user_metadata is not None:
+        question_memory.session.user_info.update(user_metadata)
+
+    await agent(
+        question_memory,
+        manager,
+    )
+
+
+async def get_state(
+    agent_id: str,
+    config: RetrievalAgentConfig,
+    internal_nua_api: str = "http://predict.learning.svc.cluster.local:8080",
+    internal_nua: bool = False,
+    local_openai: Optional[str] = None,
+    local_openai_model: Optional[str] = None,
+    external_nua_api_key: Optional[str] = None,
+    account: Optional[str] = None,
+    kbid: Optional[str] = None,
+    local_openai_model_klass: Optional[type[NuaBaseModel]] = None,
+) -> State:
+    nua: AsyncNuaClient
+    if internal_nua:
+        nua = cast(
+            AsyncNuaClient,
+            await NUAConnection.connect_internal(
+                kbid=kbid, account=account, url=internal_nua_api
+            ),
+        )
+
+    elif local_openai is not None and local_openai_model_klass is not None:
+        nua = await local_openai_model_klass.model_validate(
+            {
+                "key": external_nua_api_key,
+                "local_openai": local_openai,
+                "local_openai_model": local_openai_model,
+            }
+        ).connect()
+
+    elif external_nua_api_key is not None:
+        nua = await NUAConnection.model_validate(
+            {
+                "key": external_nua_api_key,
+            }
+        ).connect()
+
+    else:
+        logger.warning(
+            "No LLM backend configured — use a no-op client.  Agents that don't"
+            " require LLM calls (e.g. the built-in ``static`` context agent) will"
+            " work fine; any agent that actually calls NUA will raise a clear error."
+        )
+        nua = NoopNuaClient()
+
+    manager = await Manager.from_config(drivers=config.drivers, nua=nua)
+    agent = await RetrievalAgent.from_config_class(config)
+
+    return State(manager=manager, agent=agent)

hyperforge/exceptions.py

@@ -0,0 +1,14 @@
+class AutheticationException(Exception):
+    pass
+
+
+class NoAvailableAgents(Exception):
+    pass
+
+
+class MaxRetries(Exception):
+    pass
+
+
+class ModelRetry(Exception):
+    pass

hyperforge/feature_flag.py

@@ -0,0 +1,105 @@
+import hashlib
+import json
+import threading
+from functools import cached_property
+from typing import Any, Optional
+
+import mrflagly  # type: ignore
+import pydantic
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+_flag_service = threading.local()
+
+
+class FlagContext(pydantic.BaseModel):
+    """Class to hold information that can then be passed to the feature flag service to determine if a flag is enabled or not.
+
+    Having it as a class allows to simply pass this object down to any function that needs to check for a feature flag and have all the relevant information in one place, instead of having to pass multiple parameters or a dictionary.
+    """
+
+    account_id: Optional[str] = None
+    kbid: str
+
+    @pydantic.computed_field  # type: ignore[prop-decorator]
+    @cached_property
+    def account_id_md5(self) -> Optional[str]:
+        if self.account_id is not None:
+            return hashlib.md5(self.account_id.encode()).hexdigest()
+        return None
+
+
+class Features:
+    # RAO
+    FILTERED_AGENTS_FEATURE_FLAG = "rao_filtered_agents_{environment}"
+    FILTERED_DRIVERS_FEATURE_FLAG = "rao_filtered_drivers_{environment}"
+    AUDIT_RAO_ASK_ENDPOINT = "learning_audit_rao_ask"
+    RAO_ACCOUNT_ENABLED_FEATURE_FLAG = "rao_account_enabled_{account_md5}"
+
+
+class FlagService:
+    def __init__(self):
+        self.settings = Settings()
+        if self.settings.flag_settings_url is None:
+            self.flag_service = mrflagly.FlagService(data=json.dumps(DEFAULT_FLAG_DATA))
+        else:
+            self.flag_service = mrflagly.FlagService(
+                url=self.settings.flag_settings_url
+            )
+
+    def enabled(
+        self, flag_key: str, default: bool = False, context: Optional[dict] = None
+    ) -> bool:
+        if context is None:
+            context = {}
+        context["environment"] = self.settings.running_environment
+
+        return self.flag_service.enabled(flag_key, default=default, context=context)
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_parse_none_str="null")
+    running_environment: str = pydantic.Field(
+        default="local",
+        validation_alias=pydantic.AliasChoices("environment", "running_environment"),
+    )
+    flag_settings_url: str | None = (
+        "https://cdn.rag.progress.cloud/features/features-v2.json"
+    )
+
+
+DEFAULT_FLAG_DATA: dict[str, Any] = {
+    # These are just defaults to use for local dev and tests
+    Features.FILTERED_AGENTS_FEATURE_FLAG.format(environment="local"): {
+        "rollout": 0,
+        "variants": {"agents": []},
+    },
+    Features.FILTERED_DRIVERS_FEATURE_FLAG.format(environment="local"): {
+        "rollout": 0,
+        "variants": {"drivers": []},
+    },
+    Features.AUDIT_RAO_ASK_ENDPOINT: {
+        "rollout": 0,
+        "variants": {"environment": ["local"]},
+    },
+}
+
+
+def get_flag_service() -> FlagService:
+    if getattr(_flag_service, "service", None) is None:
+        _flag_service.service = FlagService()
+    return _flag_service.service
+
+
+def has_feature(
+    flag_key: str,
+    default: bool = False,
+    context: dict[str, str] | None | FlagContext = None,
+) -> bool:
+    fs = get_flag_service()
+
+    if isinstance(context, FlagContext):
+        context_dict = context.model_dump()
+    else:
+        context_dict = context or {}
+
+    return fs.enabled(flag_key, default=default, context=context_dict)