nucliadb-agentic-api 1.0.0.post156__py3-none-any.whl

nucliadb_agentic_api/db/transform.py

@@ -0,0 +1,208 @@
+from typing import Any, Dict, Tuple
+from uuid import uuid4
+
+from hyperforge.driver import DriverConfig
+from hyperforge.models import MemoryConfig, Rules
+from hyperforge.retrieval.config import RetrievalAgentConfig
+from hyperforge.workflows import WorkflowData
+from hyperforge_mcp.config import Transport
+from hyperforge_mcp.config_driver import (
+    MCPHTTPDriverConfig,
+    MCPHTTPInnerConfig,
+)
+from hyperforge_nucliadb.driver_config import (
+    NucliaDBConfig,
+    NucliaDBConnection,
+)
+from hyperforge_nucliadb_agentic.ask.model import AskRequest
+from hyperforge_nucliadb_agentic.internal_driver import (
+    InternalNucliaDBConfig,
+    InternalNucliaDBConnection,
+)
+
+from nucliadb_agentic_api.db.sources import Sources
+from nucliadb_agentic_api.models import AgenticConfigSchema
+
+
+async def transform_agentic_config(
+    agentic_config: AgenticConfigSchema,
+    source_manager: Sources,
+    account: str,
+    internal_nucliadb: bool = True,
+    internal_nucliadb_url: str | None = None,
+    external_nucliadb_url: str | None = None,
+    external_nucliadb_key: str | None = None,
+    ask_request: AskRequest | None = None,
+    kbid: str = "",
+) -> Tuple[RetrievalAgentConfig, Dict[str, DriverConfig], list[str]]:
+    drivers: Dict[str, DriverConfig] = {}
+    global_driver = []
+
+    title = agentic_config.title if agentic_config.title else "Default Agentic Config"
+
+    if agentic_config.rephrase:
+        config: Dict[str, Any] = {
+            "title": f"{title} - Retrieval",
+        }
+        if agentic_config.rephrase.model:
+            config["model"] = agentic_config.rephrase.model
+        if agentic_config.rephrase.prompt:
+            config["rules"] = [agentic_config.rephrase.prompt]
+        if agentic_config.rephrase.ask_to:
+            # Same KB with different layers
+            config["kbid"] = agentic_config.rephrase.ask_to
+        if agentic_config.rephrase.history:
+            config["history"] = agentic_config.rephrase.history
+
+        config["module"] = "rephrase"
+        preprocess = [config]
+    else:
+        preprocess = []
+
+    if agentic_config.smart_agent:
+        config = {
+            "title": f"{title} - Smart Agent",
+        }
+        if agentic_config.smart_agent.extra_prompt:
+            config["extra_prompt"] = agentic_config.smart_agent.extra_prompt
+        if agentic_config.smart_agent.models:
+            if agentic_config.smart_agent.models.context_validation:
+                config["context_validation_model"] = (
+                    agentic_config.smart_agent.models.context_validation
+                )
+            if agentic_config.smart_agent.models.planner:
+                config["planner_model"] = agentic_config.smart_agent.models.planner
+            if agentic_config.smart_agent.models.executor:
+                config["executor_model"] = agentic_config.smart_agent.models.executor
+            if agentic_config.smart_agent.history:
+                config["history"] = agentic_config.smart_agent.history
+            if agentic_config.smart_agent.extra_prompt:
+                config["extra_prompt"] = agentic_config.smart_agent.extra_prompt
+        registered_agents = []
+        for source in agentic_config.smart_agent.sources:
+            source_obj = await source_manager.get_source(account, kbid, source)
+            source_config: Dict[str, Any] = {
+                "title": f"{title} - Smart Agent - {source_obj.type.capitalize()} Source",
+            }
+
+            if source_obj.type == "nucliadb":
+                # Same KB different
+                uid = uuid4().hex
+
+                source_config["sources"] = [uid]
+                source_config["module"] = "nucliadb_agent"
+
+                ndb_driver_config: DriverConfig
+                if internal_nucliadb and internal_nucliadb_url:
+                    ndb_driver_config = InternalNucliaDBConfig(
+                        identifier=uid,
+                        name="NucliaDB",
+                        provider="nucliadb_internal",
+                        config=InternalNucliaDBConnection(
+                            url=internal_nucliadb_url,
+                            key=None,
+                            manager="",
+                            description="",
+                            kbid=kbid,
+                            filter_expression=source_obj.filter_expression,
+                            filters=source_obj.resource_filters
+                            if source_obj.resource_filters
+                            else [],
+                        ),
+                    )
+                elif external_nucliadb_url:
+                    ndb_driver_config = NucliaDBConfig(
+                        identifier=uid,
+                        name="NucliaDB",
+                        provider="nucliadb",
+                        config=NucliaDBConnection(
+                            url=external_nucliadb_url,
+                            key=external_nucliadb_key,
+                            manager="",
+                            description="",
+                            kbid=kbid,
+                            filter_expression=source_obj.filter_expression,
+                            filters=source_obj.resource_filters
+                            if source_obj.resource_filters
+                            else [],
+                        ),
+                    )
+                else:
+                    raise ValueError(
+                        "No NucliaDB URL configured for internal or external access"
+                    )
+                drivers[uid] = ndb_driver_config
+                registered_agents.append(source_config)
+            elif source_obj.type == "mcp":
+                uid = uuid4().hex
+                source_config["sources"] = [uid]
+                source_config["transport"] = Transport.HTTP
+                source_config["module"] = "mcp"
+                mcp_driver_config = MCPHTTPDriverConfig(
+                    identifier=uid,
+                    name="MCPHTTP",
+                    provider="mcphttp",
+                    config=MCPHTTPInnerConfig(
+                        uri=source_obj.uri,  # TODO: pass real URL if needed
+                        headers=source_obj.headers if source_obj.headers else {},
+                    ),
+                )  # TODO: pass real config if needed
+                drivers[uid] = mcp_driver_config
+                registered_agents.append(source_config)
+
+            elif source_obj.type == "google":
+                source_config["module"] = "google"
+                global_driver.append("google")
+                registered_agents.append(source_config)
+
+            elif source_obj.type == "perplexity":
+                source_config["module"] = "perplexity"
+                global_driver.append("perplexity")
+                registered_agents.append(source_config)
+
+        config["registered_agents"] = registered_agents
+        config["module"] = "smart"
+        context = [config]
+    else:
+        context = []
+
+    if agentic_config.summarize:
+        config = {
+            "title": f"{title} - Summarize Agent",
+        }
+        if agentic_config.summarize.model:
+            config["model"] = agentic_config.summarize.model
+        if agentic_config.summarize.user_prompt:
+            config["user_prompt"] = agentic_config.summarize.user_prompt
+        if agentic_config.summarize.system_prompt:
+            config["system_prompt"] = agentic_config.summarize.system_prompt
+        if agentic_config.summarize.conversational:
+            config["conversational"] = agentic_config.summarize.conversational
+        if agentic_config.summarize.history:
+            config["history"] = agentic_config.summarize.history
+
+        config["module"] = "summarize"
+
+        generation = [config]
+    else:
+        generation = []
+
+    return (
+        RetrievalAgentConfig(
+            drivers=[],
+            rules=Rules(),
+            memory=MemoryConfig(),
+            workflow=WorkflowData(
+                id="default",
+                name="Default Workflow",
+                description="Default workflow for agentic config transformation",
+                parameters=None,
+            ),
+            preprocess=preprocess,  # type: ignore
+            context=context,  # type: ignore
+            generation=generation,  # type: ignore
+            postprocess=[],
+        ),
+        drivers,
+        global_driver,
+    )

nucliadb_agentic_api/exceptions.py

@@ -0,0 +1,10 @@
+class Conflict(Exception):
+    """Raised when a resource already exists and cannot be created again."""
+
+
+class NotFound(Exception):
+    """Raised when a resource is not found."""
+
+
+class InvalidReference(Exception):
+    """Raised when a referenced resource does not exist."""

nucliadb_agentic_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("nucliadb_agentic_api")
+        sentry_sdk.init(
+            release=version_num,
+            environment=environment,
+            dsn=sentry_url,
+            integrations=[sentry_exception],
+        )
+        sentry_sdk.set_tag("zone", zone)

nucliadb_agentic_api/models.py

@@ -0,0 +1,321 @@
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from hyperforge.driver import DriverConfig
+from hyperforge.models import Rules
+from nucliadb_models import FilterExpression, TextFormat
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+
+class StashRoles(str, Enum):
+    # Can do anything at the stash
+    OWNER = "SOWNER"
+
+    # Can access the stash
+    MEMBER = "SMEMBER"
+
+    # Can access the stash
+    CONTRIBUTOR = "SCONTRIBUTOR"
+
+
+class AccountRoles(str, Enum):
+    OWNER = "AOWNER"
+    MEMBER = "AMEMBER"
+
+
+class AgentRole(str, Enum):
+    MEMBER = "SESSIONMEMBER"
+
+
+class NucliaDBRoles(str, Enum):
+    MANAGER = "MANAGER"
+    READER = "READER"
+    WRITER = "WRITER"
+
+
+class UserType(str, Enum):
+    ROOT = "ROOT"
+    DEALER = "DEALER"
+    USER = "USER"
+    READONLY = "READONLY"
+    MANAGER = "MANAGER"
+    SALES = "SALES"
+
+
+class AccountTypes(str, Enum):
+    TRIAL = "stash-trial"
+    STARTER = "stash-starter"
+    GROWTH = "stash-growth"
+    STARTUP = "stash-startup"
+    ENTERPRISE = "stash-enterprise"
+
+    # will be removed at some point in the near future
+    DEVELOPER = "stash-developer"
+    BUSINESS = "stash-business"
+
+    # V3 account types
+    V3_STARTER = "v3starter"
+    V3_FLY = "v3fly"
+    V3_GROWTH = "v3growth"
+    V3_PRO = "v3pro"
+    V3_ENTERPRISE = "v3enterprise"
+    COWORK = "cowork"
+
+
+class SessionData(BaseModel):
+    slug: str
+    name: str
+    summary: str
+    data: str
+    format: TextFormat
+
+
+INFO_FIELD_ID = "info"
+
+DEFAULT_RESOURCE_LIST_PAGE_SIZE = 20
+
+
+class InspectData(BaseModel):
+    contexts: List[Any]
+    driver: List[DriverConfig]
+    postprocess: List[Any]
+    preprocess: List[Any]
+    rules: Rules
+
+
+class AgentID(BaseModel):
+    id: str
+
+
+class DriverID(BaseModel):
+    id: str
+
+
+class PromptID(BaseModel):
+    id: str
+
+
+class InteractionsAuditDownloadRequest(BaseModel):
+    session_id: Optional[str] = Field(default=None, description="Filter by session ID")
+    year: Optional[int] = Field(
+        default=None,
+        description="Filter by year (e.g., 2024). If not specified, defaults to the current year.",
+    )
+    month: Optional[int] = Field(
+        default=None,
+        description="Filter by month (1-12). If not specified, defaults to the past month.",
+    )
+
+
+class DownloadStatus(BaseModel):
+    id: str
+    type: str
+    status: Literal["pending", "ready"]
+    download_url: str | None
+    query: dict[str, Any]
+
+
+class InteractionOperation(int, Enum):
+    QUESTION = 0
+    QUIT = 1
+
+
+class InteractionRequest(BaseModel):
+    question: str
+    headers: Dict[str, str] = {}
+    arguments: Dict[str, str] = {}
+    operation: InteractionOperation = InteractionOperation.QUESTION
+    streaming: bool = False
+
+
+class AgenticRephraseConfiguration(BaseModel):
+    ask_to: Optional[str] = None
+    prompt: Optional[str] = None
+    model: Optional[str] = None
+    history: Optional[bool] = None
+
+
+class AgenticSmartAgentMode(str, Enum):
+    REACTIVE = "reactive"
+    PLAN_EXECUTE = "plan_execute"
+
+
+class AgenticSmartAgentModels(BaseModel):
+    context_validation: Optional[str] = None
+    planner: Optional[str] = None
+    executor: Optional[str] = None
+
+
+class NucliaDBAgenticSource(BaseModel):
+    type: Literal["nucliadb"] = "nucliadb"
+    source_id: Optional[str] = Field(
+        default=None,
+        description="ID of an existing source in the sources table",
+    )
+    description: Optional[str] = None
+    filter_expression: Optional[str] = None
+    connection: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
+
+
+class GoogleAgenticSource(BaseModel):
+    type: Literal["google"] = "google"
+    source_id: Optional[str] = Field(
+        default=None,
+        description="ID of an existing source in the sources table",
+    )
+    description: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
+
+
+class PerplexityAgenticSource(BaseModel):
+    type: Literal["perplexity"] = "perplexity"
+    source_id: Optional[str] = Field(
+        default=None,
+        description="ID of an existing source in the sources table",
+    )
+    description: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
+
+
+class MCPAgenticSource(BaseModel):
+    type: Literal["mcp"] = "mcp"
+    source_id: Optional[str] = Field(
+        default=None,
+        description="ID of an existing source in the sources table",
+    )
+    description: Optional[str] = None
+    uri: str
+    headers: Optional[Dict[str, Any]] = None
+    tool_choice_model: Optional[str] = None
+    valid_headers: Optional[List[str]] = None
+
+
+class SyncAgenticSource(BaseModel):
+    type: Literal["sync"] = "sync"
+    source_id: Optional[str] = Field(
+        default=None,
+        description="ID of an existing source in the sources table",
+    )
+    description: Optional[str] = None
+    connection: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
+
+
+AgenticSource = Annotated[
+    NucliaDBAgenticSource
+    | GoogleAgenticSource
+    | PerplexityAgenticSource
+    | MCPAgenticSource
+    | SyncAgenticSource,
+    Field(discriminator="type"),
+]
+
+
+class AgenticSmartAgentConfiguration(BaseModel):
+    mode: AgenticSmartAgentMode = AgenticSmartAgentMode.REACTIVE
+    extra_prompt: Optional[str] = None
+    models: Optional[AgenticSmartAgentModels] = None
+    sources: List[str] = Field(default_factory=list)
+    history: Optional[bool] = None
+
+
+class AgenticSummarizeConfiguration(BaseModel):
+    user_prompt: Optional[str] = None
+    system_prompt: Optional[str] = None
+    conversational: bool = False
+    model: Optional[str] = None
+    history: Optional[bool] = None
+
+
+class AgenticConfigSchema(BaseModel):
+    title: Optional[str] = None
+    rephrase: Optional[AgenticRephraseConfiguration] = None
+    smart_agent: Optional[AgenticSmartAgentConfiguration] = None
+    summarize: Optional[AgenticSummarizeConfiguration] = None
+
+
+# ---------------------------------------------------------------------------
+# Source models
+# ---------------------------------------------------------------------------
+
+
+class GoogleTimeRange(str, Enum):
+    PAST_DAY = "past_day"
+    PAST_WEEK = "past_week"
+    PAST_MONTH = "past_month"
+    PAST_YEAR = "past_year"
+
+
+class NucliaDBSourceSchema(BaseModel):
+    type: Literal["nucliadb"] = "nucliadb"
+    description: Optional[str] = None
+
+    filter_expression: FilterExpression | None = Field(
+        default=None,
+        title="Filter expression",
+        description="Filter expression to narrow NucliaDB search results",
+    )
+    labels: List[str] | None = Field(
+        default=None,
+        description="Label filters to restrict which resources are searched",
+    )
+    resource_filters: List[str] | None = Field(
+        default=None,
+        title="Resource filters",
+        description="Additional resource-level filters passed to the NucliaDB search API",
+    )
+
+
+class PerplexitySourceSchema(BaseModel):
+    type: Literal["perplexity"] = "perplexity"
+    description: Optional[str] = None
+
+    enabled_domains: Optional[List[str]] = Field(
+        default=None,
+        description="List of domains that Perplexity is allowed to search",
+    )
+
+
+class MCPSourceSchema(BaseModel):
+    type: Literal["mcp"] = "mcp"
+    description: Optional[str] = None
+
+    uri: str = Field(description="URI of the MCP server endpoint")
+    headers: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="HTTP headers forwarded with every MCP request",
+    )
+    tool_choice_model: Optional[str] = Field(
+        default=None,
+        description="Model used for MCP tool selection",
+    )
+    valid_headers: Optional[List[str]] = Field(
+        default=None,
+        description="Allowlist of header names that may be forwarded",
+    )
+
+
+class GoogleSourceSchema(BaseModel):
+    type: Literal["google"] = "google"
+    description: Optional[str] = None
+    time_range: Optional[GoogleTimeRange] = Field(
+        default=None,
+        description="Restrict Google results to a recent time window",
+    )
+    exclude_domains: Optional[List[str]] = Field(
+        default=None,
+        description="Domains to exclude from Google search results",
+    )
+
+
+SourceSchema = Annotated[
+    Union[
+        NucliaDBSourceSchema,
+        PerplexitySourceSchema,
+        MCPSourceSchema,
+        GoogleSourceSchema,
+    ],
+    Field(discriminator="type"),
+]

nucliadb_agentic_api/server/__init__.py

@@ -0,0 +1 @@
+SERVICE_NAME = "nucliadb_agentic_api.server"

nucliadb_agentic_api/server/server.py

@@ -0,0 +1,113 @@
+import asyncio
+from importlib.metadata import version
+from typing import Any, Optional
+
+import sentry_sdk
+from hyperforge.broker.redis import RedisBroker
+from hyperforge.feature_flag import get_flag_service
+from hyperforge.server.cache import ValkeyCache
+from hyperforge.server.run import run_metrics_server
+from hyperforge_nucliadb_agentic.ask.predict import start_predict_engine
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogFormatType, LogLevel, LogSettings
+from nucliadb_telemetry.utils import setup_telemetry
+from nucliadb_utils.settings import AuditSettings
+from sentry_sdk.integrations.excepthook import ExcepthookIntegration
+
+from nucliadb_agentic_api import SERVICE_NAME, logger
+from nucliadb_agentic_api.db.agentic_configs import AgenticConfigs
+from nucliadb_agentic_api.db.settings import DataManagerSettings
+from nucliadb_agentic_api.server.session import NucliaDBAgenticSessionManager
+from nucliadb_agentic_api.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_server(
+    settings: Settings,
+    tracer: Optional[Any],
+    data_manager_settings: DataManagerSettings,
+) -> NucliaDBAgenticSessionManager:
+    if tracer:
+        await setup_telemetry(SERVICE_NAME)
+
+    await start_predict_engine()
+
+    audit_settings: AuditSettings = AuditSettings()
+    # 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 AgenticConfigs.from_settings(
+        settings=data_manager_settings,
+    )
+    await agent_manager.initialize()
+
+    session = NucliaDBAgenticSessionManager(
+        settings=settings,
+        broker=broker,
+        agent_manager=agent_manager,
+        cache=ValkeyCache(broker._client),
+        audit_settings=audit_settings,
+    )
+
+    return session
+
+
+def run():  # pragma: no cover
+    settings = Settings()
+    setup_logging(
+        settings=LogSettings(
+            log_format_type=LogFormatType.STRUCTURED
+            if not settings.debug
+            else LogFormatType.PLAIN,
+            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,
+            },
+        )
+    )
+    tracer = setup_telemetry(SERVICE_NAME)
+    if settings.sentry_url is not None:
+        set_sentry(
+            settings.zone,
+            settings.running_environment,
+            settings.sentry_url,
+        )
+
+    get_flag_service()  # precache the flag service
+
+    data_manager_settings = DataManagerSettings()  # type: ignore
+    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())
+    logger.info(f"======= Serving on Server to {settings.valkey_url} ======")
+    try:
+        loop.run_forever()
+    finally:
+        loop.run_until_complete(session.finalize())
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.close()