chained model evaluation (#58)

* multi model evaluation

* fix test

* chained model evaluation

* cache model output

* v bump

* support dict model output

* rename

* add get_model_output support

* fix

* 0.0.18-beta1

pyproject.toml

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

tests/components/business_logic/test_pinning_business_logic.py

@@ -66,6 +66,7 @@ def __init__(self):
  entity_store={},
  events=[],
  feature_map=FeatureMap(feature_map={}),
+ model_output_map={},
  ),
  )
  return await pipeline.execute(request)

tests/scenarios/test_product_ranking.py

@@ -13,11 +13,7 @@
  RealtimeFeatureComponent,
  RealtimeFeatureRequest,
 )
-from wyvern.components.models.model_component import (
- ModelComponent,
- ModelInput,
- ModelOutput,
-)
+from wyvern.components.models.model_component import ModelComponent
 from wyvern.components.pipeline_component import PipelineComponent
 from wyvern.config import settings
 from wyvern.core.compression import wyvern_encode
@@ -26,6 +22,7 @@
 from wyvern.entities.feature_entities import FeatureData, FeatureMap
 from wyvern.entities.identifier import Identifier
 from wyvern.entities.identifier_entities import ProductEntity, WyvernEntity
+from wyvern.entities.model_entities import ModelInput, ModelOutput
 from wyvern.entities.request import BaseWyvernRequest
 from wyvern.service import WyvernService
 from wyvern.wyvern_request import WyvernRequest
@@ -387,6 +384,7 @@ async def test_hydrate(mock_redis):
  json=json_input,
  headers={},
  entity_store={},
+ model_output_map={},
  events=[],
  feature_map=FeatureMap(feature_map={}),
  )
@@ -450,6 +448,7 @@ async def test_hydrate__duplicate_brand(mock_redis__duplicate_brand):
  entity_store={},
  events=[],
  feature_map=FeatureMap(feature_map={}),
+ model_output_map={},
  )
  request_context.set(test_wyvern_request)
 

wyvern/__init__.py

@@ -2,11 +2,8 @@
 from wyvern.components.features.realtime_features_component import (
  RealtimeFeatureComponent,
 )
-from wyvern.components.models.model_component import (
- ModelComponent,
- ModelInput,
- ModelOutput,
-)
+from wyvern.components.models.model_chain_component import ModelChainComponent
+from wyvern.components.models.model_component import ModelComponent
 from wyvern.components.pipeline_component import PipelineComponent
 from wyvern.components.ranking_pipeline import (
  RankingPipeline,
@@ -23,6 +20,7 @@
  WyvernDataModel,
  WyvernEntity,
 )
+from wyvern.entities.model_entities import ModelInput, ModelOutput
 from wyvern.feature_store.feature_server import generate_wyvern_store_app
 from wyvern.service import WyvernService
 from wyvern.wyvern_logging import setup_logging
@@ -41,6 +39,7 @@
  "FeatureMap",
  "Identifier",
  "IdentifierType",
+ "ModelChainComponent",
  "ModelComponent",
  "ModelInput",
  "ModelOutput",

wyvern/components/component.py

@@ -5,7 +5,7 @@
 import logging
 from enum import Enum
 from functools import cached_property
-from typing import Dict, Generic, Optional, Set
+from typing import Dict, Generic, List, Optional, Set, Union
 from uuid import uuid4
 
 from wyvern import request_context
@@ -177,3 +177,18 @@ def get_all_features(
  if not feature_data:
  return {}
  return feature_data.features
+
+ def get_model_output(
+ self,
+ model_name: str,
+ identifier: Identifier,
+ ) -> Optional[
+ Union[
+ float,
+ str,
+ List[float],
+ Dict[str, Optional[Union[float, str, list[float]]]],
+ ]
+ ]:
+ current_request = request_context.ensure_current_request()
+ return current_request.get_model_output(model_name, identifier)

wyvern/components/models/model_chain_component.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+from functools import cached_property
+from typing import Optional, Set
+
+from wyvern.components.models.model_component import ModelComponent
+from wyvern.entities.model_entities import MODEL_INPUT, MODEL_OUTPUT, ChainedModelInput
+from wyvern.exceptions import MissingModelChainOutputError
+
+
+class ModelChainComponent(ModelComponent[MODEL_INPUT, MODEL_OUTPUT]):
+ """
+ Model chaining allows you to chain models together so that the output of one model can be the input to another model
+
+ For all the models in the chain, all the request and entities in the model input are the same
+ """
+
+ def __init__(self, *upstreams: ModelComponent, name: Optional[str] = None):
+ super().__init__(*upstreams, name=name)
+ self.chain = upstreams
+
+ @cached_property
+ def manifest_feature_names(self) -> Set[str]:
+ feature_names: Set[str] = set()
+ for model in self.chain:
+ feature_names = feature_names.union(model.manifest_feature_names)
+ return feature_names
+
+ async def inference(self, input: MODEL_INPUT, **kwargs) -> MODEL_OUTPUT:
+ output = None
+ prev_model: Optional[ModelComponent] = None
+ for model in self.chain:
+ curr_input: ChainedModelInput
+ if prev_model is not None and output is not None:
+ curr_input = ChainedModelInput(
+ request=input.request,
+ entities=input.entities,
+ upstream_model_name=prev_model.name,
+ upstream_model_output=output.data,
+ )
+ else:
+ curr_input = ChainedModelInput(
+ request=input.request,
+ entities=input.entities,
+ upstream_model_name=None,
+ upstream_model_output={},
+ )
+ output = await model.execute(curr_input, **kwargs)
+ prev_model = model
+
+ if output is None:
+ raise MissingModelChainOutputError()
+
+ # TODO: do type checking to make sure the output is of the correct type
+ return output

wyvern/components/models/model_component.py

@@ -3,41 +3,19 @@
 import logging
 from datetime import datetime
 from functools import cached_property
-from typing import (
- Dict,
- Generic,
- List,
- Optional,
- Sequence,
- Set,
- Type,
- TypeVar,
- Union,
- get_args,
-)
+from typing import Dict, List, Optional, Sequence, Set, Type, Union, get_args
 
 from pydantic import BaseModel
-from pydantic.generics import GenericModel
 
 from wyvern import request_context
 from wyvern.components.component import Component
 from wyvern.components.events.events import EventType, LoggedEvent
 from wyvern.config import settings
 from wyvern.entities.identifier import Identifier
 from wyvern.entities.identifier_entities import WyvernEntity
+from wyvern.entities.model_entities import MODEL_INPUT, MODEL_OUTPUT
 from wyvern.entities.request import BaseWyvernRequest
 from wyvern.event_logging import event_logger
-from wyvern.exceptions import WyvernModelInputError
-from wyvern.wyvern_typing import GENERALIZED_WYVERN_ENTITY, REQUEST_ENTITY
-
-MODEL_OUTPUT_DATA_TYPE = TypeVar(
- "MODEL_OUTPUT_DATA_TYPE",
- bound=Union[float, str, List[float]],
-)
-"""
-MODEL_OUTPUT_DATA_TYPE is the type of the output of the model. It can be a float, a string, or a list of floats
-(e.g. a list of probabilities, embeddings, etc.)
-"""
 
 logger = logging.getLogger(__name__)
 
@@ -58,6 +36,7 @@ class ModelEventData(BaseModel):
  model_output: str
  entity_identifier: Optional[str] = None
  entity_identifier_type: Optional[str] = None
+ target: Optional[str] = None
 
 
 class ModelEvent(LoggedEvent[ModelEventData]):
@@ -71,74 +50,6 @@ class ModelEvent(LoggedEvent[ModelEventData]):
  event_type: EventType = EventType.MODEL
 
 
-class ModelOutput(GenericModel, Generic[MODEL_OUTPUT_DATA_TYPE]):
- """
- This class defines the output of a model.
-
- Args:
- data: A dictionary mapping entity identifiers to model outputs. The model outputs can also be None.
- model_name: The name of the model. This is optional.
- """
-
- data: Dict[Identifier, Optional[MODEL_OUTPUT_DATA_TYPE]]
- model_name: Optional[str] = None
-
- def get_entity_output(
- self,
- identifier: Identifier,
- ) -> Optional[MODEL_OUTPUT_DATA_TYPE]:
- """
- Get the model output for a given entity identifier.
-
- Args:
- identifier: The identifier of the entity.
-
- Returns:
- The model output for the given entity identifier. This can also be None if the model output is None.
- """
- return self.data.get(identifier)
-
-
-class ModelInput(GenericModel, Generic[GENERALIZED_WYVERN_ENTITY, REQUEST_ENTITY]):
- """
- This class defines the input to a model.
-
- Args:
- request: The request that will be used to generate the model input.
- entities: A list of entities that will be used to generate the model input.
- """
-
- request: REQUEST_ENTITY
- entities: List[GENERALIZED_WYVERN_ENTITY] = []
-
- @property
- def first_entity(self) -> GENERALIZED_WYVERN_ENTITY:
- """
- Get the first entity in the list of entities. This is useful when you know that there is only one entity.
-
- Returns:
- The first entity in the list of entities.
- """
- if not self.entities:
- raise WyvernModelInputError(model_input=self)
- return self.entities[0]
-
- @property
- def first_identifier(self) -> Identifier:
- """
- Get the identifier of the first entity in the list of entities. This is useful when you know that there is only
- one entity.
-
- Returns:
- The identifier of the first entity in the list of entities.
- """
- return self.first_entity.identifier
-
-
-MODEL_INPUT = TypeVar("MODEL_INPUT", bound=ModelInput)
-MODEL_OUTPUT = TypeVar("MODEL_OUTPUT", bound=ModelOutput)
-
-
 class ModelComponent(
  Component[
  MODEL_INPUT,
@@ -155,11 +66,14 @@ def __init__(
  self,
  *upstreams,
  name: Optional[str] = None,
+ cache_output: bool = True,
  ):
  super().__init__(*upstreams, name=name)
  self.model_input_type = self.get_type_args_simple(0)
  self.model_output_type = self.get_type_args_simple(1)
 
+ self.cache_output = cache_output
+
  @classmethod
  def get_type_args_simple(cls, index: int) -> Type:
  """
@@ -185,26 +99,51 @@ async def execute(self, input: MODEL_INPUT, **kwargs) -> MODEL_OUTPUT:
  """
  The model_name and model_score will be automatically logged
  """
- api_source = request_context.ensure_current_request().url_path
+ wyvern_request = request_context.ensure_current_request()
+ api_source = wyvern_request.url_path
  request_id = input.request.request_id
  model_output = await self.inference(input, **kwargs)
 
+ if self.cache_output:
+ wyvern_request.cache_model_output(self.name, model_output.data)
+
  def events_generator() -> List[ModelEvent]:
  timestamp = datetime.utcnow()
- return [
- ModelEvent(
- request_id=request_id,
- api_source=api_source,
- event_timestamp=timestamp,
- event_data=ModelEventData(
- model_name=model_output.model_name or self.__class__.__name__,
- model_output=str(output),
- entity_identifier=identifier.identifier,
- entity_identifier_type=identifier.identifier_type,
- ),
- )
- for identifier, output in model_output.data.items()
- ]
+ all_events: List[ModelEvent] = []
+ for identifier, output in model_output.data.items():
+ if isinstance(output, dict):
+ for key, value in output.items():
+ all_events.append(
+ ModelEvent(
+ request_id=request_id,
+ api_source=api_source,
+ event_timestamp=timestamp,
+ event_data=ModelEventData(
+ model_name=model_output.model_name
+ or self.__class__.__name__,
+ model_output=str(value),
+ target=key,
+ entity_identifier=identifier.identifier,
+ entity_identifier_type=identifier.identifier_type,
+ ),
+ ),
+ )
+ else:
+ all_events.append(
+ ModelEvent(
+ request_id=request_id,
+ api_source=api_source,
+ event_timestamp=timestamp,
+ event_data=ModelEventData(
+ model_name=model_output.model_name
+ or self.__class__.__name__,
+ model_output=str(output),
+ entity_identifier=identifier.identifier,
+ entity_identifier_type=identifier.identifier_type,
+ ),
+ ),
+ )
+ return all_events
 
  event_logger.log_events(events_generator) # type: ignore
 

wyvern/components/models/modelbit_component.py

@@ -4,15 +4,12 @@
 from functools import cached_property
 from typing import Any, Dict, List, Optional, Set, Tuple, TypeAlias, Union
 
-from wyvern.components.models.model_component import (
- MODEL_INPUT,
- MODEL_OUTPUT,
- ModelComponent,
-)
+from wyvern.components.models.model_component import ModelComponent
 from wyvern.config import settings
 from wyvern.core.http import aiohttp_client
 from wyvern.entities.identifier import Identifier
 from wyvern.entities.identifier_entities import WyvernEntity
+from wyvern.entities.model_entities import MODEL_INPUT, MODEL_OUTPUT
 from wyvern.entities.request import BaseWyvernRequest
 from wyvern.exceptions import (
  WyvernModelbitTokenMissingError,