atlas-chat 0.1.0__py3-none-any.whl

atlas/application/chat/utilities/file_processor.py

@@ -0,0 +1,613 @@
+"""
+File management utilities - pure functions for session file operations.
+
+This module provides stateless utility functions for handling files within
+chat sessions, including user uploads and tool-generated artifacts.
+"""
+
+import logging
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from atlas.modules.file_storage.content_extractor import get_content_extractor
+
+logger = logging.getLogger(__name__)
+
+# Type hint for update callback
+UpdateCallback = Callable[[Dict[str, Any]], Awaitable[None]]
+
+
+async def handle_session_files(
+    session_context: Dict[str, Any],
+    user_email: Optional[str],
+    files_map: Optional[Dict[str, Any]],
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> Dict[str, Any]:
+    """
+    Handle user file ingestion and return updated session context.
+
+    Pure function that processes files and returns new context without mutations.
+
+    Args:
+        session_context: Current session context
+        user_email: User email for file storage
+        files_map: Map of filename to file data. Can be:
+            - str: base64 content (legacy format)
+            - dict: {"content": base64, "extract": bool} (new format with extraction flag)
+        file_manager: File manager instance
+        update_callback: Optional callback for emitting updates
+
+    Returns:
+        Updated session context with file references
+    """
+    if not files_map or not file_manager or not user_email:
+        return session_context
+
+    # Work with a copy to avoid mutations
+    updated_context = dict(session_context)
+    session_files_ctx = updated_context.setdefault("files", {})
+
+    # Get content extractor
+    extractor = get_content_extractor()
+    default_extract_mode = extractor.get_default_behavior() if extractor.is_enabled() else "none"
+
+    try:
+        uploaded_refs: Dict[str, Dict[str, Any]] = {}
+        for filename, file_data in files_map.items():
+            try:
+                # Handle both legacy (string) and new (dict) formats
+                if isinstance(file_data, str):
+                    b64 = file_data
+                    extract_mode = default_extract_mode
+                else:
+                    b64 = file_data.get("content", "")
+                    # New extractMode field takes priority, then legacy extract bool
+                    if "extractMode" in file_data:
+                        extract_mode = file_data["extractMode"]
+                    elif "extract" in file_data:
+                        extract_mode = "full" if file_data["extract"] else "none"
+                    else:
+                        extract_mode = default_extract_mode
+
+                meta = await file_manager.upload_file(
+                    user_email=user_email,
+                    filename=filename,
+                    content_base64=b64,
+                    source_type="user",
+                    tags={"source": "user"}
+                )
+
+                # Store minimal reference in session context
+                file_ref = {
+                    "key": meta.get("key"),
+                    "content_type": meta.get("content_type"),
+                    "size": meta.get("size"),
+                    "source": "user",
+                    "last_modified": meta.get("last_modified"),
+                }
+
+                # Store the extraction mode for build_files_manifest
+                file_ref["extract_mode"] = extract_mode
+
+                # Attempt content extraction if enabled and mode requests it
+                if extract_mode in ("full", "preview") and extractor.is_enabled():
+                    extraction_result = await extractor.extract_content(
+                        filename=filename,
+                        content_base64=b64,
+                        mime_type=meta.get("content_type"),
+                    )
+                    if extraction_result.success:
+                        file_ref["extracted_content"] = extraction_result.content
+                        file_ref["extracted_preview"] = extraction_result.preview
+                        if extraction_result.metadata:
+                            file_ref["extraction_metadata"] = extraction_result.metadata
+                        logger.info(f"Extracted content from {filename}: {len(extraction_result.preview or '')} chars preview")
+                    else:
+                        logger.debug(f"Content extraction skipped for {filename}: {extraction_result.error}")
+
+                session_files_ctx[filename] = file_ref
+                uploaded_refs[filename] = meta
+            except Exception as e:
+                logger.error(f"Failed uploading user file {filename}: {e}")
+
+        # Emit files update if successful uploads
+        if uploaded_refs and update_callback:
+            organized = file_manager.organize_files_metadata(uploaded_refs)
+            logger.info(
+                "Emitting files_update for user uploads: total=%d",
+                len(organized.get('files', [])),
+            )
+            logger.debug("files_update details (user uploads): names=%s", list(uploaded_refs.keys()))
+            await update_callback({
+                "type": "intermediate_update",
+                "update_type": "files_update",
+                "data": organized
+            })
+
+    except Exception as e:
+        logger.error(f"Error ingesting user files: {e}", exc_info=True)
+
+    return updated_context
+
+
+async def process_tool_artifacts(
+    session_context: Dict[str, Any],
+    tool_result,
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> Dict[str, Any]:
+    """
+    Process v2 MCP artifacts produced by a tool and return updated session context.
+
+    Pure function that handles tool files without side effects on input context.
+    """
+    # Check if there's an iframe display configuration (no artifacts needed)
+    has_iframe_display = (
+        tool_result.display_config and
+        isinstance(tool_result.display_config, dict) and
+        tool_result.display_config.get("type") == "iframe" and
+        tool_result.display_config.get("url")
+    )
+
+    # Early return only if no artifacts AND no iframe display, or no file_manager
+    if (not tool_result.artifacts and not has_iframe_display) or not file_manager:
+        return session_context
+
+    # Work with a copy to avoid mutations
+    updated_context = dict(session_context)
+
+    # Process v2 artifacts (only if we have artifacts)
+    if tool_result.artifacts:
+        user_email = session_context.get("user_email")
+        if not user_email:
+            return session_context
+
+        updated_context = await ingest_v2_artifacts(
+            session_context=updated_context,
+            tool_result=tool_result,
+            user_email=user_email,
+            file_manager=file_manager,
+            update_callback=update_callback
+        )
+
+    # Handle canvas file notifications with v2 display config
+    # This handles both artifact-based displays and iframe-only displays
+    await notify_canvas_files_v2(
+        session_context=updated_context,
+        tool_result=tool_result,
+        file_manager=file_manager,
+        update_callback=update_callback
+    )
+
+    return updated_context
+
+
+async def ingest_tool_files(
+    session_context: Dict[str, Any],
+    tool_result,
+    user_email: str,
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> Dict[str, Any]:
+    """
+    Persist tool-produced files into storage and update session context.
+
+    Pure function that returns updated context without mutations.
+    """
+    if not tool_result.returned_file_names:
+        return session_context
+
+    # Work with a copy
+    updated_context = dict(session_context)
+
+    # Safety: avoid huge ingestions
+    MAX_FILES = 10
+    names = tool_result.returned_file_names[:MAX_FILES]
+    contents = tool_result.returned_file_contents[:MAX_FILES] if tool_result.returned_file_contents else []
+
+    if contents and len(contents) != len(names):
+        logger.warning(
+            "ToolResult file arrays length mismatch (names=%d, contents=%d) for tool_call_id=%s",
+            len(names), len(contents), tool_result.tool_call_id
+        )
+
+    pair_count = min(len(names), len(contents)) if contents else 0
+    session_files_ctx = updated_context.setdefault("files", {})
+    uploaded_refs: Dict[str, Dict[str, Any]] = {}
+
+    for idx, fname in enumerate(names):
+        try:
+            if idx < pair_count:
+                b64 = contents[idx]
+                meta = await file_manager.upload_file(
+                    user_email=user_email,
+                    filename=fname,
+                    content_base64=b64,
+                    source_type="tool",
+                    tags={"source": "tool"}
+                )
+                session_files_ctx[fname] = {
+                    "key": meta.get("key"),
+                    "content_type": meta.get("content_type"),
+                    "size": meta.get("size"),
+                    "source": "tool",
+                    "last_modified": meta.get("last_modified"),
+                    "tool_call_id": tool_result.tool_call_id
+                }
+                uploaded_refs[fname] = meta
+            else:
+                # Name without content – record reference placeholder only if not existing
+                if fname not in session_files_ctx:
+                    session_files_ctx[fname] = {"source": "tool", "incomplete": True}
+        except Exception as e:
+            logger.error(f"Failed uploading tool-produced file {fname}: {e}")
+
+    # Emit files update if successful uploads
+    if uploaded_refs and update_callback:
+        try:
+            organized = file_manager.organize_files_metadata(uploaded_refs)
+            logger.info(
+                "Emitting files_update for tool uploads: total=%d",
+                len(organized.get('files', [])),
+            )
+            logger.debug("files_update details (tool uploads): names=%s", list(uploaded_refs.keys()))
+            await update_callback({
+                "type": "intermediate_update",
+                "update_type": "files_update",
+                "data": organized
+            })
+        except Exception as e:
+            logger.error(f"Failed emitting tool files update: {e}")
+
+    return updated_context
+
+
+async def notify_canvas_files(
+    session_context: Dict[str, Any],
+    file_names: List[str],
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> None:
+    """
+    Send canvas files notification for tool-produced files.
+
+    Pure function with no side effects on session context.
+    """
+    if not update_callback or not file_names or not file_manager:
+        return
+
+    try:
+        uploaded_refs = {}
+        files_ctx = session_context.get("files", {})
+
+        for fname in file_names:
+            ref = files_ctx.get(fname)
+            if ref and ref.get("key"):
+                uploaded_refs[fname] = {
+                    "key": ref.get("key"),
+                    "size": ref.get("size", 0),
+                    "content_type": ref.get("content_type", "application/octet-stream"),
+                    "last_modified": ref.get("last_modified"),
+                    "tags": {"source": ref.get("source", "tool")}
+                }
+
+        if uploaded_refs:
+            canvas_files = []
+            for fname, meta in uploaded_refs.items():
+                if file_manager.should_display_in_canvas(fname):
+                    file_ext = file_manager.get_file_extension(fname).lower()
+                    canvas_files.append({
+                        "filename": fname,
+                        "type": file_manager.get_canvas_file_type(file_ext),
+                        "s3_key": meta.get("key"),
+                        "size": meta.get("size", 0),
+                    })
+
+            if canvas_files:
+                await update_callback({
+                    "type": "intermediate_update",
+                    "update_type": "canvas_files",
+                    "data": {"files": canvas_files}
+                })
+    except Exception as emit_err:
+        logger.warning(f"Non-fatal: failed to emit canvas_files update: {emit_err}")
+
+
+async def emit_files_update_from_context(
+    session_context: Dict[str, Any],
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> None:
+    """
+    Emit a files_update event based on session context files.
+
+    Pure function with no side effects.
+    """
+    if not file_manager or not update_callback:
+        return
+
+    try:
+        # Build temp structure expected by organizer
+        file_refs: Dict[str, Dict[str, Any]] = {}
+        for fname, ref in session_context.get("files", {}).items():
+            # Expand to shape similar to S3 metadata for organizer
+            file_refs[fname] = {
+                "key": ref.get("key"),
+                "size": ref.get("size", 0),
+                "content_type": ref.get("content_type", "application/octet-stream"),
+                "last_modified": ref.get("last_modified"),
+                "tags": {"source": ref.get("source", "user")}
+            }
+
+        organized = file_manager.organize_files_metadata(file_refs)
+        logger.info(
+            "Emitting files_update from context: total=%d",
+            len(organized.get('files', [])),
+        )
+        await update_callback({
+            "type": "intermediate_update",
+            "update_type": "files_update",
+            "data": organized
+        })
+    except Exception as e:
+        logger.error(f"Failed emitting files update: {e}")
+
+
+async def ingest_v2_artifacts(
+    session_context: Dict[str, Any],
+    tool_result,
+    user_email: str,
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> Dict[str, Any]:
+    """
+    Persist v2 MCP artifacts into storage and update session context.
+
+    Pure function that returns updated context without mutations.
+    """
+    if not tool_result.artifacts:
+        return session_context
+
+    # Work with a copy
+    updated_context = dict(session_context)
+
+    # Safety: avoid huge ingestions
+    MAX_ARTIFACTS = 10
+    artifacts = tool_result.artifacts[:MAX_ARTIFACTS]
+
+    try:
+        # Prepare files for upload
+        files_to_upload = []
+        for artifact in artifacts:
+            name = artifact.get("name")
+            b64_content = artifact.get("b64")
+            mime_type = artifact.get("mime")
+
+            if not name or not b64_content:
+                logger.warning("Skipping artifact with missing name or content")
+                continue
+
+            files_to_upload.append({
+                "filename": name,
+                "content": b64_content,
+                "mime_type": mime_type
+            })
+
+        if not files_to_upload:
+            return updated_context
+
+        # Upload files to storage
+        uploaded_refs = await file_manager.upload_files_from_base64(
+            files_to_upload, user_email
+        )
+
+        # Add file references to session context
+        current_files = updated_context.setdefault("files", {})
+        current_files.update(uploaded_refs)
+
+        # Emit files update if successful uploads
+        if uploaded_refs and update_callback:
+            organized = file_manager.organize_files_metadata(uploaded_refs)
+            logger.info(
+                "Emitting files_update for v2 artifacts: total=%d",
+                len(organized.get('files', [])),
+            )
+            logger.debug(
+                "files_update details (v2 artifacts): names=%s",
+                list(uploaded_refs.keys()),
+            )
+            await update_callback({
+                "type": "intermediate_update",
+                "update_type": "files_update",
+                "data": organized
+            })
+
+    except Exception as e:
+        logger.error(f"Error ingesting v2 artifacts: {e}", exc_info=True)
+
+    return updated_context
+
+
+async def notify_canvas_files_v2(
+    session_context: Dict[str, Any],
+    tool_result,
+    file_manager,
+    update_callback: Optional[UpdateCallback] = None
+) -> None:
+    """
+    Send v2 canvas files notification with display configuration.
+
+    Pure function with no side effects on session context.
+    """
+    if not update_callback:
+        return
+
+    # Check if there's an iframe display configuration (no artifacts needed)
+    has_iframe_display = (
+        tool_result.display_config and
+        isinstance(tool_result.display_config, dict) and
+        tool_result.display_config.get("type") == "iframe" and
+        tool_result.display_config.get("url")
+    )
+
+    # If no artifacts and no iframe display, nothing to show
+    if not tool_result.artifacts and not has_iframe_display:
+        return
+
+    try:
+        # Get uploaded file references from session context
+        uploaded_refs = session_context.get("files", {})
+        artifact_names = [artifact.get("name") for artifact in tool_result.artifacts if artifact.get("name")]
+
+        # Handle iframe-only display (no artifacts)
+        if has_iframe_display and not artifact_names:
+            canvas_update = {
+                "type": "intermediate_update",
+                "update_type": "canvas_files",
+                "data": {
+                    "files": [],
+                    "display": tool_result.display_config
+                }
+            }
+            logger.info("Emitting canvas_files event for iframe display")
+            logger.debug(
+                "canvas_files iframe display details: url=%s, title=%s",
+                tool_result.display_config.get("url"),
+                tool_result.display_config.get("title", "Embedded Content"),
+            )
+            await update_callback(canvas_update)
+            return
+
+        if uploaded_refs and artifact_names:
+            canvas_files = []
+            for fname in artifact_names:
+                meta = uploaded_refs.get(fname)
+                if meta and file_manager.should_display_in_canvas(fname):
+                    # Get MIME type from artifact if available
+                    artifact = next((a for a in tool_result.artifacts if a.get("name") == fname), {})
+                    mime_type = artifact.get("mime")
+
+                    file_ext = file_manager.get_file_extension(fname).lower()
+                    canvas_files.append({
+                        "filename": fname,
+                        "type": file_manager.get_canvas_file_type(file_ext),
+                        "s3_key": meta.get("key"),
+                        "size": meta.get("size", 0),
+                        "mime_type": mime_type
+                    })
+
+            if canvas_files:
+                # Reorder files to put primary_file first if provided
+                primary = None
+                if tool_result.display_config and isinstance(tool_result.display_config, dict):
+                    primary = tool_result.display_config.get("primary_file")
+                if primary:
+                    # stable reorder
+                    canvas_files = sorted(
+                        canvas_files,
+                        key=lambda f: 0 if f.get("filename") == primary else 1
+                    )
+
+                # Build canvas update with v2 display configuration
+                logger.info("Emitting canvas_files event: count=%d", len(canvas_files))
+                logger.debug(
+                    "canvas_files details: files=%s, display=%s",
+                    [f.get("filename") for f in canvas_files],
+                    tool_result.display_config,
+                )
+                canvas_update = {
+                    "type": "intermediate_update",
+                    "update_type": "canvas_files",
+                    "data": {"files": canvas_files}
+                }
+
+                # Add v2 display configuration if present
+                if tool_result.display_config:
+                    canvas_update["data"]["display"] = tool_result.display_config
+
+                await update_callback(canvas_update)
+            else:
+                logger.debug("No canvas-displayable artifacts found. artifact_names=%s", artifact_names)
+
+    except Exception as emit_err:
+        logger.warning(f"Non-fatal: failed to emit v2 canvas_files update: {emit_err}")
+
+
+def build_files_manifest(session_context: Dict[str, Any]) -> Optional[Dict[str, str]]:
+    """
+    Build ephemeral files manifest for LLM context.
+
+    Pure function that creates manifest from session context.
+    Includes extracted content previews when available.
+    """
+    files_ctx = session_context.get("files", {})
+    if not files_ctx:
+        return None
+
+    # Build file list with extracted content based on extract_mode
+    file_entries = []
+    has_full = False
+    has_preview = False
+    has_none = False
+    for name in sorted(files_ctx.keys()):
+        file_info = files_ctx[name]
+        entry = f"- {name}"
+        mode = file_info.get("extract_mode", "preview")
+
+        # Include extraction metadata if available
+        if file_info.get("extraction_metadata"):
+            meta = file_info["extraction_metadata"]
+            if meta.get("pages"):
+                entry += f" ({meta['pages']} pages)"
+
+        if mode == "full" and file_info.get("extracted_content"):
+            has_full = True
+            content = file_info["extracted_content"]
+            entry += (
+                f"\n    << content of file {name} >>\n"
+                f"    {content}\n"
+                f"    << end content of file {name} >>"
+            )
+        elif mode == "preview" and file_info.get("extracted_preview"):
+            has_preview = True
+            preview = file_info["extracted_preview"]
+            # Limit to 10 lines and 2000 characters to prevent excessive token usage
+            lines = preview.split("\n")[:10]
+            indented_preview = "\n    ".join(lines)
+            if len(indented_preview) > 2000:
+                indented_preview = indented_preview[:1997] + "..."
+            entry += f"\n    Content preview:\n    {indented_preview}"
+        else:
+            has_none = True
+
+        file_entries.append(entry)
+
+    file_list = "\n".join(file_entries)
+
+    # Build context note based on which modes were used
+    notes = []
+    if has_full:
+        notes.append(
+            "Files with full content shown above have been fully extracted. "
+            "You can reference this content directly."
+        )
+    if has_preview:
+        notes.append(
+            "Files with content previews shown above have been partially analyzed. "
+            "You can reference preview content directly."
+        )
+    if has_none:
+        notes.append(
+            "Files listed by name only can be opened or analyzed on request."
+        )
+    context_note = f"({' '.join(notes)})" if notes else ""
+
+    return {
+        "role": "system",
+        "content": (
+            "Available session files:\n"
+            f"{file_list}\n\n"
+            f"{context_note} "
+            "The user may refer to these files in their requests as session files or attachments."
+        )
+    }