fix missing set memory

Author: leehuwuj

python/llama-index-server/examples/artifact/code_workflow.py

@@ -2,8 +2,6 @@
 import time
 from typing import Any, Literal, Optional, Union
 
-from pydantic import BaseModel
-
 from llama_index.core.chat_engine.types import ChatMessage
 from llama_index.core.llms import LLM
 from llama_index.core.memory import ChatMemoryBuffer
@@ -24,6 +22,7 @@
     UIEvent,
 )
 from llama_index.server.api.utils import get_last_artifact
+from pydantic import BaseModel
 
 
 class Requirement(BaseModel):
@@ -215,6 +214,7 @@ async def planning(
                 content=f"The plan for next step: \n{response.text}",
             )
         )
+        await ctx.set("memory", memory)
         if requirement.next_step == "coding":
             return GenerateArtifactEvent(
                 requirement=requirement,

python/llama-index-server/examples/artifact/document_workflow.py

@@ -2,8 +2,6 @@
 import time
 from typing import Any, Literal, Optional
 
-from pydantic import BaseModel
-
 from llama_index.core.chat_engine.types import ChatMessage
 from llama_index.core.llms import LLM
 from llama_index.core.memory import ChatMemoryBuffer
@@ -24,6 +22,7 @@
     UIEvent,
 )
 from llama_index.server.api.utils import get_last_artifact
+from pydantic import BaseModel
 
 
 class DocumentRequirement(BaseModel):
@@ -188,6 +187,7 @@ async def planning(self, ctx: Context, event: PlanEvent) -> GenerateArtifactEven
                 content=f"Planning for the document generation: \n{response.text}",
             )
         )
+        await ctx.set("memory", memory)
         ctx.write_event_to_stream(
             UIEvent(
                 type="ui_event",