[Bug]: Unable to use ChromaDB for vector memory

[tituslhy]

Bug Description

When attempting to use chromadb as a vector memory, I received a key error on sub_dicts. I'm not sure if this is a chromadb issue because I've no errors when setting the vector_store to None and using the in-memory vector store.

Version

0.10.68

Steps to Reproduce

Assuming you already have a chromaDB instance of messages

import chromadb
from llama_index.vector_stores import ChromaVectorStore
from llama_index.core.memory import VectorMemory, Simple ComposableMemory
from llama_index.core.agent import FunctionCallingAgentWorker
from llama_index.core.tools import FunctionTool

client = chromadb.PersistentClient(path= <path to db>)
collection = client.get_or_create_collection("memory")
vector_store = ChromaVectorStore(chroma_collection = collection)
vector_memory = VectorMemory.from_defaults(
    vector_store = vector_store,
    embed_model = embed_model,
    retriever_kwargs = {"similarity_top_k": 2}
)
chat_memory_buffer = ChatMemoryBuffer.from_defaults()
composable_memory = SimpleComposableMemory.from_defaults(
    primary_memory = chat_memory_buffer,
    secondary_memory_sources=[vector_memory]
)
def add(a, b):
    """Adds 2 values given - these can be integers, floats or
    string concatenation"""
    return a + b

tools = [FunctionTool.from_defaults(add)]

agent_worker = FunctionCallingAgentWorker.from_tools(
    tools = tools,
    llm = llm,
    verbose = True
)
agent = agent_worker.as_agent(memory=composable_memory)
response = agent.chat(<Question from memory>)

Relevant Logs/Tracbacks

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/instrumentation/dispatcher.py:260, in Dispatcher.span.<locals>.wrapper(func, instance, args, kwargs)
    252 self.span_enter(
    253     id_=id_,
    254     bound_args=bound_args,
   (...)
    257     tags=tags,
    258 )
    259 try:
--> 260     result = func(*args, **kwargs)
    261 except BaseException as e:
    262     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/callbacks/utils.py:41, in trace_method.<locals>.decorator.<locals>.wrapper(self, *args, **kwargs)
     39 callback_manager = cast(CallbackManager, callback_manager)
     40 with callback_manager.as_trace(trace_id):
---> 41     return func(self, *args, **kwargs)

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/agent/runner/base.py:646, in AgentRunner.chat(self, message, chat_history, tool_choice)
    641     tool_choice = self.default_tool_choice
    642 with self.callback_manager.event(
    643     CBEventType.AGENT_STEP,
    644     payload={EventPayload.MESSAGES: [message]},
    645 ) as e:
--> 646     chat_response = self._chat(
    647         message=message,
    648         chat_history=chat_history,
    649         tool_choice=tool_choice,
    650         mode=ChatResponseMode.WAIT,
    651     )
    652     assert isinstance(chat_response, AgentChatResponse)
    653     e.on_end(payload={EventPayload.RESPONSE: chat_response})

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/instrumentation/dispatcher.py:260, in Dispatcher.span.<locals>.wrapper(func, instance, args, kwargs)
    252 self.span_enter(
    253     id_=id_,
    254     bound_args=bound_args,
   (...)
    257     tags=tags,
    258 )
    259 try:
--> 260     result = func(*args, **kwargs)
    261 except BaseException as e:
    262     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/agent/runner/base.py:578, in AgentRunner._chat(self, message, chat_history, tool_choice, mode)
    575 dispatcher.event(AgentChatWithStepStartEvent(user_msg=message))
    576 while True:
    577     # pass step queue in as argument, assume step executor is stateless
--> 578     cur_step_output = self._run_step(
    579         task.task_id, mode=mode, tool_choice=tool_choice
    580     )
    582     if cur_step_output.is_last:
    583         result_output = cur_step_output

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/instrumentation/dispatcher.py:260, in Dispatcher.span.<locals>.wrapper(func, instance, args, kwargs)
    252 self.span_enter(
    253     id_=id_,
    254     bound_args=bound_args,
   (...)
    257     tags=tags,
    258 )
    259 try:
--> 260     result = func(*args, **kwargs)
    261 except BaseException as e:
    262     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/agent/runner/base.py:412, in AgentRunner._run_step(self, task_id, step, input, mode, **kwargs)
    408 # TODO: figure out if you can dynamically swap in different step executors
    409 # not clear when you would do that by theoretically possible
    411 if mode == ChatResponseMode.WAIT:
--> 412     cur_step_output = self.agent_worker.run_step(step, task, **kwargs)
    413 elif mode == ChatResponseMode.STREAM:
    414     cur_step_output = self.agent_worker.stream_step(step, task, **kwargs)

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/instrumentation/dispatcher.py:260, in Dispatcher.span.<locals>.wrapper(func, instance, args, kwargs)
    252 self.span_enter(
    253     id_=id_,
    254     bound_args=bound_args,
   (...)
    257     tags=tags,
    258 )
    259 try:
--> 260     result = func(*args, **kwargs)
    261 except BaseException as e:
    262     self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/callbacks/utils.py:41, in trace_method.<locals>.decorator.<locals>.wrapper(self, *args, **kwargs)
     39 callback_manager = cast(CallbackManager, callback_manager)
     40 with callback_manager.as_trace(trace_id):
---> 41     return func(self, *args, **kwargs)

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/agent/function_calling/step.py:305, in FunctionCallingAgentWorker.run_step(self, step, task, **kwargs)
    299 tools = self.get_tools(task.input)
    301 # get response and tool call (if exists)
    302 response = self._llm.chat_with_tools(
    303     tools=tools,
    304     user_msg=None,
--> 305     chat_history=self.get_all_messages(task),
    306     verbose=self._verbose,
    307     allow_parallel_tool_calls=self.allow_parallel_tool_calls,
    308 )
    309 tool_calls = self._llm.get_tool_calls_from_response(
    310     response, error_on_no_tool_call=False
    311 )
    312 tool_outputs: List[ToolOutput] = []

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/agent/function_calling/step.py:195, in FunctionCallingAgentWorker.get_all_messages(self, task)
    192 def get_all_messages(self, task: Task) -> List[ChatMessage]:
    193     return (
    194         self.prefix_messages
--> 195         + task.memory.get(input=task.input)
    196         + task.extra_state["new_memory"].get_all()
    197     )

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/memory/simple_composable_memory.py:76, in SimpleComposableMemory.get(self, input, **kwargs)
     74 def get(self, input: Optional[str] = None, **kwargs: Any) -> List[ChatMessage]:
     75     """Get chat history."""
---> 76     return self._compose_message_histories(input, **kwargs)

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/memory/simple_composable_memory.py:89, in SimpleComposableMemory._compose_message_histories(self, input, **kwargs)
     87 secondary_histories = []
     88 for mem in self.secondary_memory_sources:
---> 89     secondary_history = mem.get(input, **kwargs)
     90     secondary_history = [m for m in secondary_history if m not in messages]
     92     if len(secondary_history) > 0:

File /app/tlim2/anaconda3/envs/llamaindex/lib/python3.12/site-packages/llama_index/core/memory/vector_memory.py:140, in VectorMemory.get(self, input, initial_token_count, **kwargs)
    134 nodes = retriever.retrieve(input or "")
    136 # retrieve underlying messages
    137 return [
    138     ChatMessage.parse_obj(sub_dict)
    139     for node in nodes
--> 140     for sub_dict in node.metadata["sub_dicts"]
    141 ]

KeyError: 'sub_dicts'

[dosubot[bot]]

The error you're encountering, KeyError: 'sub_dicts', indicates that the metadata attribute in the node object does not contain the expected sub_dicts key. This issue arises when using ChromaVectorStore with VectorMemory.

Potential Solution

To resolve this issue, ensure that the metadata attribute in the node object is correctly populated with the sub_dicts key. You can modify the code where the TextNode instances are created to include the sub_dicts key in the metadata dictionary.

Here's an example of how to modify the TextNode creation to include the sub_dicts key:

from llama_index.core.schema import TextNode

nodes = [
    TextNode(
        text="The Shawshank Redemption",
        metadata={
            "author": "Stephen King",
            "theme": "Friendship",
            "year": 1994,
            "sub_dicts": []  # Add this line
        },
    ),
    TextNode(
        text="The Godfather",
        metadata={
            "director": "Francis Ford Coppola",
            "theme": "Mafia",
            "year": 1972,
            "sub_dicts": []  # Add this line
        },
    ),
    # Add similar lines for other TextNode instances
]

Explanation

The metadata attribute in the node object is a dictionary that can contain various key-value pairs. In your case, it seems that the sub_dicts key is expected but not present. By adding the sub_dicts key with an appropriate value (e.g., an empty list), you can avoid the KeyError.

Additional Context

If you need more information on how the metadata attribute is populated and handled, you can refer to the relevant methods in the ChromaVectorStore class, such as the add method, which processes the metadata attribute before storing it ^[1].

By ensuring that the metadata attribute includes the sub_dicts key, you should be able to resolve the KeyError and successfully use ChromaVectorStore with VectorMemory.

To continue talking to Dosu, mention @dosu.