dara-core 1.17.6__py3-none-any.whl → 1.18.1__py3-none-any.whl

dara/core/interactivity/any_data_variable.py

@@ -18,7 +18,8 @@ limitations under the License.
 import abc
 import io
 import os
-from typing import Any, Awaitable, Callable, Literal, Optional, TypedDict, Union, cast
+from collections.abc import Awaitable
+from typing import Any, Callable, Literal, Optional, TypedDict, Union, cast
 
 import pandas
 from fastapi import UploadFile
@@ -105,8 +106,8 @@ async def upload(data: UploadFile, data_uid: Optional[str] = None, resolver_id:
     if data_uid is not None:
         try:
             variable = await registry_mgr.get(data_variable_registry, data_uid)
-        except KeyError:
-            raise ValueError(f'Data Variable {data_uid} does not exist')
+        except KeyError as e:
+            raise ValueError(f'Data Variable {data_uid} does not exist') from e
 
         if variable.type == 'derived':
             raise ValueError('Cannot upload data to DerivedDataVariable')
@@ -126,12 +127,12 @@ async def upload(data: UploadFile, data_uid: Optional[str] = None, resolver_id:
     elif file_type == '.xlsx':
         file_object_xlsx = io.BytesIO(content)
         content = pandas.read_excel(file_object_xlsx, index_col=None)
-        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)   # type: ignore
+        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)  # type: ignore
     else:
         # default to csv
         file_object_csv = io.StringIO(content.decode('utf-8'))
         content = pandas.read_csv(file_object_csv, index_col=0)
-        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)   # type: ignore
+        content.columns = content.columns.str.replace('Unnamed: *', 'column_', regex=True)  # type: ignore
 
     # If a data variable is provided, update it with the new content
     if variable:

dara/core/interactivity/any_variable.py

@@ -30,9 +30,8 @@ from fastapi.encoders import jsonable_encoder
 from pydantic import ConfigDict
 
 from dara.core.auth.definitions import SESSION_ID, USER, UserData
-from dara.core.base_definitions import BaseTask
+from dara.core.base_definitions import BaseTask, PendingTask
 from dara.core.base_definitions import DaraBaseModel as BaseModel
-from dara.core.base_definitions import PendingTask
 from dara.core.interactivity.condition import Condition, Operator
 from dara.core.internal.cache_store import CacheStore
 from dara.core.internal.tasks import TaskManager
@@ -204,7 +203,6 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
 
         for session, channels in session_channels.items():
             for ws in channels:
-
                 raw_result = raw_results[ws]
                 # Skip values from clients where the variable is not registered
                 if raw_result == NOT_REGISTERED:
@@ -238,12 +236,12 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
 
         # If we're returning multiple values, in Jupyter environments print an explainer
         try:
-            from IPython import get_ipython
+            from IPython import get_ipython  # pyright: ignore[reportMissingImports]
         except ImportError:
             pass
         else:
             if get_ipython() is not None:
-                from IPython.display import HTML, display
+                from IPython.display import HTML, display  # pyright: ignore[reportMissingImports]
 
                 display(
                     HTML(
@@ -272,7 +270,7 @@ async def get_current_value(variable: dict, timeout: float = 3, raw: bool = Fals
         return results
 
 
-class AnyVariable(BaseModel, abc.ABC):
+class AnyVariable(BaseModel, abc.ABC):  # noqa: PLW1641 # we override equals to create conditions, otherwise we should define hash
     """
     Base class for all variables. Used for typing to specify that any variable can be provided.
     """
@@ -282,7 +280,6 @@ class AnyVariable(BaseModel, abc.ABC):
     uid: str
 
     def __init__(self, uid: Optional[str] = None, **kwargs) -> None:
-
         new_uid = uid
         if new_uid is None:
             new_uid = str(uuid.uuid4())

dara/core/interactivity/data_variable.py

@@ -312,7 +312,7 @@ class DataVariable(AnyDataVariable):
     def ser_model(self, nxt: SerializerFunctionWrapHandler) -> dict:
         parent_dict = nxt(self)
         if 'data' in parent_dict:
-            parent_dict.pop('data')   # make sure data is not included in the serialised dict
+            parent_dict.pop('data')  # make sure data is not included in the serialised dict
         return {**parent_dict, '__typename': 'DataVariable', 'uid': str(parent_dict['uid'])}
 
 

dara/core/interactivity/derived_data_variable.py

@@ -18,7 +18,8 @@ limitations under the License.
 from __future__ import annotations
 
 import asyncio
-from typing import Any, Callable, Coroutine, List, Optional, Union, cast
+from collections.abc import Coroutine
+from typing import Any, Callable, List, Optional, Union, cast
 from uuid import uuid4
 
 from pandas import DataFrame
@@ -185,7 +186,7 @@ class DerivedDataVariable(AnyDataVariable, DerivedVariable):
         store: CacheStore,
         task_mgr: TaskManager,
         args: List[Any],
-        force: bool = False,
+        force_key: Optional[str] = None,
     ) -> DerivedVariableResult:
         """
         Update the underlying derived variable.
@@ -201,7 +202,7 @@ class DerivedDataVariable(AnyDataVariable, DerivedVariable):
         eng_logger.info(
             f'Derived Data Variable {_uid_short} calling superclass get_value', {'uid': var_entry.uid, 'args': args}
         )
-        value = await super().get_value(var_entry, store, task_mgr, args, force)
+        value = await super().get_value(var_entry, store, task_mgr, args, force_key)
 
         # Pin the value in the store until it's read by get data
         await asyncio.gather(
@@ -339,8 +340,8 @@ class DerivedDataVariable(AnyDataVariable, DerivedVariable):
         store: CacheStore,
         task_mgr: TaskManager,
         args: List[Any],
-        force: bool,
         filters: Optional[Union[FilterQuery, dict]] = None,
+        force_key: Optional[str] = None,
     ):
         """
         Helper method to resolve the filtered value of a derived data variable.
@@ -351,11 +352,11 @@ class DerivedDataVariable(AnyDataVariable, DerivedVariable):
         :param store: the store instance to check for cached values
         :param task_mgr: task manager instance
         :param args: the arguments to call the underlying function with
-        :param force: whether to ignore cache
         :param filters: the filters to apply to the data
+        :param force_key: unique key for forced execution, if provided forces cache bypass
         :param pagination: the pagination to apply to the data
         """
-        dv_result = await cls.get_value(dv_entry, store, task_mgr, args, force)
+        dv_result = await cls.get_value(dv_entry, store, task_mgr, args, force_key)
 
         # If the intermediate result was a task/metatask, we need to run it
         # get_data will then pick up the result from the pending task for it

dara/core/interactivity/derived_variable.py

@@ -19,10 +19,10 @@ from __future__ import annotations
 
 import json
 import uuid
+from collections.abc import Awaitable
 from inspect import Parameter, signature
 from typing import (
     Any,
-    Awaitable,
     Callable,
     Generic,
     List,
@@ -32,6 +32,7 @@ from typing import (
     cast,
 )
 
+from cachetools import LRUCache
 from pydantic import (
     ConfigDict,
     Field,
@@ -63,6 +64,16 @@ from dara.core.metrics import RUNTIME_METRICS_TRACKER
 
 VariableType = TypeVar('VariableType')
 
+# Global set to track force keys that have been encountered
+# LRU with 2048 entries should be sufficient to not drop in-progress force keys
+# but also not have to worry about memory leaks
+_force_keys_seen: LRUCache[str, bool] = LRUCache(maxsize=2048)
+
+VALUE_MISSING = object()
+"""
+Sentinel value to indicate that a value is missing from the cache
+"""
+
 
 class DerivedVariableResult(TypedDict):
     cache_key: str
@@ -146,7 +157,12 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
                     raise RuntimeError('run_as_task is not supported within a Jupyter environment')
 
         super().__init__(
-            cache=cache, uid=uid, variables=variables, polling_interval=polling_interval, deps=deps, nested=nested
+            cache=cache,
+            uid=uid,
+            variables=variables,
+            polling_interval=polling_interval,
+            deps=deps,
+            nested=nested,
         )
 
         # Import the registry of variables and register the function at import
@@ -208,15 +224,17 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
         :param uid: uid of a DerivedVariable
         :param deps: list of indexes of dependencies
         """
+        from dara.core.internal.dependency_resolution import clean_force_key
+
         key = f'{uid}'
 
         filtered_args = [arg for idx, arg in enumerate(args) if idx in deps] if deps is not None else args
 
-        for arg in filtered_args:
-            if isinstance(arg, dict):
-                key = f'{key}:{json.dumps(arg, sort_keys=True, default=str)}'
-            else:
-                key = f'{key}:{arg}'
+        for raw_arg in filtered_args:
+            # remove force keys from the arg to not cause extra cache misses
+            arg = clean_force_key(raw_arg)
+
+            key = f'{key}:{json.dumps(arg, sort_keys=True, default=str)}' if isinstance(arg, dict) else f'{key}:{arg}'
         return key
 
     @staticmethod
@@ -269,7 +287,8 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
         if not latest_value_registry.has(var_entry.uid):
             # Keep latest entry per scope (user,session); if cache_type is None, use GLOBAL
             reg_entry = LatestValueRegistryEntry(
-                uid=var_entry.uid, cache=Cache.Policy.MostRecent(cache_type=cache_type or Cache.Type.GLOBAL)
+                uid=var_entry.uid,
+                cache=Cache.Policy.MostRecent(cache_type=cache_type or Cache.Type.GLOBAL),
             )
             latest_value_registry.register(var_entry.uid, reg_entry)
         else:
@@ -285,7 +304,7 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
         store: CacheStore,
         task_mgr: TaskManager,
         args: List[Any],
-        force: bool = False,
+        force_key: Optional[str] = None,
     ) -> DerivedVariableResult:
         """
         Get the value of this DerivedVariable. This method will check the main app store for an appropriate response
@@ -296,7 +315,7 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
         :param store: the store instance to check for cached values
         :param task_mgr: task manager instance
         :param args: the arguments to call the underlying function with
-        :param force: whether to ignore cache
+        :param force_key: unique key for forced execution, if provided forces cache bypass
         """
         assert var_entry.func is not None, 'DerivedVariable function is not defined'
 
@@ -318,30 +337,32 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
             values = []
 
             # dynamic import due to circular import
-            from dara.core.internal.dependency_resolution import (
-                is_resolved_derived_variable,
-                resolve_dependency,
-            )
+            from dara.core.internal.dependency_resolution import resolve_dependency
 
-            eng_logger.info(f'Derived Variable {_uid_short} get_value', {'uid': var_entry.uid, 'args': args})
+            eng_logger.info(
+                f'Derived Variable {_uid_short} get_value',
+                {'uid': var_entry.uid, 'args': args},
+            )
 
             for val in args:
-                # Don't force nested DVs
-                if is_resolved_derived_variable(val):
-                    val['force'] = False
-
                 var_value = await resolve_dependency(val, store, task_mgr)
                 values.append(var_value)
 
             eng_logger.debug(
-                f'DerivedVariable {_uid_short}', 'resolved arguments', {'values': values, 'uid': var_entry.uid}
+                f'DerivedVariable {_uid_short}',
+                'resolved arguments',
+                {'values': values, 'uid': var_entry.uid},
             )
 
             # Loop over the passed arguments and if the expected type is a BaseModel and arg is a dict then convert the dict
             # to an instance of the BaseModel class.
             parsed_args = DerivedVariable._restore_pydantic_models(var_entry.func, *values)
 
-            dev_logger.debug(f'DerivedVariable {_uid_short}', 'executing', {'args': parsed_args, 'uid': var_entry.uid})
+            dev_logger.debug(
+                f'DerivedVariable {_uid_short}',
+                'executing',
+                {'args': parsed_args, 'uid': var_entry.uid},
+            )
 
             # Check if there are any Tasks to be run in the args
             has_tasks = any(isinstance(arg, BaseTask) for arg in parsed_args)
@@ -351,6 +372,24 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
 
             cache_type = var_entry.cache
 
+            # Handle force key tracking to prevent double execution
+            effective_force = force_key is not None
+            if force_key is not None:
+                if force_key in _force_keys_seen:
+                    # This force key has been seen before, don't force again
+                    effective_force = False
+                    eng_logger.debug(
+                        f'DerivedVariable {_uid_short} force key already seen, using cached value',
+                        extra={'uid': var_entry.uid, 'force_key': force_key},
+                    )
+                else:
+                    # First time seeing this force key, add it to the set
+                    _force_keys_seen[force_key] = True
+                    eng_logger.debug(
+                        f'DerivedVariable {_uid_short} new force key, will force recalculation',
+                        extra={'uid': var_entry.uid, 'force_key': force_key},
+                    )
+
             # If deps is not None, force session use
             # Note: this is temporarily commented out as no tests were broken by removing it;
             # once we find what scenario this fixes, we should add a test to cover that scenario and move this snippet
@@ -364,18 +403,33 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
                 {'uid': var_entry.uid},
             )
 
+            # Start with a sentinel value to indicate that the value is missing
+            # from cache, this lets us distinguish between a cache miss and a
+            # value that is None
+            value = VALUE_MISSING
+
             ignore_cache = (
                 var_entry.cache is None
                 or var_entry.polling_interval
                 or DerivedVariable.check_polling(var_entry.variables)
+                or effective_force
             )
-            value = await store.get(var_entry, key=cache_key) if not ignore_cache else None
-
-            eng_logger.debug(
-                f'DerivedVariable {_uid_short}',
-                'retrieved value from cache',
-                {'uid': var_entry.uid, 'cached_value': value},
-            )
+            if not ignore_cache:
+                try:
+                    value = await store.get(var_entry, key=cache_key, raise_for_missing=True)
+                    eng_logger.debug(
+                        f'DerivedVariable {_uid_short}',
+                        'retrieved value from cache',
+                        {'uid': var_entry.uid, 'cached_value': value},
+                    )
+                except KeyError:
+                    eng_logger.debug(
+                        f'DerivedVariable {_uid_short}',
+                        'no value found in cache',
+                        {'uid': var_entry.uid},
+                    )
+                    # key error means no entry found;
+                    # this lets us distinguish from a None value stored and not found
 
             # If it's a PendingTask then return that task so it can be awaited later by a MetaTask
             if isinstance(value, PendingTask):
@@ -392,11 +446,13 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
                     f'DerivedVariable {_uid_short} waiting for pending value',
                     {'uid': var_entry.uid, 'pending_value': value},
                 )
-                return {'cache_key': cache_key, 'value': await store.get_or_wait(var_entry, key=cache_key)}
+                return {
+                    'cache_key': cache_key,
+                    'value': await store.get_or_wait(var_entry, key=cache_key),
+                }
 
-            # If there is a value that is not pending then we have the result so return it
-            # If force is True, don't return even if value is found and recalculate
-            if not force and value is not None:
+            # We retrieved an actual value from the cache, return it
+            if not ignore_cache and value is not VALUE_MISSING:
                 eng_logger.info(
                     f'DerivedVariable {_uid_short} returning cached value directly',
                     {'uid': var_entry.uid, 'cached_value': value},
@@ -492,7 +548,11 @@ class DerivedVariable(NonDataVariable, Generic[VariableType]):
     @model_serializer(mode='wrap')
     def ser_model(self, nxt: SerializerFunctionWrapHandler) -> dict:
         parent_dict = nxt(self)
-        return {**parent_dict, '__typename': 'DerivedVariable', 'uid': str(parent_dict['uid'])}
+        return {
+            **parent_dict,
+            '__typename': 'DerivedVariable',
+            'uid': str(parent_dict['uid']),
+        }
 
 
 class DerivedVariableRegistryEntry(CachedRegistryEntry):

dara/core/interactivity/filtering.py

@@ -20,10 +20,10 @@ from __future__ import annotations
 import re
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union, cast
 
 import numpy
-from pandas import DataFrame, Series  # pylint: disable=unused-import
+from pandas import DataFrame, Series  # noqa: F401
 
 from dara.core.base_definitions import DaraBaseModel as BaseModel
 from dara.core.logging import dev_logger
@@ -157,11 +157,11 @@ def infer_column_type(series: Series) -> ColumnType:
         return ColumnType.CATEGORICAL
 
 
-def _filter_to_series(data: DataFrame, column: str, operator: QueryOperator, value: Any) -> Optional['Series[bool]']:
+def _filter_to_series(data: DataFrame, column: str, operator: QueryOperator, value: Any) -> Optional[Series]:
     """
     Convert a single filter to a Series[bool] for filtering
     """
-    series = data[column]
+    series = cast(Series, data[column])
 
     # Contains is a special case, we always treat the column as a string
     if operator == QueryOperator.CONTAINS:
@@ -175,19 +175,15 @@ def _filter_to_series(data: DataFrame, column: str, operator: QueryOperator, val
             return series.isin(value)
         # Converts date passed from frontend to the right format to compare with pandas
         if col_type == ColumnType.DATETIME:
-            if isinstance(value, List):
-                value = [parseISO(value[0]), parseISO(value[1])]
-            else:
-                value = parseISO(value)
+            value = [parseISO(value[0]), parseISO(value[1])] if isinstance(value, List) else parseISO(value)
         elif col_type == ColumnType.CATEGORICAL:
             value = str(value)
+        elif isinstance(value, List):
+            lower_bound = float(value[0]) if '.' in str(value[0]) else int(value[0])
+            upper_bound = float(value[1]) if '.' in str(value[1]) else int(value[1])
+            value = [lower_bound, upper_bound]
         else:
-            if isinstance(value, List):
-                lower_bound = float(value[0]) if '.' in str(value[0]) else int(value[0])
-                upper_bound = float(value[1]) if '.' in str(value[1]) else int(value[1])
-                value = [lower_bound, upper_bound]
-            else:
-                value = float(value) if '.' in str(value) else int(value)
+            value = float(value) if '.' in str(value) else int(value)
 
         if operator == QueryOperator.GT:
             return series > value
@@ -208,12 +204,14 @@ def _filter_to_series(data: DataFrame, column: str, operator: QueryOperator, val
         return None
 
 
-def _resolve_filter_query(data: DataFrame, query: FilterQuery) -> 'Optional[Series[bool]]':
+def _resolve_filter_query(data: DataFrame, query: FilterQuery) -> Optional[Series]:
     """
     Resolve a FilterQuery to a Series[bool] for filtering. Strips the internal column index from the query.
     """
     if isinstance(query, ValueQuery):
-        return _filter_to_series(data, re.sub(COLUMN_PREFIX_REGEX, '', query.column, 1), query.operator, query.value)
+        return _filter_to_series(
+            data, re.sub(COLUMN_PREFIX_REGEX, repl='', string=query.column, count=1), query.operator, query.value
+        )
     elif isinstance(query, ClauseQuery):
         filters = None
 
@@ -222,15 +220,9 @@ def _resolve_filter_query(data: DataFrame, query: FilterQuery) -> 'Optional[Seri
 
             if resolved_clause is not None:
                 if query.combinator == QueryCombinator.AND:
-                    if filters is None:
-                        filters = resolved_clause
-                    else:
-                        filters = filters & resolved_clause
+                    filters = resolved_clause if filters is None else filters & resolved_clause
                 elif query.combinator == QueryCombinator.OR:
-                    if filters is None:
-                        filters = resolved_clause
-                    else:
-                        filters = filters | resolved_clause
+                    filters = resolved_clause if filters is None else filters | resolved_clause
                 else:
                     raise ValueError(f'Unknown combinator {query.combinator}')
 
@@ -262,7 +254,7 @@ def apply_filters(
     if pagination is not None:
         # ON FETCHING SPECIFIC ROW
         if pagination.index is not None:
-            return data[int(pagination.index) : int(pagination.index) + 1], total_count
+            return cast(DataFrame, data[int(pagination.index) : int(pagination.index) + 1]), total_count
 
         # SORT
         if pagination.orderBy is not None:
@@ -278,7 +270,7 @@ def apply_filters(
             if col == 'index':
                 new_data = new_data.sort_index(ascending=ascending, inplace=False)
             else:
-                new_data = new_data.sort_values(by=col, ascending=ascending, inplace=False)
+                new_data = new_data.sort_values(by=col, ascending=ascending, inplace=False)  # type: ignore
 
         # PAGINATE
         start_index = pagination.offset if pagination.offset is not None else 0
@@ -286,4 +278,4 @@ def apply_filters(
 
         new_data = new_data.iloc[start_index:stop_index]
 
-    return new_data, total_count
+    return cast(DataFrame, new_data), total_count

dara/core/interactivity/plain_variable.py

@@ -42,6 +42,7 @@ VARIABLE_INIT_OVERRIDE = ContextVar[Optional[Callable[[dict], dict]]]('VARIABLE_
 VariableType = TypeVar('VariableType')
 PersistenceStoreType_co = TypeVar('PersistenceStoreType_co', bound=PersistenceStore, covariant=True)
 
+
 # TODO: once Python supports a default value for a generic type properly we can make PersistenceStoreType a second generic param
 class Variable(NonDataVariable, Generic[VariableType]):
     """
@@ -84,7 +85,7 @@ class Variable(NonDataVariable, Generic[VariableType]):
             # TODO: this is temporary, persist_value will eventually become a type of store
             raise ValueError('Cannot provide a Variable with both a store and persist_value set to True')
 
-        super().__init__(**kwargs)   # type: ignore
+        super().__init__(**kwargs)  # type: ignore
 
         if self.store:
             call_async(self.store.init, self)
@@ -257,7 +258,7 @@ class Variable(NonDataVariable, Generic[VariableType]):
                 'Cannot create a Variable from a DerivedDataVariable, only standard DerivedVariables are allowed'
             )
 
-        return cls(default=other)   # type: ignore
+        return cls(default=other)  # type: ignore
 
     @model_serializer(mode='wrap')
     def ser_model(self, nxt: SerializerFunctionWrapHandler) -> dict:

dara/core/interactivity/switch_variable.py

@@ -216,8 +216,8 @@ class SwitchVariable(NonDataVariable):
     Key Serialization:
         When using mappings with SwitchVariable, be aware that JavaScript object keys
         are always strings. The system automatically converts lookup keys to strings:
-        - Python: {True: 'admin', False: 'user'}
-        - JavaScript: {"true": "admin", "false": "user"}
+        - Python: `{True: 'admin', False: 'user'}`
+        - JavaScript: `{"true": "admin", "false": "user"}`
         - Boolean values are converted to lowercase strings ("true"/"false")
         - Other values use standard string conversion to match JavaScript's String() behavior
     """
@@ -274,7 +274,7 @@ class SwitchVariable(NonDataVariable):
         true_value: Union[Any, NonDataVariable],
         false_value: Union[Any, NonDataVariable],
         uid: Optional[str] = None,
-    ) -> 'SwitchVariable':
+    ) -> SwitchVariable:
         """
         Create a SwitchVariable for boolean conditions.
 
@@ -350,7 +350,7 @@ class SwitchVariable(NonDataVariable):
         mapping: Union[Dict[Any, Any], NonDataVariable],
         default: Optional[Union[Any, NonDataVariable]] = None,
         uid: Optional[str] = None,
-    ) -> 'SwitchVariable':
+    ) -> SwitchVariable:
         """
         Create a SwitchVariable with a custom mapping.
 

dara/core/internal/cache_store/base_impl.py

@@ -19,12 +19,13 @@ class CacheStoreImpl(abc.ABC, Generic[PolicyT]):
         """
 
     @abc.abstractmethod
-    async def get(self, key: str, unpin: bool = False) -> Any:
+    async def get(self, key: str, unpin: bool = False, raise_for_missing: bool = False) -> Any:
         """
         Retrieve an entry from the cache.
 
         :param key: The key of the entry to retrieve.
         :param unpin: If true, the entry will be unpinned if it is pinned.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         """
 
     @abc.abstractmethod

dara/core/internal/cache_store/cache_store.py

@@ -23,7 +23,7 @@ def cache_impl_for_policy(policy: PolicyT) -> CacheStoreImpl[PolicyT]:
     """
     impl: Optional[CacheStoreImpl] = None
 
-    if isinstance(policy, LruCachePolicy) or isinstance(policy, MostRecentCachePolicy):
+    if isinstance(policy, (LruCachePolicy, MostRecentCachePolicy)):
         impl = LRUCache(policy)
     elif isinstance(policy, TTLCachePolicy):
         impl = TTLCache(policy)
@@ -62,21 +62,24 @@ class CacheScopeStore(Generic[PolicyT]):
 
         return await cache.delete(key)
 
-    async def get(self, key: str, unpin: bool = False) -> Optional[Any]:
+    async def get(self, key: str, unpin: bool = False, raise_for_missing: bool = False) -> Optional[Any]:
         """
         Retrieve an entry from the cache.
 
         :param key: The key of the entry to retrieve.
         :param unpin: If true, the entry will be unpinned if it is pinned.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         """
         scope = get_cache_scope(self.policy.cache_type)
         cache = self.caches.get(scope)
 
         # No cache for this scope yet
         if cache is None:
+            if raise_for_missing:
+                raise KeyError(f'No cache found for {scope}')
             return None
 
-        return await cache.get(key, unpin=unpin)
+        return await cache.get(key, unpin=unpin, raise_for_missing=raise_for_missing)
 
     async def set(self, key: str, value: Any, pin: bool = False):
         """
@@ -150,21 +153,30 @@ class CacheStore:
 
         return prev_entry
 
-    async def get(self, registry_entry: CachedRegistryEntry, key: str, unpin: bool = False) -> Optional[Any]:
+    async def get(
+        self,
+        registry_entry: CachedRegistryEntry,
+        key: str,
+        unpin: bool = False,
+        raise_for_missing: bool = False,
+    ) -> Optional[Any]:
         """
         Retrieve an entry from the cache for the given registry entry and cache key.
 
         :param registry_entry: The registry entry to retrieve the value for.
         :param key: The key of the entry to retrieve.
         :param unpin: If true, the entry will be unpinned if it is pinned.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         """
         registry_store = self.registry_stores.get(registry_entry.to_store_key())
 
         # No store for this entry yet
         if registry_store is None:
+            if raise_for_missing:
+                raise KeyError(f'No cache store found for {registry_entry.to_store_key()}')
             return None
 
-        return await registry_store.get(key, unpin=unpin)
+        return await registry_store.get(key, unpin=unpin, raise_for_missing=raise_for_missing)
 
     async def get_or_wait(self, registry_entry: CachedRegistryEntry, key: str):
         """

dara/core/internal/cache_store/keep_all.py

@@ -42,18 +42,21 @@ class KeepAllCache(CacheStoreImpl[KeepAllCachePolicy]):
             del self.cache[key]
             return entry
 
-    async def get(self, key: str, unpin: bool = False) -> Optional[Any]:
+    async def get(self, key: str, unpin: bool = False, raise_for_missing: bool = False) -> Optional[Any]:
         """
         Retrieve a value from the cache.
 
         :param key: The key of the value to retrieve.
         :param unpin: This parameter is ignored in KeepAllCache as entries are never evicted.
+        :param raise_for_missing: If true, an exception will be raised if the entry is not found
         :return: The value associated with the key, or None if the key is not in the cache.
         """
         async with self.lock:
             entry = self.cache.get(key)
 
             if entry is None:
+                if raise_for_missing:
+                    raise KeyError(f'No cache entry found for {key}')
                 return None
 
             if unpin: