Refactor memory layer to be used across flow and chat agent #1707

arjunkumargiri · 2023-11-29T00:08:31Z

Description

Refactored conversation creation to MLAgentExecutor so that the same memory ID could be used across multiple agents. Generated memory ID will be passed to agents as part of input parameters. Integrated memory with flow agent, with output of the flow agent being appended to additional info field of interaction.

Request:

{
  "parameters": {
    "question": "How many indexes do I have in my cluster?",
    "verbose": false
  }
}

Response:

{
    "inference_results": [
        {
            "output": [
                {
                    "name": "response",
                    "dataAsMap": {
                        "response": "There are 7 indexes in the cluster."
                    }
                },
                {
                    "name": "QuestionSuggestor",
                    "result": " Here are some follow up questions the Human may ask:\n\n[question1: 'How many shards are there per index?',\nquestion2: 'What is the health status of the indexes?',  \nquestion3: 'What are the names of the indexes?']"
                }
            ]
        }
    ]
}

Conversation history:

{
    "interactions": [
        {
            "conversation_id": "EY9sGIwBQ8497oqLl2Np",
            "interaction_id": "Eo9sGIwBQ8497oqLnWMT",
            "create_time": "2023-11-29T00:13:39.188239Z",
            "input": "How many indexes do I have in my cluster?",
            "prompt_template": null,
            "response": "There are 7 indexes in the cluster.",
            "origin": null,
            "additional_info": {
                "LLMResponseGenerator.output": "ModelTensorOutput(mlModelOutputs=[ModelTensors(mlModelTensors=[ModelTensor(name=response, data=null, shape=null, dataType=null, byteBuffer=null, result=null, dataAsMap={response=There are 7 indexes in the cluster.})], statusCode=null)])",
                "QuestionSuggestor.output": " Here are some follow up questions the Human may ask:\n\n[question1: 'How many shards are there per index?',\nquestion2: 'What is the health status of the indexes?',  \nquestion3: 'What are the names of the indexes?']",
                "None": "None"
            }
        }
    ]
}

TODO:

Add unit test while merging in main branch
Add output of each flow agent step to memory
Instead of agent updating final answer, agent executor should update final answer interaction in memory

Issues Resolved

[List any issues this PR will resolve]

Check List

New functionality includes testing.
- All tests pass
New functionality has been documented.
- New functionality has javadoc added
Commits are signed per the DCO using --signoff

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.

codecov · 2023-11-29T00:32:57Z

Codecov Report

Attention: 146 lines in your changes are missing coverage. Please review.

Comparison is base (bb86968) 68.60% compared to head (371b685) 68.94%.
Report is 2 commits behind head on feature/agent_framework_dev.

Files	Patch %	Lines
...ch/ml/engine/algorithms/agent/MLAgentExecutor.java	0.00%	87 Missing ⚠️
.../ml/engine/algorithms/agent/MLFlowAgentRunner.java	31.25%	18 Missing and 4 partials ⚠️
...nsearch/ml/engine/tools/AbstractRetrieverTool.java	74.19%	14 Missing and 2 partials ⚠️
...a/org/opensearch/ml/engine/tools/VectorDBTool.java	0.00%	11 Missing ⚠️
.../ml/engine/algorithms/agent/MLChatAgentRunner.java	89.13%	4 Missing and 1 partial ⚠️
...g/opensearch/ml/engine/tools/NeuralSparseTool.java	86.48%	3 Missing and 2 partials ⚠️

Additional details and impacted files

@@                        Coverage Diff                        @@
##             feature/agent_framework_dev    #1707      +/-   ##
=================================================================
+ Coverage                          68.60%   68.94%   +0.34%     
- Complexity                          2591     2611      +20     
=================================================================
  Files                                239      241       +2     
  Lines                              12757    12875     +118     
  Branches                            1284     1291       +7     
=================================================================
+ Hits                                8752     8877     +125     
+ Misses                              3404     3392      -12     
- Partials                             601      606       +5

Flag	Coverage Δ
ml-commons	`68.94% <48.77%> (+0.34%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

arjunkumargiri · 2023-11-29T03:51:27Z

@jngz-es @ylwu-amzn please review

Hailong-am · 2023-11-30T05:35:07Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java

+            .getFinalInteractions(memoryId, PREVIOUS_INTERACTION, ActionListener.<List<Interaction>>wrap(interactions -> {
+                if (interactions.size() == 0) {
+                    throw new IllegalStateException("No existing interactions to update");
+                }


Few questions here

if flow agent don't have conversational agent as its tools, there should not have any interactions behind the memory. Should we throw exception for this case?

Get last 1 interaction may not work when race condition happens, there may have new interaction before previous interaction have final answer

Good question:

Will update to create an empty interaction in case there are no existing interaction.

As per my understanding getFinalInteraction will return previous interaction with final answer. @Zhangxunmt please confirm if getFinalInteractions can result in race condition.

getFinalInteractions returns all interactions that are not traces. The final interaction itself is always created when an interaction starts, with an empty response and it waits until all the traces are complete to obtain the final answer. So it's possible that new interaction is created before the previous one receives the final response.

Hailong-am · 2023-11-30T06:51:27Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLAgentExecutor.java

+                            String title = inputDataSet.getParameters().get(QUESTION);
+
+                            ConversationIndexMemory.Factory conversationIndexMemoryFactory =
+                                (ConversationIndexMemory.Factory) memoryFactoryMap.get(memoryType);


does this mean we can only use conversation_index as memory type for all agents?

In my opinion, memory should not be tied to an agent and should be common to all agents. We could have different implementations for memory that could be dynamically chosen but it should reusable across all agents.

This memoryFactoryMap is supposed to support different memory types. Any memory that implements the "Memory" class can be registered in this Map at initialization. So any agent can fetch the memory based on their needs using a memory type. ATM, we only implemented ConversationIndexMemory but there will be more on the way.

@jngz-es @ylwu-amzn have question about Memory interface, is memory instance a single memory object or it's a memory container?

BufferedMemory implements like a container for all sessions

ConversationIndexMemory have a property conversation_id associate with it, it's more like a single memory object

Please ignore BufferedMemory for now, we don't use it right now as it is not P0 I think.
Yeah, ConversationIndexMemory is an object which is allocated for each request.

Zhangxunmt

This PR name is a little confusing. I thought its a big refactor for all the memory layer :).

Hailong-am · 2023-12-04T09:19:38Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLFlowAgentRunner.java

-                    }
+                    String outputResponse = parseResponse(output);
+                    params.put(outputKey, outputResponse);
+                    additionalInfo.put(outputKey, outputResponse);


Using previous step output as chat history may not sufficient for generating suggestions. At least we should include question, that could include in prompt template. Further more, chat history in conversational agent should used.

Parameters field already contains question as part of the map.

Adding historical chat context to flow parameters will be added as part of a different PR.

Hailong-am · 2023-12-05T07:21:49Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

+                                .updateInteraction(
+                                    parentInteractionId,
+                                    ImmutableMap.of(AI_RESPONSE_FIELD, finalAnswer1),
+                                    ActionListener.<UpdateResponse>wrap(updateResponse -> {


why ADDITIONAL_INFO_FIELD, additionalInfo been removed ?

Resolved merge conflict

Signed-off-by: Arjun kumar Giri <[email protected]>

jngz-es · 2023-12-06T18:26:05Z

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

@@ -299,17 +298,7 @@ private void runReAct(
        String maxIteration = Optional.ofNullable(tmpParameters.get("max_iteration")).orElse("3");

        // Create root interaction.
-        StepListener<CreateInteractionResponse> createRootItListener = new StepListener<>();


Got it, thanks!

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:08 — with GitHub Actions Failure

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:08 — with GitHub Actions Error

arjunkumargiri force-pushed the flow_memory branch from 3d43a3c to 0e2c00d Compare November 29, 2023 00:21

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:21 — with GitHub Actions Error

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:21 — with GitHub Actions Failure

arjunkumargiri force-pushed the flow_memory branch from 0e2c00d to c9765fa Compare November 29, 2023 00:25

arjunkumargiri temporarily deployed to ml-commons-cicd-env November 29, 2023 00:25 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:25 — with GitHub Actions Error

arjunkumargiri temporarily deployed to ml-commons-cicd-env November 29, 2023 00:25 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env November 29, 2023 00:25 — with GitHub Actions Failure

arjunkumargiri marked this pull request as ready for review November 29, 2023 00:35

Hailong-am reviewed Nov 30, 2023

View reviewed changes

Zhangxunmt reviewed Nov 30, 2023

View reviewed changes

arjunkumargiri marked this pull request as draft December 1, 2023 18:18

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 1, 2023 18:18 — with GitHub Actions Failure

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 1, 2023 18:18 — with GitHub Actions Error

arjunkumargiri force-pushed the flow_memory branch from f2215ec to 584dfbc Compare December 1, 2023 18:19

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 1, 2023 18:19 — with GitHub Actions Failure

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 1, 2023 18:19 — with GitHub Actions Error

arjunkumargiri force-pushed the flow_memory branch from 584dfbc to b4dbe23 Compare December 1, 2023 18:38

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 4, 2023 04:07 — with GitHub Actions Inactive

Hailong-am reviewed Dec 4, 2023

View reviewed changes

Hailong-am reviewed Dec 5, 2023

View reviewed changes

arjunkumargiri force-pushed the flow_memory branch from c29dee1 to 6dd5242 Compare December 5, 2023 21:45

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 21:45 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 21:45 — with GitHub Actions Error

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 21:46 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 21:46 — with GitHub Actions Failure

arjunkumargiri marked this pull request as draft December 5, 2023 22:00

arjunkumargiri added 4 commits December 5, 2023 14:04

Refactor memory layer to be used across flow and chat agent

9cc059f

Signed-off-by: Arjun kumar Giri <[email protected]>

Create parent interaction ID in MLAgentExecutor

97e80c3

Signed-off-by: Arjun kumar Giri <[email protected]>

Do not convert object to string before toJson

fbc1096

Signed-off-by: Arjun kumar Giri <[email protected]>

Rebased and fixed escape JSON

15d1f72

Signed-off-by: Arjun kumar Giri <[email protected]>

arjunkumargiri force-pushed the flow_memory branch from 6dd5242 to a7bb564 Compare December 5, 2023 22:05

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 22:05 — with GitHub Actions Error

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 22:05 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 22:05 — with GitHub Actions Failure

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 22:05 — with GitHub Actions Inactive

Fix merge conflict

371b685

Signed-off-by: Arjun kumar Giri <[email protected]>

arjunkumargiri force-pushed the flow_memory branch from a7bb564 to 371b685 Compare December 5, 2023 22:12

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 22:12 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 22:12 — with GitHub Actions Failure

arjunkumargiri temporarily deployed to ml-commons-cicd-env December 5, 2023 22:12 — with GitHub Actions Inactive

arjunkumargiri had a problem deploying to ml-commons-cicd-env December 5, 2023 22:12 — with GitHub Actions Error

arjunkumargiri marked this pull request as ready for review December 5, 2023 22:47

jngz-es approved these changes Dec 6, 2023

View reviewed changes

ylwu-amzn merged commit fcbfc7e into opensearch-project:feature/agent_framework_dev Dec 6, 2023
4 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor memory layer to be used across flow and chat agent #1707

Refactor memory layer to be used across flow and chat agent #1707

arjunkumargiri commented Nov 29, 2023 •

edited

Loading

codecov bot commented Nov 29, 2023 •

edited

Loading

arjunkumargiri commented Nov 29, 2023

Hailong-am Nov 30, 2023

arjunkumargiri Nov 30, 2023

Zhangxunmt Nov 30, 2023

Hailong-am Nov 30, 2023

arjunkumargiri Nov 30, 2023

Zhangxunmt Nov 30, 2023

Hailong-am Dec 1, 2023

jngz-es Dec 1, 2023

Zhangxunmt left a comment

Hailong-am Dec 4, 2023

arjunkumargiri Dec 5, 2023

Hailong-am Dec 5, 2023

arjunkumargiri Dec 5, 2023

jngz-es Dec 6, 2023

Refactor memory layer to be used across flow and chat agent #1707

Refactor memory layer to be used across flow and chat agent #1707

Conversation

arjunkumargiri commented Nov 29, 2023 • edited Loading

Description

Issues Resolved

Check List

codecov bot commented Nov 29, 2023 • edited Loading

Codecov Report

arjunkumargiri commented Nov 29, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Zhangxunmt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

arjunkumargiri commented Nov 29, 2023 •

edited

Loading

codecov bot commented Nov 29, 2023 •

edited

Loading