-
Notifications
You must be signed in to change notification settings - Fork 130
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Migrate RAG pipeline to async processing. #2345
Conversation
Signed-off-by: Austin Lee <austin@aryn.ai>
@msfroh Can you take a look at my code changes? I have a really basic question about exception handling in |
} | ||
final int timeout = t; | ||
log.info("Timeout for this request: {} seconds.", timeout); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Seems unnecessary to print this on info level. Move to debug level?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed to debug.
if (conversationId != null && !Strings.hasText(conversationId)) { | ||
throw new IllegalArgumentException("Empty conversation_id is not allowed."); | ||
} | ||
// log.info("LLM question: {}, LLM model {}, conversation id: {}", llmQuestion, llmModel, conversationId); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove this line?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Removed.
List<Interaction> chatHistory = (conversationId == null) | ||
? Collections.emptyList() | ||
: memoryClient.getInteractions(conversationId, interactionSize); | ||
log.info("Using interaction size of {}", interactionSize); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Move to debug level?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes
try { | ||
ChatCompletionOutput output = llm | ||
.doChatCompletion( | ||
// log.info("system_prompt: {}", systemPrompt); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove these two lines?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes
} else { | ||
final Instant memoryStart = Instant.now(); | ||
memoryClient.getInteractions(conversationId, interactionSize, ActionListener.wrap(r -> { | ||
log.info("getInteractions complete. ({})", getDuration(memoryStart)); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove this line or move to debug level?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed to debug.
llm.doChatCompletion(input, new ActionListener<>() { | ||
@Override | ||
public void onResponse(ChatCompletionOutput output) { | ||
log.info("doChatCompletion complete. ({})", getDuration(chatStart)); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove this line or move to debug level?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed to debug.
.getMlModelTensors() | ||
.get(0) | ||
.getDataAsMap(); | ||
// log.info("dataAsMap: {}", dataAsMap.toString()); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
remove this line?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
SearchResponse response, | ||
PipelineProcessingContext requestContext, | ||
ActionListener<SearchResponse> responseListener | ||
) { | ||
log.info("Entering processResponse."); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This log can be removed right?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed to debug.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A lot of "onFailure(Exception e)" implementations in the ActionListeners do not include log.error(). Should we add more logs for errors? This was brought up in the earlier security review too.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the quick fix
* Migrate RAG pipeline to async processing. Signed-off-by: Austin Lee <austin@aryn.ai> * Address reviewer comments. Signed-off-by: Austin Lee <austin@aryn.ai> --------- Signed-off-by: Austin Lee <austin@aryn.ai> (cherry picked from commit 4b26ebf)
* Migrate RAG pipeline to async processing. Signed-off-by: Austin Lee <austin@aryn.ai> * Address reviewer comments. Signed-off-by: Austin Lee <austin@aryn.ai> --------- Signed-off-by: Austin Lee <austin@aryn.ai> (cherry picked from commit 4b26ebf)
@@ -128,14 +141,15 @@ public SearchResponse processResponse(SearchRequest request, SearchResponse resp | |||
} | |||
String conversationId = params.getConversationId(); | |||
|
|||
if (conversationId != null && !Strings.hasText(conversationId)) { | |||
throw new IllegalArgumentException("Empty conversation_id is not allowed."); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Have you managed to test this?
You should probably invoke responseListener.onFailure()
. Otherwise, the current thread may throw and the listener would sit there waiting for a response.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, I have a test for this, but it does not go through the REST layer. I may need an IT test.
log.error("Context " + contextField + " not found in search hit " + hits[i]); | ||
// TODO throw a more meaningful error here? | ||
throw new RuntimeException(); | ||
throw new RuntimeException("Context " + contextField + " not found in search hit " + hits[i]); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Similarly, you need to make sure that this exception gets propagated to the listener. (I don't remember if that's covered by ActionListener.wrap()
. Maybe?)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'll check and also test.
Shoot -- I didn't see your question before this got merged. (I spent most of yesterday traveling to a conference.) The general contract in the async world is that every possible code path needs to notify the listener exactly once or else it will wait indefinitely.
|
Description
Use the async version of the search pipeline process to avoid blocking (remote) calls.
Original bug - opensearch-project/OpenSearch#10248.
Issues Resolved
Check List
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.