review: cli fixes

examples/document-search/basic.py

@@ -32,11 +32,14 @@
 
 import asyncio
 
+from ragbits.core import audit
 from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
 
+audit.set_trace_handlers("cli")
+
 documents = [
     DocumentMeta.create_text_document_from_literal(
         """

examples/document-search/chroma.py

@@ -35,12 +35,15 @@
 
 from chromadb import EphemeralClient
 
+from ragbits.core import audit
 from ragbits.core.embeddings.litellm import LiteLLMEmbeddings, LiteLLMEmbeddingsOptions
 from ragbits.core.vector_stores import VectorStoreOptions
 from ragbits.core.vector_stores.chroma import ChromaVectorStore
 from ragbits.document_search import DocumentSearch, SearchConfig
 from ragbits.document_search.documents.document import DocumentMeta
 
+audit.set_trace_handlers("cli")
+
 documents = [
     DocumentMeta.create_text_document_from_literal(
         """

examples/document-search/distributed.py

@@ -31,12 +31,15 @@
 
 import asyncio
 
+from ragbits.core import audit
 from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
 from ragbits.document_search.ingestion.processor_strategies import DistributedProcessing
 
+audit.set_trace_handlers("cli")
+
 documents = [
     DocumentMeta.create_text_document_from_literal(
         """

examples/document-search/from_config.py

@@ -30,9 +30,12 @@ class to rephrase the query.
 
 import asyncio
 
+from ragbits.core import audit
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
 
+audit.set_trace_handlers("cli")
+
 documents = [
     DocumentMeta.create_text_document_from_literal(
         """

examples/document-search/multimodal.py

@@ -33,6 +33,7 @@
 import asyncio
 from pathlib import Path
 
+from ragbits.core import audit
 from ragbits.core.embeddings.vertex_multimodal import VertexAIMultimodelEmbeddings
 from ragbits.core.vector_stores.in_memory import InMemoryVectorStore
 from ragbits.document_search import DocumentSearch
@@ -41,6 +42,8 @@
 from ragbits.document_search.ingestion.document_processor import DocumentProcessorRouter
 from ragbits.document_search.ingestion.providers.dummy import DummyImageProvider
 
+audit.set_trace_handlers("cli")
+
 IMAGES_PATH = Path(__file__).parent / "images"
 
 

examples/document-search/qdrant.py

@@ -35,11 +35,14 @@
 
 from qdrant_client import AsyncQdrantClient
 
+from ragbits.core import audit
 from ragbits.core.embeddings.litellm import LiteLLMEmbeddings
 from ragbits.core.vector_stores.qdrant import QdrantVectorStore
 from ragbits.document_search import DocumentSearch, SearchConfig
 from ragbits.document_search.documents.document import DocumentMeta
 
+audit.set_trace_handlers("cli")
+
 documents = [
     DocumentMeta.create_text_document_from_literal(
         """

packages/ragbits-core/src/ragbits/core/audit/__init__.py

@@ -64,6 +64,7 @@ def clear_event_handlers() -> None:
 
     _trace_handlers.clear()
 
+
 @contextmanager
 def trace(name: str | None = None, **inputs: Any) -> Iterator[SimpleNamespace]:  # noqa: ANN401
     """

packages/ragbits-core/src/ragbits/core/audit/base.py

@@ -86,3 +86,36 @@ def trace(self, name: str, **inputs: Any) -> Iterator[SimpleNamespace]:  # noqa:
 
         span = self._spans.get().pop()
         self.stop(outputs=vars(outputs), current_span=span)
+
+
+def format_attributes(data: dict, prefix: str | None = None) -> dict:
+    """
+    Format attributes for CLI.
+
+    Args:
+        data: The data to format.
+        prefix: The prefix to use for the keys.
+
+    Returns:
+        The formatted attributes.
+    """
+    flattened = {}
+
+    for key, value in data.items():
+        current_key = f"{prefix}.{key}" if prefix else key
+
+        if isinstance(value, dict):
+            flattened.update(format_attributes(value, current_key))
+        elif isinstance(value, list | tuple):
+            flattened[current_key] = repr(
+                [
+                    item if isinstance(item, str | float | int | bool) else repr(item)
+                    for item in value  # type: ignore
+                ]
+            )
+        elif isinstance(value, str | float | int | bool):
+            flattened[current_key] = value
+        else:
+            flattened[current_key] = repr(value)
+
+    return flattened