From b994546fc737cd29c623120d43fc7e352e648251 Mon Sep 17 00:00:00 2001 From: Matteo Rosati Date: Thu, 23 Apr 2026 15:53:14 +0200 Subject: [PATCH] use typer for better syntax highlight --- chromy/utilities.py | 28 +++++++++++++++++++++------- uv.lock | 2 +- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/chromy/utilities.py b/chromy/utilities.py index 9b39b15..117491b 100644 --- a/chromy/utilities.py +++ b/chromy/utilities.py @@ -1,5 +1,9 @@ from __future__ import annotations +from rich.text import Text +from rich.rule import Rule +from rich.console import Console + from collections.abc import Mapping, Sequence from chromadb import QueryResult @@ -8,10 +12,12 @@ from chromy.chroma_functions import add_data, query_data from chromy.chunk_functions import chunk_file from chromy.embed import embed +CONSOLE = Console() + def print_lines(lines: Sequence[str]) -> None: for line in lines: - print(line) + CONSOLE.print(line) def ingest_file(collection_name: str, file_path: str) -> int: @@ -39,14 +45,19 @@ def format_query_result(result: QueryResult) -> list[str]: if not first_ids: return ["No results found."] - lines = ["Query results:"] + lines = [Rule(title="Query results")] for index, document_id in enumerate(first_ids, start=1): - lines.append(f"{index}.\tid: {document_id}") + # lines.append(f"{index}.\tid: {document_id}") + lines.append( + Text.from_markup(f"[bold]{index}[/].\t[green]id[/]\t\t{document_id}") + ) i = index - 1 if i < len(first_distances): - lines.append(f"\tdistance: {first_distances[i]}") + lines.append( + Text.from_markup(f"\t[green]distance[/]\t{first_distances[i]}") + ) if i < len(first_metadatas): metadata = first_metadatas[i] @@ -55,12 +66,15 @@ def format_query_result(result: QueryResult) -> list[str]: file_name = metadata.get("file_name") if file_name: - lines.append(f"\tfile_name: {file_name}") + lines.append( + Text.from_markup(f"\t[green]file_name[/]\t{file_name}") + ) if i < len(first_documents): - lines.append(f"\tdocument: {first_documents[i]}") + lines.append(Text.from_markup("\n[bold green]Retrieved contents[/]\n")) + lines.append(first_documents[i]) # Print a separator between documents - lines.append(60 * "-") + lines.append(Rule()) return lines diff --git a/uv.lock b/uv.lock index c949307..f0535fd 100644 --- a/uv.lock +++ b/uv.lock @@ -254,7 +254,7 @@ wheels = [ [[package]] name = "chromy" -version = "1.0.0" +version = "1.1.0" source = { editable = "." } dependencies = [ { name = "chromadb" },