Update main.py

main.py

@@ -39,8 +39,6 @@ async def favicon():
     favicon_path = Path(__file__).parent / "favicon.ico"
     return FileResponse(favicon_path, media_type="image/x-icon")
     
-import requests
-import json
 
 def generate_search(query: str, stream: bool = True) -> str:
     headers = {"User-Agent": ""}
@@ -48,7 +46,6 @@ def generate_search(query: str, stream: bool = True) -> str:
         {"role": "user", "content": query},
     ]
     
-    # Insert the system prompt at the beginning of the conversation history
     prompt.insert(0, {"content": "Be Helpful and Friendly", "role": "system"})
     
     payload = {
@@ -58,30 +55,28 @@ def generate_search(query: str, stream: bool = True) -> str:
         "user_input": prompt[-1]["content"],
     }
     
-    # Use the newly added SECRET_API_ENDPOINT_3 for the search API call
     chat_endpoint = secret_api_endpoint_3
     response = requests.post(chat_endpoint, headers=headers, json=payload, stream=True)
     
-    # Collect streamed text content
     streaming_text = ""
     for value in response.iter_lines(decode_unicode=True):
-        # Ensure the value starts with 'data: ' and process it
         if value.startswith("data: "):  
             try:
-                # Parse JSON and remove 'data: ' prefix
                 json_modified_value = json.loads(value[6:])
-                content = json_modified_value.get("content", "")  # Get the actual content
-                
-                if content.strip():  # Ensure non-empty content
-                    if stream:
-                        yield content  # Pass content directly, removing 'data: ' and JSON overhead
-                    streaming_text += content  # Accumulate content for non-streaming case
+                content = json_modified_value.get("choices", [{}])[0].get("delta", {}).get("content", "")
+                if stream:
+                    yield f"data: {json.dumps({'choices': [{'delta': {'content': content}}]})}\n\n"
+                streaming_text += content
             except json.JSONDecodeError:
-                continue  # Skip lines that are not valid JSON
+                continue
     
-    # If not streaming, yield the full collected content
     if not stream:
-        yield streaming_text
+        yield json.dumps({
+            "choices": [{
+                "message": {"content": streaming_text},
+                "finish_reason": "stop"
+            }]
+        })
 @app.get("/searchgpt")
 async def search_gpt(q: str, stream: Optional[bool] = False):
     if not q: