Update main.py

main.py

@@ -57,18 +57,20 @@ def generate_search(query: str, stream: bool = True) -> str:
     
     # Collect streamed text content
     streaming_text = ""
-    for value in response.iter_lines(decode_unicode=True, chunk_size=12):
-        modified_value = re.sub("data:", "", value)
-        if modified_value:
+    for value in response.iter_lines(decode_unicode=True):
+        # Ensure the value is clean and processable
+        if value.startswith("data: "):  # Ensure the prefix matches expected format
             try:
-                json_modified_value = json.loads(modified_value)
-                content = json_modified_value["choices"][0]["delta"]["content"]
-                if stream:
-                    yield f"data: {content}\n\n"
-                streaming_text += content
-            except:
-                continue
+                json_modified_value = json.loads(value[6:])  # Remove 'data: ' prefix
+                content = json_modified_value.get("choices", [{}])[0].get("delta", {}).get("content", "")
+                if content.strip():  # Avoid empty content from delta
+                    if stream:
+                        yield f"data: {json.dumps({'response': content})}\n\n"
+                    streaming_text += content
+            except json.JSONDecodeError:
+                continue  # Skip lines that are not valid JSON
     
+    # If not streaming, yield the full text
     if not stream:
         yield streaming_text