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

nucliadb_agentic_api/server/session.py

@@ -0,0 +1,276 @@
+import os
+from asyncio import Task, create_task, timeout  # type: ignore
+from functools import partial
+
+import nucliadb_telemetry.context
+import nucliadb_telemetry.metrics
+import prometheus_client
+from hyperforge.broker import Broker
+from hyperforge.configure import GLOBAL_REGISTRY, load_all_configurations, scan
+from hyperforge.engine import State, get_state
+from hyperforge.interaction import AnswerOperation, AragAnswer, ARAGException
+from hyperforge.memory import QuestionMemory
+from hyperforge.pubsub import AgentDone, StartInteraction
+from hyperforge.server.cache import Cache
+from hyperforge.server.session import SessionManager
+from hyperforge.server.utils import get_memory
+from hyperforge_nucliadb_agentic.ask.audit import (
+    start_audit_utility,
+    stop_audit_utility,
+)
+from hyperforge_nucliadb_agentic.ask.model import AskRequest
+from lru import LRU
+from nucliadb_telemetry import errors
+from nucliadb_telemetry.utils import get_telemetry
+from nucliadb_utils.settings import AuditSettings
+from opentelemetry import trace
+
+from nucliadb_agentic_api import logger
+from nucliadb_agentic_api.db.agentic_configs import AgenticConfigs
+from nucliadb_agentic_api.server import SERVICE_NAME
+from nucliadb_agentic_api.server.settings import Settings as ServerSettings
+
+HOSTNAME = os.environ.get("HOSTNAME", "nucliadb-agentic-api-server").encode()
+
+answer_observer = nucliadb_telemetry.metrics.Observer("nucliadb_agentic_api_answer")
+activation_observer = nucliadb_telemetry.metrics.Observer(
+    "nucliadb_agentic_api_activation"
+)
+answer_running = prometheus_client.Gauge(
+    "nucliadb_agentic_api_running_answers_count",
+    "Number of answering processess currently running",
+)
+
+
+def tracer():
+    provider = get_telemetry(SERVICE_NAME)
+    if provider:
+        return provider.get_tracer(__name__)
+    else:
+        return trace.NoOpTracer()
+
+
+class NucliaDBAgenticSessionManager(SessionManager):
+    agent_manager: AgenticConfigs  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        settings: ServerSettings,
+        audit_settings: AuditSettings,
+        broker: Broker,
+        agent_manager: AgenticConfigs,
+        cache: Cache,
+    ):
+        self.settings = settings
+        self.audit_settings = audit_settings
+        self.agent_manager = agent_manager
+        self.broker = broker
+        self.memory: LRU = LRU(800)
+        self.activation_task: Task | None = None
+        self.tasks = []
+        self.cache = cache
+
+    async def initialize(self, health_check: bool = True):
+        await super().initialize(health_check)
+
+        await start_audit_utility(SERVICE_NAME, self.audit_settings)
+
+        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 finalize(self):
+        await super().finalize()
+        await stop_audit_utility()
+        GLOBAL_REGISTRY.clear()
+
+    async def activate(self, message: StartInteraction):
+        topic = None
+
+        ask_request_json = message.arguments.get("ask_request")
+        ask_request = None
+        if ask_request_json:
+            ask_request = AskRequest.model_validate_json(ask_request_json)
+
+        logger.info("Activation message received: %s", message)
+        observation = activation_observer()
+        observation.start()
+        try:
+            nucliadb_telemetry.context.add_context(
+                {
+                    "agent_id": message.agent_id,
+                    "session_id": message.session,
+                    "question_id": message.question_id,
+                }
+            )
+
+            topic = self.question_topic(
+                message.account,
+                message.agent_id,
+                message.session,
+                message.question_id,
+                message.workflow_id,
+            )
+
+            # Get or load session
+            config = await self.agent_manager.get_agent_config(
+                account=message.account,
+                kbid=message.agent_id,
+                internal_nucliadb_url=self.settings.internal_nucliadb_url,
+                internal_nucliadb=self.settings.internal_nucliadb,
+                external_nucliadb_url=self.settings.external_nucliadb_url,
+                external_nucliadb_key=self.settings.external_nucliadb_key,
+                workflow_id=message.workflow_id,
+                ask_request=ask_request,
+            )
+
+            state = await get_state(
+                agent_id=message.agent_id,
+                config=config,
+                internal_nua_api=self.settings.internal_nua_api,
+                internal_nua=self.settings.internal_nua,
+                local_openai=self.settings.local_openai,
+                external_nua_api_key=self.settings.external_nua_api_key,
+                account=message.account,
+                kbid=None if self.settings.standalone else message.agent_id,
+            )
+
+            if message.session not in self.memory:
+                memory = await get_memory(
+                    settings=self.settings,
+                    session=message.session,
+                    cache=self.cache,
+                    config=config.memory,
+                    agent=message.agent_id,
+                    workflow_id=message.workflow_id,
+                )
+                self.memory[message.session] = memory
+            else:
+                memory = self.memory[message.session]
+
+            memory.rules = config.rules.rules
+
+            question = memory.start_question(
+                message.question,
+                question_id=message.question_id,
+                headers=message.headers,
+                arguments=message.arguments,
+                streaming=message.streaming,
+                chat_history=message.chat_history,
+            )
+
+            task = create_task(
+                self.answer(
+                    message.account,
+                    message.agent_id,
+                    message.workflow_id,
+                    topic,
+                    state,
+                    question,
+                )
+            )
+            task.add_done_callback(self._remove_task)
+            self.tasks.append(task)
+
+        except Exception as e:
+            logger.exception("Activation exception")
+            errors.capture_exception(e)
+            observation.set_status("error")
+            if topic:
+                await self.callback(
+                    topic,
+                    AragAnswer(
+                        exception=ARAGException(detail="Unable to start agent"),
+                        operation=AnswerOperation.ERROR,
+                    ),
+                )
+                await self.send_message(topic, AgentDone())
+
+        observation.end()
+
+    async def answer(
+        self,
+        account_id: str,
+        agent_id: str,
+        workflow_id: str,
+        topic: str,
+        state: State,
+        question_memory: QuestionMemory,
+    ):
+        error = None
+
+        keepalive = create_task(self.keep_alive(topic))
+        observation = answer_observer()
+        observation.start()
+        answer_running.inc()
+
+        try:
+            callback = partial(self.callback, topic)
+            question_memory.set_callback_fn(callback)
+
+            feedback = partial(self.feedback, topic)
+            question_memory.set_feedback_fn(feedback)
+
+            oauth = partial(self.oauth, topic)
+            question_memory.set_oauth_fn(oauth)
+
+            oauth_callback = partial(
+                self.get_oauth_callback,
+                account_id,
+                agent_id,
+                question_memory.session.id,
+                workflow_id,
+            )
+            question_memory.set_oauth_callback_fn(oauth_callback)
+
+            await self.callback(
+                topic,
+                AragAnswer(operation=AnswerOperation.START),
+            )
+
+            async with timeout(self.settings.question_timeout_seconds):
+                await state.agent(question_memory, state.manager)
+
+        except Exception as e:
+            logger.exception("Answering exception")
+            errors.capture_exception(e)
+            error = ARAGException(detail=str(e))
+            observation.set_status("error")
+
+        observation.end()
+        answer_running.dec()
+        keepalive.cancel()
+
+        await self.callback(
+            topic,
+            AragAnswer(
+                exception=error,
+                answer=question_memory.final_answer,
+                answer_citations=question_memory.final_answer_citations,
+                answer_urls=question_memory.final_answer_urls,
+                operation=AnswerOperation.ERROR
+                if error is not None
+                else AnswerOperation.ANSWER,
+                data_visualizations=question_memory.data_visualizations
+                if question_memory.data_visualizations
+                else None,
+            ),
+        )
+        await self.send_message(
+            topic,
+            AgentDone(),
+        )
+
+        try:
+            await question_memory.save()
+            self.process_event(
+                "memory_saved",
+                {"account_id": account_id, "question_memory": question_memory},
+            )
+        except Exception as e:
+            # Log memory errors but don't report them to the user
+            logger.exception("Error saving memory")
+            errors.capture_exception(e)

nucliadb_agentic_api/server/settings.py

@@ -0,0 +1,22 @@
+from hyperforge.server.settings import Settings as HyperforgeServerSettings
+
+MODULES = [
+    "hyperforge_rephrase",
+    "hyperforge_nucliadb",
+    "hyperforge_nucliadb_agentic",
+    "hyperforge_summarize",
+    "hyperforge_smart",
+    "hyperforge_mcp",
+    "hyperforge_google",
+    "hyperforge_perplexity",
+    "hyperforge_nucliadb_agentic.internal_driver",
+]
+
+
+class Settings(HyperforgeServerSettings):
+    load_modules: list[str] = MODULES
+    activate_subject: str = "ndb_agentic.activate"
+    answers_subject: str = (
+        "ndb_agentic.{account}.{agent_id}.{workflow_id}.{session}.{question}.answer"
+    )
+    oauth_subject: str = "ndb_agentic.{account}.{agent_id}.{workflow_id}.{session}.{question}.oauth.{oauth_uuid}"

nucliadb_agentic_api/settings.py

@@ -0,0 +1,39 @@
+from typing import Optional
+
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    debug: bool = False
+    log_level: str = "ERROR"
+    http_host: str = "0.0.0.0"
+    http_port: int = 8080
+
+    sentry_url: Optional[str] = None
+    running_environment: str = "stage"
+    zone: str = "stashify"
+    grpc_port: int = 8030
+
+    valkey_url: str = "redis://ndb_agentic-valkey-cluster"
+    valkey_cluster_mode: bool = False
+    answers_subject: str = (
+        "ndb_agentic.{account}.{agent_id}.{workflow_id}.{session}.{question}.answer"
+    )
+    oauth_subject: str = "ndb_agentic.{account}.{agent_id}.{workflow_id}.{session}.{question}.oauth.{oauth_uuid}"
+    activate_subject: str = "ndb_agentic.activate"
+    pubsub_keepalive_seconds: float = 20
+
+    # Hydra settings for MCP Oauth
+    hydra_public_url: str = "https://oauth.progress.cloud"
+    hydra_scopes_supported: list[str] = ["offline_access", "openid"]
+
+    hyperforge_google_credentials: Optional[str] = None
+    hyperforge_google_location: Optional[str] = None
+    hyperforge_perplexity_key: Optional[str] = None
+    internal_nua_api: str = "http://predict.learning.svc.cluster.local:8080"
+
+    internal_nucliadb: bool = False
+    internal_nucliadb_url: Optional[str] = None
+
+    external_nucliadb_key: Optional[str] = None
+    external_nucliadb_url: Optional[str] = None

nucliadb_agentic_api/utils.py

@@ -0,0 +1,284 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, List, TypeVar
+
+import pydantic
+from typing_extensions import TypeGuard
+
+_T = TypeVar("_T")
+
+FLOW_PROPERTIES = [
+    "next_agent",
+    "fallback",
+    "then",
+    "else_",
+    "agents",
+    "registered_agents",
+    "agent",
+]
+
+
+def to_strict_json_schema(
+    model: type[pydantic.BaseModel] | pydantic.TypeAdapter[Any],
+    exclude_properties: List[str] = [],
+    exclude_defs: List[str] = [],
+) -> dict[str, Any]:
+    if inspect.isclass(model) and is_basemodel_type(model):
+        schema = model.model_json_schema()
+    elif isinstance(model, pydantic.TypeAdapter):
+        schema = model.json_schema()
+    else:
+        raise TypeError(
+            f"Non BaseModel types are only supported with Pydantic v2 - {model}"
+        )
+
+    for exclude in exclude_properties:
+        if "properties" in schema and exclude in schema["properties"]:
+            del schema["properties"][exclude]
+    for exclude in exclude_defs:
+        if "$defs" in schema and exclude in schema["$defs"]:
+            del schema["$defs"][exclude]
+    return _ensure_strict_json_schema(schema, path=(), root=schema)
+
+
+def _ensure_strict_json_schema(
+    json_schema: object,
+    *,
+    path: tuple[str, ...],
+    root: dict[str, object],
+) -> dict[str, Any]:
+    """Mutates the given JSON schema to ensure it conforms to the `strict` standard
+    that the API expects.
+    """
+    if not is_dict(json_schema):
+        raise TypeError(f"Expected {json_schema} to be a dictionary; path={path}")
+
+    # defs = json_schema.get("$defs")
+    # if is_dict(defs):
+    #     for def_name, def_schema in defs.items():
+    #         _ensure_strict_json_schema(
+    #             def_schema, path=(*path, "$defs", def_name), root=root
+    #         )
+
+    definitions = json_schema.get("definitions")
+    if is_dict(definitions):
+        for definition_name, definition_schema in definitions.items():
+            _ensure_strict_json_schema(
+                definition_schema,
+                path=(*path, "definitions", definition_name),
+                root=root,
+            )
+
+    typ = json_schema.get("type")
+    if typ == "object" and "additionalProperties" not in json_schema:
+        json_schema["additionalProperties"] = False
+
+    # object types
+    # { 'type': 'object', 'properties': { 'a':  {...} } }
+    properties = json_schema.get("properties")
+    if is_dict(properties):
+        json_schema["required"] = [prop for prop in properties.keys()]
+        json_schema["properties"] = {
+            key: _ensure_strict_json_schema(
+                prop_schema, path=(*path, "properties", key), root=root
+            )
+            for key, prop_schema in properties.items()
+        }
+
+    # arrays
+    # { 'type': 'array', 'items': {...} }
+    items = json_schema.get("items")
+    if is_dict(items):
+        json_schema["items"] = _ensure_strict_json_schema(
+            items, path=(*path, "items"), root=root
+        )
+
+    # unions
+    one_of = json_schema.get("oneOf")
+    if isinstance(one_of, list):
+        json_schema["oneOf"] = [
+            _ensure_strict_json_schema(
+                variant, path=(*path, "oneOf", str(i)), root=root
+            )
+            for i, variant in enumerate(one_of)
+        ]
+
+    # unions
+    any_of = json_schema.get("anyOf")
+    if isinstance(any_of, list):
+        json_schema["anyOf"] = [
+            _ensure_strict_json_schema(
+                variant, path=(*path, "anyOf", str(i)), root=root
+            )
+            for i, variant in enumerate(any_of)
+        ]
+
+    # intersections
+    all_of = json_schema.get("allOf")
+    if isinstance(all_of, list):
+        if len(all_of) == 1:
+            json_schema.update(
+                _ensure_strict_json_schema(
+                    all_of[0], path=(*path, "allOf", "0"), root=root
+                )
+            )
+            json_schema.pop("allOf")
+        else:
+            json_schema["allOf"] = [
+                _ensure_strict_json_schema(
+                    entry, path=(*path, "allOf", str(i)), root=root
+                )
+                for i, entry in enumerate(all_of)
+            ]
+
+    # strip `None` defaults as there's no meaningful distinction here
+    # the schema will still be `nullable` and the model will default
+    # to using `None` anyway
+    if json_schema.get("default", "") is None:
+        json_schema.pop("default")
+
+    # we can't use `$ref`s if there are also other properties defined, e.g.
+    # `{"$ref": "...", "description": "my description"}`
+    #
+    # so we unravel the ref
+    # `{"type": "string", "description": "my description"}`
+    ref = json_schema.get("$ref")
+    if ref is not None:
+        assert isinstance(ref, str), f"Received non-string $ref - {ref}"
+
+        resolved = resolve_ref(root=root, ref=ref)
+        if not is_dict(resolved):
+            raise ValueError(
+                f"Expected `$ref: {ref}` to resolved to a dictionary but got {resolved}"
+            )
+
+        # properties from the json schema take priority over the ones on the `$ref`
+        json_schema.update({**resolved, **json_schema})
+        json_schema.pop("$ref")
+        # Since the schema expanded from `$ref` might not have `additionalProperties: false` applied,
+        # we call `_ensure_strict_json_schema` again to fix the inlined schema and ensure it's valid.
+        return _ensure_strict_json_schema(json_schema, path=path, root=root)
+
+    return json_schema
+
+
+def resolve_ref(*, root: dict[str, object], ref: str) -> object:
+    if not ref.startswith("#/"):
+        raise ValueError(f"Unexpected $ref format {ref!r}; Does not start with #/")
+
+    path = ref[2:].split("/")
+    resolved = root
+    for key in path:
+        value = resolved[key]
+        assert is_dict(value), (
+            f"encountered non-dictionary entry while resolving {ref} - {resolved}"
+        )
+        resolved = value
+
+    return resolved
+
+
+def is_basemodel_type(typ: type) -> TypeGuard[type[pydantic.BaseModel]]:
+    if not inspect.isclass(typ):
+        return False
+    return issubclass(typ, pydantic.BaseModel)
+
+
+def is_dataclass_like_type(typ: type) -> bool:
+    """Returns True if the given type likely used `@pydantic.dataclass`"""
+    return hasattr(typ, "__pydantic_config__")
+
+
+def is_dict(obj: object) -> TypeGuard[dict[str, object]]:
+    # just pretend that we know there are only `str` keys
+    # as that check is not worth the performance cost
+    return isinstance(obj, dict)
+
+
+async def clean_up_items(items: dict[str, Any], filtered: list[str]) -> dict[str, Any]:
+    """Cleans up the items section of the schema by removing filtered agents and drivers from the references."""
+    if "discriminator" in items and "mapping" in items["discriminator"]:
+        for name, module in list(items["discriminator"]["mapping"].items()):
+            if filtered and module.split("/")[-1] in filtered:
+                del items["discriminator"]["mapping"][name]
+    if "oneOf" in items:
+        for item in list(items["oneOf"]):
+            if "$ref" in item:
+                module_name = item["$ref"].split("/")[-1]
+                if module_name in filtered:
+                    items["oneOf"].remove(item)
+    return items
+
+
+async def cleanup_anyof(anyof: list[dict[str, Any]], filtered: list[str]):
+    """Cleans up the anyof section of the schema by removing filtered agents and drivers from the references."""
+    mapping = {}
+    if "discriminator" in anyof[0]:
+        anyof_mapping = anyof[0]["discriminator"].get("mapping", {})
+        for name, module in anyof_mapping.items():
+            module_name = module.split("/")[-1]
+            if all(module_name not in fa for fa in filtered):
+                mapping[name] = module
+        anyof[0]["discriminator"]["mapping"] = mapping
+    if "oneOf" in anyof[0]:
+        for module in anyof[0]["oneOf"]:
+            if "$ref" in module:
+                module_name = module["$ref"].split("/")[-1]
+                if module_name in filtered:
+                    anyof[0]["oneOf"].remove(module)
+    return anyof
+
+
+async def cleanup_properties(
+    properties: dict[str, Any], filtered_agents: list[str], filtered_drivers: list[str]
+) -> dict[str, Any]:
+    """Cleans up the properties section of the schema by removing filtered agents and drivers from the references."""
+    filtered = filtered_agents + filtered_drivers
+    steps = ["preprocess", "context", "generation", "postprocess", "drivers"]
+    for step in steps:
+        if step not in properties:
+            continue
+        items = properties[step].get("items")
+        if items:
+            properties[step]["items"] = await clean_up_items(items, filtered)
+    return properties
+
+
+async def cleanup_definitions(
+    definitions: dict[str, Any], filtered_agents: list[str]
+) -> dict[str, Any]:
+    """Cleans up the definitions section of the schema by removing filtered agents from the definitions."""
+    for item, item_schema in definitions.items():
+        if item.endswith("AgentConfig"):
+            if item in filtered_agents:
+                del definitions[item]
+                continue
+            else:
+                properties = item_schema.get("properties", {})
+                for property in FLOW_PROPERTIES:
+                    if property in item_schema["properties"]:
+                        if "anyOf" in properties[property]:
+                            anyof = properties[property]["anyOf"]
+                            if "items" in anyof[0]:
+                                items = anyof[0]["items"]
+                                properties[property]["anyOf"][0][
+                                    "items"
+                                ] = await clean_up_items(items, filtered_agents)
+                            else:
+                                properties[property]["anyOf"] = await cleanup_anyof(
+                                    anyof, filtered_agents
+                                )
+                        elif "items" in properties[property]:
+                            items = properties[property]["items"]
+                            properties[property]["items"] = await clean_up_items(
+                                items, filtered_agents
+                            )
+                        elif "discriminator" in properties[property]:
+                            properties[property] = await clean_up_items(
+                                properties[property], filtered_agents
+                            )
+                item_schema["properties"] = properties
+        definitions[item] = item_schema
+
+    return definitions

nucliadb_agentic_api/v1/__init__.py

@@ -0,0 +1,19 @@
+from nucliadb_agentic_api.v1 import (
+    agentic_configs,
+    ask,
+    ask_websocket,
+    mcp_nucliadb,
+    oauth,
+    sources,
+)
+from nucliadb_agentic_api.v1.router import router
+
+__all__ = [
+    "agentic_configs",
+    "ask_websocket",
+    "ask",
+    "mcp_nucliadb",
+    "oauth",
+    "sources",
+    "router",
+]