ai-pipeline-core 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/{debug/content.py → observability/_debug/_content.py}

@@ -1,6 +1,7 @@
 """Content writing and extraction for trace debugging V3.
 
 Uses hash-based artifact storage with automatic deduplication.
+Handles Document attachments by externalizing large binary/text attachments to the artifact store.
 """
 
 import base64
@@ -10,13 +11,13 @@ import re
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 from uuid import UUID
 
 import yaml
 from pydantic import BaseModel, ConfigDict, SecretStr
 
-from .config import TraceDebugConfig
+from ._config import TraceDebugConfig
 
 
 class ContentRef(BaseModel):
@@ -54,9 +55,7 @@ class ArtifactStore:
             return self._known_hashes[content_hash]
 
         # Create sharded path: ab/cd/abcdef...1234.txt
-        file_path = (
-            self._artifacts_path / content_hash[:2] / content_hash[2:4] / f"{content_hash}.txt"
-        )
+        file_path = self._artifacts_path / content_hash[:2] / content_hash[2:4] / f"{content_hash}.txt"
         file_path.parent.mkdir(parents=True, exist_ok=True)
 
         if not file_path.exists():
@@ -85,13 +84,12 @@ class ArtifactStore:
             "image/png": ".png",
             "image/jpeg": ".jpg",
             "image/gif": ".gif",
+            "image/webp": ".webp",
             "application/pdf": ".pdf",
         }
         ext = ext_map.get(mime_type, ".bin")
 
-        file_path = (
-            self._artifacts_path / content_hash[:2] / content_hash[2:4] / f"{content_hash}{ext}"
-        )
+        file_path = self._artifacts_path / content_hash[:2] / content_hash[2:4] / f"{content_hash}{ext}"
         file_path.parent.mkdir(parents=True, exist_ok=True)
 
         if not file_path.exists():
@@ -162,9 +160,7 @@ class ContentWriter:
         if size > self._config.max_file_bytes:
             # Reduce preview sizes to fit under limit
             structured = self._reduce_previews(structured)
-            serialized = yaml.dump(
-                structured, default_flow_style=False, allow_unicode=True, sort_keys=False
-            )
+            serialized = yaml.dump(structured, default_flow_style=False, allow_unicode=True, sort_keys=False)
             serialized = self._redact(serialized)
             size = len(serialized.encode("utf-8"))
 
@@ -172,9 +168,7 @@ class ContentWriter:
             if size > self._config.max_file_bytes:
                 serialized = serialized[: self._config.max_file_bytes]
                 max_bytes = self._config.max_file_bytes
-                serialized += (
-                    f"\n\n# [TRUNCATED: original {size} bytes exceeded {max_bytes} limit]\n"
-                )
+                serialized += f"\n\n# [TRUNCATED: original {size} bytes exceeded {max_bytes} limit]\n"
                 size = len(serialized.encode("utf-8"))
 
         # Write file
@@ -192,10 +186,9 @@ class ContentWriter:
         """Convert raw content to structured YAML-ready format."""
         if self._is_llm_messages(content):
             return self._structure_llm_messages(content)
-        elif self._is_document_list(content):
+        if self._is_document_list(content):
             return self._structure_documents(content)
-        else:
-            return self._structure_generic(content)
+        return self._structure_generic(content)
 
     def _is_llm_messages(self, content: Any) -> bool:
         """Check if content looks like LLM messages."""
@@ -203,21 +196,21 @@ class ContentWriter:
             return False
         if not content:
             return False
-        first = content[0]
+        first = cast(Any, content[0])
         if not isinstance(first, dict):
             return False
         return "role" in first and "content" in first
 
     def _is_document_list(self, content: Any) -> bool:
-        """Check if content looks like a DocumentList."""
+        """Check if content looks like a list of serialized documents."""
         if not isinstance(content, list):
             return False
         if not content:
             return False
-        first = content[0]
+        first = cast(Any, content[0])
         if not isinstance(first, dict):
             return False
-        return "base_type" in first and "content" in first
+        return "class_name" in first and "content" in first
 
     def _structure_llm_messages(self, messages: list[Any]) -> dict[str, Any]:
         """Structure LLM messages preserving ALL parts losslessly."""
@@ -238,16 +231,18 @@ class ContentWriter:
 
             if isinstance(content, list):
                 # Multimodal: preserve each part separately
-                msg_entry["parts"] = []
-                for j, part in enumerate(content):
+                content_parts = cast(list[Any], content)
+                msg_parts: list[dict[str, Any]] = []
+                msg_entry["parts"] = msg_parts
+                for j, part in enumerate(content_parts):
                     structured_part, part_bytes = self._structure_message_part(part, j)
-                    msg_entry["parts"].append(structured_part)
+                    msg_parts.append(structured_part)
                     part_type = structured_part.get("type", "")
                     if part_type == "text":
                         total_text_bytes += part_bytes
                     elif part_type == "image":
                         total_image_bytes += part_bytes
-                    elif part_type in ("tool_use", "tool_result"):
+                    elif part_type in {"tool_use", "tool_result"}:
                         total_tool_bytes += part_bytes
             elif isinstance(content, str):
                 # Simple text message
@@ -284,9 +279,7 @@ class ContentWriter:
             "size_bytes": total_text_bytes + total_image_bytes + total_tool_bytes,
         }
 
-    def _structure_message_part(
-        self, part: dict[str, Any], sequence: int
-    ) -> tuple[dict[str, Any], int]:
+    def _structure_message_part(self, part: dict[str, Any], sequence: int) -> tuple[dict[str, Any], int]:
         """Structure a single message part losslessly.
 
         Returns:
@@ -297,13 +290,13 @@ class ContentWriter:
         if part_type == "text":
             entry = self._structure_text_element(part.get("text", ""), sequence)
             return entry, entry.get("size_bytes", 0)
-        elif part_type == "image_url":
+        if part_type == "image_url":
             entry = self._structure_image_openai(part, sequence)
             return entry, entry.get("size_bytes", 0)
-        elif part_type == "image":
+        if part_type == "image":
             entry = self._structure_image_anthropic(part, sequence)
             return entry, entry.get("size_bytes", 0)
-        elif part_type == "tool_use":
+        if part_type == "tool_use":
             input_str = json.dumps(part.get("input", {}))
             size = len(input_str.encode("utf-8"))
             return {
@@ -313,7 +306,7 @@ class ContentWriter:
                 "name": part.get("name"),
                 "input": self._convert_types(part.get("input")),
             }, size
-        elif part_type == "tool_result":
+        if part_type == "tool_result":
             result_content = part.get("content")
             entry: dict[str, Any] = {
                 "type": "tool_result",
@@ -327,25 +320,26 @@ class ContentWriter:
                 entry["content"] = text_entry
                 size = text_entry.get("size_bytes", 0)
             elif isinstance(result_content, list):
-                entry["content"] = []
-                for k, p in enumerate(result_content):
+                result_parts = cast(list[Any], result_content)
+                content_list: list[dict[str, Any]] = []
+                entry["content"] = content_list
+                for k, p in enumerate(result_parts):
                     part_entry, part_size = self._structure_message_part(p, k)
-                    entry["content"].append(part_entry)
+                    content_list.append(part_entry)
                     size += part_size
             else:
                 entry["content"] = self._convert_types(result_content)
             return entry, size
-        else:
-            # Unknown type — preserve raw data, never drop
-            raw = self._convert_types(part)
-            raw_str = json.dumps(raw)
-            size = len(raw_str.encode("utf-8"))
-            return {
-                "type": "unknown",
-                "sequence": sequence,
-                "original_type": part_type,
-                "raw_data": raw,
-            }, size
+        # Unknown type — preserve raw data, never drop
+        raw = self._convert_types(part)
+        raw_str = json.dumps(raw)
+        size = len(raw_str.encode("utf-8"))
+        return {
+            "type": "unknown",
+            "sequence": sequence,
+            "original_type": part_type,
+            "raw_data": raw,
+        }, size
 
     def _structure_text_element(self, text: str, sequence: int) -> dict[str, Any]:
         """Structure a text element, optionally externalizing large content."""
@@ -369,9 +363,7 @@ class ContentWriter:
                     "mime_type": ref.mime_type,
                     "encoding": ref.encoding,
                 }
-                entry["excerpt"] = (
-                    text[:excerpt_len] + "\n[TRUNCATED - see artifact for full content]"
-                )
+                entry["excerpt"] = text[:excerpt_len] + "\n[TRUNCATED - see artifact for full content]"
             else:
                 # No artifact store — truncate with marker
                 entry["content"] = text[: self._config.max_element_bytes]
@@ -486,20 +478,20 @@ class ContentWriter:
 
         return entry
 
-    def _structure_documents(self, docs: list[Any]) -> dict[str, Any]:
-        """Structure document list."""
+    def _structure_documents(self, docs: list[Any]) -> dict[str, Any]:  # noqa: PLR0914
+        """Structure document list with attachment externalization."""
         doc_entries: list[dict[str, Any]] = []
 
         for i, doc in enumerate(docs):
             doc_name = doc.get("name", f"doc_{i}")
-            base_type = doc.get("base_type", "unknown")
+            class_name = doc.get("class_name", "Document")
             content = doc.get("content", "")
             content_encoding = doc.get("content_encoding", "utf-8")
 
             doc_entry: dict[str, Any] = {
                 "index": i,
                 "name": doc_name,
-                "base_type": base_type,
+                "class_name": class_name,
             }
 
             if content_encoding == "base64":
@@ -541,12 +533,75 @@ class ContentWriter:
                         "mime_type": ref.mime_type,
                         "encoding": ref.encoding,
                     }
-                    doc_entry["excerpt"] = (
-                        text[:excerpt_len] + "\n[TRUNCATED - see artifact for full content]"
-                    )
+                    doc_entry["excerpt"] = text[:excerpt_len] + "\n[TRUNCATED - see artifact for full content]"
                 else:
                     doc_entry["content"] = text
 
+            # Structure attachments if present
+            raw_attachments = doc.get("attachments")
+            if isinstance(raw_attachments, list) and raw_attachments:
+                att_entries: list[dict[str, Any]] = []
+                attachments_list = cast(list[Any], raw_attachments)
+                for j, att in enumerate(attachments_list):
+                    if not isinstance(att, dict):
+                        continue
+                    att_dict = cast(dict[str, Any], att)
+                    att_name = att_dict.get("name", f"attachment_{j}")
+                    att_encoding = att_dict.get("content_encoding", "utf-8")
+                    att_content = att_dict.get("content", "")
+
+                    att_entry: dict[str, Any] = {
+                        "index": j,
+                        "name": att_name,
+                    }
+                    if att_dict.get("description"):
+                        att_entry["description"] = att_dict["description"]
+
+                    if att_encoding == "base64":
+                        try:
+                            binary_data = base64.b64decode(att_content)
+                            size = len(binary_data)
+                            att_entry["size_bytes"] = size
+                            att_entry["encoding"] = "base64"
+                            mime_type = att_dict.get("mime_type", "application/octet-stream")
+
+                            if size > self._config.max_element_bytes and self._artifact_store:
+                                ref = self._artifact_store.store_binary(binary_data, mime_type)
+                                att_entry["content_ref"] = {
+                                    "hash": ref.hash,
+                                    "path": ref.path,
+                                    "mime_type": ref.mime_type,
+                                    "encoding": ref.encoding,
+                                }
+                                att_entry["preview"] = f"[Binary attachment, {size} bytes]"
+                            else:
+                                att_entry["content"] = att_content
+                        except Exception:
+                            att_entry["content"] = "[binary content - decode failed]"
+                            att_entry["size_bytes"] = 0
+                    else:
+                        text = self._redact(str(att_content))
+                        text_bytes = len(text.encode("utf-8"))
+                        att_entry["size_bytes"] = text_bytes
+
+                        if text_bytes > self._config.max_element_bytes and self._artifact_store:
+                            ref = self._artifact_store.store_text(text)
+                            excerpt_len = self._config.element_excerpt_bytes
+                            att_entry["content_ref"] = {
+                                "hash": ref.hash,
+                                "path": ref.path,
+                                "mime_type": ref.mime_type,
+                                "encoding": ref.encoding,
+                            }
+                            att_entry["excerpt"] = text[:excerpt_len] + "\n[TRUNCATED - see artifact for full content]"
+                        else:
+                            att_entry["content"] = text
+
+                    att_entries.append(att_entry)
+
+                doc_entry["attachment_count"] = len(att_entries)
+                doc_entry["attachments"] = att_entries
+
             doc_entries.append(doc_entry)
 
         return {
@@ -583,11 +638,10 @@ class ContentWriter:
                 "image_bytes": metadata.get("total_image_bytes", 0),
                 "tool_bytes": metadata.get("total_tool_bytes", 0),
             }
-        elif "size_bytes" in structured:
+        if "size_bytes" in structured:
             return {"total_bytes": structured["size_bytes"]}
-        else:
-            serialized = json.dumps(self._convert_types(structured))
-            return {"total_bytes": len(serialized.encode("utf-8"))}
+        serialized = json.dumps(self._convert_types(structured))
+        return {"total_bytes": len(serialized.encode("utf-8"))}
 
     def _reduce_previews(self, structured: dict[str, Any]) -> dict[str, Any]:
         """Reduce preview/excerpt sizes to fit file under max_file_bytes."""
@@ -606,7 +660,7 @@ class ContentWriter:
             text = pattern.sub("[REDACTED]", text)
         return text
 
-    def _convert_types(self, value: Any, seen: set[int] | None = None) -> Any:
+    def _convert_types(self, value: Any, seen: set[int] | None = None) -> Any:  # noqa: PLR0911
         """Convert non-serializable types recursively with cycle detection."""
         # Cycle detection
         if seen is None:
@@ -634,7 +688,7 @@ class ContentWriter:
             case Enum():
                 return value.value
             case set() | frozenset():
-                return sorted(str(x) for x in value)
+                return sorted(str(x) for x in cast(set[Any] | frozenset[Any], value))
             case BaseModel():
                 try:
                     return value.model_dump(mode="json")
@@ -642,12 +696,14 @@ class ContentWriter:
                     return str(value)
             case dict():
                 seen.add(obj_id)
-                result = {str(k): self._convert_types(v, seen) for k, v in value.items()}
+                typed_dict = cast(dict[Any, Any], value)
+                result = {str(k): self._convert_types(v, seen) for k, v in typed_dict.items()}
                 seen.discard(obj_id)
                 return result
             case list() | tuple():
                 seen.add(obj_id)
-                result = [self._convert_types(x, seen) for x in value]
+                typed_seq = cast(list[Any] | tuple[Any, ...], value)
+                result = [self._convert_types(x, seen) for x in typed_seq]
                 seen.discard(obj_id)
                 return result
             case _:
@@ -680,26 +736,29 @@ def reconstruct_span_content(trace_root: Path, span_dir: Path, content_type: str
 def _rehydrate(obj: Any, trace_root: Path) -> Any:
     """Recursively replace content_ref entries with actual content."""
     if isinstance(obj, dict):
-        if "content_ref" in obj:
+        obj_dict = cast(dict[str, Any], obj)
+        if "content_ref" in obj_dict:
             # This is an artifact reference - load the full content
-            ref = obj["content_ref"]
-            artifact_path = trace_root / ref["path"]
+            ref: dict[str, Any] = obj_dict["content_ref"]
+            artifact_path: Path = trace_root / ref["path"]
 
+            full_content: str | bytes
             if ref.get("encoding") == "utf-8":
                 full_content = artifact_path.read_text(encoding="utf-8")
             else:
                 full_content = artifact_path.read_bytes()
 
             # Replace ref with full content
-            obj = obj.copy()
-            obj["content"] = full_content
-            del obj["content_ref"]
-            if "excerpt" in obj:
-                del obj["excerpt"]
+            obj_dict = obj_dict.copy()
+            obj_dict["content"] = full_content
+            del obj_dict["content_ref"]
+            if "excerpt" in obj_dict:
+                del obj_dict["excerpt"]
 
-        return {k: _rehydrate(v, trace_root) for k, v in obj.items()}
+        return {k: _rehydrate(v, trace_root) for k, v in obj_dict.items()}
 
-    elif isinstance(obj, list):
-        return [_rehydrate(v, trace_root) for v in obj]
+    if isinstance(obj, list):
+        obj_list = cast(list[Any], obj)
+        return [_rehydrate(v, trace_root) for v in obj_list]
 
     return obj

ai_pipeline_core/{debug/processor.py → observability/_debug/_processor.py}

@@ -1,10 +1,12 @@
 """OpenTelemetry SpanProcessor for local trace debugging."""
 
+import contextlib
+
 from opentelemetry.context import Context
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.trace import StatusCode
 
-from .writer import LocalTraceWriter, WriteJob
+from ._writer import LocalTraceWriter, WriteJob
 
 
 class LocalDebugSpanProcessor(SpanProcessor):
@@ -29,7 +31,7 @@ class LocalDebugSpanProcessor(SpanProcessor):
         Creates the span directory early so we can see "running" spans.
         Input/output data is not available yet - will be captured in on_end().
         """
-        try:
+        with contextlib.suppress(Exception):
             if span.context is None:
                 return
             trace_id = format(span.context.trace_id, "032x")
@@ -37,9 +39,6 @@ class LocalDebugSpanProcessor(SpanProcessor):
             parent_id = self._get_parent_span_id(span)
 
             self._writer.on_span_start(trace_id, span_id, parent_id, span.name)
-        except Exception:
-            # Never fail the actual span - debug tracing should be transparent
-            pass
 
     def on_end(self, span: ReadableSpan) -> None:
         """Handle span end - queue full span data for background write.
@@ -47,7 +46,7 @@ class LocalDebugSpanProcessor(SpanProcessor):
         All data (input, output, attributes, events) is captured here because
         Laminar sets these attributes after span start.
         """
-        try:
+        with contextlib.suppress(Exception):
             if span.context is None or span.start_time is None or span.end_time is None:
                 return
             job = WriteJob(
@@ -63,19 +62,18 @@ class LocalDebugSpanProcessor(SpanProcessor):
                 end_time_ns=span.end_time,
             )
             self._writer.on_span_end(job)
-        except Exception:
-            # Never fail the actual span
-            pass
 
     def shutdown(self) -> None:
         """Shutdown the processor and writer."""
         self._writer.shutdown()
 
-    def force_flush(self, timeout_millis: int = 30000) -> bool:
+    def force_flush(self, timeout_millis: int = 30000) -> bool:  # noqa: PLR6301
         """Force flush is not needed for this processor."""
+        _ = timeout_millis
         return True
 
-    def _get_parent_span_id(self, span: Span) -> str | None:
+    @staticmethod
+    def _get_parent_span_id(span: Span) -> str | None:
         """Extract parent span ID from a writable Span."""
         if hasattr(span, "parent") and span.parent:
             parent_ctx = span.parent
@@ -83,17 +81,18 @@ class LocalDebugSpanProcessor(SpanProcessor):
                 return format(parent_ctx.span_id, "016x")
         return None
 
-    def _get_parent_span_id_from_readable(self, span: ReadableSpan) -> str | None:
+    @staticmethod
+    def _get_parent_span_id_from_readable(span: ReadableSpan) -> str | None:
         """Extract parent span ID from a ReadableSpan."""
-        if span.parent:
-            if hasattr(span.parent, "span_id") and span.parent.span_id:
-                return format(span.parent.span_id, "016x")
+        if span.parent and hasattr(span.parent, "span_id") and span.parent.span_id:
+            return format(span.parent.span_id, "016x")
         return None
 
-    def _get_status_code(self, span: ReadableSpan) -> str:
+    @staticmethod
+    def _get_status_code(span: ReadableSpan) -> str:
         """Get status code as string."""
         if span.status.status_code == StatusCode.OK:
             return "OK"
-        elif span.status.status_code == StatusCode.ERROR:
+        if span.status.status_code == StatusCode.ERROR:
             return "ERROR"
         return "UNSET"