get the run_id from the http header and log run_id by default for eve…

…ry event (#77) * wyvern run_id from the header * add run_id to all the events * string for the log event run_id * v0.0.23
Skyvern-AI · Oct 3, 2023 · 5cfc424 · 5cfc424
1 parent 0db764d
commit 5cfc424
Show file tree

Hide file tree

Showing 10 changed files with 42 additions and 7 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wyvern-ai"
-version = "0.0.22"
+version = "0.0.23"
 description = ""
 authors = ["Wyvern AI <[email protected]>"]
 readme = "README.md"

diff --git a/wyvern/components/candidates/candidate_logger.py b/wyvern/components/candidates/candidate_logger.py
@@ -57,13 +57,16 @@ async def execute(
         current_span = tracer.current_span()
         if current_span:
             current_span.set_tag("candidate_size", len(input.scored_candidates))
-        url_path = request_context.ensure_current_request().url_path
+        wyvern_request = request_context.ensure_current_request()
+        url_path = wyvern_request.url_path
+        run_id = wyvern_request.run_id
 
         def candidate_events_generator() -> List[CandidateEvent]:
             timestamp = datetime.utcnow()
             candidate_events = [
                 CandidateEvent(
                     request_id=input.request.request_id,
+                    run_id=run_id,
                     api_source=url_path,
                     event_timestamp=timestamp,
                     event_data=CandidateEventData(

diff --git a/wyvern/components/events/events.py b/wyvern/components/events/events.py
@@ -3,12 +3,21 @@
 from enum import Enum
 from typing import Generic, Optional, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from pydantic.generics import GenericModel
 
+from wyvern import request_context
+
 EVENT_DATA = TypeVar("EVENT_DATA", bound=BaseModel)
 
 
+def generate_run_id() -> str:
+    curr_wyvern_request = request_context.current()
+    if curr_wyvern_request is None:
+        return str(0)
+    return str(curr_wyvern_request.run_id)
+
+
 class EventType(str, Enum):
     """Enum for the different types of events that can be logged."""
 
@@ -37,6 +46,7 @@ class LoggedEvent(GenericModel, Generic[EVENT_DATA]):
     event_timestamp: Optional[datetime]
     event_type: EventType
     event_data: EVENT_DATA
+    run_id: str = Field(default_factory=generate_run_id)
 
 
 class EntityEventData(BaseModel):

diff --git a/wyvern/components/features/feature_logger.py b/wyvern/components/features/feature_logger.py
@@ -64,7 +64,9 @@ async def execute(
         self, input: FeatureEventLoggingRequest[REQUEST_ENTITY], **kwargs
     ) -> None:
         """Logs feature events."""
-        url_path = request_context.ensure_current_request().url_path
+        wyvern_request = request_context.ensure_current_request()
+        url_path = wyvern_request.url_path
+        run_id = wyvern_request.run_id
 
         def feature_event_generator():
             """Generates feature events. This is a generator function that's called by the event logger. It's never called directly.
@@ -76,6 +78,7 @@ def feature_event_generator():
             return [
                 FeatureEvent(
                     request_id=input.request.request_id,
+                    run_id=run_id,
                     api_source=url_path,
                     event_timestamp=timestamp,
                     event_data=FeatureLogEventData(

diff --git a/wyvern/components/impressions/impression_logger.py b/wyvern/components/impressions/impression_logger.py
@@ -86,13 +86,16 @@ async def execute(
         current_span = tracer.current_span()
         if current_span:
             current_span.set_tag("impression_size", len(input.scored_impressions))
-        url_path = request_context.ensure_current_request().url_path
+        wyvern_request = request_context.ensure_current_request()
+        url_path = wyvern_request.url_path
+        run_id = wyvern_request.run_id
 
         def impression_events_generator() -> List[ImpressionEvent]:
             timestamp = datetime.utcnow()
             impression_events = [
                 ImpressionEvent(
                     request_id=input.request.request_id,
+                    run_id=run_id,
                     api_source=url_path,
                     event_timestamp=timestamp,
                     event_data=ImpressionEventData(

diff --git a/wyvern/components/models/model_component.py b/wyvern/components/models/model_component.py
@@ -105,6 +105,7 @@ async def execute(self, input: INPUT_TYPE, **kwargs) -> MODEL_OUTPUT:
         api_source = wyvern_request.url_path
         request_id = self._get_request_id(input)
         model_output = await self.inference(input, **kwargs)
+        run_id = wyvern_request.run_id
 
         if self.cache_output:
             wyvern_request.cache_model_output(self.name, model_output.data)
@@ -119,6 +120,7 @@ def events_generator() -> List[ModelEvent]:
                         all_events.append(
                             ModelEvent(
                                 request_id=request_id,
+                                run_id=run_id,
                                 api_source=api_source,
                                 event_timestamp=timestamp,
                                 event_data=ModelEventData(
@@ -134,6 +136,7 @@ def events_generator() -> List[ModelEvent]:
                     all_events.append(
                         ModelEvent(
                             request_id=request_id,
+                            run_id=run_id,
                             api_source=api_source,
                             event_timestamp=timestamp,
                             event_data=ModelEventData(

diff --git a/wyvern/event_logging/event_logger.py b/wyvern/event_logging/event_logger.py
@@ -50,12 +50,14 @@ def log_custom_events(events: List[ENTITY_EVENT_DATA_TYPE]) -> None:
     request = request_context.ensure_current_request()
     api_source = request.url_path
     request_id = request.request_id
+    run_id = request.run_id
 
     def event_generator() -> List[LoggedEvent[Any]]:
         timestamp = datetime.utcnow()
         return [
             CustomEvent(
                 request_id=request_id,
+                run_id=run_id,
                 api_source=api_source,
                 event_timestamp=timestamp,
                 event_data=event,

diff --git a/wyvern/experimentation/providers/eppo_provider.py b/wyvern/experimentation/providers/eppo_provider.py
@@ -90,10 +90,12 @@ def event_generator() -> List[LoggedEvent[ExperimentationEventData]]:
             timestamp = datetime.utcnow()
             request_id = request.request_id
             api_source = request.url_path
+            run_id = request.run_id
 
             return [
                 ExperimentationEvent(
                     request_id=request_id,
+                    run_id=run_id,
                     api_source=api_source,
                     event_timestamp=timestamp,
                     event_data=ExperimentationEventData(

diff --git a/wyvern/web_frameworks/fastapi.py b/wyvern/web_frameworks/fastapi.py
@@ -2,10 +2,10 @@
 import logging
 import time
 from contextlib import asynccontextmanager
-from typing import Dict, Type, Union
+from typing import Annotated, Dict, Type, Union
 
 import uvicorn
-from fastapi import BackgroundTasks, FastAPI, HTTPException, Request
+from fastapi import BackgroundTasks, FastAPI, Header, HTTPException, Request
 from fastapi.responses import JSONResponse
 from pydantic import ValidationError
 
@@ -131,6 +131,10 @@ async def post(
             data: root_component.REQUEST_SCHEMA_CLASS,  # type: ignore
             fastapi_request: Request,
             background_tasks: BackgroundTasks,
+            x_wyvern_run_id: Annotated[
+                int,
+                Header(),
+            ] = 0,
         ) -> root_component.RESPONSE_SCHEMA_CLASS:  # type: ignore
             """
             The main entrypoint for the route component. This will parse the request payload, set the WyvernRequest in
@@ -153,10 +157,12 @@ async def post(
                 request_id = None
                 if isinstance(data, BaseWyvernRequest):
                     request_id = data.request_id
+
                 wyvern_req = WyvernRequest.parse_fastapi_request(
                     json=data,
                     req=fastapi_request,
                     request_id=request_id,
+                    run_id=str(x_wyvern_run_id),
                 )
                 request_context.set(wyvern_req)
 

diff --git a/wyvern/wyvern_request.py b/wyvern/wyvern_request.py
@@ -61,6 +61,7 @@ class WyvernRequest:
     ]
 
     request_id: Optional[str] = None
+    run_id: str = "0"
 
     # TODO: params
 
@@ -69,6 +70,7 @@ def parse_fastapi_request(
         cls,
         json: BaseModel,
         req: fastapi.Request,
+        run_id: str = "0",
         request_id: Optional[str] = None,
     ) -> WyvernRequest:
         """
@@ -93,6 +95,7 @@ def parse_fastapi_request(
             feature_map=FeatureMap(feature_map={}),
             model_output_map={},
             request_id=request_id,
+            run_id=run_id,
         )
 
     def cache_model_output(