hyperforge-nucliadb-agentic 1.0.0.post78__tar.gz → 1.0.0.post83__tar.gz

{hyperforge_nucliadb_agentic-1.0.0.post78 → hyperforge_nucliadb_agentic-1.0.0.post83}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyperforge_nucliadb_agentic
-Version: 1.0.0.post78
+Version: 1.0.0.post83
 Summary: NucliaDB Hyperforge agent
 Author: Nuclia
 Author-email: Nuclia <nucliadb@nuclia.com>
@@ -11,10 +11,10 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Dist: hyperforge>=1.0.0.post76
+Requires-Dist: hyperforge>=1.0.0.post184
 Requires-Dist: lingua-language-detector
 Requires-Dist: nucliadb-sdk
-Requires-Python: >=3.10, <4
+Requires-Python: >=3.12, <4
 Project-URL: Homepage, https://progress.com
 Project-URL: Repository, https://github.com/nuclia/nucliadb_agentic_api
 Project-URL: Changelog, https://github.com/nuclia/nucliadb_agentic_api/blob/main/agents/nucliadb/CHANGELOG.md

hyperforge_nucliadb_agentic-1.0.0.post83/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["uv_build>=0.11.19,<0.12"]
+build-backend = "uv_build"
+
+[project]
+name = "hyperforge_nucliadb_agentic"
+version = "1.0.0.post83"
+license = "Apache-2.0"
+description = "NucliaDB Hyperforge agent"
+readme = "README.md"
+classifiers = [
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+requires-python = ">=3.12, <4"
+dependencies = [
+    "hyperforge>=1.0.0.post184",
+    "lingua-language-detector",
+    "nucliadb_sdk",
+]
+
+[[project.authors]]
+name = "Nuclia"
+email = "nucliadb@nuclia.com"
+
+[project.urls]
+Homepage = "https://progress.com"
+Repository = "https://github.com/nuclia/nucliadb_agentic_api"
+Changelog = "https://github.com/nuclia/nucliadb_agentic_api/blob/main/agents/nucliadb/CHANGELOG.md"
+
+[dependency-groups]
+dev = [
+    "pytest",
+    "pytest-benchmark",
+    "pytest-docker-fixtures>=1.4.2",
+    "pytest-lazy-fixtures",
+    "hyperforge_rephrase",
+]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
+
+[tool.uv.build-backend]
+source-exclude = ["tests/**/*"]

hyperforge_nucliadb_agentic-1.0.0.post78/pyproject.toml → hyperforge_nucliadb_agentic-1.0.0.post83/pyproject.toml.orig

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "hyperforge_nucliadb_agentic"
-version = "1.0.0.post78"
+version = "1.0.0.post83"
 license = "Apache-2.0"
 description = "NucliaDB Hyperforge agent"
 authors = [{ name = "Nuclia", email = "nucliadb@nuclia.com" }]
@@ -17,9 +17,9 @@ classifiers = [
     "Programming Language :: Python :: 3 :: Only",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-requires-python = ">=3.10, <4"
+requires-python = ">=3.12, <4"
 dependencies = [
-    "hyperforge>=1.0.0.post76",
+    "hyperforge>=1.0.0.post184",
     "lingua-language-detector",
     "nucliadb_sdk",
 ]

{hyperforge_nucliadb_agentic-1.0.0.post78 → hyperforge_nucliadb_agentic-1.0.0.post83}/src/hyperforge_nucliadb_agentic/ask/audit.py

@@ -9,6 +9,7 @@ import mmh3
 import nats
 from fastapi import Request
 from hyperforge.feature_flag import Features, has_feature
+from hyperforge.models import ExternalUsage, Step
 from nucliadb_models.retrieval import RawQuery, RetrievalRequest
 from nucliadb_models.search import (
     NucliaDBClientType,
@@ -19,9 +20,19 @@ from nucliadb_protos.audit_pb2 import (
     ChatContext,
     RetrievedContext,
 )
+from nucliadb_protos.kb_usage_pb2 import (
+    ActivityLogMatch,
+    ActivityLogMatchType,
+    ClientType as KbUsageClientType,
+    KBSource,
+    Predict,
+    PredictType,
+    Service,
+)
 from nucliadb_telemetry.jetstream import get_traced_jetstream, get_traced_nats_client
 from nucliadb_utils import logger
-from nucliadb_utils.settings import AuditSettings
+from nucliadb_utils.nuclia_usage.utils.kb_usage_report import KbUsageReportUtility
+from nucliadb_utils.settings import AuditSettings, usage_settings
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 from opentelemetry.trace import INVALID_SPAN, format_trace_id, get_current_span
 from starlette.background import BackgroundTask
@@ -38,6 +49,21 @@ from hyperforge_nucliadb_agentic.ask.predict import AnswerStatusCode
 from hyperforge_nucliadb_agentic.ask.utils.proto import client_type
 
 
+def external_usage_to_predict(
+    event: ExternalUsage, ndb_client_type: NucliaDBClientType
+) -> Predict:
+    return Predict(
+        client=KbUsageClientType.Value(ndb_client_type.name),
+        type=PredictType.QUESTION_ANSWER,
+        model=f"{event.provider}/{event.model}",
+        input=event.input_tokens,
+        output=event.output_tokens,
+        num_predicts=event.requests,
+        image=event.image,
+        customer_key=False,
+    )
+
+
 class RequestContext:
     def __init__(self: "RequestContext"):
         self.audit_request: AuditRequest = AuditRequest()
@@ -120,6 +146,7 @@ class StreamAuditStorage:
         self.service = service
         self.task = None
         self.initialized = False
+        self.kb_usage_utility: KbUsageReportUtility | None = None
 
     def get_partition(self, kbid: str):
         return mmh3.hash(kbid, self.seed, signed=False) % self.partitions
@@ -155,10 +182,17 @@ class StreamAuditStorage:
 
         self.js = get_traced_jetstream(self.nc, self.service)
         self.task = asyncio.create_task(self.run())
+        self.kb_usage_utility = KbUsageReportUtility(
+            nats_stream=self.js,
+            nats_subject=cast(str, usage_settings.usage_jetstream_subject),
+        )
+        await self.kb_usage_utility.initialize()
 
         self.initialized = True
 
     async def finalize(self):
+        if self.kb_usage_utility is not None:
+            await self.kb_usage_utility.finalize()
         if self.task is not None:
             self.task.cancel()
         if self.nc:
@@ -166,6 +200,35 @@ class StreamAuditStorage:
             await self.nc.close()
             self.nc = None
 
+    def report_step_usage(
+        self,
+        *,
+        account_id: str,
+        kbid: str,
+        client_type: NucliaDBClientType,
+        step: Step,
+        trace_id: str | None = None,
+    ) -> None:
+        predicts = [
+            external_usage_to_predict(event, client_type)
+            for event in step.external_usage or []
+        ]
+        if not predicts or self.kb_usage_utility is None:
+            return
+
+        self.kb_usage_utility.send_kb_usage(
+            service=Service.RAO,
+            account_id=account_id,
+            kb_id=kbid,
+            kb_source=KBSource.HOSTED,
+            predicts=predicts,
+            activity_log_match=(
+                ActivityLogMatch(id=trace_id, type=ActivityLogMatchType.TRACE_ID)
+                if trace_id
+                else None
+            ),
+        )
+
     async def run(self):
         while True:
             item_dequeued = False