Upload ask.py

ask.py

@@ -19,7 +19,7 @@ script_dir = os.path.dirname(os.path.abspath(__file__))
 default_env_file = os.path.abspath(os.path.join(script_dir, ".env"))
 
 
-def _get_logger(log_level: str) -> logging.Logger:
+def get_logger(log_level: str) -> logging.Logger:
     logger = logging.getLogger(__name__)
     logger.setLevel(log_level)
     handler = logging.StreamHandler()
@@ -29,20 +29,6 @@ def _get_logger(log_level: str) -> logging.Logger:
     return logger
 
 
-def _read_url_list(url_list_file: str) -> str:
-    if url_list_file is None:
-        return None
-
-    with open(url_list_file, "r") as f:
-        links = f.readlines()
-    links = [
-        link.strip()
-        for link in links
-        if link.strip() != "" and not link.startswith("#")
-    ]
-    return "\n".join(links)
-
-
 class Ask:
 
     def __init__(self, logger: Optional[logging.Logger] = None):
@@ -51,7 +37,7 @@ class Ask:
         if logger is not None:
             self.logger = logger
         else:
-            self.logger = _get_logger("INFO")
+            self.logger = get_logger("INFO")
 
         self.table_name = "document_chunks"
         self.db_con = duckdb.connect(":memory:")
@@ -397,66 +383,84 @@ Here is the context:
         response_str = completion.choices[0].message.content
         return response_str
 
-    def run_query(
-        self,
-        query: str,
-        date_restrict: int,
-        target_site: str,
-        output_language: str,
-        output_length: int,
-        url_list_str: str,
-        model_name: str,
-    ) -> str:
-        logger = self.logger
-        if url_list_str is None or url_list_str.strip() == "":
-            logger.info("Searching the web ...")
-            links = self.search_web(query, date_restrict, target_site)
-            logger.info(f"✅ Found {len(links)} links for query: {query}")
-            for i, link in enumerate(links):
-                logger.debug(f"{i+1}. {link}")
-        else:
-            links = url_list_str.split("\n")
-
-        logger.info("Scraping the URLs ...")
-        scrape_results = self.scrape_urls(links)
-        logger.info(f"✅ Scraped {len(scrape_results)} URLs.")
-
-        logger.info("Chunking the text ...")
-        chunking_results = self.chunk_results(scrape_results, 1000, 100)
-        total_chunks = 0
-        for url, chunks in chunking_results.items():
-            logger.debug(f"URL: {url}")
-            total_chunks += len(chunks)
-            for i, chunk in enumerate(chunks):
-                logger.debug(f"Chunk {i+1}: {chunk}")
-        logger.info(f"✅ Generated {total_chunks} chunks ...")
-
-        logger.info(f"Saving {total_chunks} chunks to DB ...")
-        self.save_to_db(chunking_results)
-        logger.info(f"✅ Successfully embedded and saved chunks to DB.")
-
-        logger.info("Querying the vector DB to get context ...")
-        matched_chunks = self.vector_search(query)
-        for i, result in enumerate(matched_chunks):
-            logger.debug(f"{i+1}. {result}")
-        logger.info(f"✅ Got {len(matched_chunks)} matched chunks.")
-
-        logger.info("Running inference with context ...")
-        answer = self.run_inference(
-            query=query,
-            model_name=model_name,
-            matched_chunks=matched_chunks,
-            output_language=output_language,
-            output_length=output_length,
-        )
-        logger.info("✅ Finished inference API call.")
-        logger.info("generateing output ...")
 
-        answer = f"# Answer\n\n{answer}\n"
-        references = "\n".join(
-            [f"[{i+1}] {result['url']}" for i, result in enumerate(matched_chunks)]
-        )
-        return f"{answer}\n\n# References\n\n{references}"
+def _read_url_list(url_list_file: str) -> str:
+    if url_list_file is None:
+        return None
+
+    with open(url_list_file, "r") as f:
+        links = f.readlines()
+    links = [
+        link.strip()
+        for link in links
+        if link.strip() != "" and not link.startswith("#")
+    ]
+    return "\n".join(links)
+
+
+def _run_query(
+    query: str,
+    date_restrict: int,
+    target_site: str,
+    output_language: str,
+    output_length: int,
+    url_list_str: str,
+    model_name: str,
+    log_level: str,
+) -> str:
+    logger = get_logger(log_level)
+
+    ask = Ask(logger=logger)
+
+    if url_list_str is None or url_list_str.strip() == "":
+        logger.info("Searching the web ...")
+        links = ask.search_web(query, date_restrict, target_site)
+        logger.info(f"✅ Found {len(links)} links for query: {query}")
+        for i, link in enumerate(links):
+            logger.debug(f"{i+1}. {link}")
+    else:
+        links = url_list_str.split("\n")
+
+    logger.info("Scraping the URLs ...")
+    scrape_results = ask.scrape_urls(links)
+    logger.info(f"✅ Scraped {len(scrape_results)} URLs.")
+
+    logger.info("Chunking the text ...")
+    chunking_results = ask.chunk_results(scrape_results, 1000, 100)
+    total_chunks = 0
+    for url, chunks in chunking_results.items():
+        logger.debug(f"URL: {url}")
+        total_chunks += len(chunks)
+        for i, chunk in enumerate(chunks):
+            logger.debug(f"Chunk {i+1}: {chunk}")
+    logger.info(f"✅ Generated {total_chunks} chunks ...")
+
+    logger.info(f"Saving {total_chunks} chunks to DB ...")
+    ask.save_to_db(chunking_results)
+    logger.info(f"✅ Successfully embedded and saved chunks to DB.")
+
+    logger.info("Querying the vector DB to get context ...")
+    matched_chunks = ask.vector_search(query)
+    for i, result in enumerate(matched_chunks):
+        logger.debug(f"{i+1}. {result}")
+    logger.info(f"✅ Got {len(matched_chunks)} matched chunks.")
+
+    logger.info("Running inference with context ...")
+    answer = ask.run_inference(
+        query=query,
+        model_name=model_name,
+        matched_chunks=matched_chunks,
+        output_language=output_language,
+        output_length=output_length,
+    )
+    logger.info("✅ Finished inference API call.")
+    logger.info("generateing output ...")
+
+    answer = f"# Answer\n\n{answer}\n"
+    references = "\n".join(
+        [f"[{i+1}] {result['url']}" for i, result in enumerate(matched_chunks)]
+    )
+    return f"{answer}\n\n# References\n\n{references}"
 
 
 def launch_gradio(
@@ -467,12 +471,11 @@ def launch_gradio(
     output_length: int,
     url_list_str: str,
     model_name: str,
+    log_level: str,
     share_ui: bool,
-    logger: logging.Logger,
 ) -> None:
-    ask = Ask(logger=logger)
     iface = gr.Interface(
-        fn=ask.run_query,
+        fn=_run_query,
         inputs=[
             gr.Textbox(label="Query", value=query),
             gr.Number(
@@ -500,6 +503,7 @@ def launch_gradio(
         ],
         additional_inputs=[
             gr.Textbox(label="Model Name", value=model_name),
+            gr.Textbox(label="Log Level", value=log_level),
         ],
         outputs="text",
         show_progress=True,
@@ -582,7 +586,6 @@ def search_extract_summarize(
     log_level: str,
 ):
     load_dotenv(dotenv_path=default_env_file, override=False)
-    logger = _get_logger(log_level)
 
     if web_ui or os.environ.get("RUN_GRADIO_UI", "false").lower() != "false":
         if os.environ.get("SHARE_GRADIO_UI", "false").lower() == "true":
@@ -597,14 +600,14 @@ def search_extract_summarize(
             output_length=output_length,
             url_list_str=_read_url_list(url_list_file),
             model_name=model_name,
+            log_level=log_level,
             share_ui=share_ui,
-            logger=logger,
         )
     else:
         if query is None:
             raise Exception("Query is required for the command line mode")
-        ask = Ask(logger=logger)
-        result = ask.run_query(
+
+        result = _run_query(
             query=query,
             date_restrict=date_restrict,
             target_site=target_site,
@@ -612,6 +615,7 @@ def search_extract_summarize(
             output_length=output_length,
             url_list_str=_read_url_list(url_list_file),
             model_name=model_name,
+            log_level=log_level,
         )
         click.echo(result)