Handle server {STOP} token

jithinraj · Aug 20, 2024 · de35005 · de35005
1 parent bebb0b5
commit de35005
Showing 1 changed file with 36 additions and 38 deletions.
diff --git a/interpreter/core/async_core.py b/interpreter/core/async_core.py
@@ -701,47 +701,45 @@ async def openai_compatible_generator():
     async def chat_completion(request: ChatCompletionRequest):
         # Convert to LMC
 
-        user_messages = []
-        for message in reversed(request.messages):
-            if message.role == "user":
-                user_messages.append(message)
-            else:
-                break
-        user_messages.reverse()
-
-        for message in user_messages:
-            if type(message.content) == str:
-                async_interpreter.messages.append(
-                    {"role": "user", "type": "message", "content": message.content}
-                )
-            if type(message.content) == list:
-                for content in message.content:
-                    if content["type"] == "text":
-                        async_interpreter.messages.append(
-                            {"role": "user", "type": "message", "content": content}
+        last_message = request.messages[-1]
+
+        if last_message.role != "user":
+            raise ValueError("Last message must be from the user.")
+
+        if last_message.content == "{STOP}":
+            # Handle special STOP token
+            return
+
+        if type(last_message.content) == str:
+            async_interpreter.messages.append(last_message)
+        if type(last_message.content) == list:
+            for content in last_message.content:
+                if content["type"] == "text":
+                    async_interpreter.messages.append(
+                        {"role": "user", "type": "message", "content": content}
+                    )
+                elif content["type"] == "image_url":
+                    if "url" not in content["image_url"]:
+                        raise Exception("`url` must be in `image_url`.")
+                    url = content["image_url"]["url"]
+                    print(url[:100])
+                    if "base64," not in url:
+                        raise Exception(
+                            '''Image must be in the format: "data:image/jpeg;base64,{base64_image}"'''
                         )
-                    elif content["type"] == "image_url":
-                        if "url" not in content["image_url"]:
-                            raise Exception("`url` must be in `image_url`.")
-                        url = content["image_url"]["url"]
-                        print(url[:100])
-                        if "base64," not in url:
-                            raise Exception(
-                                '''Image must be in the format: "data:image/jpeg;base64,{base64_image}"'''
-                            )
 
-                        # data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAA6oA...
+                    # data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAA6oA...
 
-                        data = url.split("base64,")[1]
-                        format = "base64." + url.split(";")[0].split("/")[1]
-                        async_interpreter.messages.append(
-                            {
-                                "role": "user",
-                                "type": "image",
-                                "format": format,
-                                "content": data,
-                            }
-                        )
+                    data = url.split("base64,")[1]
+                    format = "base64." + url.split(";")[0].split("/")[1]
+                    async_interpreter.messages.append(
+                        {
+                            "role": "user",
+                            "type": "image",
+                            "format": format,
+                            "content": data,
+                        }
+                    )
 
         if request.stream:
             return StreamingResponse(