fix modelbit response bug + making feature store value available in realtime feature component

pyproject.toml

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

wyvern/components/features/feature_retrieval_pipeline.py

@@ -6,6 +6,7 @@
 from ddtrace import tracer
 from pydantic.generics import GenericModel
 
+from wyvern import request_context
 from wyvern.components.component import Component
 from wyvern.components.features.feature_logger import (
     FeatureEventLoggingComponent,
@@ -22,7 +23,7 @@
 )
 from wyvern.entities.candidate_entities import CandidateSetEntity
 from wyvern.entities.feature_entities import FeatureData, FeatureMap
-from wyvern.entities.feature_entity_helpers import feature_map_create, feature_map_join
+from wyvern.entities.feature_entity_helpers import feature_map_create
 from wyvern.entities.identifier_entities import WyvernEntity
 from wyvern.wyvern_typing import REQUEST_ENTITY
 
@@ -143,6 +144,9 @@ async def execute(
 
         all_entities = input.request.get_all_entities(cached=True)
         all_identifiers = input.request.get_all_identifiers(cached=True)
+
+        current_request = request_context.ensure_current_request()
+
         # TODO (suchintan): Pass in the feature retrieval features here so they can leverage them
         feature_retrieval_request = FeatureStoreRetrievalRequest(
             identifiers=all_identifiers,
@@ -156,6 +160,7 @@ async def execute(
                 **kwargs,
             )
         )
+        current_request.feature_map = feature_retrieval_response
 
         """
         TODO (suchintan):
@@ -307,6 +312,7 @@ async def execute(
                 *real_time_candidate_features,
                 *real_time_candidate_combination_features,
             )
+            current_request.extend_feature_map(real_time_feature_responses)
 
         with tracer.trace("FeatureRetrievalPipeline.create_feature_response"):
             await self.feature_logger_component.execute(
@@ -316,10 +322,4 @@ async def execute(
                 ),
             )
 
-            # TODO (suchintan): Improve performance of this
-            feature_responses = feature_map_join(
-                feature_retrieval_response,
-                real_time_feature_responses,
-            )
-
-        return feature_responses
+        return current_request.feature_map

wyvern/components/models/modelbit_component.py

@@ -134,12 +134,12 @@ async def inference(self, input: MODEL_INPUT, **kwargs) -> MODEL_OUTPUT:
         output_data: Dict[Identifier, Optional[Union[float, str, List[float]]]] = {}
 
         for batch_idx, resp in enumerate(responses):
-            if resp.status_code != 200:
+            if resp.status != 200:
                 logger.warning(f"Modelbit inference failed: {resp.text}")
                 continue
-            resp_list: List[
-                List[Union[float, str, List[float], None]]
-            ] = resp.json().get(
+            resp_list: List[List[Union[float, str, List[float], None]]] = (
+                await resp.json()
+            ).get(
                 "data",
                 [],
             )

wyvern/components/pipeline_component.py

@@ -2,7 +2,6 @@
 from functools import cached_property
 from typing import Optional, Set, Type
 
-from wyvern import request_context
 from wyvern.components.api_route_component import APIRouteComponent
 from wyvern.components.component import Component
 from wyvern.components.features.feature_retrieval_pipeline import (
@@ -67,11 +66,9 @@ async def retrieve_features(self, request: REQUEST_ENTITY) -> None:
             requested_feature_names=self.feature_names,
             feature_overrides=self.realtime_features_overrides,
         )
-        feature_map = await self.feature_retrieval_pipeline.execute(
+        await self.feature_retrieval_pipeline.execute(
             feature_request,
         )
-        current_request = request_context.ensure_current_request()
-        current_request.feature_map = feature_map
 
     async def warm_up(self, input: REQUEST_ENTITY) -> None:
         await super().warm_up(input)

wyvern/wyvern_request.py

@@ -77,3 +77,6 @@ def parse_fastapi_request(
             feature_map=FeatureMap(feature_map={}),
             request_id=request_id,
         )
+
+    def extend_feature_map(self, feature_map: FeatureMap) -> None:
+        self.feature_map.feature_map.update(feature_map.feature_map)