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

{hyperforge_nucliadb_agentic-1.0.0.post64 → 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.post64
+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.post64/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.post64"
+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.post64 → hyperforge_nucliadb_agentic-1.0.0.post83}/src/hyperforge_nucliadb_agentic/agent.py

@@ -1056,7 +1056,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
             )
         )
 
-        paragraphs = await paragraphs_result.model_dump()
+        paragraphs = await paragraphs_result.to_sync_response()
         answer = None
         input_tokens = (
             paragraphs.consumption.normalized_tokens.input

{hyperforge_nucliadb_agentic-1.0.0.post64 → 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

{hyperforge_nucliadb_agentic-1.0.0.post64 → hyperforge_nucliadb_agentic-1.0.0.post83}/src/hyperforge_nucliadb_agentic/ask/predict.py

@@ -296,13 +296,18 @@ class PredictEngine:
         headers = self.get_predict_headers(kbid)
         headers["Accept"] = "application/x-ndjson"
 
-        resp = await self.make_stream_request(
-            "POST",
-            url=self.get_predict_url(CHAT, kbid),
-            json=item.model_dump(),
-            headers={**headers, **(extra_headers or {})},
-            timeout=httpx.Timeout(180.0, read=120.0),
-        )
+        try:
+            resp = await self.make_stream_request(
+                "POST",
+                url=self.get_predict_url(CHAT, kbid),
+                json=item.model_dump(),
+                headers={**headers, **(extra_headers or {})},
+                timeout=httpx.Timeout(30.0, read=None),
+            )
+        except httpx.TimeoutException as exc:
+            raise ProxiedPredictAPIError(
+                status=504, detail="Timeout while connecting to predict"
+            ) from exc
         await self.check_response(kbid, resp, expected_status=200)
         ident = resp.headers.get(NUCLIA_LEARNING_ID_HEADER) or "unknown"
         model = resp.headers.get(NUCLIA_LEARNING_MODEL_HEADER) or "unknown"

{hyperforge_nucliadb_agentic-1.0.0.post64 → hyperforge_nucliadb_agentic-1.0.0.post83}/src/hyperforge_nucliadb_agentic/ask/search/ask.py

@@ -378,7 +378,7 @@ class AskResult:
                 metrics=self.metrics.dump(),
             )
 
-    async def model_dump(self) -> SyncAskResponse:
+    async def to_sync_response(self) -> SyncAskResponse:
         # Run the stream to completion to ensure all data is available in memory
         # First, run the stream in memory to get all the data in memory
         async for _ in self._stream():
@@ -458,7 +458,7 @@ class AskResult:
         return response
 
     async def json(self) -> str:
-        response = await self.model_dump()
+        response = await self.to_sync_response()
         return response.model_dump_json(exclude_none=True, by_alias=True)
 
     async def get_relations_results(self) -> Relations:
@@ -551,7 +551,7 @@ class NotEnoughContextAskResult(AskResult):
                 )
             )
 
-    async def json(self) -> str:
+    async def to_sync_response(self) -> SyncAskResponse:
         prequeries = (
             {
                 prequery.id or f"prequery_{index}": prequery_result
@@ -567,7 +567,10 @@ class NotEnoughContextAskResult(AskResult):
             retrieval_results=self.main_results,
             prequeries=prequeries,
             status=AnswerStatusCode.NO_RETRIEVAL_DATA.prettify(),
-        ).model_dump_json()
+        )
+
+    async def json(self) -> str:
+        return (await self.to_sync_response()).model_dump_json()
 
 
 async def rephrase_query(