Context Matters: Using Chat History to Retrieve More Relevant Memories for AI Agent
Posted by Shreyas Shinde, AI Engineer of Gaudiy Inc.
Introduction
In the rapidly evolving field of artificial intelligence, ensuring that AI agents can provide relevant and accurate responses is crucial. One of the key challenges in AI conversations is handling inappropriate or ambiguous user queries. This blog explores the importance of query contextualization using chat history to improve agents’ memory retrieval and overall conversation quality. It also goes deep into how to systematically tune a prompt to achieve query contextualization.
At Gaudiy, we provide “Gaudiy Fanlink,” a community platform for multiple IPs (intellectual properties). It has a Beta feature that allows users to chat with AI avatars of specific characters (such as actual celebrities or anime characters) in collaboration with specific IP holders. We use vector cosine similarity search to retrieve the Character AI’s memories (base information about the character and summaries of chat history with a user) relevant to the user’s message query from a Vector Store (using Vertex AI) and then add them to the LLM’s prompt to initiate a reply generation process. As the user’s message doesn’t always contain enough context to retrieve relevant memories it becomes essential to first contextualize the user’s message using recent chat history and then use it as a query.
Problem Example
Imagine a conversation between a User and Agent below:
User: Have you read any new books recently?
Agent: Yes, I recently read ‘Harry Potter’.
User: That's great. Which volume are you reading?
Agent: I've started reading the first volume. It's very interesting.
User: Nice!When we directly use the user’s last message Nice! as a query for the vector db we would get results which would not be related to what that particular Nice! was about. It would miss out on getting information about key topics and nuance in the current conversation i.e Harry Potter, books, first volume & interesting . If our Agent has memories which are relevant to these topics within the context of Nice! , it would miss out on the chance to utilize them to have a much richer & personalized conversation.
Literature Survey
We investigated 3 different possible approaches:
1.
Prompt an LLM to rewrite the user’s message based on the chat history such that it can be understood standalone. This has been documented in LangChain’s docs.
We quickly found out that although the simple prompt provided by LangChain performs great when used with a powerful but expensive LLM like gpt4o, there were several issues when used with a cost effective LLM like gpt3.5, gemini-1.0, gemini-1.5-flash, claude-sonnet-3 depending upon the model chosen.
Some Issues faced:
- Replies to the user’s last message instead of creating a contextualized message
- Only partially contextualizes the message
- Adds hallucination in the contextualized message
The choice for this LLM needed to be a cost effective one for it be used in our product. We found gemini-1.5-flash to be the most promising among the cheaper models.
2.
To counter the issue of incomplete contextualization mentioned in 1, instead of only 1 generate 3 standalone queries(statement or question) from user message and chat history and retrieve per query.
But this approach has the additional overhead of post-retrieval reranking and filtering the set of memories retrieved by 3 queries. We could use an approach like reciprocal rank fusion or llm filter to do so. We would also need some prompt engineering to improve the relevance of the 3 queries and to avoid adding hallucinations.
3.
HypotheticalDocumentEmbedding: This approach first generates an Agent reply without using any memories and then uses the reply to retrieve memories and then regenerate the reply. As our reply generation process is multi-step this approach leads to doubling our cost so it is not very cost effective.
💡 Decision:
We chose to go ahead with the first approach above as it was judged that by prompt tuning we would be able to recognize and resolve issues and use it in our product cost effectively.
Prompt Tuning for Query Contextualization
In order to improve the simple prompt from LangChain in a structured way we took below steps:
- Generate a synthetic dataset of various types of conversations between a user and agent for evaluating the prompt against using chatgpt4. Add ground truth contextualized queries per example. These are the so called ideal queries which would lead to the retrieval of most relevant memories Upload the dataset to LangSmith.
- Write an evaluator function using LangSmith which calculates the cosine similarity distance between the embeddings of the generated contextualized query and the ground truth contextualized query. Shorter distance would mean the prompt is producing a contextualized query for an example that is similar to the ground truth and leads to similar relevant memories retrieval as would with the ground truth query. We used the langsmith helper library we recently released as OSS to help us with this evaluation.
- Execute the query contextualiser llm per example and evaluate the generated contextualized query. Calculate average score across categories of examples. Identify troublesome examples and update the prompt in an attempt to resolve them. Rerun the evaluation to confirm if they are resolved. This cycle of evaluation not only gave us a direction for improving the prompt on specific category of examples but as we were also computing an average evaluation score during each iteration it also made sure that the prompt was also performing well on the overall dataset.
- Run it against actual production queries. Add few shot examples and set the LLM’s temperature to 0 for consistent outputs as it reduces the LLM’s creativity.
1. Generating Synthetic Dataset using chatgpt4
By chatting with chatgpt4 we created a high quality synthetic dataset of 75 conversations by varying the types of chat_history and user_last_message. It is important to consider the example types deeply for better coverage of actual scenarios.
We also generated ground truth contextualized queries for evaluation.
Get Gaudiy Lab’s stories in your inbox
Join Medium for free to get updates from this writer.
Example:
"chat_history": "L: The SPK has valuable intel.
Shreyas: Should we contact them?
L: Yes, but cautiously.
Shreyas: What's the best approach?
L: A secure line.",
"user_last_message": "Should I set up the call?",
"ground_truth_standalone_message": "Should I set up a secure call to contact the SPK for their valuable intel?",
"example_type": "Multiple turns, question, context from multiple messages"Uploaded to LangSmith:
What was the reason for Human in the loop style of synthetic dataset generation using chatgpt4?
💡 There were clear merits of doing this task in a chat way rather than one prompt and script. Chat way gave more control over the dataset being generated (I first asked it to generate 10 examples with the given prompt and then verify what mistakes it was doing). As we can see how it is failing to follow the instructions we can immediately ask it to correct it, verify again and then produce more examples. We decide the conditions of how we want to vary the examples (in my case I wanted example types like user_last_message is statement or question, chat_history has single message or multiple messages etc). It also allows us to flexibly make changes to the schema mid-way without having to be perfect about it in the first go in case it is needed to change it. As it has a long context window all the history is being used to keep following all instructions in previous messages. Chatgpt easily outputs files in csv, json formats. CSV file can be directly uploaded from LangSmith’s UI to create a new dataset! It also helped me finally translate the file into Japanese because first I had to verify the data in English. This time we only needed to create 75 examples. In case we need 1000 examples the idea will be that after performing this chat and we can satisfactorily and consistently produce examples we ask it to summarize the conversation and create a prompt. And then use that prompt by script and api to generate 1000 examples.
2. Evaluator function using LangSmith
We used the LangSmith’s off the shelf embedding_distance evaluator
from typing import Any
from langchain_google_vertexai import VertexAIEmbeddings
from langsmith.evaluation import LangChainStringEvaluator
from langsmith.evaluation.integrations import _langchain
from langsmith.schemas import Example, Run
from gaudiy_python_kit.ai_testing.tests.utils import get_child_run, get_output
def prepare_data_for_similarity_eval_single_contextualized_query_and_ground_truth( # noqa: D103
run: Run, example: Example | None
) -> _langchain.SingleEvaluatorInput:
prompt_run = get_child_run(run, "prompt")
relevant_memories_run = get_child_run(prompt_run, "get_relevant_memories")
pre_retrieval_query_translation_run = get_child_run(relevant_memories_run, "_pre_retrieval_query_translation")
contextualize_query_run = get_child_run(pre_retrieval_query_translation_run, "contextualize_query")
return {
"prediction": get_output(contextualize_query_run, "output")[0],
"reference": get_output(example, "ground_truth_standalone_message"),
"input": "", # added because mypy complains about missing key
}
evaluators = [
LangChainStringEvaluator(
"embedding_distance",
config={
"embeddings": VertexAIEmbeddings(model_name="textembedding-gecko-multilingual@latest"),
"distance_metric": "cosine",
},
prepare_data=prepare_data_for_similarity_eval_single_contextualized_query_and_ground_truth,
),
]
summary_evaluators: list[Any] = []3. Evaluating the query contextualizer’s prompt on the dataset and prompt tuning
- First of all in order to create a benchmark we ran the evaluator without query contextualization. The score was 0.11 (average of embedding distances between user_last_message and ground_truth_standalone_message).
- Next we ran evaluation on following configuration:
a. v1 of the prompt (langchain’s simple prompt)
b. model: gemini-1.5-flash
The score was 0.08. As the average distance reduced we could conclude that query contextualization makes the user_last_message more cosine similar to the ground_truth query. - We then checked the type of examples which had largest embedding distance as those were the clear targets for prompt tuning. After 22 such prompt revisions we finally reached a score of 0.04 and a good enough prompt that performed well on all the different types of examples. We made sure to version the prompts in code in order to be able to back track and compare variations effectively. Again chatgpt4 helped here in improving the prompts by just giving it the current prompt and a list of examples where the prompt was performing poorly!
💡 We can see from the above graph that we keep tuning the prompt in a manner that the distance keeps reducing!
4. Evaluating against actual production chat log
Running only against synthetic dataset is not enough to confirm the prompt works well in actual practice. It is because synthetic datasets suffer with the problem of textual diversity even if we have made sure to create a variety of examples. For better confidence we tested the prompt against a few production queries and chat history and quickly found some improvements that needed to be done to the prompt. We also found that adding few shot examples to the prompt and setting the LLM’s temperature to 0 leads to much more consistent results. As gpt4o-mini was recently released we also did comparative evaluation using it and found it performed better so we switched to using it.
Results
1. Some actual outputs (in Japanese)
#Chat History:
shreyas: 最近何か新しい本を読みましたか？
AI: はい、最近「ハリーポッター」を読みました。
shreyas: それは素晴らしいですね。どの巻を読んでいますか？
AI: 最初の巻を読み始めました。とても面白いです。
#Last user message:
いいね！
#Contextualized query:
最近「ハリーポッター」の最初の巻を読み始めたことに対して、いいね！と言います。
#Chat History:
shreyas: ピカチュウって、ほんとにあるのかな？
#Last user message:
あるよ！
#Contextualized query:
ピカチュウは実在するキャラクターで、ポケットモンスターの一部として知られていますよ！
#Chat History:
shreyas: 明日、アイドルフェスティバルでライブ配信があるんだ
#Last user message:
それって何ですか？
#Contextualized query:
明日行われるアイドルフェスティバルについて詳しく教えてもらえますか？2. Prompt and model config
Model: gpt-4o-mini
Temperature: 0
Few shots: Yes (ref: using LangChain)
Prompt:
Examine the provided chat history and the last user message to determine its clarity as a complete query or statement.
The chat history is formatted as a continuous string where each message is on a new line and includes the sender's name, while the last user message is given separately without a username identifier.
If the last message is already a clear, complete query or statement that can be understood independently, return it as is.
If it lacks sufficient clarity or requires contextual details to be fully understood on its own, carefully examine all previous interactions in chat history to gather every relevant piece of information.
If the last message implicitly or explicitly references any aspect discussed earlier in the chat history that is essential for understanding then Rewrite the last message to include all necessary details discussed earlier, ensuring it clearly & completely integrates those pertinent details to stand alone confidently and directly connects the context from chat history to the user's statement.
The rewritten message in japanese should maintain the original message's intent, either as a question or a statement, without explaining or answering itself.Note: update the prompt if you want output in English!
Conclusion
In conclusion, the approach of contextualizing user queries using chat history has proven to be a critical enhancement for improving memory retrieval in AI conversational agents. By systematically tuning prompts and leveraging synthetic datasets for evaluation, we were able to significantly reduce the embedding distance between the original and contextualized queries. This improvement ensures that AI agents can retrieve more relevant memories, leading to richer, more personalized conversations. The iterative process of prompt tuning, coupled with real-world testing, has resulted in a robust system capable of handling diverse conversation scenarios effectively. This method not only enhances the user experience but also optimizes the performance of AI agents in practical applications. The success of this approach underscores the importance of context in AI conversations and demonstrates the value of carefully designed prompt engineering and evaluation strategies in developing more capable AI systems.
Gaudiy will continue to advance research and development in AI agents.
More Articles:
About Company:
