Streaming custom LLM with NemoGuardrails

[nicolemzh]

What are the requirements to stream a custom LLM with NemoGuardrails?

The custom LLM is working with NemoGuardrails without the streaming function. However, I'm having an issue implementing streaming with stream_async or the StreamingHandler. I was wondering if anyone has tried this task and if there are extra steps to stream with a custom LLM compared to GPT or LangChain through NemoGuardrails.

[drazvan]

Hi @nicolemzh ! The NeMo LLM provider is a good example that also supports streaming: https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/nemoguardrails/llm/providers/nemollm.py. In a custom LLM, it needs to be implemented similarly. Let me know if you need more guidance.

[nicolemzh]

Thanks @drazvan, appreciate your help! To give more context, we want to implement streaming with mixtral-8x7b-32768 through the Groq provider. Since we tried using Groq using Langchain (https://python.langchain.com/docs/integrations/chat/groq/) but it is not a listed provider, we decided to register ChatGroq as a "custom LLM" with the same methods, and also taking in a BaseChatModel (https://github.com/langchain-ai/langchain/blob/3856dedff469b2ce08e3d416e9e351e2b8950e27/libs/partners/groq/langchain_groq/chat_models.py):

1. We added "streaming: True" and specified our model in the .yml file
2. We registered ChatGroq as a custom LLM
3. implemented streaming the same way as before (was working with GPT)

However, when I send a request using the streaming implementation, I get a response, but it is not streaming and is returned as the whole message.

I was wondering why this is because we were running into the same issue when streaming GPT with NemoGuardrails in the past. The request would stream perfectly fine using gpt-3.5-turbo-instruct but would not stream when we changed the model to gpt-4 in the .yml file.

[drazvan]

@nicolemzh : Sorry for the delayed follow-up on this. Can you double check that the streaming flag is set correctly here: https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/nemoguardrails/rails/llm/llmrails.py#L358. (add a breakpoint/print statement there).

Also, is the behavior with streaming with gpt-3.5-turbo-instruct vs. gpt-4 still reproducible? It should work with both.