Hey! I'm currently working on a chat app. App is based on websockets (fastapi lib).I have

   async def get_dog_name_streaming(llm, messages):
       async for chunk in await llm.astream_chat(messages):
           print(chunk.delta, end="", flush=True)

   asyncio.run(get_dog_name_streaming(llm, messages))
   

   Settings.context_window = 2048
   

   memory = ChatMemoryBuffer.from_defaults(token_limit=4096)  # or any higher value
   

from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.chat_engine import SimpleChatEngine

memory = ChatMemoryBuffer.from_defaults(token_limit=4096)
chat_engine = SimpleChatEngine.from_defaults(memory=memory)

# Now use chat_engine.astream_chat()

from fastapi.responses import StreamingResponse

@app.post("/chat")
async def chat(...):
  response = await chat_engine.astream_chat(...)
  async def gen():
    async for token in response.async_response_gen():
      # fastapi typically needs a newline to deliminate chunks
      yeild token + "\n" 
  
  return StreamingResponse(gen())

chat_memory = ChatMemoryBuffer.from_defaults(
  token_limit=4096,
  chat_history=self.chat_history
)

self.chat_engine: ContextChatEngine = index.as_chat_engine(
   similarity_top_k=self.similarity_top_k,
   chat_mode='context',
   llm=self.llm,
   context_template=formatted_context_template,
   verbose=True,
   system_prompt=formatted_system_prompt,
   streaming=False,
   memory=chat_memory
)

'ElasticsearchStore' object has no attribute '__pydantic_private__'

In this example:
- The `run_ollama` step is decorated with `@step(num_workers=4)`, allowing up to 4 concurrent executions.
- The `start` step emits multiple events, which will be processed in parallel by `run_ollama`.

To use this workflow:

This approach should allow you to run multiple Ollama instances concurrently without blocking.

Additionally, ensure that you're using the asynchronous methods correctly in your FastAPI application. Here's a basic example of how you might structure your FastAPI route:

# initialize the embedding model in order to speed up the response time
Settings.embed_model = HuggingFaceEmbedding(
    model_name="google-bert/bert-base-multilingual-cased"
)

# initialize router
router = APIRouter()

sessions = {} # stores the chat sessions

# define routes - WebSocket
@router.websocket("/{uuid}") # uuid - user unique identifier
async def chat(websocket: WebSocket, uuid: str):
...
  try:
        while True:
            # parse received data
            data = await websocket.receive_text()
            payload = json.loads(data)
            ...
            message_handler = MessageHandler(llm_instance=chat_LLM, websocket=websocket)
            # generate response from the model
                            message_task = asyncio.to_thread(await message_handler.handle_RAG_message(message=payload['message']), sessions[current_user.get_uuid()]["chat_session"].get_current_conversation().get_conversation_id())
                            assistant_generated_message = await message_task

async def handle_RAG_message(self, message: str) -> str:
        """
        Handles the message from the client
        Returns the generated assistant_generated_message
        """

        # token generation loop
        await self.websocket.send_text(WSResponse(action="generate-state", state="started", error=False, error_message="", timestamp=int(time.time())).to_json())
        
        assistant_generated_message = ""
        chat_engine = self.llm_instance.get_chat_engine()

        print(f"chat_engine: {chat_engine}")

        response = await chat_engine.astream_chat(message=message)
        async for token in response.async_response_gen():
            assistant_generated_message += token
            await self.websocket.send_text(WSResponse(action="generate", partial=token, error=False, error_message="", timestamp=int(time.time())).to_json())
        await self.websocket.send_text(WSResponse(action="generate-state", state="finished", error=False, error_message="", timestamp=int(time.time())).to_json())

        return assistant_generated_message

from llama_index.core import MockEmbedding

embed_model = MockEmbedding(embed_dim=1234)

# token generation loop
        await self.websocket.send_text(WSResponse(action="generate-state", state="started", error=False, error_message="", timestamp=int(time.time())).to_json())
        
        assistant_generated_message = ""
        chat_engine = self.llm_instance.get_chat_engine()

        response = await chat_engine.astream_chat(message=message)
        async def gen():
            async for token in response.async_response_gen():
                yield token

        async for token in gen():
            assistant_generated_message += token
            await self.websocket.send_text(WSResponse(action="generate", partial=token, error=False, error_message="", timestamp=int(time.time())).to_json())

        await self.websocket.send_text(WSResponse(action="generate-state", state="finished", error=False, error_message="", timestamp=int(time.time())).to_json())

response = await chat_engine.astream_chat(message=message)
        async def gen():
            async for token in response.async_response_gen():
                yield token

        async for token in gen():
            assistant_generated_message += token
            await # send websocket response

        await # send websocket response

async def stream_chat(message: str):
    response = await chat_engine.astream_chat(message=message)
    assistant_generated_message = ""
    
    async for token in response.async_response_gen():
        assistant_generated_message += token
        # Yield each token immediately
        yield token
        # Optionally, you can add a small delay to ensure the event loop can process other tasks
        await asyncio.sleep(0)

# In your websocket handler
async for token in stream_chat(message):
    await websocket.send_text(token)

from fastapi.responses import StreamingResponse

@app.get("/stream")
async def stream_endpoint(message: str):
    return StreamingResponse(stream_chat(message), media_type="text/plain")

@app.websocket("/chat")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    
    while True:
        message = await websocket.receive_text()
        
        response = await chat_engine.astream_chat(message=message)
        
        async for token in response.async_response_gen():
            await websocket.send_text(token)
        
        # Optionally send a message to indicate the end of the stream
        await websocket.send_text("[END_OF_STREAM]")

   async for token in response.async_response_gen():
       await websocket.send_text(token)
       await asyncio.sleep(0.01)  # Small delay to prevent buffering
   

async for token in response.async_response_gen():
  print(token, end="", flush=True)