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

nucliadb_agentic_api/commands.py

@@ -0,0 +1,66 @@
+import uvicorn
+from hyperforge import openapi
+from hyperforge.feature_flag import get_flag_service
+from nucliadb_telemetry.fastapi import instrument_app
+from nucliadb_telemetry.logs import setup_logging
+from nucliadb_telemetry.settings import LogFormatType, LogLevel, LogSettings
+from nucliadb_telemetry.utils import get_telemetry, setup_telemetry
+from nucliadb_utils.settings import AuditSettings
+
+from nucliadb_agentic_api import SERVICE_NAME
+from nucliadb_agentic_api.app import HTTPApplication
+from nucliadb_agentic_api.db.settings import DataManagerSettings
+from nucliadb_agentic_api.logging import set_sentry
+from nucliadb_agentic_api.settings import Settings
+from nucliadb_agentic_api.v1.router import router
+
+
+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,
+            },
+        )
+    )
+    setup_telemetry(SERVICE_NAME)  # type: ignore
+    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
+    audit_settings = AuditSettings()
+    app = HTTPApplication(
+        settings,
+        data_manager_settings=data_manager_settings,
+        audit_settings=audit_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(
+        "nucliadb_agentic_api", "NucliaDB Agentic API", router
+    )

nucliadb_agentic_api/db/agentic_configs.py

@@ -0,0 +1,296 @@
+import datetime
+import json
+from time import time
+from typing import Dict
+
+import databases
+import sqlalchemy as sa
+from hyperforge.database import metadata
+from hyperforge.retrieval.config import RetrievalAgentConfig
+from hyperforge_google.config import GoogleDriverConfig, GoogleInnerConfig
+from hyperforge_nucliadb_agentic.ask.model import AskRequest
+from hyperforge_perplexity.config import PerplexityDriverConfig, PerplexityInnerConfig
+from lru import LRU
+from nucliadb_telemetry.utils import get_telemetry, init_telemetry
+from sqlalchemy.dialects.postgresql import JSONB
+
+from nucliadb_agentic_api import exceptions
+from nucliadb_agentic_api.db.settings import DataManagerSettings
+from nucliadb_agentic_api.db.sources import Sources
+from nucliadb_agentic_api.db.transform import transform_agentic_config
+from nucliadb_agentic_api.models import AgenticConfigSchema
+
+
+# Imported lazily in methods to avoid any load-order sensitivity between the two
+# table modules (both share the same `hyperforge.database.metadata`).
+def _get_sources_table():  # pragma: no cover
+    from nucliadb_agentic_api.db.sources import sources_table
+
+    return sources_table
+
+
+SERVICE_NAME = "AGENTIC_CONFIGS_DB"
+
+
+def utc_now() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
+agentic_config_table = sa.Table(
+    "agentic_config_table",
+    metadata,
+    sa.Column("account", sa.String, primary_key=True, nullable=False, index=True),
+    sa.Column("kbid", sa.String, primary_key=True, nullable=False, index=True),  # KBID
+    sa.Column("agentic_id", sa.String, primary_key=True, nullable=False),  # Agentic ID
+    sa.Column("created", sa.DateTime, default=sa.func.now()),
+    sa.Column("modified", sa.DateTime, onupdate=sa.func.now()),
+    sa.Column("title", sa.String, nullable=True),
+    sa.Column("config", JSONB, nullable=False),
+)
+
+
+CACHE: Dict[str, AgenticConfigSchema] = LRU(size=1024)  # type: ignore
+
+
+def _cache_key(account: str, kbid: str, agentic_id: str) -> str:
+    return f"{account}:{kbid}:{agentic_id}:{int(time()) // 5}"
+
+
+def _collect_source_ids(config: AgenticConfigSchema) -> list[str]:
+    """Return all non-None source_id values declared in smart_agent sources."""
+    if not config.smart_agent:
+        return []
+    return [source for source in config.smart_agent.sources if source is not None]
+
+
+def _serialize_config(config: AgenticConfigSchema) -> dict:
+    return config.model_dump(mode="json")
+
+
+def _config_from_row(row) -> AgenticConfigSchema:
+    return AgenticConfigSchema.model_validate(row["config"])
+
+
+class AgenticConfigs:
+    settings: DataManagerSettings
+
+    def __init__(
+        self,
+        database: databases.Database,
+        settings: DataManagerSettings,
+    ):
+        self.database = database
+        self.settings = settings
+
+    @classmethod
+    async def from_settings(
+        cls,
+        settings: DataManagerSettings,
+    ):
+        tracer_provider = get_telemetry(SERVICE_NAME)
+        if tracer_provider:
+            await init_telemetry(tracer_provider)
+
+        database = databases.Database(settings.postgresql_dsn)
+
+        return cls(database=database, settings=settings)
+
+    async def initialize(self):
+        await self.database.connect()
+
+    async def finalize(self):
+        await self.database.disconnect()
+
+    async def patch_agentic_config(
+        self, account: str, kbid: str, agentic_id: str, config: AgenticConfigSchema
+    ):
+        source_ids = _collect_source_ids(config)
+        if source_ids:
+            await self._validate_sources_exist(account, kbid, source_ids)
+
+        query = (
+            sa.update(agentic_config_table)
+            .values(
+                title=config.title,
+                config=_serialize_config(config),
+            )
+            .where(
+                agentic_config_table.c.account == account,
+                agentic_config_table.c.kbid == kbid,
+                agentic_config_table.c.agentic_id == agentic_id,
+            )
+            .returning(agentic_config_table.c.agentic_id)
+        )
+        updated = await self.database.fetch_one(query)
+        if updated is None:
+            raise exceptions.NotFound("Agentic configuration not found")
+        CACHE[_cache_key(account, kbid, agentic_id)] = config
+
+    async def create_agentic_config(
+        self, account: str, kbid: str, agentic_id: str, config: AgenticConfigSchema
+    ):
+        source_ids = _collect_source_ids(config)
+        if source_ids:
+            await self._validate_sources_exist(account, kbid, source_ids)
+
+        query = sa.select(agentic_config_table.c.agentic_id).where(
+            agentic_config_table.c.account == account,
+            agentic_config_table.c.kbid == kbid,
+            agentic_config_table.c.agentic_id == agentic_id,
+        )
+        existing = await self.database.fetch_one(query)
+        if existing is not None:
+            raise exceptions.Conflict("Agentic configuration already exists")
+
+        insert_query = sa.insert(agentic_config_table).values(
+            account=account,
+            kbid=kbid,
+            agentic_id=agentic_id,
+            title=config.title,
+            config=_serialize_config(config),
+        )
+        await self.database.execute(insert_query)
+        CACHE[_cache_key(account, kbid, agentic_id)] = config
+
+    async def _validate_sources_exist(
+        self, account: str, kbid: str, source_ids: list[str]
+    ) -> None:
+        """Raise InvalidReference if any source_id is not present in sources_table."""
+        sources_table = _get_sources_table()
+        query = sa.select(sources_table.c.source_id).where(
+            sources_table.c.account == account,
+            sources_table.c.kbid == kbid,
+            sources_table.c.source_id.in_(source_ids),
+        )
+        rows = await self.database.fetch_all(query)
+        found = {row["source_id"] for row in rows}
+        missing = sorted(set(source_ids) - found)
+        if missing:
+            raise exceptions.InvalidReference(
+                f"Source(s) not found: {', '.join(missing)}"
+            )
+
+    async def delete_configs_referencing_source(
+        self, account: str, kbid: str, source_id: str
+    ) -> int:
+        """Delete every agentic config that references source_id in its smart_agent
+        sources list.  Returns the number of deleted configs.
+
+        Uses the JSONB containment operator (@>) to find matching rows:
+            config->'smart_agent'->'sources' @> '[{"source_id": "<id>"}]'
+        """
+        query = (
+            sa.delete(agentic_config_table)
+            .where(
+                agentic_config_table.c.account == account,
+                agentic_config_table.c.kbid == kbid,
+                agentic_config_table.c.config["smart_agent"]["sources"].op("@>")(
+                    sa.cast(
+                        json.dumps([{"source_id": source_id}]),
+                        JSONB,
+                    )
+                ),
+            )
+            .returning(agentic_config_table.c.agentic_id)
+        )
+        deleted_rows = await self.database.fetch_all(query)
+        for row in deleted_rows:
+            key = _cache_key(account, kbid, row["agentic_id"])
+            if key in CACHE:
+                del CACHE[key]
+        return len(deleted_rows)
+
+    async def get_agentic_config(
+        self, account: str, kbid: str, agentic_id: str
+    ) -> AgenticConfigSchema:
+        key = _cache_key(account, kbid, agentic_id)
+        if key not in CACHE:
+            query = sa.select(agentic_config_table).where(
+                agentic_config_table.c.account == account,
+                agentic_config_table.c.kbid == kbid,
+                agentic_config_table.c.agentic_id == agentic_id,
+            )
+            row = await self.database.fetch_one(query)
+            if not row:
+                raise exceptions.NotFound("Agentic configuration not found")
+
+            config = _config_from_row(row)
+
+            CACHE[key] = config
+        else:
+            config = CACHE[key]
+        return config
+
+    async def list_agentic_configs(
+        self, account: str, kbid: str
+    ) -> dict[str, AgenticConfigSchema]:
+        query = sa.select(agentic_config_table).where(
+            agentic_config_table.c.account == account,
+            agentic_config_table.c.kbid == kbid,
+        )
+        rows = await self.database.fetch_all(query)
+        return {row["agentic_id"]: _config_from_row(row) for row in rows}
+
+    async def get_agent_config(
+        self,
+        account: str,
+        kbid: str,
+        internal_nucliadb_url: str | None = None,
+        internal_nucliadb: bool = True,
+        external_nucliadb_key: str | None = None,
+        external_nucliadb_url: str | None = None,
+        default_memory: bool = False,
+        workflow_id: str = "default",
+        ask_request: AskRequest | None = None,
+    ) -> RetrievalAgentConfig:
+        # For now, we only support one config per KB, so we ignore agent_id and workflow_id, but in the future we can extend this method to support multiple configs per KB and select
+        # the right one based on these parameters
+        retrieval_config: RetrievalAgentConfig
+        agentic_config = await self.get_agentic_config(account, kbid, workflow_id)
+
+        source_manager = Sources(self.database, self.settings)
+        retrieval_config, drivers, global_drivers = await transform_agentic_config(
+            agentic_config,
+            source_manager,
+            account,
+            internal_nucliadb=internal_nucliadb,
+            internal_nucliadb_url=internal_nucliadb_url,
+            external_nucliadb_url=external_nucliadb_url,
+            external_nucliadb_key=external_nucliadb_key,
+            ask_request=ask_request,
+            kbid=kbid,
+        )
+
+        for driver_config in drivers.values():
+            retrieval_config.drivers.append(driver_config)
+
+        if "google" in global_drivers:
+            retrieval_config.drivers.append(
+                GoogleDriverConfig(
+                    identifier="google",
+                    name="google",
+                    config=GoogleInnerConfig(
+                        credentials=self.settings.hyperforge_google_credentials,
+                        location=self.settings.hyperforge_google_location,
+                        vertexai=True,
+                    ),
+                )
+            )
+
+        if "perplexity" in global_drivers:
+            retrieval_config.drivers.append(
+                PerplexityDriverConfig(
+                    identifier="perplexity",
+                    name="perplexity",
+                    provider="perplexity",
+                    config=PerplexityInnerConfig(
+                        key=self.settings.hyperforge_perplexity_key
+                    ),
+                )
+            )
+        return retrieval_config
+
+    async def ensure_workflow_active(self, account: str, kbid: str, workflow_id: str):
+        # For now, we only support one config per KB, so we ignore agent_id and workflow_id, but in the future we can extend this method to support multiple configs per KB and select
+        # the right one based on these parameters
+        await self.get_agentic_config(account, kbid, workflow_id)

nucliadb_agentic_api/db/settings.py

@@ -0,0 +1,11 @@
+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
+
+    hyperforge_google_credentials: str = ""
+    hyperforge_google_location: str = ""
+    hyperforge_perplexity_key: str = ""

nucliadb_agentic_api/db/sources.py

@@ -0,0 +1,202 @@
+import datetime
+from time import time
+
+import databases
+import sqlalchemy as sa
+from hyperforge.database import metadata
+from lru import LRU
+from nucliadb_telemetry.utils import get_telemetry, init_telemetry
+from pydantic import TypeAdapter
+from sqlalchemy.dialects.postgresql import JSONB
+
+from nucliadb_agentic_api import exceptions
+from nucliadb_agentic_api.db.settings import DataManagerSettings
+from nucliadb_agentic_api.models import SourceSchema
+
+SERVICE_NAME = "SOURCES_DB"
+
+_source_adapter: TypeAdapter = TypeAdapter(SourceSchema)
+
+
+def utc_now() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
+sources_table = sa.Table(
+    "sources_table",
+    metadata,
+    sa.Column("account", sa.String, primary_key=True, nullable=False, index=True),
+    sa.Column("kbid", sa.String, primary_key=True, nullable=False, index=True),
+    sa.Column("source_id", sa.String, primary_key=True, nullable=False),
+    sa.Column("agentic_config_id", sa.String, nullable=True, index=True),
+    sa.Column("created", sa.DateTime, default=sa.func.now()),
+    sa.Column("modified", sa.DateTime, onupdate=sa.func.now()),
+    sa.Column("description", sa.String, nullable=True),
+    sa.Column("type", sa.String, nullable=False),
+    sa.Column("config", JSONB, nullable=False),
+    sa.ForeignKeyConstraint(
+        ["account", "kbid", "agentic_config_id"],
+        [
+            "agentic_config_table.account",
+            "agentic_config_table.kbid",
+            "agentic_config_table.agentic_id",
+        ],
+        ondelete="CASCADE",
+    ),
+)
+
+
+CACHE: LRU = LRU(size=1024)
+
+
+def _cache_key(account: str, kbid: str, source_id: str) -> str:
+    return f"src:{account}:{kbid}:{source_id}:{int(time()) // 5}"
+
+
+def _serialize_source(source: SourceSchema) -> dict:  # type: ignore[valid-type]
+    return source.model_dump(mode="json")  # type: ignore[union-attr]
+
+
+def _source_from_row(row) -> SourceSchema:  # type: ignore[valid-type]
+    return _source_adapter.validate_python(row["config"])
+
+
+class Sources:
+    settings: DataManagerSettings
+
+    def __init__(
+        self,
+        database: databases.Database,
+        settings: DataManagerSettings,
+    ):
+        self.database = database
+        self.settings = settings
+
+    @classmethod
+    async def from_settings(cls, settings: DataManagerSettings) -> "Sources":
+        tracer_provider = get_telemetry(SERVICE_NAME)
+        if tracer_provider:
+            await init_telemetry(tracer_provider)
+
+        database = databases.Database(settings.postgresql_dsn)
+        return cls(database=database, settings=settings)
+
+    async def initialize(self) -> None:
+        await self.database.connect()
+
+    async def finalize(self) -> None:
+        await self.database.disconnect()
+
+    # ------------------------------------------------------------------
+    # CRUD
+    # ------------------------------------------------------------------
+
+    async def create_source(
+        self,
+        account: str,
+        kbid: str,
+        source_id: str,
+        source: SourceSchema,  # type: ignore[valid-type]
+    ) -> None:
+        query = sa.select(sources_table.c.source_id).where(
+            sources_table.c.account == account,
+            sources_table.c.kbid == kbid,
+            sources_table.c.source_id == source_id,
+        )
+        existing = await self.database.fetch_one(query)
+        if existing is not None:
+            raise exceptions.Conflict("Source already exists")
+
+        serialized = _serialize_source(source)
+        insert_query = sa.insert(sources_table).values(
+            account=account,
+            kbid=kbid,
+            source_id=source_id,
+            description=serialized.get("description"),
+            type=serialized["type"],
+            config=serialized,
+        )
+        await self.database.execute(insert_query)
+        CACHE[_cache_key(account, kbid, source_id)] = source
+
+    async def get_source(
+        self,
+        account: str,
+        kbid: str,
+        source_id: str,
+    ) -> SourceSchema:  # type: ignore[valid-type]
+        key = _cache_key(account, kbid, source_id)
+        if key not in CACHE:
+            query = sa.select(sources_table).where(
+                sources_table.c.account == account,
+                sources_table.c.kbid == kbid,
+                sources_table.c.source_id == source_id,
+            )
+            row = await self.database.fetch_one(query)
+            if not row:
+                raise exceptions.NotFound("Source not found")
+
+            source = _source_from_row(row)
+            CACHE[key] = source
+        else:
+            source = CACHE[key]
+        return source
+
+    async def patch_source(
+        self,
+        account: str,
+        kbid: str,
+        source_id: str,
+        source: SourceSchema,  # type: ignore[valid-type]
+    ) -> None:
+        serialized = _serialize_source(source)
+        query = (
+            sa.update(sources_table)
+            .values(
+                description=serialized.get("description"),
+                type=serialized["type"],
+                config=serialized,
+            )
+            .where(
+                sources_table.c.account == account,
+                sources_table.c.kbid == kbid,
+                sources_table.c.source_id == source_id,
+            )
+            .returning(sources_table.c.source_id)
+        )
+        updated = await self.database.fetch_one(query)
+        if updated is None:
+            raise exceptions.NotFound("Source not found")
+        CACHE[_cache_key(account, kbid, source_id)] = source
+
+    async def delete_source(
+        self,
+        account: str,
+        kbid: str,
+        source_id: str,
+    ) -> None:
+        query = (
+            sa.delete(sources_table)
+            .where(
+                sources_table.c.account == account,
+                sources_table.c.kbid == kbid,
+                sources_table.c.source_id == source_id,
+            )
+            .returning(sources_table.c.source_id)
+        )
+        deleted = await self.database.fetch_one(query)
+        if deleted is None:
+            raise exceptions.NotFound("Source not found")
+        CACHE.pop(_cache_key(account, kbid, source_id), None)
+
+    async def list_sources(
+        self,
+        account: str,
+        kbid: str,
+    ) -> dict[str, SourceSchema]:  # type: ignore[valid-type]
+        query = sa.select(sources_table).where(
+            sources_table.c.account == account,
+            sources_table.c.kbid == kbid,
+        )
+        rows = await self.database.fetch_all(query)
+        return {row["source_id"]: _source_from_row(row) for row in rows}