From 02c3178b6bd3e413bd3b7f7791b9e64fa153d125 Mon Sep 17 00:00:00 2001
From: clean99 <xff9924@gmail.com>
Date: Tue, 21 Nov 2023 21:14:45 +0800
Subject: [PATCH] feat: update api to support image comparison

---
 backend/main.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/backend/main.py b/backend/main.py
index 4eb09fa..2fa9e9f 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -96,7 +96,10 @@ async def stream_code_test(websocket: WebSocket):
     async def process_chunk(content):
         await websocket.send_json({"type": "chunk", "value": content})
 
-    prompt_messages = assemble_prompt(params["image"])
+    if params.get("resultImg") and params["resultImg"]:
+        prompt_messages = assemble_prompt(params["image"], params["resultImg"])
+    else:
+        prompt_messages = assemble_prompt(params["image"])
 
     # Image cache for updates so that we don't have to regenerate images
     image_cache = {}