[Bug]: ValueError Triggered for Missing Key in Context with Default = None

[nguyenph17]

Bug Description

I followed your code in https://docs.llamaindex.ai/en/stable/understanding/workflows/state/ and got the error with key some_database not in the context llama-index-core/llama_index/core/workflow/context.py

Version

0.11.3

Steps to Reproduce

Mimic your guideline in https://docs.llamaindex.ai/en/stable/understanding/workflows/state/ Also, adding await for ctx.get and ctx.set funcions

Relevant Logs/Tracbacks

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[9], line 2
      1 w = StatefulFlow(timeout=10, verbose=True, disable_validation=True)
----> 2 result = await w.run(query="Some query")
      3 print(result)

File ~/miniconda3/envs/pipeline/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py:291, in Dispatcher.span.<locals>.async_wrapper(func, instance, args, kwargs)
    283 self.span_enter(
    284     id_=id_,
    285     bound_args=bound_args,
   (...)
    288     tags=tags,
    289 )
    290 try:
--> 291     result = await func(*args, **kwargs)
    292 except BaseException as e:
    293     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

File ~/miniconda3/envs/pipeline/lib/python3.11/site-packages/llama_index/core/workflow/workflow.py:252, in Workflow.run(self, **kwargs)
    249 if exception_raised:
    250     # Make sure to stop streaming, in case the workflow terminated abnormally
    251     ctx.write_event_to_stream(StopEvent())
--> 252     raise exception_raised
    254 # Raise WorkflowTimeoutError if the workflow timed out
    255 if not we_done:

File ~/miniconda3/envs/pipeline/lib/python3.11/site-packages/llama_index/core/workflow/workflow.py:152, in Workflow._start.<locals>._task(name, queue, step, config)
    149 instrumented_step = dispatcher.span(step)
    151 if asyncio.iscoroutinefunction(step):
--> 152     new_ev = await instrumented_step(**kwargs)
    153 else:
    154     run_task = functools.partial(instrumented_step, **kwargs)

File ~/miniconda3/envs/pipeline/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py:291, in Dispatcher.span.<locals>.async_wrapper(func, instance, args, kwargs)
    283 self.span_enter(
    284     id_=id_,
    285     bound_args=bound_args,
   (...)
    288     tags=tags,
    289 )
    290 try:
--> 291     result = await func(*args, **kwargs)
    292 except BaseException as e:
    293     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

Cell In[8], line 14
     10 @step
     11 async def start(
     12     self, ctx: Context, ev: StartEvent
     13 ) -> SetupEvent | StepTwoEvent:
---> 14     db = await ctx.get("some_database", default=None)
     15     if db is None or db == 0:
     16         print("Need to load data")

File ~/miniconda3/envs/pipeline/lib/python3.11/site-packages/llama_index/core/workflow/context.py:77, in Context.get(self, key, default)
     74         return default
     76 msg = f"Key '{key}' not found in Context"
---> 77 raise ValueError(msg)

ValueError: Key 'some_database' not found in Context

[dosubot[bot]]

@nguyenph17 Hello! I'm here to help you with any issues or questions you have. If you need assistance with anything related to the repository, feel free to ask.

To resolve the ValueError triggered by a missing key 'some_database' in the Context class when using the StatefulFlow workflow, you should ensure that the value for "some_database" is set in the context before trying to access it. Here is an example of how you can modify your code to handle the missing key:

from llama_index.core.workflow import (
    StartEvent,
    StopEvent,
    Workflow,
    step,
    Event,
    Context,
)

class SetupEvent(Event):
    query: str

class StepTwoEvent(Event):
    query: str

class StatefulFlow(Workflow):
    @step
    async def start(
        self, ctx: Context, ev: StartEvent
    ) -> SetupEvent | StepTwoEvent:
        db = await ctx.get("some_database", default=None)
        if db is None:
            print("Need to load data")
            return SetupEvent(query=ev.query)

        # do something with the query
        return StepTwoEvent(query=ev.query)

    @step
    async def setup(self, ctx: Context, ev: SetupEvent) -> StartEvent:
        # load data
        await ctx.set("some_database", [1, 2, 3])
        return StartEvent(query=ev.query)

    @step
    async def step_two(self, ctx: Context, ev: StepTwoEvent) -> StopEvent:
        # do something with the data
        print("Data is ", await ctx.get("some_database"))

        return StopEvent(result=(await ctx.get("some_database"))[1])

w = StatefulFlow(timeout=10, verbose=False)
result = await w.run(query="Some query")
print(result)

By providing a default value in the get method, you can avoid the ValueError and handle the missing key scenario gracefully. This approach ensures that your workflow can proceed even if the key 'some_database' is not initially present in the context ^[1][2].

To continue talking to Dosu, mention @dosu.

[nguyenph17]

@nguyenph17 Hello! I'm here to help you with any issues or questions you have. If you need assistance with anything related to the repository, feel free to ask.

To resolve the ValueError triggered by a missing key 'some_database' in the Context class when using the StatefulFlow workflow, you should ensure that the value for "some_database" is set in the context before trying to access it. Here is an example of how you can modify your code to handle the missing key:

from llama_index.core.workflow import (
    StartEvent,
    StopEvent,
    Workflow,
    step,
    Event,
    Context,
)

class SetupEvent(Event):
    query: str

class StepTwoEvent(Event):
    query: str

class StatefulFlow(Workflow):
    @step
    async def start(
        self, ctx: Context, ev: StartEvent
    ) -> SetupEvent | StepTwoEvent:
        db = await ctx.get("some_database", default=None)
        if db is None:
            print("Need to load data")
            return SetupEvent(query=ev.query)

        # do something with the query
        return StepTwoEvent(query=ev.query)

    @step
    async def setup(self, ctx: Context, ev: SetupEvent) -> StartEvent:
        # load data
        await ctx.set("some_database", [1, 2, 3])
        return StartEvent(query=ev.query)

    @step
    async def step_two(self, ctx: Context, ev: StepTwoEvent) -> StopEvent:
        # do something with the data
        print("Data is ", await ctx.get("some_database"))

        return StopEvent(result=(await ctx.get("some_database"))[1])

w = StatefulFlow(timeout=10, verbose=False)
result = await w.run(query="Some query")
print(result)

By providing a default value in the get method, you can avoid the ValueError and handle the missing key scenario gracefully. This approach ensures that your workflow can proceed even if the key 'some_database' is not initially present in the context [1][2].

To continue talking to Dosu, mention @dosu.

Yes, but I think the function ctx.get should return the default value, even if the default value is None, so that it can trigger SetupEvent

[logan-markewich]

Yea this seems like a bug, will make a fix for this

[huhu415]

same bug

ValueError Traceback (most recent call last)
Cell In[11], [line 28](vscode-notebook-cell:?execution_count=11&line=28)
     [24](vscode-notebook-cell:?execution_count=11&line=24) return StopEvent(result=await ctx.get("some_database")[1])
     [27](vscode-notebook-cell:?execution_count=11&line=27) w = StatefulFlow(timeout=10, verbose=False)
---> [28](vscode-notebook-cell:?execution_count=11&line=28) result = await w.run(query="Some query")
     [29](vscode-notebook-cell:?execution_count=11&line=29) print(result)

File c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:80, in Handle._run(self)
     [78](c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:78) def _run(self):
     [79](c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:79) try:
---> [80](c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:80) self._context.run(self._callback, *self._args)
     [81](c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:81) except (SystemExit, KeyboardInterrupt):
     [82](c:\ProgramData\miniconda3\envs\zzy10\lib\asyncio\events.py:82) raise

File c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:281, in Dispatcher.span.<locals>.wrapper.<locals>.handle_future_result(future, span_id, bound_args, instance, context)
    [273](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:273) def handle_future_result(
    [274](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:274) future: asyncio.Future,
    [275](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:275) span_id: str,
   (...)
    [278](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:278) context: Context,
    [279](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:279) ) -> None:
    [280](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:280) try:
--> [281](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:281) result = future.result()
    [282](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\instrumentation\dispatcher.py:282) self.span_exit(
...
    [173](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\workflow\context.py:173) return default
    [175](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\workflow\context.py:175) msg = f"Key '{key}' not found in Context"
--> [176](c:\ProgramData\miniconda3\envs\zzy10\lib\site-packages\llama_index\core\workflow\context.py:176) raise ValueError(msg)

ValueError: Key 'some_database' not found in Context

fromhere, its not work

[sharathe100111]

This issue should be fixed once my pull-request is completed by maintainers at Fixed issue with default value set as None for get method