Lookii125/image_editing_agent
0
1import os2import json3import re4import time5import uuid6import inspect7from functools import lru_cache8from typing import TypedDict, Annotated, List, Dict, Any, Optional, Tuple9from dotenv import load_dotenv10 11load_dotenv()12 13from langgraph.graph import StateGraph, START, END14from langgraph.prebuilt import ToolNode, InjectedState15from langgraph.graph.message import add_messages16from langchain_huggingface import ChatHuggingFace17from langchain_huggingface.llms import HuggingFaceEndpoint18from langchain_core.messages import (19 BaseMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage, RemoveMessage,20)21from langchain_core.messages.tool import ToolCall22from langchain_core.tools import tool23 24class AgentState(TypedDict):25 input: str26 chat_history: List[BaseMessage]27 messages: Annotated[List[BaseMessage], add_messages]28 parsed_action: Dict[str, Any]29 output: str30 img_state_ref: str31 parse_error_count: int32 error: Optional[str]33 step_count: int34 retry_hint: Optional[str]35 36MAX_STEPS = 2037MAX_PARSE_RETRIES = 238MAX_LLM_RETRIES = 339LLM_RETRY_DELAY = 1.040 41_IMG_STATE_REGISTRY: Dict[str, Any] = {}42 43def register_img_state(session_key: str, img_state_obj: Any) -> None:44 _IMG_STATE_REGISTRY[session_key] = img_state_obj45 46def _get_img_state(state: dict) -> Any:47 key = state.get("img_state_ref")48 if not isinstance(key, str):49 raise RuntimeError("img_state_ref must be a session key string.")50 obj = _IMG_STATE_REGISTRY.get(key)51 if obj is None:52 raise RuntimeError(f"No ImageState registered under key '{key}'.")53 return obj54 55@lru_cache(maxsize=4)56def _make_llm(hf_token: str, max_new_tokens: int = 512) -> ChatHuggingFace:57 llm = HuggingFaceEndpoint(58 repo_id="Qwen/Qwen2.5-72B-Instruct",59 huggingfacehub_api_token=hf_token,60 task="text-generation",61 max_new_tokens=max_new_tokens,62 temperature=0.01,63 do_sample=False,64 timeout=45,65 )66 return ChatHuggingFace(llm=llm)67 68BANNED_INPUT_PATTERNS = [69 re.compile(r"\bzignoruj\b", re.IGNORECASE),70 re.compile(r"\bzapomnij\s+o\b", re.IGNORECASE),71 re.compile(r"\boverrid(e|ing|den)\b", re.IGNORECASE),72 re.compile(r"\bsystem\s*prompt\b", re.IGNORECASE),73 re.compile(r"\byou\s+are\s+now\b", re.IGNORECASE),74 re.compile(75 r"\bignore\s+(all|any|previous|prior|the)\s+(instructions?|rules?|prompts?|guidelines?)\b",76 re.IGNORECASE,77 ),78]79 80def validate_input(state: AgentState) -> Dict[str, Any]:81 text = (state.get("input") or "").strip()82 83 if not text:84 return {"error": "Proszę opisać, co chcesz zrobić z obrazem.", "output": ""}85 if len(text) > 800:86 return {"error": "Twoja instrukcja jest zbyt długa. Maksymalnie 800 znaków.", "output": ""}87 if any(p.search(text) for p in BANNED_INPUT_PATTERNS):88 return {"error": "Wykryto zapytanie niezwiązane z edycją obrazów.", "output": ""}89 90 try:91 img_state = _get_img_state(state)92 if img_state.current is None:93 return {"error": "Brak załadowanego obrazu. Najpierw prześlij plik.", "output": ""}94 except RuntimeError:95 return {"error": "Brak załadowanego obrazu. Najpierw prześlij plik.", "output": ""}96 97 return {"error": None, "parse_error_count": 0, "step_count": 0, "retry_hint": None}98 99def route_after_validation(state: AgentState) -> str:100 return "abort" if state.get("error") else "call_model"101 102def extract_json_from_response(text: str) -> Dict[str, Any]:103 cleaned = (text or "").strip()104 105 if not cleaned:106 return {"action": "__parse_failed__", "raw": text, "reason": "empty"}107 108 if cleaned.startswith("```"):109 cleaned = re.sub(r"^```[a-z]*\n?", "", cleaned)110 cleaned = re.sub(r"\n?```$", "", cleaned)111 112 cleaned = cleaned.strip()113 cleaned = re.sub(r'\bNone\b', 'null', cleaned)114 cleaned = re.sub(r'\bTrue\b', 'true', cleaned)115 cleaned = re.sub(r'\bFalse\b', 'false', cleaned)116 117 if cleaned.startswith("{") and not cleaned.endswith("}"):118 if cleaned[-1] != '"' and cleaned.count('"') % 2 != 0:119 cleaned += '"'120 cleaned += "\n}"121 122 try:123 return json.loads(cleaned)124 except json.JSONDecodeError:125 match = re.search(r"\{.*\}", cleaned, re.DOTALL)126 if match:127 try:128 return json.loads(match.group(0))129 except json.JSONDecodeError:130 pass131 132 return {"action": "__parse_failed__", "raw": text}133 134def format_mod_history(modifications_list: list) -> str:135 if not modifications_list:136 return "No modifications applied yet."137 return "\n".join(f" {i+1}. {d}" for i, d in enumerate(modifications_list))138 139def safe_truncate_history(messages: List[BaseMessage], max_messages: int) -> List[BaseMessage]:140 if len(messages) <= max_messages:141 return messages142 truncated = list(messages[-max_messages:])143 while truncated:144 first = truncated[0]145 if isinstance(first, AIMessage) and getattr(first, "tool_calls", None):146 if len(truncated) < 2 or not isinstance(truncated[1], ToolMessage):147 truncated = truncated[1:]148 continue149 break150 return truncated151 152@lru_cache(maxsize=1)153def _get_yolo_model():154 from ultralytics import YOLO155 return YOLO("yolov8n.pt") 156 157@tool158def detect_elements_tool(159 state: Annotated[dict, InjectedState],160 conf_threshold: float = 0.15,161) -> str:162 """163 Detects salient visual elements in the image using a YOLO object detection model. 164 Always call this before any element-indexed tool.165 Args:166 conf_threshold: The minimum confidence threshold for a valid detection. Default 0.15.167 """168 img_state = _get_img_state(state)169 safe_conf = float(conf_threshold) if conf_threshold is not None else 0.15170 model = _get_yolo_model()171 172 elements = img_state.detect_elements_op(173 model=model,174 conf_threshold=safe_conf,175 )176 177 if not elements:178 return "No elements detected with the given parameters. Try reducing conf_threshold."179 180 h, w = img_state.current.shape[:2]181 lines = [f"Image size: {w}x{h} px. Detected {len(elements)} element(s):\n"]182 for e in elements:183 lines.append(e.summary(w, h))184 return "\n".join(lines)185 186 187@tool188def get_element_properties_tool(189 element_index: int,190 state: Annotated[dict, InjectedState],191) -> str:192 """193 Returns the full geometric and colour properties of a single detected element.194 Args:195 element_index: The integer index of the specific element to inspect, as discovered by detect_elements_tool.196 """197 img_state = _get_img_state(state)198 elem = img_state.get_element(int(element_index))199 return json.dumps(elem.to_dict(), indent=2)200 201 202@tool203def analyze_scene_tool(204 question: str,205 state: Annotated[dict, InjectedState],206) -> str:207 """208 Asks a reasoning question about the detected elements and receives a plain-text analytical answer.209 Useful for resolving spatial clues (e.g., "Which element is on the left?").210 Args:211 question: A specific query about the scene, positions, or colours of the detected elements.212 """213 hf_token = os.getenv("HF_TOKEN")214 if not hf_token:215 return "Error: HF_TOKEN not set."216 217 img_state = _get_img_state(state)218 219 if not img_state.detections:220 return (221 "No detections available. Call detect_elements_tool first, "222 "then call analyze_scene_tool with your question."223 )224 225 h, w = img_state.current.shape[:2]226 detection_text = "\n".join(e.summary(w, h) for e in img_state.detections)227 228 user_input = state.get("input", "No specific user input provided.")229 230 reasoning_prompt = (231 f"You are a computer vision analyst. "232 f"You have measurements of visual elements detected in an image.\n\n"233 f"USER'S CURRENT INSTRUCTION: \"{user_input}\"\n"234 f"(Pay close attention to any spatial hints like 'left', 'right', 'top', 'bottom', or explicit element indices mentioned by the user.)\n\n"235 f"IMAGE SIZE: {w}x{h} pixels\n\n"236 f"DETECTED ELEMENTS (sorted largest-first by area):\n{detection_text}\n\n"237 f"QUESTION FROM ORCHESTRATOR: {question}\n\n"238 f"INSTRUCTIONS:\n"239 f"1. Answer concisely and precisely using only the data above.\n"240 f"2. If the user's instruction provides spatial hints, use the element position (e.g., 'left', 'right') and centroid data to match them.\n"241 f"3. If the user explicitly identifies an element by index in their instruction, treat that assignment as ground truth.\n"242 f"4. Refer to elements by their index number.\n"243 f"5. If the question asks which element index to use for an operation, state the index explicitly at the start of your answer."244 )245 246 last_exc = None247 for attempt in range(MAX_LLM_RETRIES):248 try:249 chat_model = _make_llm(hf_token, max_new_tokens=300)250 response = chat_model.invoke([HumanMessage(content=reasoning_prompt)])251 raw_content = response.content or ""252 content = raw_content.strip()253 254 if not content:255 raise ValueError("Reasoning sub-model returned an empty completion.")256 return content257 258 except Exception as exc:259 last_exc = exc260 time.sleep(LLM_RETRY_DELAY * (2 ** attempt))261 262 return f"Reasoning sub-model failed after {MAX_LLM_RETRIES} retries: {last_exc}"263 264@tool265def blur_element_tool(266 element_index: int,267 state: Annotated[dict, InjectedState],268 kernel_size: int = 15,269) -> str:270 """271 Applies a Gaussian blur to the bounding box of a specific detected element.272 Args:273 element_index: The index of the target element from detect_elements_tool. Required.274 kernel_size: The intensity of the blur. Must be an odd positive integer. Optional, default 15.275 """276 _get_img_state(state).blur_element(int(element_index), int(kernel_size))277 return f"Blurred element #{element_index} with kernel size {kernel_size}."278 279@tool280def crop_to_element_tool(281 element_index: int,282 state: Annotated[dict, InjectedState],283 padding: int = 0,284) -> str:285 """286 Crops the canvas to exactly the bounding box of a specified detected element.287 Args:288 element_index: The index of the target element. Required.289 padding: Extra pixels to leave around the element's bounding box. Optional, default 0 (tight crop).290 """291 _get_img_state(state).crop_to_element(int(element_index), int(padding))292 return f"Cropped canvas to element #{element_index} (padding={padding}px)."293 294@tool295def brightness_element_tool(296 element_index: int,297 factor: float,298 state: Annotated[dict, InjectedState],299) -> str:300 """301 Adjusts the brightness exclusively inside the bounding box of a specific detected element.302 Args:303 element_index: The index of the target element.304 factor: Brightness multiplier (e.g., 1.5 increases brightness, 0.5 decreases it).305 """306 _get_img_state(state).brightness_element(int(element_index), float(factor))307 return f"Adjusted brightness of element #{element_index} by factor {factor}."308 309@tool310def grayscale_tool(311 state: Annotated[dict, InjectedState],312 x: Optional[int] = None,313 y: Optional[int] = None,314 width: Optional[int] = None,315 height: Optional[int] = None316) -> str:317 """318 Converts the image to grayscale. If coordinates (x, y, width, height) are left empty/None,319 it will convert the entire image. If provided, it will apply only to that localized region.320 All four coordinate arguments are OPTIONAL — omit all of them for a whole-image conversion.321 Args:322 x: Top-left x-coordinate of a localized region (optional).323 y: Top-left y-coordinate of a localized region (optional).324 width: Width of the region (optional).325 height: Height of the region (optional).326 """327 _get_img_state(state).to_grayscale(x, y, width, height)328 return "Image converted to grayscale successfully."329 330@tool331def blur_tool(332 kernel_size: int,333 state: Annotated[dict, InjectedState],334 x: Optional[int] = None,335 y: Optional[int] = None,336 width: Optional[int] = None,337 height: Optional[int] = None338) -> str:339 """340 Applies a Gaussian blur. If coordinates are left empty/None, it blurs the entire image.341 If provided, it blurs only the localized bounding box. Coordinates are OPTIONAL.342 Args:343 kernel_size: The intensity of the blur. Must be an odd positive integer (e.g., 3, 9, 15). Required.344 x: Top-left x-coordinate of a localized region (optional).345 y: Top-left y-coordinate of a localized region (optional).346 width: Width of the localized region (optional).347 height: Height of the localized region (optional).348 """349 _get_img_state(state).blur(int(kernel_size), x, y, width, height)350 return f"Gaussian blur applied with kernel size {kernel_size}."351 352@tool353def resize_tool(354 width: int,355 height: int,356 state: Annotated[dict, InjectedState]357) -> str:358 """359 Resizes the entire canvas to exactly the requested width and height.360 Args:361 width: The new target width in pixels.362 height: The new target height in pixels.363 """364 _get_img_state(state).resize(int(width), int(height))365 return f"Image resized to {width}x{height}."366 367@tool368def flip_tool(369 axis: str,370 state: Annotated[dict, InjectedState]371) -> str:372 """373 Flips the entire canvas along a specified axis.374 Args:375 axis: Must be one of 'horizontal', 'vertical', or 'both'.376 """377 _get_img_state(state).flip(axis)378 return f"Flipped image ({axis})."379 380@tool381def rotate_tool(382 degrees: float,383 state: Annotated[dict, InjectedState],384 expand: bool = True,385) -> str:386 """387 Rotates the entire canvas about its center.388 DIRECTION (this is the most common mistake — read carefully):389 - POSITIVE degrees -> counter-clockwise (the top of the image swings toward the LEFT).390 - NEGATIVE degrees -> clockwise (the top of the image swings toward the RIGHT).391 - "rotate/turn right", "rotate clockwise" -> use a NEGATIVE value.392 - "rotate/turn left", "rotate counter-clockwise" -> use a POSITIVE value.393 - "upside down" / "flip it around" -> degrees = 180 (sign doesn't matter at 180).394 MAGNITUDE — only use the amount actually requested, do not round up:395 - "straighten a bit" / "small tilt" / "nudge" -> roughly 3 to 10 degrees.396 - "quarter turn" / "turn it on its side" -> 90 degrees.397 - "upside down" / "half turn" -> 180 degrees.398 - If the user gives an explicit number, use that exact number.399 - Never call this tool more than once for a single user request — rotations400 accumulate on the canvas, so calling it twice doubles the effect.401 Args:402 degrees: Amount to rotate, in degrees. Positive = counter-clockwise, negative = clockwise.403 expand: Optional, default True. True expands the canvas so nothing gets clipped;404 False keeps the original canvas size and clips the corners.405 """406 _get_img_state(state).rotate(float(degrees), bool(expand))407 direction = "counter-clockwise" if degrees >= 0 else "clockwise"408 return f"Rotated {abs(degrees)}\u00b0 {direction} (expand={expand})."409 410@tool411def crop_tool(412 x: int,413 y: int,414 width: int,415 height: int,416 state: Annotated[dict, InjectedState]417) -> str:418 """419 Crops the canvas to a specified rectangle. Both width and height MUST be positive non-zero integers.420 CALCULATING CANVAS DIMENSIONS:421 - To trim edges from an image, calculate final remaining dimensions relative to current size:422 Example: If image is 640x426 px and user asks to "crop 250px off the right side",423 new_width = 640 - 250 = 390. Height remains 426.424 Call: crop_tool(x=0, y=0, width=390, height=426).425 - NEVER pass 0, None, or negative values for width or height.426 Args:427 x: Top-left x-coordinate of the box.428 y: Top-left y-coordinate of the box.429 width: Target width in pixels (must be > 0).430 height: Target height in pixels (must be > 0).431 """432 _get_img_state(state).crop(int(x), int(y), int(width), int(height))433 return f"Cropped to ({x},{y}) with size {width}x{height}."434 435@tool436def brightness_tool(437 factor: float,438 state: Annotated[dict, InjectedState],439 x: Optional[int] = None,440 y: Optional[int] = None,441 width: Optional[int] = None,442 height: Optional[int] = None443) -> str:444 """445 Multiplies pixel brightness by a factor. If coordinates are omitted, the entire image is affected.446 Coordinates are OPTIONAL.447 Args:448 factor: Multiplier for brightness. >1.0 lightens, <1.0 darkens. Required.449 x: Top-left x-coordinate of a localized region (optional).450 y: Top-left y-coordinate of a localized region (optional).451 width: Width of the localized region (optional).452 height: Height of the localized region (optional).453 """454 _get_img_state(state).adjust_brightness(float(factor), x, y, width, height)455 return f"Brightness adjusted by factor {factor}."456 457@tool458def contrast_tool(459 alpha: float,460 beta: float,461 state: Annotated[dict, InjectedState],462 x: Optional[int] = None,463 y: Optional[int] = None,464 width: Optional[int] = None,465 height: Optional[int] = None466) -> str:467 """468 Adjusts image contrast. If coordinates are omitted, the entire image is adjusted. Coordinates are OPTIONAL.469 Args:470 alpha: Simple contrast control (e.g., 1.0 to 3.0). Required.471 beta: Brightness control applied during contrast adjustment (-100 to 100). Required.472 x: Top-left x-coordinate of a localized region (optional).473 y: Top-left y-coordinate of a localized region (optional).474 width: Width of the localized region (optional).475 height: Height of the localized region (optional).476 """477 _get_img_state(state).adjust_contrast(float(alpha), float(beta), x, y, width, height)478 return f"Contrast adjusted (alpha={alpha}, beta={beta})."479 480@tool481def canny_tool(482 threshold1: int,483 threshold2: int,484 state: Annotated[dict, InjectedState],485 x: Optional[int] = None,486 y: Optional[int] = None,487 width: Optional[int] = None,488 height: Optional[int] = None489) -> str:490 """491 Applies Canny edge detection. If coordinates are omitted, detects edges on the entire image.492 Coordinates are OPTIONAL.493 Args:494 threshold1: First threshold for the hysteresis procedure (e.g., 100). Required.495 threshold2: Second threshold for the hysteresis procedure (e.g., 200). Required.496 x: Top-left x-coordinate of a localized region (optional).497 y: Top-left y-coordinate of a localized region (optional).498 width: Width of the localized region (optional).499 height: Height of the localized region (optional).500 """501 _get_img_state(state).canny(int(threshold1), int(threshold2), x, y, width, height)502 return f"Canny edge detection applied (th1={threshold1}, th2={threshold2})."503 504@tool505def threshold_tool(506 threshold_value: int,507 max_value: int,508 state: Annotated[dict, InjectedState],509 x: Optional[int] = None,510 y: Optional[int] = None,511 width: Optional[int] = None,512 height: Optional[int] = None513) -> str:514 """515 Applies a binary threshold filter. If coordinates are omitted, applies to the entire image.516 Coordinates are OPTIONAL.517 Args:518 threshold_value: The value (0-255) to threshold against. Required.519 max_value: The value given if a pixel is greater than the threshold (usually 255). Required.520 x: Top-left x-coordinate of a localized region (optional).521 y: Top-left y-coordinate of a localized region (optional).522 width: Width of the localized region (optional).523 height: Height of the localized region (optional).524 """525 _get_img_state(state).threshold(int(threshold_value), int(max_value), x, y, width, height)526 return f"Binary threshold applied (val={threshold_value}, max={max_value})."527 528@tool529def sharpen_tool(530 state: Annotated[dict, InjectedState],531 x: Optional[int] = None,532 y: Optional[int] = None,533 width: Optional[int] = None,534 height: Optional[int] = None535) -> str:536 """537 Applies an unsharp-mask sharpening filter. If coordinates are omitted, applies to the entire image.538 Coordinates are OPTIONAL.539 Args:540 x: Top-left x-coordinate of a localized region (optional).541 y: Top-left y-coordinate of a localized region (optional).542 width: Width of the localized region (optional).543 height: Height of the localized region (optional).544 """545 _get_img_state(state).sharpen(x, y, width, height)546 return "Canvas sharpened."547 548image_tools_catalog = [549 detect_elements_tool,550 get_element_properties_tool,551 analyze_scene_tool,552 blur_element_tool,553 crop_to_element_tool,554 brightness_element_tool,555 grayscale_tool, blur_tool, resize_tool, flip_tool, rotate_tool,556 crop_tool, brightness_tool, contrast_tool, canny_tool, threshold_tool, sharpen_tool,557]558 559def _tool_arg_info(t) -> Tuple[set, set]:560 try:561 schema_model = t.args_schema562 if hasattr(schema_model, "model_json_schema"):563 schema = schema_model.model_json_schema()564 else:565 schema = schema_model.schema()566 all_args = {k for k in schema.get("properties", {}).keys() if k != "state"}567 required_args = {k for k in schema.get("required", []) if k != "state"}568 return all_args, required_args569 except Exception:570 all_args = {k for k in t.args.keys() if k != "state"}571 return all_args, set(all_args)572 573TOOL_ALL_ARGS: Dict[str, set] = {}574TOOL_REQUIRED_ARGS: Dict[str, set] = {}575for _t in image_tools_catalog:576 _all, _req = _tool_arg_info(_t)577 TOOL_ALL_ARGS[_t.name] = _all578 TOOL_REQUIRED_ARGS[_t.name] = _req579 580def _tool_signature_line(t) -> str:581 try:582 schema_model = t.args_schema583 schema = (584 schema_model.model_json_schema()585 if hasattr(schema_model, "model_json_schema")586 else schema_model.schema()587 )588 props = schema.get("properties", {})589 required = set(schema.get("required", []))590 parts = []591 for name, spec in props.items():592 if name == "state":593 continue594 if name in required:595 parts.append(name)596 else:597 parts.append(f"{name}={spec.get('default')!r}")598 return f"{t.name}({', '.join(parts)})"599 except Exception:600 return f"{t.name}(...)"601 602def _tool_docstring(t) -> str:603 func = getattr(t, "func", None) or getattr(t, "coroutine", None)604 doc = inspect.getdoc(func) if func else None605 if not doc:606 doc = (t.description or "").strip()607 return doc.strip()608 609def _build_tool_catalog_text(tools) -> str:610 blocks = []611 missing_docs = []612 for t in tools:613 sig = _tool_signature_line(t)614 doc = _tool_docstring(t)615 if not doc:616 missing_docs.append(t.name)617 blocks.append(f" {sig}")618 continue619 indented = "\n".join(620 (" " + line if line.strip() else "") for line in doc.splitlines()621 )622 blocks.append(f" {sig}\n{indented}")623 if missing_docs:624 print(f"[agent_graph] WARNING: tools with no docstring (model will get no guidance for these): {missing_docs}")625 return "\n".join(blocks)626 627TOOL_CATALOG_TEXT = _build_tool_catalog_text(image_tools_catalog)628 629def call_model(state: AgentState) -> Dict[str, Any]:630 hf_token = os.getenv("HF_TOKEN")631 if not hf_token:632 return {"error": "Missing HF_TOKEN credential.", "output": ""}633 634 img_state = _get_img_state(state)635 h, w = img_state.current.shape[:2]636 live_mod_history = format_mod_history(img_state.modifications_list)637 638 if img_state.detections and not img_state.detections_are_stale:639 det_lines = [f" {e.summary(w, h)}" for e in img_state.detections]640 detection_context = (641 f"Current detections ({len(img_state.detections)} elements, image {w}x{h}px):\n"642 + "\n".join(det_lines)643 )644 elif img_state.detections and img_state.detections_are_stale:645 detection_context = (646 "Detections exist but are STALE (canvas was modified after detection). "647 "Re-run detect_elements_tool if element-level operations are needed."648 )649 else:650 detection_context = (651 f"No detections yet. Image is {w}x{h}px. "652 "Call detect_elements_tool first if element-level operations are needed."653 )654 655 extra_context = ""656 hint = state.get("retry_hint")657 if hint:658 extra_context = f"\n\nIMPORTANT: {hint}"659 elif state.get("parse_error_count", 0) > 0:660 extra_context = (661 "\n\nIMPORTANT: Your previous response could not be parsed as valid JSON. "662 "Reply with ONLY the raw JSON object — no prose, no markdown fences, "663 "and never leave your response empty."664 )665 666 system_prompt = f"""You are an expert image editing orchestrator. You plan and execute image operations step by step.667CRITICAL GROUND TRUTH RULES:6681. Base your understanding of the current canvas ONLY on the "SCENE STATE" and "MODIFICATION HISTORY" sections below.6692. If an action is NOT listed in MODIFICATION HISTORY, it is NOT applied (for example, if the user performed an 'Undo').6703. Ignore past assistant chat messages if they claim an action was done that isn't reflected in MODIFICATION HISTORY.671AVAILABLE TOOLS:672{TOOL_CATALOG_TEXT}673CANVAS & SCENE STATE:674- Canvas Dimensions: {w}x{h} px675- Active History Step: {img_state.history_index}676{detection_context}677APPLIED MODIFICATIONS (Source of Truth):678{live_mod_history}679RESPONSE FORMAT — output exactly ONE of these two JSON structures per turn, nothing else:680Format 1 — call a tool:681{{682 "thought": "Step-by-step reasoning.",683 "action": "call_tool",684 "tool_name": "EXACT_TOOL_NAME",685 "parameters": {{ ... }}686}}687Format 2 — final answer or clarification:688{{689 "thought": "Why the task is complete, or why you need to ask the user something.",690 "action": "final_response",691 "response": "Natural language message to the user."692}}693RULES:694- Output exactly one JSON block per turn, nothing else. Never leave the JSON — or the "response" field in Format 2 — empty or blank.695- One tool call per turn; wait for the result before the next step.696- Never call an element-indexed tool without a preceding detect_elements_tool call.697- Do not invent pixel coordinates or other numeric values you are not confident about. If a tool's required parameter cannot be determined from the user's instruction or the scene state, ask a short, specific clarifying question using Format 2 instead of guessing.{extra_context}"""698 699 MAX_HISTORY_MESSAGES = 8700 truncated_history = safe_truncate_history(state["chat_history"], MAX_HISTORY_MESSAGES)701 payload = [SystemMessage(content=system_prompt)] + truncated_history + state["messages"]702 703 last_exc = None704 for attempt in range(MAX_LLM_RETRIES):705 try:706 chat_model = _make_llm(hf_token, max_new_tokens=400)707 response = chat_model.invoke(payload)708 raw_content = response.content or ""709 content = raw_content.strip() 710 711 if not content:712 print(f"[call_model] WARNING: empty completion on attempt {attempt + 1}/{MAX_LLM_RETRIES}, retrying...")713 raise ValueError("Model returned an empty completion.")714 715 return {716 "messages": [AIMessage(content=response.content)],717 "step_count": state.get("step_count", 0) + 1,718 "retry_hint": None,719 }720 except Exception as exc:721 last_exc = exc722 time.sleep(LLM_RETRY_DELAY * (2 ** attempt))723 724 return {"error": f"LLM unreachable or returned empty output after {MAX_LLM_RETRIES} retries: {last_exc}", "output": ""}725 726def route_after_model(state: AgentState) -> str:727 return "abort" if state.get("error") else "custom_parser"728 729SUSPICIOUS_OUTPUT_PATTERNS = [730 re.compile(r"\bprzepis\w*\b", re.IGNORECASE),731 re.compile(r"\bszarlotk\w*\b", re.IGNORECASE),732 re.compile(r"\bciast[ao]\w*\b", re.IGNORECASE),733 re.compile(r"\bmąk\w*\b", re.IGNORECASE),734 re.compile(r"\bskładnik\w*\b", re.IGNORECASE),735 re.compile(r"\bupiecz\w*\b", re.IGNORECASE),736 re.compile(r"\bbaking\b", re.IGNORECASE),737 re.compile(r"\brecipe\b", re.IGNORECASE),738 re.compile(r"\bingredients?\b", re.IGNORECASE),739]740 741def custom_parser(state: AgentState) -> Dict[str, Any]:742 raw_content = state["messages"][-1].content743 parsed = extract_json_from_response(raw_content)744 745 print("\n[parser] Raw:\n", raw_content)746 parsed = extract_json_from_response(raw_content)747 print("\n[parser] Parsed:", json.dumps(parsed, indent=2))748 749 if parsed.get("action") == "__parse_failed__":750 count = state.get("parse_error_count", 0) + 1751 if count >= MAX_PARSE_RETRIES:752 return {753 "error": "I'm sorry, I encountered an error during execution. Could you ask again?",754 "parse_error_count": count,755 }756 757 hint = (758 "Your last response was empty or was not valid JSON. Reply with exactly one JSON "759 "object (Format 1 or Format 2), with no prose before or after it, and make sure it "760 "is not empty."761 )762 763 failure_feedback = HumanMessage(764 content="[SYSTEM NOTICE]: Your previous JSON response was malformed and COULD NOT be executed. "765 "No changes were made to the image canvas. Please re-issue your tool call using valid JSON."766 )767 768 return {769 "parse_error_count": count, 770 "parsed_action": parsed, 771 "retry_hint": hint,772 "messages": [failure_feedback]773 }774 775 if parsed.get("action") == "call_tool":776 tool_name = parsed.get("tool_name", "")777 params = parsed.get("parameters", {}) or {}778 779 if tool_name not in TOOL_ALL_ARGS:780 count = state.get("parse_error_count", 0) + 1781 if count >= MAX_PARSE_RETRIES:782 return {783 "error": f"Nie udało się wykonać żądania — nieznane narzędzie '{tool_name}'.",784 "parse_error_count": count,785 }786 hint = (787 f"'{tool_name}' is not a valid tool name. Choose one of the exact tool names "788 f"listed in AVAILABLE TOOLS."789 )790 return {"parse_error_count": count, "parsed_action": parsed, "retry_hint": hint}791 792 missing_args = TOOL_REQUIRED_ARGS[tool_name] - set(params.keys())793 if missing_args:794 count = state.get("parse_error_count", 0) + 1795 if count >= MAX_PARSE_RETRIES:796 return {797 "error": (798 f"Nie udało się wykonać narzędzia '{tool_name}' — brakuje wymaganych "799 f"parametrów: {sorted(missing_args)}."800 ),801 "parse_error_count": count,802 }803 hint = (804 f"Your call to '{tool_name}' was missing required parameter(s): "805 f"{sorted(missing_args)}. Either provide concrete values for all of them, or — "806 f"if the user's instruction doesn't give you enough information to determine "807 f"them — respond using Format 2 and ask the user a short clarifying question "808 f"instead of guessing."809 )810 return {"parse_error_count": count, "parsed_action": parsed, "retry_hint": hint}811 812 filtered_params = {k: v for k, v in params.items() if k in TOOL_ALL_ARGS[tool_name]}813 814 tool_call = ToolCall(name=tool_name, args=filtered_params, id=f"call_{uuid.uuid4().hex[:8]}")815 816 echo_content = json.dumps({817 "thought": parsed.get("thought", ""),818 "action": "call_tool",819 "tool_name": tool_name,820 "parameters": filtered_params,821 })822 823 return {824 "messages": [AIMessage(content=echo_content, tool_calls=[tool_call])],825 "parsed_action": parsed,826 "parse_error_count": 0,827 "retry_hint": None,828 }829 830 if parsed.get("action") == "final_response":831 response_text = (parsed.get("response") or "").strip()832 if not response_text:833 count = state.get("parse_error_count", 0) + 1834 hint = "Your final_response 'response' field was empty. Provide a non-empty message to the user."835 if count >= MAX_PARSE_RETRIES:836 return {837 "error": "I'm sorry, I encountered an error during execution. Could you ask again?",838 "parse_error_count": count,839 }840 return {"parse_error_count": count, "parsed_action": parsed, "retry_hint": hint}841 842 if any(p.search(response_text) for p in SUSPICIOUS_OUTPUT_PATTERNS):843 return {844 "parsed_action": parsed,845 "output": "Odpowiedź zablokowana: wykryto treść niezwiązaną z operacjami graficznymi.",846 "parse_error_count": 0,847 "retry_hint": None,848 }849 850 return {"parsed_action": parsed, "output": response_text, "parse_error_count": 0, "retry_hint": None}851 852 count = state.get("parse_error_count", 0) + 1853 hint = (854 'Your "action" field must be exactly "call_tool" or "final_response". '855 "Re-read the two response formats and reply with one of them."856 )857 if count >= MAX_PARSE_RETRIES:858 return {859 "error": "I'm sorry, I encountered an error during execution. Could you ask again?",860 "parse_error_count": count,861 }862 return {"parse_error_count": count, "parsed_action": parsed, "retry_hint": hint}863 864def route_after_parser(state: AgentState) -> str:865 if state.get("error"):866 return "abort"867 parsed = state.get("parsed_action", {})868 if parsed.get("action") in ("__parse_failed__", None):869 return "call_model"870 if state.get("output"):871 return "end_workflow"872 latest_msg = state["messages"][-1]873 if hasattr(latest_msg, "tool_calls") and latest_msg.tool_calls:874 return "tools"875 return "call_model"876 877def handle_tool_error(state: AgentState) -> Dict[str, Any]:878 last_msg = state["messages"][-1]879 error_detail = getattr(last_msg, "content", None) or state.get("error") or "Unknown tool execution fault."880 881 tool_call_id = getattr(last_msg, "tool_call_id", None)882 if tool_call_id is None:883 for msg in reversed(state["messages"]):884 if hasattr(msg, "tool_calls") and msg.tool_calls:885 tool_call_id = msg.tool_calls[0]["id"]886 break887 tool_call_id = tool_call_id or "unknown"888 889 messages_out: List[BaseMessage] = []890 msg_id = getattr(last_msg, "id", None)891 if isinstance(last_msg, ToolMessage) and msg_id:892 messages_out.append(RemoveMessage(id=msg_id))893 894 messages_out.append(895 ToolMessage(896 content=f"Tool failed: {error_detail}. Re-evaluate parameters or choose a different tool/approach.",897 tool_call_id=tool_call_id,898 )899 )900 return {"messages": messages_out, "error": None}901 902def route_after_tools(state: AgentState) -> str:903 last = state["messages"][-1]904 is_error = False905 if isinstance(last, ToolMessage):906 status = getattr(last, "status", None)907 if status == "error":908 is_error = True909 elif status is None and isinstance(last.content, str) and last.content.startswith("Error"):910 is_error = True911 912 if is_error:913 return "handle_tool_error"914 if state.get("step_count", 0) >= MAX_STEPS:915 return "abort"916 return "call_model"917 918def abort_node(state: AgentState) -> Dict[str, Any]:919 return {"output": state.get("error", "Pipeline fault caused termination.")}920 921def build_agent_graph():922 workflow = StateGraph(AgentState)923 924 workflow.add_node("validate_input", validate_input)925 workflow.add_node("call_model", call_model)926 workflow.add_node("custom_parser", custom_parser)927 workflow.add_node("tools", ToolNode(image_tools_catalog, handle_tool_errors=True))928 workflow.add_node("handle_tool_error", handle_tool_error)929 workflow.add_node("abort", abort_node)930 931 workflow.add_edge(START, "validate_input")932 workflow.add_conditional_edges("validate_input", route_after_validation, {"call_model": "call_model", "abort": "abort"})933 workflow.add_conditional_edges("call_model", route_after_model, {"custom_parser": "custom_parser", "abort": "abort"})934 workflow.add_conditional_edges("custom_parser", route_after_parser, {935 "tools": "tools",936 "call_model": "call_model",937 "end_workflow": END,938 "abort": "abort",939 })940 workflow.add_conditional_edges("tools", route_after_tools, {941 "handle_tool_error": "handle_tool_error",942 "call_model": "call_model",943 "abort": "abort",944 })945 workflow.add_edge("handle_tool_error", "call_model")946 workflow.add_edge("abort", END)947 948 return workflow.compile()949 950compiled_agent = build_agent_graph()