import json import time from concurrent.futures import ThreadPoolExecutor, TimeoutError from langgraph.graph import StateGraph, END from langgraph.checkpoint.memory import MemorySaver from src.config import config as app_config, resolve_ai_client from src.state import AgentState from src.schemas import ValidationFeedbackSchema, DesignReviewSchema from src.retrieval.retriever import Retriever from src.vision.analyzer import analyze_asset from src.utils.logging import get_logger, log_latency from src.utils.model_caller import call_with_fallback from langchain_core.runnables import RunnableConfig from openai import OpenAI import PIL.Image import io import base64 import os logger = get_logger(__name__) TIMEOUT_SECONDS = 120.0 def run_with_timeout(func, timeout_seconds, *args, **kwargs): with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(func, *args, **kwargs) return future.result(timeout=timeout_seconds) def append_error(state: AgentState, error_msg: str): errors = state.get("errors", []) errors.append(error_msg) return {"errors": errors} def retrieval_node(state: AgentState): query = state.get("query", "") project_id = state.get("project_id", "default") retriever = Retriever(project_id) # pad the query with brand-related keywords for better vector search hits augmented_query = f"{query} brand guidelines color palette hex typography fonts layout rules".strip() context = retriever.search(augmented_query, k=8) return {"retrieved_context": context} def vision_node(state: AgentState): """Run vision analysis on the uploaded design.""" image_path = state.get("image_path") if not image_path: return {"visual_analysis": {"error": "No image provided"}} analysis_result = analyze_asset(image_path).model_dump() return {"visual_analysis": analysis_result} def synthesis_node(state: AgentState, config: RunnableConfig): """Main review node -- builds the 3-part report (consistency, collision, enhancements).""" query = state.get("query", "") context = state.get("retrieved_context", []) vision = state.get("visual_analysis", {}) count = state.get("recursion_count", 0) image_path = state.get("image_path") project_id = state.get("project_id", "default") project_name = state.get("project_name", "Unknown Brand") q = config.get("configurable", {}).get("queue") # Resolve AI provider from request overrides ai_provider = config.get("configurable", {}).get("ai_provider") ai_api_key = config.get("configurable", {}).get("ai_api_key") user_client, user_model = resolve_ai_client(ai_provider, ai_api_key) # stitch together the retrieved guideline chunks into a prompt block guidelines_summary = "" if context: guidelines_summary = "\n".join([ f"- Source: {c.get('source_file', 'unknown')}, Section: {c.get('section_title', 'N/A')}\n Content: {c.get('content', '')[:500]}" for c in context[:5] ]) else: guidelines_summary = ( "NO BRAND GUIDELINES HAVE BEEN INGESTED FOR THIS PROJECT. " "You have NO reference material to compare against. " "Because of this, the consistency_score MUST be capped at 50 maximum. " "You MUST add a finding with element 'Missing Guidelines' and status 'needs_attention' " "explaining that no brand guidelines were uploaded, so a reliable consistency check is impossible. " "Provide general best-practice feedback only." ) vision_summary = json.dumps(vision, indent=2) if vision else "No visual analysis available." parent_review = state.get("design_review") parent_review_summary = "" if parent_review: parent_review_summary = f""" ## PREVIOUS REVIEW (V1) The designer is uploading a revised version (V2). Here was the previous review: {json.dumps(parent_review, indent=2)} **CRITICAL INSTRUCTION**: Since this is a revision, you MUST generate a "Delta Report" in your findings. Compare V1 to V2. Explicitly state what improved, what regressed, and if previous feedback was addressed. """ competitors = [] base_dir = os.path.join("data", "projects", project_id) if project_id and project_id != "default" else "" if base_dir and os.path.exists(os.path.join(base_dir, "competitors.json")): try: with open(os.path.join(base_dir, "competitors.json"), "r") as f: competitors = json.load(f) except Exception: pass competitors_summary = "" if competitors: competitors_summary = "## KNOWN COMPETITORS TO AVOID\nWe have identified the following direct competitors. Ensure the design does NOT resemble their brand kits:\n" for c in competitors: competitors_summary += f"- {c.get('name', 'Unknown')}: Primary Colors: {c.get('primary_colors', [])}, Typography: {c.get('typography', [])}\n" system_prompt = f"""You are Artitude, an expert AI Brand Consistency & Design Enhancement Co-Pilot. ## THE BRAND YOU ARE PROTECTING You are reviewing a design submitted under the project: **"{project_name}"**. Your ONLY job is to evaluate whether this uploaded design is consistent with **"{project_name}"**'s brand guidelines listed below. **CRITICAL RULE — BRAND IDENTITY MISMATCH**: If the uploaded design clearly depicts, contains, or IS the logo/visual identity of a DIFFERENT brand (e.g., Nike, Adidas, Apple, Coca-Cola, or any other recognizable third-party brand), then: - The `consistency_score` MUST be between 0-15. It is fundamentally NOT "{project_name}"'s brand. - Every finding related to that foreign brand identity must be marked "inconsistent". - Add a finding with element "Brand Identity" explaining that this design belongs to a different brand entirely. - The Market Collision section MUST flag this as a "high" risk collision. Do NOT evaluate the design against the foreign brand's own guidelines. You are ONLY checking against "{project_name}". ## DESIGN ASSET VISUAL ANALYSIS {vision_summary} ## "{project_name}" BRAND GUIDELINES (from RAG) {guidelines_summary} {parent_review_summary} {competitors_summary} ## DESIGNER'S QUESTION {query if query else "No specific question — provide a general design review."} ## YOUR TASK Analyze this design and produce a structured review with these three sections: ### 1. Brand Consistency (against "{project_name}" ONLY) Compare the design's colors, typography, layout density, and visual theme against the "{project_name}" brand guidelines above. Assign a consistency_score (0-100). List specific findings for each element (color palette, fonts, visual tone, etc.) and mark each as "consistent", "inconsistent", or "needs_attention". Remember: if this design is clearly from another brand, the score must be near 0. **IMPORTANT**: For each finding, if it refers to a specific visual element on the image, provide a `bounding_box` as `[x_percentage, y_percentage, width_percentage, height_percentage]` relative percentages (0.0 to 100.0). If it's a general finding, leave it null. ### 2. Market Collision Check Using your broad knowledge of famous global brands AND the explicit "KNOWN COMPETITORS TO AVOID" list above, check if this design accidentally resembles OR directly copies any visual identity. Consider color schemes, layout patterns, logo styles, wordmarks, and overall aesthetic. - If the design IS another brand's asset (e.g., an Adidas logo uploaded to a non-Adidas project), flag it as "high" risk. - CRITICAL: If the uploaded design's dominant colors or typography match the extracted KNOWN COMPETITORS data, you MUST heavily penalize the design and explicitly generate a "Market Collision Warning" in your findings with a "high" risk level. - If it merely resembles another brand, flag appropriately as low/medium/high. Flag any matches with the brand name, what's similar, and a risk level. If no collisions found, return an empty list. ### 3. World-Class Enhancements Provide actionable, professional-grade suggestions to elevate this design. Cover areas like: - Typography (modern font pairings, weight hierarchy) - Color harmony (complementary palettes, contrast ratios) - Layout (whitespace, visual hierarchy, grid alignment) - Modern design trends that could apply Rate each suggestion's impact as low/medium/high. **IMPORTANT**: For each enhancement, if it refers to a specific area on the image, provide a `bounding_box` as `[x_percentage, y_percentage, width_percentage, height_percentage]` relative percentages (0.0 to 100.0). For this step, write out your synthesis and Delta Report as a professional design narrative. Do not output JSON yet. Just write the text. """ pil_image = None b64_img = None if image_path: if os.path.exists(image_path): pil_image = PIL.Image.open(image_path) pil_image.load() # read into memory so the file handle gets freed buffered = io.BytesIO() if pil_image.mode != 'RGB': pil_image = pil_image.convert('RGB') pil_image.save(buffered, format="JPEG") b64_img = base64.b64encode(buffered.getvalue()).decode("utf-8") start_time = time.time() try: messages = [{"role": "system", "content": system_prompt}] if b64_img: messages.append({ "role": "user", "content": [ { "type": "image_url", "image_url": { "url": f"data:image/jpeg;base64,{b64_img}" } } ] }) # stream the narrative text first using the resolved client narrative_stream = user_client.chat.completions.create( model=user_model, messages=messages, stream=True ) raw_synthesis = "" for chunk in narrative_stream: if hasattr(chunk, "choices") and chunk.choices and len(chunk.choices) > 0: delta = chunk.choices[0].delta content = getattr(delta, "content", None) if content: raw_synthesis += content if q: q.put({"type": "token", "data": content}) # now ask for structured json messages.append({"role": "assistant", "content": raw_synthesis}) messages.append({"role": "user", "content": "Now, output the exact corresponding structured JSON for this analysis matching the required schema."}) def call_api(): return call_with_fallback( messages=messages, response_format=DesignReviewSchema, primary_model_name=app_config.ROUTER_MODEL_NAME, fallback_model_names=app_config.ROUTER_MODEL_FALLBACKS, client_override=user_client, model_override=user_model, ) response_text = run_with_timeout(call_api, TIMEOUT_SECONDS) review_data = json.loads(response_text) # raw_synthesis was built up from the stream above except TimeoutError: logger.warning("Synthesis node timed out", extra={"custom_fields": {"node": "synthesis"}}) return { "raw_synthesis": "Error: Design review timed out.", "design_review": {}, "recursion_count": count + 1, "errors": state.get("errors", []) + ["Synthesis node timed out"] } except Exception as e: logger.error(f"Synthesis node error: {e}", extra={"custom_fields": {"node": "synthesis"}}) return { "raw_synthesis": f"Error: {str(e)}", "design_review": {}, "recursion_count": count + 1, "errors": state.get("errors", []) + [f"Synthesis error: {str(e)}"] } latency = (time.time() - start_time) * 1000 log_latency(logger, "synthesis_node", latency) return { "raw_synthesis": raw_synthesis, "design_review": review_data, "recursion_count": count + 1 } def _format_review_summary(review: dict) -> str: """Turn the structured review dict into readable text.""" parts = [] bc = review.get("brand_consistency", {}) score = bc.get("consistency_score", "N/A") parts.append(f"## Brand Consistency Score: {score}/100") for f in bc.get("findings", []): parts.append(f"- **{f.get('element', '')}** [{f.get('status', '')}]: {f.get('detail', '')}") mc = review.get("market_collision", {}) collisions = mc.get("collisions", []) if collisions: parts.append(f"\n## Market Collision Alerts ({len(collisions)} found)") for c in collisions: parts.append(f"- WARNING: **{c.get('brand_name', '')}** ({c.get('risk_level', '')} risk): {c.get('detail', '')}") else: parts.append("\n## Market Collision Check: No collisions detected") en = review.get("enhancements", {}) suggestions = en.get("suggestions", []) if suggestions: parts.append(f"\n## Enhancement Suggestions ({len(suggestions)})") for s in suggestions: parts.append(f"- **{s.get('category', '')}** [{s.get('impact', '')} impact]: {s.get('suggestion', '')}") return "\n".join(parts) def validator_node(state: AgentState): """Quick check that all three sections came back from the model.""" review = state.get("design_review", {}) # make sure each section is present and not empty has_consistency = bool(review.get("brand_consistency", {}).get("findings")) has_collision = "collisions" in review.get("market_collision", {}) has_enhancements = bool(review.get("enhancements", {}).get("suggestions")) is_valid = has_consistency and has_collision and has_enhancements if not is_valid and not review: # empty review = synthesis blew up, no point retrying return {"validation_passed": False} return {"validation_passed": is_valid} def build_graph(): workflow = StateGraph(AgentState) workflow.add_node("retrieval", retrieval_node) workflow.add_node("vision", vision_node) workflow.add_node("synthesis", synthesis_node) workflow.add_node("validator", validator_node) workflow.set_entry_point("retrieval") def retrieval_conditional(state: AgentState): if state.get("image_path"): return "vision" return "synthesis" workflow.add_conditional_edges( "retrieval", retrieval_conditional, { "vision": "vision", "synthesis": "synthesis" } ) workflow.add_edge("vision", "synthesis") workflow.add_edge("synthesis", "validator") def validator_conditional(state: AgentState): # bail after 3 tries, or if it passed, or if the review is empty (hard error) if state.get("validation_passed") or state.get("recursion_count", 0) >= 3 or not state.get("design_review"): return END return "synthesis" workflow.add_conditional_edges( "validator", validator_conditional, { END: END, "synthesis": "synthesis" } ) return workflow.compile(checkpointer=MemorySaver()) graph = build_graph()