Update app.py

app.py

@@ -221,17 +221,12 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
 
     try:
         if use_web_search:
-            for main_content, sources, search_results in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
-                if sources:
-                    formatted_sources = match_and_format_sources(sources, search_results)
-                    response = f"{main_content}\n\n{formatted_sources}"
-                else:
-                    response = main_content
+            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+                response = f"{main_content}\n\n{sources}"
                 first_line = response.split('\n')[0] if response else ''
                 logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
-            # PDF search logic (unchanged)
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
                 embed = get_embeddings()
@@ -243,11 +238,10 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                 else:
                     context_str = "No documents available."
                 
-                for main_content, sources in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    response = f"{main_content}\n\n{sources}" if sources else main_content
-                    first_line = response.split('\n')[0] if response else ''
+                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
                     logging.info(f"Generated Response (first line): {first_line}")
-                    yield response
+                    yield partial_response
             else:
                 # Use Hugging Face API
                 for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
@@ -305,49 +299,39 @@ 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
-                                    # Attempt to split the response into main content and sources
-                                    main_content, sources = split_content_and_sources(full_response)
-                                    yield main_content, sources
+                                    yield full_response
                             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.", ""
-
-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, ""
+        yield "I apologize, but I couldn't generate a response at this time. Please try again later."
 
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
-    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['title']}" 
+    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
     
-    prompt = f"""Using the following context from web search results:
+    prompt = f"""Using the following context:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
 After writing the document, please provide a list of sources used in your response."""
 
     if model == "@cf/meta/llama-3.1-8b-instruct":
         # Use Cloudflare API
-        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, search_results
+        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
     else:
-        # Use Hugging Face API for other models
+        # Use Hugging Face API
         client = InferenceClient(model, token=huggingface_token)
         
-        full_response = ""
+        main_content = ""
         for i in range(num_calls):
             for message in client.chat_completion(
                 messages=[{"role": "user", "content": prompt}],
@@ -357,16 +341,8 @@ After writing the document, please provide a list of sources used in your respon
             ):
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                     chunk = message.choices[0].delta.content
-                    full_response += chunk
-                    main_content, sources = split_content_and_sources(full_response)
-                    yield main_content, sources, search_results
-
-def split_content_and_sources(text):
-    parts = text.split("Sources:", 1)
-    if len(parts) > 1:
-        return parts[0].strip(), "Sources:" + parts[1].strip()
-    else:
-        return text.strip(), ""
+                    main_content += chunk
+                    yield main_content, ""  # Yield partial main content without sources
 
 def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
     embed = get_embeddings()
@@ -483,4 +459,4 @@ with demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)