This commit is contained in:
+21
-7
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user