feat: save context in runtime

Author: GabrielVasilescu04

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/runtime.py

@@ -143,10 +143,7 @@ async def _run_workflow(
         workflow_input = UiPathChatMessagesMapper.map_input(input or {})
         is_resuming = bool(options and options.resume)
 
-        if is_resuming:
-            self._context = await self._load_context()
-        else:
-            self._context = Context(self.workflow)
+        self._context = await self._load_context()
 
         # Make the Context discoverable from inside steps
         if self.debug_mode and self._context is not None:
@@ -241,7 +238,9 @@ async def _run_workflow(
             else:
                 yield self._create_suspended_result(suspended_event)
         else:
-            yield self._create_success_result(await handler)
+            final_result = await handler
+            await self._save_context()
+            yield self._create_success_result(final_result)
 
     async def get_schema(self) -> UiPathRuntimeSchema:
         """Get schema for this LlamaIndex runtime."""