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,6 +29,20 @@ 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):
@@ -37,7 +51,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:")
@@ -383,84 +397,66 @@ 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 ...")
 
-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}"
+        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(
@@ -471,11 +467,12 @@ 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=_run_query,
+        fn=ask.run_query,
         inputs=[
             gr.Textbox(label="Query", value=query),
             gr.Number(
@@ -503,7 +500,6 @@ 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,
@@ -586,6 +582,7 @@ 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":
@@ -600,14 +597,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")
-
-        result = _run_query(
+        ask = Ask(logger=logger)
+        result = ask.run_query(
             query=query,
             date_restrict=date_restrict,
             target_site=target_site,
@@ -615,7 +612,6 @@ 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)