cadSet/CADDesigner-Code/web_interface/utils.py

"""
Utility function module.
"""

import time
import uuid
import base64
from pathlib import Path
import re
from typing import Tuple, Optional, Any, Union, Dict
from fastapi import HTTPException
from fastapi.responses import JSONResponse
from SimpleLLMFunc.type import ImgPath, ImgUrl, Text

from .models import (
    ChatCompletionRequest,
    Usage,
    ChatCompletionResponse,
    ChatMessage,
    ChatChoice,
)
from context.conversation_manager import ConversationManager, Conversation
from SimpleLLMFunc.logger import (
    app_log,
    push_warning,
    push_error,
    get_current_context_attribute,
)
from react_stream import extract_output_text, is_response_yield
from observability import propagate_conversation_session


DATA_URL_PATTERN = re.compile(r"^data:(image/[a-zA-Z0-9.+-]+);base64,(.*)$", re.DOTALL)
IMAGE_MIME_EXTENSIONS = {
    "image/jpeg": ".jpg",
    "image/png": ".png",
    "image/gif": ".gif",
    "image/bmp": ".bmp",
    "image/webp": ".webp",
}


def _content_item_type(item: Any) -> Optional[str]:
    if isinstance(item, dict):
        value = item.get("type")
        return value if isinstance(value, str) else None
    value = getattr(item, "type", None)
    return value if isinstance(value, str) else None


def _content_item_text(item: Any) -> Optional[str]:
    if isinstance(item, dict):
        value = item.get("text")
        return value if isinstance(value, str) else None
    value = getattr(item, "text", None)
    return value if isinstance(value, str) else None


def _content_item_image_url_payload(item: Any) -> Any:
    if isinstance(item, dict):
        return item.get("image_url")
    return getattr(item, "image_url", None)


def _image_payload_url(image_payload: Any) -> Optional[str]:
    if isinstance(image_payload, dict):
        value = image_payload.get("url")
        return value if isinstance(value, str) else None
    value = getattr(image_payload, "url", None)
    return value if isinstance(value, str) else None


def _image_payload_detail(image_payload: Any) -> str:
    if isinstance(image_payload, dict):
        value = image_payload.get("detail")
        return value if isinstance(value, str) else "auto"
    value = getattr(image_payload, "detail", None)
    return value if isinstance(value, str) else "auto"


def _image_payload_local_path(image_payload: Any) -> Optional[str]:
    if isinstance(image_payload, dict):
        value = image_payload.get("local_path")
        return value if isinstance(value, str) else None
    value = getattr(image_payload, "local_path", None)
    return value if isinstance(value, str) else None


def _set_image_payload_local_path(image_payload: Any, local_path: str) -> None:
    if isinstance(image_payload, dict):
        image_payload["local_path"] = local_path
        return
    if hasattr(image_payload, "local_path"):
        image_payload.local_path = local_path


def get_agent_for_model(model_name: str, agent_registry) -> Any:
    """
    Get an Agent instance by model name.

    Args:
        model_name: Model name
        agent_registry: Agent registry instance

    Returns:
        Agent instance

    Raises:
        HTTPException: If the model does not exist or creation fails
    """
    if not agent_registry:
        raise HTTPException(status_code=500, detail="Agent registry not initialized")

    # First try to get an existing Agent instance.
    agent = agent_registry.get_agent(model_name)
    if agent:
        return agent

    # If it does not exist, try to create a new Agent instance.
    try:
        agent = agent_registry.get_or_create_agent(
            model_name,
            name=f"Agent for {model_name}",
            description=f"Agent instance for model {model_name}",
        )
        return agent
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Error: {e} was caught. Maybe due to: [Unknown model: {model_name}]",
        )


# Removed create_error_response because it has been moved to error_handlers.py.


def _has_nonempty_user_content(content: Any) -> bool:
    if isinstance(content, str):
        return bool(content.strip())

    if isinstance(content, list):
        for item in content:
            item_type = _content_item_type(item)
            if item_type == "text":
                text_value = _content_item_text(item)
                if isinstance(text_value, str) and text_value.strip():
                    return True
            if item_type == "image_url":
                image_url = _content_item_image_url_payload(item)
                if _image_payload_url(image_url):
                    return True
        return False

    return False


def persist_request_images(
    request: ChatCompletionRequest,
    conversation_id: str,
    workspace_root: Optional[str | Path] = None,
) -> None:
    """Persist inline image uploads into a deterministic workspace folder.

    Images sent as data URLs are written under `./uploads/<conversation_id>/` so tools that
    expect a local image path can reference stable files.
    """

    root = Path(workspace_root).resolve() if workspace_root else Path.cwd().resolve()
    uploads_dir = root / "uploads" / conversation_id
    upload_index = 0

    for message in request.messages:
        if message.role != "user" or not isinstance(message.content, list):
            continue

        for item in message.content:
            item_type = _content_item_type(item)
            if item_type != "image_url":
                continue

            image_url = _content_item_image_url_payload(item)
            if image_url is None:
                continue

            local_path = _image_payload_local_path(image_url)
            if isinstance(local_path, str) and local_path.strip():
                continue

            url = _image_payload_url(image_url)
            if not isinstance(url, str):
                continue

            match = DATA_URL_PATTERN.match(url)
            if not match:
                continue

            mime_type, encoded_data = match.groups()
            extension = IMAGE_MIME_EXTENSIONS.get(mime_type, ".jpg")
            uploads_dir.mkdir(parents=True, exist_ok=True)
            upload_index += 1
            image_path = uploads_dir / f"query_image_{upload_index:03d}{extension}"
            image_bytes = base64.b64decode(encoded_data)
            image_path.write_bytes(image_bytes)

            resolved_path = str(image_path)
            _set_image_payload_local_path(image_url, resolved_path)


def _convert_chat_content_to_agent_query(content: Any) -> Any:
    if isinstance(content, str):
        return content

    if not isinstance(content, list):
        return ""

    parts: list[Any] = []
    image_path_notes: list[str] = []
    for item in content:
        item_type = _content_item_type(item)

        if item_type == "text":
            text_value = _content_item_text(item)
            if isinstance(text_value, str) and text_value.strip():
                parts.append(Text(text_value))
            continue

        if item_type == "image_url":
            image_url = _content_item_image_url_payload(item)
            local_path = _image_payload_local_path(image_url)
            if isinstance(local_path, str) and local_path.strip():
                cleaned_local_path = local_path.strip()
                parts.append(ImgPath(cleaned_local_path, detail="high"))
                image_path_notes.append(
                    "Reference image saved at: "
                    f"{cleaned_local_path}. You can pass this path to tools like "
                    "`make_user_query_more_detailed(query_image_path=...)` or "
                    "`get_visual_feedback(query_image_path=...)` when needed."
                )
                continue
            image_url_value = _image_payload_url(image_url)
            if isinstance(image_url_value, str):
                parts.append(
                    ImgUrl(image_url_value, detail=_image_payload_detail(image_url))
                )

    for note in image_path_notes:
        parts.append(Text(note))

    if not parts:
        return ""
    if len(parts) == 1 and isinstance(parts[0], Text):
        return str(parts[0])
    return parts


def validate_chat_request(request: ChatCompletionRequest) -> Tuple[Any, str, Any]:
    """
    Validate the chat request and extract the user message.

    Returns:
        (query_for_agent, request_id, raw_user_content)
    """
    if not request.messages:
        raise HTTPException(
            status_code=400, detail="Missing required parameter: messages"
        )

    # Get the user's last message.
    user_messages = [msg for msg in request.messages if msg.role == "user"]
    if not user_messages:
        raise HTTPException(
            status_code=400, detail="No user message found in conversation"
        )

    last_user_message = user_messages[-1]

    if not _has_nonempty_user_content(last_user_message.content):
        raise HTTPException(
            status_code=400, detail="User message content cannot be empty"
        )

    query = _convert_chat_content_to_agent_query(last_user_message.content)
    request_id = f"chatcmpl-{uuid.uuid4().hex[:29]}"
    return query, request_id, last_user_message.content


def get_or_create_conversation(
    conversation_id: Optional[str], conversation_manager: ConversationManager
) -> Tuple[Conversation, str]:
    """
    Get or create a conversation.

    If conversation_id is None, create a new conversation.
    If conversation_id exists but the corresponding conversation does not exist, create a new conversation.
    Otherwise, return the existing conversation.

    Args:
        conversation_id: Optional conversation ID
        conversation_manager: ConversationManager instance

    Returns:
        (conversation, conversation_id)
    """
    if not conversation_id:
        conversation = conversation_manager.create_conversation()
        conversation_id = conversation.uuid
    else:
        conversation = conversation_manager.get_conversation(conversation_id)  # type: ignore
        if conversation is None:
            # If the conversation does not exist, create a new one.
            conversation = conversation_manager.create_conversation(
                conversation_id=conversation_id
            )

    return conversation, conversation_id


def _extract_tokens_from_chunk(chunk: Any) -> Tuple[Optional[int], Optional[int]]:
    """Extract token statistics from multiple possible chunk structures.

    Return (prompt_tokens, completion_tokens). Either value is None if absent.
    """
    # 1) Pydantic model: chunk.usage.prompt_tokens
    try:
        usage = getattr(chunk, "usage", None)
        if usage is not None:
            pt = getattr(usage, "prompt_tokens", None)
            ct = getattr(usage, "completion_tokens", None)
            if isinstance(pt, int) or isinstance(ct, int):
                return (
                    int(pt) if isinstance(pt, int) else None,
                    int(ct) if isinstance(ct, int) else None,
                )
    except Exception:
        pass

    # 2) Dictionary: {"usage": {"prompt_tokens": x, "completion_tokens": y}}
    try:
        if isinstance(chunk, dict):
            u = chunk.get("usage")
            if isinstance(u, dict):
                pt = u.get("prompt_tokens")
                ct = u.get("completion_tokens")
                pt_v = int(pt) if isinstance(pt, (int, float)) else None
                ct_v = int(ct) if isinstance(ct, (int, float)) else None
                if pt_v is not None or ct_v is not None:
                    return (pt_v, ct_v)
    except Exception:
        pass

    # 3) Flat dictionary: {"prompt_tokens": x, "completion_tokens": y}
    try:
        if isinstance(chunk, dict):
            pt = chunk.get("prompt_tokens")
            ct = chunk.get("completion_tokens")
            pt_v = int(pt) if isinstance(pt, (int, float)) else None
            ct_v = int(ct) if isinstance(ct, (int, float)) else None
            if pt_v is not None or ct_v is not None:
                return (pt_v, ct_v)
    except Exception:
        pass

    return (None, None)


async def process_agent_response(
    query: Any,
    conversation: Conversation,
    agent: Any,
    raw_user_content: Any = None,
) -> Tuple[str, Optional[int], Optional[int]]:
    """Process the Agent response and return (full_text, prompt_tokens, completion_tokens)."""
    full_response: str = ""
    prompt_tokens: Optional[int] = None
    completion_tokens: Optional[int] = None
    try:
        with propagate_conversation_session(
            conversation_id=conversation.uuid,
            metadata={
                "model": getattr(agent, "model_name", None),
                "agent_name": getattr(agent, "name", None),
                "transport": "non_stream",
            },
            tags=["cadagent", "non_stream"],
        ):
            with conversation:
                async for output in agent.run(query, raw_user_content=raw_user_content):
                    delta = extract_output_text(output, "agent_non_stream")
                    if delta:
                        full_response += delta
                    if (
                        prompt_tokens is None or completion_tokens is None
                    ) and is_response_yield(output):
                        pt, ct = _extract_tokens_from_chunk(output.response)
                        if pt is not None:
                            prompt_tokens = pt
                        if ct is not None:
                            completion_tokens = ct

                # Persist the conversation immediately after completion.
                try:
                    await conversation.context.persist()
                    # Directly call the sketch_pad persist method.
                    conversation.sketch_pad.persist()
                    app_log(
                        f"✅ Auto-saved conversation {conversation.uuid} after agent response"
                    )
                except Exception as save_error:
                    # Save failure should not affect the response, but it should be logged.
                    push_warning(
                        f"⚠️ Warning: Failed to save conversation {conversation.uuid}: {save_error}"
                    )

        return full_response.strip(), prompt_tokens, completion_tokens
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Agent processing error: {str(e)}")


def create_chat_response(
    request_id: str,
    model: str,
    full_response: str,
    prompt_tokens: Optional[int] = None,
    completion_tokens: Optional[int] = None,
) -> ChatCompletionResponse:
    """Create a chat response."""
    created_time = int(time.time())

    response_message = ChatMessage(
        role="assistant",
        content=full_response,
        name=None,
        tool_calls=None,
        tool_call_id=None,
    )

    choice = ChatChoice(index=0, message=response_message, finish_reason="stop")

    # Prefer upstream token statistics; otherwise fall back to context statistics; finally use 0.
    if prompt_tokens is None:
        _in = get_current_context_attribute("input_tokens")
        try:
            prompt_tokens = int(_in) if _in is not None else 0
        except Exception:
            prompt_tokens = 0
    if completion_tokens is None:
        _out = get_current_context_attribute("output_tokens")
        try:
            completion_tokens = int(_out) if _out is not None else 0
        except Exception:
            completion_tokens = 0

    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=(prompt_tokens or 0) + (completion_tokens or 0),
    )

    return ChatCompletionResponse(
        id=request_id,
        object="chat.completion",
        created=created_time,
        model=model,
        choices=[choice],
        usage=usage,
        system_fingerprint=None,
    )