Update app.py

app.py

@@ -272,13 +272,23 @@ def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temp
     return final_response
 
 class SimpleDDGSearch:
-    def search(self, query: str, model: str = "claude-3-haiku") -> str:
-        result = DDGS().chat(
-            keywords=query,
-            model=model,
-            timeout=30
-        )
-        return result
+    def search(self, query: str, use_chat: bool = False, model: str = "claude-3-haiku", num_results: int = 10):
+        if use_chat:
+            result = DDGS().chat(
+                keywords=query,
+                model=model,
+                timeout=30
+            )
+            return [{"content": result, "url": "AI-generated response"}]
+        else:
+            results = DDGS().news(
+                keywords=query,
+                region='wt-wt',
+                safesearch='off',
+                timelimit='m',
+                max_results=num_results
+            )
+            return [{"content": res["body"], "url": res["url"]} for res in results]
 
 class TrafilaturaWebCrawler:
     def get_website_content_from_url(self, url: str) -> str:
@@ -506,13 +516,16 @@ def create_web_search_vectors(search_results):
 
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     searcher = SimpleDDGSearch()
-    search_results = searcher.search(query, num_results=10)
+    use_chat = model == "@cf/meta/llama-3.1-8b-instruct"  # Or however you want to determine when to use chat
+    search_results = searcher.search(query, use_chat=use_chat, model="claude-3-haiku" if use_chat else None, num_results=10)
     
-    crawler = TrafilaturaWebCrawler()
     context = ""
-    
-    for url in search_results:
-        context += crawler.get_website_content_from_url(url) + "\n"
+    for result in search_results:
+        if use_chat:
+            context += result['content']
+        else:
+            crawler = TrafilaturaWebCrawler()
+            context += crawler.get_website_content_from_url(result['url']) + "\n"
     
     prompt = f"""You are an expert AI named Sentinel and have been given a task to create a detailed and complete research article using the following context from web search results:
 {context} that fulfills the following user request: '{query}'