File size: 5,329 Bytes
399944f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
from typing import Any, Dict, Optional, Sequence
import rich

from kbdebugger.compat.langchain import Document
from kbdebugger.utils.json import write_json
from kbdebugger.utils.time import now_utc_compact, now_utc_human
from .types import Qualities, SourceKind, DecomposeMode

def build_chunked_documents_payload(
    *,
    docs: Sequence[Document],
    created_at: str | None = None,
) -> dict[str, Any]:
    """
    Build a clean, JSON-serializable payload for chunked paragraph Documents.

    Design goals
    ------------
    - Avoid logging massive Docling internals (dl_meta.doc_items, bbox, prov, etc.)
    - Keep only what is useful for debugging + UI traceability:
        * top-level `source` once (if available)
        * per-doc `page_content`
        * per-doc `headings` (if available): metadata["dl_meta"]["headings"]

    Parameters
    ----------
    docs:
        Chunked LangChain Documents.
    created_at:
        Optional timestamp string; defaults to now_utc_human().

    Returns
    -------
    dict[str, Any]
        Clean payload ready for write_json(...) or API responses.
    """
    created_at = created_at or now_utc_human()

    # Source is repeated in every doc metadata; log it once if possible.
    source: str | None = None
    if docs:
        md0 = dict(getattr(docs[0], "metadata", {}) or {})
        source = md0.get("source")

    cleaned_docs: list[dict[str, Any]] = []
    for doc in docs:
        md = dict(getattr(doc, "metadata", {}) or {})
        dl_meta = md.get("dl_meta") or {}

        headings = None
        if isinstance(dl_meta, dict):
            headings_val = dl_meta.get("headings")
            if isinstance(headings_val, list) and headings_val:
                headings = headings_val

        cleaned_docs.append(
            {
                "page_content": getattr(doc, "page_content", "") or "",
                # Only include headings if we actually have them
                **({"headings": headings} if headings is not None else {}),
            }
        )

    payload: dict[str, Any] = {
        "source": source,
        "num_docs": len(docs),
        "docs": cleaned_docs,
        "created_at": created_at,
    }

    return payload


def save_chunked_documents_json(
    *,
    docs: list[Document],
    source_kind: SourceKind,
) -> dict[str, Any]:
    """
    Save chunked Documents to a clean JSON log for debugging/demo purposes.

    What we store
    -------------
    - source (once, top-level)
    - num_docs
    - docs: [{ page_content, headings? }]
    - created_at (human-friendly)

    Returns
    -------
    dict[str, Any]
        The payload (useful if caller wants to reuse it without rebuilding).
    """
    payload = build_chunked_documents_payload(docs=docs)

    # Keep compact time ONLY for filenames.
    created_at_compact = now_utc_compact()
    path = f"logs/01.1_chunker_output_docs_[{source_kind}]_{created_at_compact}.json"
    write_json(path, payload)

    rich.print(f"\n[INFO] Wrote chunker output log to {path}")
    return payload


def build_decomposer_payload(
    *,
    qualities: Qualities,
    mode: DecomposeMode,
    num_input_docs: int,
    use_batch_decomposer: bool,
    batch_size: Optional[int],
    num_batches: Optional[int],
    parallel: bool,
    max_workers: Optional[int],
    created_at: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Build a clean, JSON-serializable payload for the Decomposer stage.

    Notes
    -----
    - `created_at` is human-readable (for payloads/UI).
    - Filename timestamp stays compact (handled by logger).
    """
    created_at = created_at or now_utc_human()

    payload: Dict[str, Any] = {
        "created_at": created_at,
        # "mode": str(mode),  # or mode.value if it's an Enum
        "mode": mode.value if hasattr(mode, "value") else str(mode),
        "num_input_docs": num_input_docs,
        "num_output_qualities": len(qualities),
        "use_batch_decomposer": use_batch_decomposer,
        "batch_size": batch_size,
        "num_batches": num_batches,
        "parallel": parallel,
        "max_workers": max_workers if parallel else None,
        "qualities": list(qualities),
    }

    # Keep payload clean: drop None fields
    return {k: v for k, v in payload.items() if v is not None}


def save_qualities_json(
    *,
    qualities: Qualities,
    mode: DecomposeMode,
    num_input_docs: int,
    use_batch_decomposer: bool,
    batch_size: Optional[int] = None,
    num_batches: Optional[int] = None,
    parallel: bool = False,
    max_workers: Optional[int] = None,
    output_dir: str = "logs",
) -> Dict[str, Any]:
    """
    Save Decomposer output qualities to JSON and return the written payload.

    - Payload uses human-readable created_at.
    - Filename uses compact timestamp.
    """
    payload = build_decomposer_payload(
        qualities=qualities,
        mode=mode,
        num_input_docs=num_input_docs,
        use_batch_decomposer=use_batch_decomposer,
        batch_size=batch_size,
        num_batches=num_batches,
        parallel=parallel,
        max_workers=max_workers,
    )

    ts = now_utc_compact()
    path = f"{output_dir}/01.2_decomposer_qualities_{mode}_{ts}.json"
    write_json(path, payload)

    rich.print(f"\n[INFO] Wrote decomposer qualities log to {path}")
    return payload