Ai assistant features

.gitignore

@@ -41,4 +41,4 @@ node_modules
 .env.*
 
 # But include .env.example files
-!.env.example
+!.env.example*.tgz

backend/app.py

@@ -28,7 +28,7 @@
 from routes.admin import admin_bp
 from routes.frontend import frontend_bp
 from routes.analytics import analytics_bp
-from routes.export import export_bp
+from routes.ai_assistant import ai_assistant_bp
 from services.db import redis_client
 from services.canvas_counter import get_canvas_draw_count
 from services.graphql_service import commit_transaction_via_graphql
@@ -170,33 +170,32 @@ def handle_all_exceptions(e):
 app.register_blueprint(undo_redo_bp)
 app.register_blueprint(metrics_bp)
 app.register_blueprint(auth_bp)
+app.register_blueprint(ai_assistant_bp)
 app.register_blueprint(rooms_bp)
 app.register_blueprint(submit_room_line_bp)
 app.register_blueprint(admin_bp)
-app.register_blueprint(export_bp)
 
 # Register versioned API v1 blueprints for external applications
 from api_v1.auth import auth_v1_bp
 from api_v1.canvases import canvases_v1_bp
 from api_v1.collaborations import collaborations_v1_bp
 from api_v1.notifications import notifications_v1_bp
 from api_v1.users import users_v1_bp
-from routes.stamps import stamps_bp
 from api_v1.templates import templates_v1_bp
 
 app.register_blueprint(auth_v1_bp)
 app.register_blueprint(canvases_v1_bp)
 app.register_blueprint(collaborations_v1_bp)
 app.register_blueprint(notifications_v1_bp)
 app.register_blueprint(users_v1_bp)
-app.register_blueprint(stamps_bp, url_prefix='/api')
 app.register_blueprint(templates_v1_bp)
 
 # Frontend serving must be last to avoid route conflicts
 app.register_blueprint(frontend_bp)
 app.register_blueprint(analytics_bp)
 
 if __name__ == '__main__':
+    print(SIGNER_PUBLIC_KEY, SIGNER_PRIVATE_KEY, RECIPIENT_PUBLIC_KEY)
     if not redis_client.exists('res-canvas-draw-count'):
         init_count = {"id": "res-canvas-draw-count", "value": 0}
         logger = __import__('logging').getLogger(__name__)

backend/config.py

@@ -24,6 +24,7 @@
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 ANALYTICS_COLLECTION_NAME = os.getenv("ANALYTICS_COLLECTION_NAME", "analytics_events")
 ANALYTICS_AGGREGATES_COLLECTION = os.getenv("ANALYTICS_AGGREGATES_COLLECTION", "analytics_aggregates")
+HUGGINGFACE_API_KEY=os.getenv("HUGGINGFACE_API_KEY")
 
 JWT_SECRET = os.getenv("JWT_SECRET", "dev-insecure-change-me")
 JWT_ISSUER = "rescanvas"

backend/middleware/validators.py

@@ -487,6 +487,10 @@ def validate_stroke_payload(value) -> Tuple[bool, str]:
 
     if "pathData" not in stroke:
         return False, "Stroke must have pathData"
+
+    path_data = stroke.get("pathData")
+    if not isinstance(path_data, (dict, list)):
+        return False, "Stroke pathData must be an object or array"
 
     is_valid, error = validate_color(stroke.get("color"))
     if not is_valid:
@@ -498,6 +502,21 @@ def validate_stroke_payload(value) -> Tuple[bool, str]:
             return False, "Line width must be between 1 and 100"
     except (TypeError, ValueError):
         return False, "Line width must be a number"
+
+    timestamp = stroke.get("timestamp")
+    if timestamp is not None:
+        try:
+            int(timestamp)
+        except (TypeError, ValueError):
+            return False, "Timestamp must be a number"
+
+    drawing_id = stroke.get("drawingId")
+    if drawing_id is not None and not isinstance(drawing_id, str):
+        return False, "drawingId must be a string"
+
+    stroke_id = stroke.get("id")
+    if stroke_id is not None and not isinstance(stroke_id, str):
+        return False, "id must be a string"
 
     # Validate optional signature fields (will be enforced for secure rooms in handler)
     signature = value.get("signature")

backend/routes/ai_assistant.py

@@ -0,0 +1,149 @@
+from flask import Blueprint, request, jsonify
+from services.llm_service import prompt_to_drawings, complete_shape_from_canvas, beautify_canvas_state
+# from services.image_generation_service import (
+#     text_to_image as img_text_to_image,
+# )
+import logging
+import base64
+import io
+
+ai_assistant_bp = Blueprint('ai_assistant', __name__)
+logger = logging.getLogger(__name__)
+
+
+@ai_assistant_bp.route('/api/ai_assistant/drawing', methods=['POST'])
+def text_to_drawings():
+    """
+    Body: { "prompt": "<natural language description>", "canvasState": { ... } }
+    Returns: stroke JSON or an error payload.
+    """
+    try:
+        payload = request.get_json(silent=True) or {}
+        prompt = payload.get("prompt")
+        canvas_state = payload.get("canvasState") or {}
+
+        if not isinstance(prompt, str) or not prompt.strip():
+            return jsonify({
+                "error": "bad_request",
+                "detail": "Missing or invalid 'prompt' (string)."
+            }), 400
+
+        if not isinstance(canvas_state, dict):
+            return jsonify({
+                "error": "bad_request",
+                "detail": "Invalid 'canvasState' (object)."
+            }), 400
+
+        logger.info("AI drawing requested: route entered")
+        logger.info("Calling prompt_to_drawings now")
+        result = prompt_to_drawings(prompt.strip(), canvas_state)
+        logger.info("prompt_to_drawings returned")
+
+        if isinstance(result, dict) and "error" in result:
+            logger.warning("AI drawing failed: %s", result)
+            return jsonify({
+                "error": "upstream_model_error",
+                "detail": result
+            }), 502
+
+        logger.info("AI drawing generated successfully")
+        logger.debug("AI drawing result: %r", result)
+        return jsonify(result), 200
+
+    except Exception as e:
+        logger.exception("Unhandled error in /drawing")
+        return jsonify({"error": "server_error", "detail": str(e)}), 500
+
+
+@ai_assistant_bp.route('/api/ai_assistant/complete', methods=['POST'])
+def shape_completion():
+    """
+    Body: { "canvasState": { ... } }
+    Returns: { complete, confidence, object{ color, lineWidth, pathData{...} } } or an error payload.
+    """
+    try:
+        payload = request.get_json(silent=True) or {}
+        canvas_state = payload.get("canvasState")
+        if not isinstance(canvas_state, dict):
+            return jsonify({"error": "bad_request", "detail": "Missing or invalid 'canvas_state' (object)."}), 400
+
+        logger.info("AI shape completion requested")
+        suggestion = complete_shape_from_canvas(canvas_state)
+
+        if not isinstance(canvas_state, dict):
+            return jsonify({
+                "error": "bad_request",
+                "detail": "Missing or invalid 'canvasState' (object)."
+            }), 400
+
+        return jsonify(suggestion), 200
+    except Exception as e:
+        logger.exception("Unhandled error in /complete")
+        return jsonify({"error": "server_error", "detail": str(e)}), 500
+
+
+@ai_assistant_bp.route('/api/ai_assistant/image', methods=['POST'])
+def text_to_image():
+    """
+    TODO: To be implemented
+    Body: { "prompt": "<string>", "width"?: int, "height"?: int, "style"?: str }
+    Returns: { "imageDataUrl": "data:image/png;base64,..." }
+    """
+    try:
+        payload = request.get_json(silent=True) or {}
+        prompt = payload.get("prompt", "")
+        width = payload.get("width") or 512
+        height = payload.get("height") or 512
+        style = payload.get("style") or "default"
+
+        if not isinstance(prompt, str) or not prompt.strip():
+            return jsonify({
+                "error": "bad_request",
+                "detail": "Missing or invalid 'prompt' (string)."
+            }), 400
+
+        logger.info("AI text-to-image requested")
+
+        pil_image = []      # img_text_to_image(prompt.strip(), width=width, height=height, style=style) 
+
+        buf = io.BytesIO()
+        pil_image.save(buf, format="PNG")
+        buf.seek(0)
+        encoded = base64.b64encode(buf.read()).decode("utf-8")
+        data_url = f"data:image/png;base64,{encoded}"
+
+        return jsonify({"imageDataUrl": data_url}), 200
+
+    except Exception as e:
+        logger.exception("Unhandled error in /image")
+        return jsonify({"error": "server_error", "detail": str(e)}), 500
+
+
+@ai_assistant_bp.route("/api/ai_assistant/beautify", methods=["POST"])
+def beautify_sketch():
+    try:
+        payload = request.get_json(silent=True) or {}
+        canvas_state = payload.get("canvasState")
+
+        if not isinstance(canvas_state, dict):
+            return jsonify({
+                "error": "bad_request",
+                "detail": "Missing or invalid 'canvasState' (object)."
+            }), 400
+
+        result = beautify_canvas_state(canvas_state)
+        # print("\n\ncanvas_state!!!", canvas_state, "\n\n")
+        # print("\n\nResult!!!", result, "\n\n")
+
+        if not isinstance(result, dict) or "objects" not in result:
+            logger.warning("Beautify returned invalid payload: %r", result)
+            return jsonify({
+                "error": "upstream_model_error",
+                "detail": "Beautify model returned invalid payload."
+            }), 502
+
+        return jsonify(result), 200
+
+    except Exception as e:
+        logger.exception("Unhandled error in /beautify")
+        return jsonify({"error": "server_error", "detail": str(e)}), 500

backend/routes/rooms.py

@@ -699,7 +699,10 @@ def admin_fill_wrapped_key(roomId):
 @require_room_access(room_id_param="roomId")
 @limiter.limit(f"{RATE_LIMIT_STROKE_MINUTE}/minute")
 @validate_request_data({
-    "stroke": {"validator": lambda v: (isinstance(v, dict), "Stroke must be an object") if not isinstance(v, dict) else (True, None), "required": True},
+    "stroke": {
+        "validator": lambda v: validate_stroke_payload(request.get_json() or {}),
+        "required": True
+    },
     "signature": {"validator": validate_optional_string(max_length=1000), "required": False},
     "signerPubKey": {"validator": validate_optional_string(max_length=1000), "required": False}
 })
@@ -732,6 +735,19 @@ def post_stroke(roomId):
         parent_paste_id = stroke.get("parentPasteId", "NOT SET")
         logger.warning(f"POST STROKE DEBUG - roomId={roomId}, brushType={brush_type}, brushParams={brush_params}, parentPasteId={parent_paste_id}")
         logger.warning(f"POST STROKE DEBUG - Full stroke object: {json.dumps(stroke, default=str)}")
+        path_data = stroke.get("pathData")
+        is_ai_batch_marker = isinstance(path_data, dict) and path_data.get("tool") == "paste" and path_data.get("aiGenerated") is True
+        is_ai_child_stroke = bool(stroke.get("parentPasteId")) or (
+            isinstance(path_data, dict) and bool(path_data.get("parentPasteId"))
+        )
+        if is_ai_batch_marker or is_ai_child_stroke:
+            logger.info(
+                "AI drawing submitted: roomId=%s drawingId=%s parentPasteId=%s batchMarker=%s",
+                roomId,
+                stroke.get("drawingId") or stroke.get("id"),
+                stroke.get("parentPasteId") or (path_data.get("parentPasteId") if isinstance(path_data, dict) else None),
+                is_ai_batch_marker,
+            )
     except Exception as e:
         logger.error(f"POST STROKE DEBUG - Error logging stroke: {e}")
 
@@ -827,6 +843,20 @@ def post_stroke(roomId):
         logger.warning(f"STORING FULL STROKE: {json.dumps(stroke, default=str)[:500]}...")
 
         strokes_coll.insert_one({"roomId": roomId, "ts": stroke["ts"], "stroke": stroke})
+        try:
+            path_data = stroke.get("pathData")
+            if (
+                stroke.get("parentPasteId") or
+                (isinstance(path_data, dict) and (path_data.get("parentPasteId") or path_data.get("aiGenerated")))
+            ):
+                logger.info(
+                    "AI drawing stored: roomId=%s strokeId=%s parentPasteId=%s",
+                    roomId,
+                    stroke.get("id") or stroke.get("drawingId"),
+                    stroke.get("parentPasteId") or (path_data.get("parentPasteId") if isinstance(path_data, dict) else None),
+                )
+        except Exception:
+            logger.exception("Failed to emit AI storage log for room %s", roomId)
 
         rooms_coll.update_one({"_id": room["_id"]}, {"$set": {"updatedAt": datetime.utcnow()}})
 
@@ -1340,6 +1370,14 @@ def get_strokes(roomId):
                         "filterParams": stroke_data["filterParams"],
                     }
 
+                    if parent_paste_id or (isinstance(stroke_data.get("pathData"), dict) and stroke_data["pathData"].get("aiGenerated")):
+                        logger.info(
+                            "AI drawing restored: roomId=%s strokeId=%s parentPasteId=%s",
+                            roomId,
+                            stroke_id,
+                            parent_paste_id,
+                        )
+
                     filtered_strokes.append(stroke_data)
                     if stroke_id:
                         seen_stroke_ids.add(stroke_id)

backend/services/db.py

@@ -2,6 +2,7 @@
 
 import threading
 import redis
+import os
 from pymongo import MongoClient
 from pymongo.server_api import ServerApi
 import logging
@@ -65,7 +66,22 @@
 
 settings_coll = mongo_client[DB_NAME]["settings"]
 
-redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0)
+DISABLE_REDIS = os.getenv("DISABLE_REDIS", "false").lower() == "true"
+
+if DISABLE_REDIS:
+    try:
+        import fakeredis
+        redis_client = fakeredis.FakeStrictRedis(decode_responses=False)
+        print("Using fakeredis fallback (no real Redis server).")
+    except ImportError:
+        raise RuntimeError(
+            "DISABLE_REDIS=true but fakeredis is not installed. "
+            "Run: python -m pip install fakeredis"
+        )
+else:
+    REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
+    REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
+    redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0)
 
 lock = threading.Lock()