""" 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//` 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, )