Quick Chroma Test

chroma_db_test.py

import chromadb
from pathlib import Path

def get_project_root() -> Path:
  return Path(__file__).parent


if __name__ == '__main__':
  project_root = get_project_root()

  markdown_files = [f for f in project_root.iterdir() if f.is_file() and str(f)[-2:] == "md"]

  chroma_client = chromadb.Client()

  collection = chroma_client.create_collection(name='test_collection')

  ids = []
  documents = []
  for file_path in markdown_files:
    ids.append(str(file_path))
    with open(file=file_path, mode='r', encoding='utf8') as file:
      content = file.read()
      documents.append(content)

  collection.add(
    documents=documents,
    ids=ids
  )

  food_results = collection.query(
    query_texts=["Food"],
    n_results=1
  )

  print(f'File: {food_results['ids'][0][0]}')
  print(f'Distance: {food_results['distances'][0][0]}')
  print(f'Peek: """{food_results['documents'][0][0][:120]}"""')
  print('\n\n')


  tech_results = collection.query(
    query_texts=["tech"],
    n_results=1
  )

  print(f'File: {tech_results['ids'][0][0]}')
  print(f'Distance: {tech_results['distances'][0][0]}')
  print(f'Peek: """{tech_results['documents'][0][0][:120]}"""')
  print('\n\n')


  money_results = collection.query(
    query_texts=["money"],
    n_results=1
  )

  print(f'File: {money_results['ids'][0][0]}')
  print(f'Distance: {money_results['distances'][0][0]}')
  print(f'Peek: """{money_results['documents'][0][0][:120]}"""')
  print('\n\n')