Skip to content

Commit

Permalink
Fix #2960 by checking if the values are a list of lists. (#2971)
Browse files Browse the repository at this point in the history
* Fix #2960 by checking values are list of list

* Reduce dictionary look up overhead
  • Loading branch information
thinkall authored and victordibia committed Jul 30, 2024
1 parent 60b1fd4 commit 668146e
Showing 1 changed file with 9 additions and 4 deletions.
13 changes: 9 additions & 4 deletions autogen/agentchat/contrib/vectordb/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,15 +99,20 @@ def chroma_results_to_query_results(data_dict: Dict[str, List[List[Any]]], speci
]
"""

keys = [key for key in data_dict if key != special_key]
keys = [
key
for key in data_dict
if key != special_key and data_dict[key] is not None and isinstance(data_dict[key][0], list)
]
result = []
data_special_key = data_dict[special_key]

for i in range(len(data_dict[special_key])):
for i in range(len(data_special_key)):
sub_result = []
for j, distance in enumerate(data_dict[special_key][i]):
for j, distance in enumerate(data_special_key[i]):
sub_dict = {}
for key in keys:
if data_dict[key] is not None and len(data_dict[key]) > i:
if len(data_dict[key]) > i:
sub_dict[key[:-1]] = data_dict[key][i][j] # remove 's' in the end from key
sub_result.append((sub_dict, distance))
result.append(sub_result)
Expand Down

0 comments on commit 668146e

Please sign in to comment.