fix(monitoring): inject context into langfuse (#889)

* fix streaming bug

* revert

* Update unit coverage badge

---------

Co-authored-by: Benjamin PILIA <benjamin.pilia@protonmail.com>
Co-authored-by: leoguillaume <leoguillaume@users.noreply.github.com>

Author: 3 people

.github/badges/coverage.json

@@ -1 +1 @@
-{"schemaVersion":1,"label":"coverage","message":"56.29%","color":"red"}
+{"schemaVersion":1,"label":"coverage","message":"56.26%","color":"red"}

api/clients/model/_basemodelprovider.py

@@ -285,9 +285,11 @@ async def _log_performance_metric(self, redis_client: AsyncRedis, ttft: int | No
             logger.error(f"Failed to log request metrics (latency) in redis (id: {self.id})", exc_info=True)
             await safe_redis_reset(redis_client)
 
-    def _start_langfuse_observation(self, request_content: RequestContent) -> Any | None:
+    def _start_langfuse_observation(self, request_content: RequestContent, ctx=None) -> Any | None:
         langfuse_obs = None
         if global_context.langfuse_client is not None:
+            if ctx is None:
+                ctx = request_context.get()
             try:
                 langfuse_obs = global_context.langfuse_client.start_observation(
                     as_type="generation",
@@ -299,9 +301,10 @@ def _start_langfuse_observation(self, request_content: RequestContent) -> Any |
 
         return langfuse_obs
 
-    def _end_langfuse_observation(self, langfuse_obs: Any, latency: int | None, ttft: int | None = None):
+    def _end_langfuse_observation(self, langfuse_obs: Any, latency: int | None, ttft: int | None = None, ctx=None):
         if global_context.langfuse_client is not None and langfuse_obs is not None:
-            ctx = request_context.get()
+            if ctx is None:
+                ctx = request_context.get()
             if ctx.usage is not None:
                 global_context.langfuse_client.update_observation(
                     langfuse_obs,
@@ -429,7 +432,8 @@ async def forward_stream(self, request_content: RequestContent, redis_client: As
         url = urljoin(base=self.url, url=self.ENDPOINT_TABLE.get_endpoint(endpoint=request_content.endpoint).lstrip("/"))
         request_content = self._format_request(request_content=request_content)
 
-        langfuse_obs = self._start_langfuse_observation(request_content=request_content)
+        ctx = request_context.get()
+        langfuse_obs = self._start_langfuse_observation(request_content=request_content, ctx=ctx)
         inflight_key = f"{PREFIX__REDIS_METRIC_GAUGE}:{Metric.INFLIGHT.value}:{self.id}"
         inflight_incremented = False
 
@@ -481,7 +485,7 @@ async def forward_stream(self, request_content: RequestContent, redis_client: As
                             latency = self._elapsed_ms(start_time=start_time)
                             extra_chunk = self._get_extra_stream_chunk(request_content=request_content, buffer=buffer, latency=latency)
 
-                            self._end_langfuse_observation(langfuse_obs=langfuse_obs, latency=latency, ttft=ttft)
+                            self._end_langfuse_observation(langfuse_obs=langfuse_obs, latency=latency, ttft=ttft, ctx=ctx)
 
                             if extra_chunk is not None:
                                 yield f"data: {dumps(extra_chunk)}\n\n", response.status_code

api/endpoints/ocr.py

@@ -23,7 +23,6 @@
 router = APIRouter(prefix="/v1", tags=[RouterName.OCR.upper()])
 
 
-@hooks
 @router.post(
     path=EndpointRoute.OCR,
     dependencies=[Security(dependency=AccessController())],
@@ -35,6 +34,7 @@
     },
     response_model=OCR,
 )
+@hooks
 async def ocr(
     request: Request,
     body: CreateOCR,