Kz/287 feat nested fields as columns (#308)

kzamlynska · web-flow · commit 885bc68bab01 · 2025-01-29T15:03:42.000+01:00
diff --git a/packages/ragbits-cli/src/ragbits/cli/state.py b/packages/ragbits-cli/src/ragbits/cli/state.py
@@ -2,10 +2,12 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from enum import Enum
-from typing import TypeVar
+from types import UnionType
+from typing import Optional, TypeVar, Union, get_args, get_origin
 
 import typer
 from pydantic import BaseModel
+from pydantic.fields import FieldInfo
 from rich.console import Console
 from rich.table import Column, Table
 
@@ -48,39 +50,72 @@ def print_output_table(
         console.print("No results")
         return
 
-    fields = {**data[0].model_fields, **data[0].model_computed_fields}
-
-    # Human readable titles for columns
-    titles = {
-        key: value.get("title", key)
-        for key, value in data[0].model_json_schema(mode="serialization")["properties"].items()
-    }
+    base_fields = {**data[0].model_fields, **data[0].model_computed_fields}
 
     # Normalize the list of columns
     if columns is None:
-        columns = {key: Column() for key in fields}
+        columns = {key: Column() for key in base_fields}
     elif isinstance(columns, str):
         columns = {key: Column() for key in columns.split(",")}
     elif isinstance(columns, Sequence):
         columns = {key: Column() for key in columns}
 
-    # Add headers to columns if not provided
-    for key in columns:
-        if key not in fields:
-            Console(stderr=True).print(f"Unknown column: {key}")
-            raise typer.Exit(1)
-
-        column = columns[key]
+    # check if columns are correct
+    for column_name in columns:
+        field = _get_nested_field(column_name, base_fields)
+        column = columns[column_name]
         if column.header == "":
-            column.header = titles.get(key, key)
+            column.header = field.title if field.title else column_name.replace("_", " ").replace(".", " ").title()
 
     # Create and print the table
     table = Table(*columns.values(), show_header=True, header_style="bold magenta")
+
     for row in data:
-        table.add_row(*[str(getattr(row, key)) for key in columns])
+        row_to_add = []
+        for key in columns:
+            *path_fragments, field_name = key.strip().split(".")
+            base_row = row
+            for fragment in path_fragments:
+                base_row = getattr(base_row, fragment)
+            z = getattr(base_row, field_name)
+            row_to_add.append(str(z))
+        table.add_row(*row_to_add)
+
     console.print(table)
 
 
+def _get_nested_field(column_name: str, base_fields: dict) -> FieldInfo:
+    """
+    Check if column name exists in the model schema.
+
+    Args:
+        column_name: name of the column to check
+        base_fields: model fields
+    Returns:
+        field: nested field
+    """
+    fields = base_fields
+    *path_fragments, field_name = column_name.strip().split(".")
+    for fragment in path_fragments:
+        if fragment not in fields:
+            Console(stderr=True).print(
+                f"Unknown column: {'.'.join(path_fragments + [field_name])} ({fragment} not found)"
+            )
+            raise typer.Exit(1)
+        model_class = fields[fragment].annotation
+        if get_origin(model_class) in [UnionType, Optional, Union]:
+            types = get_args(model_class)
+            model_class = next((t for t in types if t is not type(None)), None)
+        if model_class and issubclass(model_class, BaseModel):
+            fields = {**model_class.model_fields, **model_class.model_computed_fields}
+    if field_name not in fields:
+        Console(stderr=True).print(
+            f"Unknown column: {'.'.join(path_fragments + [field_name])} ({field_name} not found)"
+        )
+        raise typer.Exit(1)
+    return fields[field_name]
+
+
 def print_output_json(data: Sequence[ModelT]) -> None:
     """
     Display data from Pydantic models in a JSON format.
diff --git a/packages/ragbits-cli/tests/unit/test_state.py b/packages/ragbits-cli/tests/unit/test_state.py
@@ -0,0 +1,121 @@
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import typer
+from pydantic import BaseModel
+from pydantic.fields import Field, FieldInfo
+from rich.table import Column, Table
+
+from ragbits.cli.state import OutputType, _get_nested_field, print_output, print_output_table
+from ragbits.document_search.documents.sources import LocalFileSource
+
+
+class InnerTestModel(BaseModel):
+    id: int
+    name: str = Field(title="Name of the inner model", description="Name of the inner model")
+    location: LocalFileSource
+
+
+class OtherTestModel(BaseModel):
+    id: int
+    name: str
+    location: InnerTestModel
+
+
+class MainTestModel(BaseModel):
+    id: int
+    name: str
+    model: OtherTestModel | None
+
+
+data = [
+    MainTestModel(
+        id=1,
+        name="A",
+        model=OtherTestModel(
+            id=11,
+            name="aa",
+            location=InnerTestModel(id=111, name="aa1", location=LocalFileSource(path=Path("folder_1"))),
+        ),
+    ),
+    MainTestModel(
+        id=2,
+        name="B",
+        model=OtherTestModel(
+            id=22,
+            name="bb",
+            location=InnerTestModel(id=222, name="aa2", location=LocalFileSource(path=Path("folder_2"))),
+        ),
+    ),
+]
+
+
+@patch("ragbits.cli.state.print_output_table")
+@patch("ragbits.cli.state.print_output_json")
+def test_print_output_text(mock_print_output_json: MagicMock, mock_print_output_table: MagicMock):
+    with patch("ragbits.cli.state.cli_state") as mock_cli_state:
+        mock_cli_state.output_type = OutputType.text
+        columns = {"id": Column(), "name": Column()}
+        print_output(data, columns=columns)
+        mock_print_output_table.assert_called_once_with(data, columns)
+        mock_print_output_json.assert_not_called()
+
+
+@patch("ragbits.cli.state.print_output_table")
+@patch("ragbits.cli.state.print_output_json")
+def test_print_output_json(mock_print_output_json: MagicMock, mock_print_output_table: MagicMock):
+    with patch("ragbits.cli.state.cli_state") as mock_cli_state:
+        mock_cli_state.output_type = OutputType.json
+        print_output(data)
+        mock_print_output_table.assert_not_called()
+        mock_print_output_json.assert_called_once_with(data)
+
+
+def test_print_output_unsupported_output_type():
+    with patch("ragbits.cli.state.cli_state") as mock_cli_state:
+        mock_cli_state.output_type = "unsupported_type"
+        with pytest.raises(ValueError, match="Unsupported output type: unsupported_type"):
+            print_output(data)
+
+
+def test_print_output_table():
+    with patch("rich.console.Console.print") as mock_print:
+        columns = {"id": Column(), "model.location.location.path": Column(), "model.location.name": Column()}
+        print_output_table(data, columns)
+        mock_print.assert_called_once()
+        args, _ = mock_print.call_args_list[0]
+        printed_table = args[0]
+        assert isinstance(printed_table, Table)
+        assert printed_table.columns[0].header == "Id"
+        assert printed_table.columns[1].header == "Model Location Location Path"
+        assert printed_table.columns[2].header == "Name of the inner model"
+        assert printed_table.row_count == 2
+
+
+def test_get_nested_field():
+    column = "model.location.location.path"
+    fields = {"name": FieldInfo(annotation=str), "model": FieldInfo(annotation=OtherTestModel)}
+
+    try:
+        result = _get_nested_field(column, fields)
+        assert result.annotation == Path
+    except typer.Exit:
+        pytest.fail("typer.Exit was raised unexpectedly")
+
+
+def test_get_nested_field_wrong_field():
+    column_names = [
+        ("model.location.wrong_field", "wrong_field"),
+        ("model.wrong_path.location.path", "wrong_path"),
+        ("wrong_path.location.location.path", "wrong_path"),
+        ("model.location.path", "path"),
+        ("model.location.location.path.additional_field", "additional_field"),
+    ]
+    fields = {"name": FieldInfo(annotation=str), "model": FieldInfo(annotation=OtherTestModel)}
+
+    for wrong_column, wrong_fragment in column_names:
+        with patch("rich.console.Console.print") as mock_print:
+            with pytest.raises(typer.Exit, match="1"):
+                _get_nested_field(wrong_column, fields)
+            mock_print.assert_called_once_with(f"Unknown column: {wrong_column} ({wrong_fragment} not found)")
diff --git a/packages/ragbits-core/tests/unit/utils/test_config_handling.py b/packages/ragbits-core/tests/unit/utils/test_config_handling.py
@@ -32,7 +32,7 @@ def example_factory() -> ExampleClassWithConfigMixin:
     return ExampleSubclass("aligator", 42)
 
 
-def test_defacult_from_config():
+def test_default_from_config():
     config = {"foo": "foo", "bar": 1}
     instance = ExampleClassWithConfigMixin.from_config(config)
     assert instance.foo == "foo"
@@ -113,7 +113,6 @@ def test_subclass_from_defaults_instance_yaml():
         subproject="core",
         current_dir=projects_dir / "project_with_instances_yaml",
     )
-    print(config)
     instance = ExampleClassWithConfigMixin.subclass_from_defaults(config)
     assert isinstance(instance, ExampleSubclass)
     assert instance.foo == "I am a foo"
diff --git a/packages/ragbits-document-search/tests/cli/test_search.py b/packages/ragbits-document-search/tests/cli/test_search.py
@@ -69,6 +69,29 @@ def test_search_columns():
     assert "<DocumentType.TXT: 'txt'>" in result.stdout
 
 
+def test_search_nested_columns():
+    runner = CliRunner(mix_stderr=False)
+    result = runner.invoke(
+        ds_app,
+        [
+            "--factory-path",
+            factory_path,
+            "search",
+            "example query",
+            "--columns",
+            "location,location.coordinates,location.page_number",
+        ],
+    )
+    assert result.exit_code == 0
+    print(result.stdout)
+    assert "Foo document" not in result.stdout
+    assert "Bar document" not in result.stdout
+    assert "Baz document" not in result.stdout
+    assert "Location" in result.stdout
+    assert "Location Coordinates" in result.stdout
+    assert "Location Page Number" in result.stdout
+
+
 def test_search_columns_non_existent():
     runner = CliRunner(mix_stderr=False)
     result = runner.invoke(
@@ -86,6 +109,23 @@ def test_search_columns_non_existent():
     assert "Unknown column: non_existent" in result.stderr
 
 
+def test_search_nested_columns_non_existent():
+    runner = CliRunner(mix_stderr=False)
+    result = runner.invoke(
+        ds_app,
+        [
+            "--factory-path",
+            factory_path,
+            "search",
+            "example query",
+            "--columns",
+            "document_meta,location,location.non_existent",
+        ],
+    )
+    assert result.exit_code == 1
+    assert "Unknown column: location.non_existent" in result.stderr
+
+
 def test_search_json():
     autoregister()
     runner = CliRunner(mix_stderr=False)
@@ -101,7 +141,6 @@ def test_search_json():
             "example query",
         ],
     )
-    print(result.stderr)
     assert result.exit_code == 0
     elements = json.loads(result.stdout)
     assert len(elements) == 3
