# pip install chromadb

import chromadb

client = chromadb.Client()

collection = client.get_or_create_collection("code_snippets")

collection.add(

ids=["1", "2", "3"],

documents=[

"def add(a, b): return a + b",

"numbers_doubled = list(map(lambda x: x * 2, numbers))",

"print('Hello, world!')"

]

)

func_regex = r"^\s*def\s+\w+\("

matched = collection.get(where_document={"$regex": func_regex})

print(f"Snippets with a function definition: {matched['documents']}")

non_matched = collection.get(where_document={"$not_regex": func_regex})

print(f"Snippets without a function definition: {non_matched['documents']}")