AIS-308 Optimize Chat History #24
Conversation
| def add_doc(self, doc: dict[str, Any]) -> None: | ||
| """Add a dict of message to the chat history.""" | ||
| self._db.insert_document(self._collection_name, doc) | ||
|
|
There was a problem hiding this comment.
Let's consider introducing a function called add_qa_message which is used specifically for ArangoGraphQAChain purposes.
When we do this, we can the role property set to "qa, e.g:
{
"_key": "...",
"role": "qa",
"user_input": user_input,
"aql_query": aql_query,
"result": result.content
}
There was a problem hiding this comment.
This would allow us to slightly optimize the AQL query to look like this:
FOR doc IN {collection_name}
FILTER doc.role == "qa"
SORT doc._key DESC
LIMIT @n
RETURN doc
There was a problem hiding this comment.
Idea: Maybe use type instead of role? This would require modifying the other message insertions to include a new type property...
There was a problem hiding this comment.
Idea: Better yet, we can introduce a get_messages_by_role method in ChatHistory that accepts n and role as the parameters
| aql = f""" | ||
| FOR doc IN {collection_name} | ||
| SORT doc._key DESC | ||
| LIMIT @n | ||
| RETURN UNSET(doc, ["_id", "_key", "_rev", "session_id"]) | ||
| """ |
There was a problem hiding this comment.
Small adjustment to this query, could we wrap this in def get_messages_by_role and store it in the ChatHistory class?
| aql = f""" | |
| FOR doc IN {collection_name} | |
| SORT doc._key DESC | |
| LIMIT @n | |
| RETURN UNSET(doc, ["_id", "_key", "_rev", "session_id"]) | |
| """ | |
| aql = f""" | |
| FOR doc IN {collection_name} | |
| FILTER doc.session_id == @session_id | |
| AND doc.role == @role | |
| SORT doc._key DESC | |
| LIMIT @n | |
| RETURN UNSET(doc, ["_id", "_key", "_rev", "session_id"]) | |
| """ |
if include_history:
chat_history.extend(self.chat_history_store.get_messages_by_role("qa"))There was a problem hiding this comment.
I just realized the reason I didn't wrap this in get_messages_by_role() here is because we want to include all messages i.e. both qa and human message(feedback). Then I thought of wrapping this with messages() but it wouldn't recognize qa message as it's not BaseMessage instance.
Therefore, I add get_messages() that optionally takes a role argument. If role is given, it retrieves messages by role otherwise it returns all messages. What do you think?
There was a problem hiding this comment.
This makes sense agreed! Thanks
| }, | ||
| ) | ||
|
|
||
| def add_qa_message(self, user_input: str, aql_query: str, result: str) -> None: |
There was a problem hiding this comment.
for both add_message and add_qa_message, can we introduce a new property added to all documents called time?
import time
....
self._db.collection(self._collection_name).insert(
{
...,
"time": time.time()
}
)
| " Use the 'add_messages' instead." | ||
| ) | ||
|
|
||
| def get_messages(self, role: Optional[str] = None, n_messages: int = 10) -> list: |
There was a problem hiding this comment.
thank you good call - slight adjustment to the signature and the AQL query:
def get_messages(self, role: Optional[str] = None, n_messages: int = 10, excluded_fields: list[str] = ["_id", "_key", "_rev", "session_id"]) -> list:
"""Retrieve messages from ArangoDB, optionally filtered by role."""
query = f"""
FOR doc IN @@col
FILTER doc.session_id == @session_id
{"AND doc.role == @role" if role else ""}
SORT doc.time DESC
LIMIT @n
RETURN UNSET(doc, @excluded_fields)
"""
bind_vars = {
"@col": self._collection_name,
"session_id": self._session_id,
"role": role,
"n": n_messages,
"excluded_fields": excluded_fields,
}
| for msg in self.chat_history_store.messages[-max_history_messages:]: | ||
| cls = HumanMessage if msg.type == "human" else AIMessage | ||
| chat_history.append(cls(content=msg.content)) | ||
| chat_history.extend( |
There was a problem hiding this comment.
good catch on not passing role here since we also want to include human messages
| ) | ||
|
|
||
| def add_qa_message(self, user_input: str, aql_query: str, result: str) -> None: | ||
| """Add a QA message to the chat history.""" |
There was a problem hiding this comment.
can we add the full docstring to this function, i.e :param and :type statements
| n_messages: int = 10, | ||
| excluded_fields: list[str] = ["_id", "_key", "_rev", "session_id", "time"], | ||
| ) -> list: | ||
| """Retrieve messages from ArangoDB, optionally filtered by role.""" |
aMahanna
left a comment
aMahanna
left a comment
There was a problem hiding this comment.
LGTM! Small adjustment to the new docstrings, but feel free to merge when those are implemented and the CI passes
Description
Type of Change
Complexity
How Has This Been Tested?
Checklist