[WIP] Streaming Chat with Referenced Documents PoC

[TamiTakamiya]

Description

Streaming Chat with Referenced Documents PoC. Requires llama-stack 0.2.7.

Screenshare.-.2025-05-19.9_24_25.PM.mp4

Type of change

• Refactor
• New feature
• Bug fix
• CVE fix
• Optimization
• Documentation Update
• Configuration Update
• Bump-up dependent library
• Bump-up library or tool used for development (does not change the final image)
• CI configuration change
• Konflux configuration change

Related Tickets & Documents

• Related Issue #
• Closes #

Checklist before requesting a review

• I have performed a self-review of my code.
• PR has passed all pre-merge test jobs.
• If it is a core feature, I have added thorough tests.

Testing

• Please provide detailed steps to perform tests related to this code change.
• How were the fix/results from this change verified? Please provide relevant screenshots or results.

[manstis]

@TamiTakamiya From what I understand this effectively adds post processing of the response turns to add formatting and most importantly extract the "Referenced Documents" from tool_call="query_from_memory"? Since, as we discussed yesterday, there will be no runtime "lightspeed-stack" component and "lightspeed-stack" will in essence be an "image builder"/configuration tool the code you propose in this PR should probably be elsewhere.

I am wondering whether it's possible to move all of this to a custom agent overriding the run method (or another, but this seems a plausible candidate?) to handle the post processing?

We'd then configure the llama-stack to use our custom external provider:

providers: 
  agent: 
    provider_type: inline::our-super-agent
...

At best this code would have to live in ansible-ai-connect-service but we should favour trying to write our own agent.

WDYT?

[manstis]

retrieve_response would be in ansible-ai-connect-service as it's using LlamaStackClient to invoke llama-stack etc.

response_processing_wrapper would live in our custom Agent but instead of post-processing LlamaStackClient objects it would post-process the llama-stack server equivalents.

[TamiTakamiya]

@manstis Yes, it's a PoC that is intended to show the possibility to implement the road-core/service /streaming-chat endpoint using the llama-stack builtin::rag/knowledge_search. Initially I thought opening a PR for ansible-ai-connect-service, but after I saw the /query endpint in lightspeed-stack, I decided to creating a PR here. This can be anywhere in between llama-stack and the chatbot client.

It is true that the logic is a bit redundant because the post-processing code parses the streaming output using a regex, stores the results, and then the results are sent at the end of streaming. If we write our own agent or tool, implementations could be simplified.

The _yield_printable_events method of the RAGTurnStreamEventPrinterclass was copied from theTurnStreamEventPrinter` class and contains lots of duplicated lines. I want to reduce the amount of duplicated lines.

[manstis]

@TamiTakamiya Was you wanting to keep this open given your PR for ansible-ai-connect-service? Perhaps this is now better refactoring to a README or something?

[umago]

Thanks Tami for this change! I would love to see a streaming endpoint here. But, I do think that this change needs some refactor and make both /query and /streaming_query compatible. I see a lot of potential in sharing code between these two endpoints I think we should do it otherwise we risk making them incompatible in the future.

[umago]

Let's keep models in the models/ directory.

Also, there's now a QueryRequest model that I think we should use instead of introducing a new one:

lightspeed-stack/src/models/requests.py

Line 55 in 0405417

class QueryRequest(BaseModel):

[umago]

Can we use the QueryResponse for this ? I don't see why introducing a new StreamingResponse model.

[umago]

Let's keep compatibility with the /query endpoint, which supports passing a model/provider in the request

[umago]

Ditto, keeping compatibility with the /query endpoint which supports passing a system_prompt in the request

[umago]

A lot of this code is duplicated with /query. I think we can refactor these two endpoints and have a common base of code. Otherwise we will endpoint creating a lot of inconsistencies between the two.

[TamiTakamiya]

@umago @manstis /streaming_query endpoint was implemented with #126. I will close this and open a new PR only for enabling referenced documents support.