Browse Source

update api endpoint and response format

david 1 month ago
parent
commit
dec7c7620e
1 changed files with 30 additions and 29 deletions
  1. 30 29
      main.py

+ 30 - 29
main.py

@@ -58,7 +58,7 @@ def extract_text_from_links(links, timeout=5):
     return extracted_texts
 
 
-def summarize_individual_texts(texts_and_urls, query, model, ollama_url="http://localhost:11434/api/generate"):
+def summarize_individual_texts(texts_and_urls, query, model, api_url="http://localhost:8000/api/v1/completions"):
     summaries = []
     for url, text in texts_and_urls:
         prompt = f"Extract the relevant information from the following text with regards to the original \
@@ -74,19 +74,21 @@ def summarize_individual_texts(texts_and_urls, query, model, ollama_url="http://
         }
         
         try:
-            response = requests.post(ollama_url, json=payload)
+            response = requests.post(api_url, json=payload)
             if response.status_code == 200:
-                result = json.loads(response.text)["response"]
+                #result = json.loads(response.text)["response"]
+                result_json = json.loads(response.text)
+                result = result_json["choices"][0]["text"].strip()                
                 summaries.append((url, result))
             else:
-                print(f"Failed to get summary from Ollama server for {url}. Status code: {response.status_code}")
+                print(f"Failed to get summary from server for {url}. Status code: {response.status_code}")
         except requests.RequestException as e:
-            print(f"An error occurred while sending request to Ollama server for {url}: {e}")
+            print(f"An error occurred while sending request to server for {url}: {e}")
     
     return summaries
 
 
-def summarize_with_ollama(texts_and_urls, query, model, ollama_url="http://localhost:11434/api/generate"):
+def summarize(texts_and_urls, query, model, api_url="http://localhost:8000/api/v1/completions"):
     # Prepare the context and prompt
     context = "\n".join([f"URL: {url}\nText: {text}" for url, text in texts_and_urls])
     prompt = f"Summarize the following search results with regards to the original query: '{query}' \
@@ -104,25 +106,27 @@ def summarize_with_ollama(texts_and_urls, query, model, ollama_url="http://local
 	}
     }
     
-    # Send the POST request to the Ollama server
+    # Send the POST request to the server
     try:
         print("Processing")
-        response = requests.post(ollama_url, json=payload)
+        response = requests.post(api_url, json=payload)
         if response.status_code == 200:
-            result = json.loads(response.text)["response"]
-            return result
+            #result = json.loads(response.text)["response"]
+            #return result            
+            result = json.loads(response.text)            
+            return(result["choices"][0]["text"].strip())
         else:
-            print(f"Failed to get summary from Ollama server. Status code: {response.status_code}")
+            print(f"Failed to get summary from the server. Status code: {response.status_code}")
             return None
     except requests.RequestException as e:
-        print(f"An error occurred while sending request to Ollama server: {e}")
+        print(f"An error occurred while sending request to the server: {e}")
         return None
 
 
-def optimize_search_query(query, query_model, ollama_url="http://localhost:11434/api/generate"):
+def optimize_search_query(query, query_model, api_url="http://localhost:8000/api/v1/completions"):
     # Prepare the prompt for optimizing the search query
     prompt = f"Optimize the following natural language query to improve its effectiveness in a web search.\
-        Make it very concise. Return just the optimized query no explanations. Query: '{query}'"
+        Make it very concise. Return only the optimized query text no additional texts, quotations or thoughts. Query: '{query}'"
     
     # Create the payload for the POST request
     payload = {        
@@ -132,18 +136,18 @@ def optimize_search_query(query, query_model, ollama_url="http://localhost:11434
         "max_tokens": 50
     }
     
-    # Send the POST request to the Ollama server
+    # Send the POST request to the server
     try:
         print("Optimizing search query")
-        response = requests.post(ollama_url, json=payload)
+        response = requests.post(api_url, json=payload)
         if response.status_code == 200:
-            result = json.loads(response.text)["response"].strip()
-            return result.strip('"')
+            result = json.loads(response.text)            
+            return(result["choices"][0]["text"].strip())
         else:
-            print(f"Failed to optimize search query from Ollama server. Status code: {response.status_code}")
+            print(f"Failed to optimize search query from the server. Status code: {response.status_code}")
             return query
     except requests.RequestException as e:
-        print(f"An error occurred while sending request to Ollama server for optimizing the search query: {e}")
+        print(f"An error occurred while sending request to the server for optimizing the search query: {e}")
         return query
 
 
@@ -155,7 +159,7 @@ def pretty_print_markdown(markdown_text):
 
 if __name__ == "__main__":
     # Set up argument parser
-    parser = argparse.ArgumentParser(description="Search DuckDuckGo, extract text from results, and summarize with Ollama.")
+    parser = argparse.ArgumentParser(description="Search DuckDuckGo, extract text from results, and summarize with LLM.")
     parser.add_argument("query", type=str, help="The search query to use on DuckDuckGo")
     parser.add_argument("--num_results", type=int, default=5, help="Number of search results to process (default: 5)")
     
@@ -163,10 +167,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
     
     original_query = args.query
-    model = "command-r"
-    #model = "qwq"
-    #model = "qwen2.5:32b"
-    query_model = model
+    query_model = "Gemma-3-4b-it-GGUF"    
 
     # Optimize the search query
     optimized_query = optimize_search_query(original_query, query_model)
@@ -180,11 +181,11 @@ if __name__ == "__main__":
     for i, link in enumerate(links, start=1):
         print(f"{i}. {link}")
     
-    texts_and_urls = extract_text_from_links(links)
-    
+    texts_and_urls = extract_text_from_links(links)    
+    summary_model = "Gemma-3-4b-it-GGUF"
     print("Summarizing individual search results")
-    intermediate_summaries = summarize_individual_texts(texts_and_urls, original_query, model)
-    final_summary = summarize_with_ollama(intermediate_summaries, original_query, model)
+    intermediate_summaries = summarize_individual_texts(texts_and_urls, original_query, summary_model)
+    final_summary = summarize(intermediate_summaries, original_query, summary_model)
     
     if final_summary:
         print("\nFinal Summary of search results:\n")