Update app.py

app.py

@@ -231,6 +231,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                 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()
@@ -330,10 +331,10 @@ def split_content_and_sources(text):
 
 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['href']}\n" 
+    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['title']}" 
                         for result in search_results if 'body' in result)
     
-    prompt = f"""Using the following context:
+    prompt = f"""Using the following context from web search results:
 {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."""
@@ -341,12 +342,12 @@ 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 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
+            yield main_content, sources, search_results
     else:
-        # Use Hugging Face API
+        # Use Hugging Face API for other models
         client = InferenceClient(model, token=huggingface_token)
         
-        main_content = ""
+        full_response = ""
         for i in range(num_calls):
             for message in client.chat_completion(
                 messages=[{"role": "user", "content": prompt}],
@@ -356,8 +357,16 @@ 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
-                    main_content += chunk
-                    yield main_content, ""  # Yield partial main content without sources
+                    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(), ""
 
 def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
     embed = get_embeddings()