Build support for shadow requests (#89)

* Build support for shadow requests

* Implement ShadowRequest class

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wyvern-ai"
-version = "0.0.27"
+version = "0.0.28"
 description = ""
 authors = ["Wyvern AI <info@wyvern.ai>"]
 readme = "README.md"

wyvern/web_frameworks/fastapi.py

@@ -159,6 +159,10 @@ async def post(
                 await root_component.warm_up(data)
                 output = await root_component.execute(data)
 
+                background_tasks.add_task(
+                    wyvern_req.execute_shadow_requests,
+                )
+
                 background_tasks.add_task(
                     wyvern_kinesis_firehose.put_record_batch_callable,
                     KinesisFirehoseStream.EVENT_STREAM,

wyvern/wyvern_request.py

@@ -3,18 +3,32 @@
 
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Coroutine, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 import fastapi
 from pydantic import BaseModel
 
+from wyvern import request_context
 from wyvern.components.events.events import LoggedEvent
 from wyvern.entities.feature_entities import FeatureDataFrame
 from wyvern.entities.identifier import Identifier
 from wyvern.exceptions import WyvernLoggingOriginalIdentifierMissingError
 
 
+class ShadowRequest:
+    callable: Callable[..., Coroutine[Any, Any, Any]]
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+
+    def __init__(
+        self, shadow_callable: Callable[..., Coroutine[Any, Any, Any]], *args, **kwargs
+    ):
+        self.callable = shadow_callable
+        self.args = args
+        self.kwargs = kwargs
+
+
 @dataclass
 class WyvernRequest:
     """
@@ -70,6 +84,8 @@ class WyvernRequest:
     request_id: Optional[str] = None
     run_id: str = "0"
 
+    shadow_requests: Optional[List[ShadowRequest]] = None
+
     # TODO: params
 
     @classmethod
@@ -163,3 +179,24 @@ def get_original_identifier(
                 identifier=primary_identifier_key,
                 feature_name=feature_name,
             )
+
+    def add_shadow_request_call(
+        self,
+        shadow_request: ShadowRequest,
+    ):
+        if self.shadow_requests is None:
+            self.shadow_requests = []
+
+        self.shadow_requests.append(shadow_request)
+
+    async def execute_shadow_requests(self):
+        if self.shadow_requests is None:
+            return
+        try:
+            request_context.set(self)
+            for shadow_request in self.shadow_requests:
+                await shadow_request.callable(
+                    *shadow_request.args, **shadow_request.kwargs
+                )
+        finally:
+            request_context.reset()