468 lines
16 KiB
Python
468 lines
16 KiB
Python
"""
|
|
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,
|
|
)
|