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)