Update app.py

app.py

@@ -219,10 +219,10 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
 
-    try:
+   try:
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
-                response = f"{main_content}\n\n{sources}"
+                response = f"{main_content}\n\n{sources}" if sources else main_content
                 first_line = response.split('\n')[0] if response else ''
                 logging.info(f"Generated Response (first line): {first_line}")
                 yield response
@@ -299,19 +299,29 @@ After writing the document, please provide a list of sources used in your respon
                                 if 'response' in json_response:
                                     chunk = json_response['response']
                                     full_response += chunk
-                                    yield full_response
+                                    # Attempt to split the response into main content and sources
+                                    main_content, sources = split_content_and_sources(full_response)
+                                    yield main_content, sources
                             except (json.JSONDecodeError, IndexError) as e:
                                 logging.error(f"Error parsing streaming response: {str(e)}")
                                 continue
                 else:
                     logging.error(f"HTTP Error: {response.status_code}, Response: {response.text}")
-                    yield f"I apologize, but I encountered an HTTP error: {response.status_code}. Please try again later."
+                    yield f"I apologize, but I encountered an HTTP error: {response.status_code}. Please try again later.", ""
         except Exception as e:
             logging.error(f"Error in generating response from Cloudflare: {str(e)}")
-            yield f"I apologize, but an error occurred: {str(e)}. Please try again later."
+            yield f"I apologize, but an error occurred: {str(e)}. Please try again later.", ""
     
     if not full_response:
-        yield "I apologize, but I couldn't generate a response at this time. Please try again later."
+        yield "I apologize, but I couldn't generate a response at this time. Please try again later.", ""
+
+def split_content_and_sources(text):
+    # Attempt to split the text into main content and sources
+    parts = text.split("Sources:", 1)
+    if len(parts) > 1:
+        return parts[0].strip(), "Sources:" + parts[1]
+    else:
+        return text, ""
 
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
@@ -325,8 +335,8 @@ After writing the document, please provide a list of sources used in your respon
 
     if model == "@cf/meta/llama-3.1-8b-instruct":
         # Use Cloudflare API
-        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
-            yield response, ""  # Yield streaming response without sources
+        for main_content, sources in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
+            yield main_content, sources
     else:
         # Use Hugging Face API
         client = InferenceClient(model, token=huggingface_token)