Update local_knowledge_base.py

local_knowledge_base.py

@@ -1,50 +1,50 @@
-from phi.vectordb.lancedb import LanceDb
-from phi.knowledge.pdf import PDFKnowledgeBase, PDFReader
-from phi.embedder.google import GeminiEmbedder
-from phi.vectordb.search import SearchType
-from phi.utils.log import logger
-import os
-
-def load_knowledge_base():
-    """
-    Loads or creates a knowledge base from PDF documents in the 'knowledge' folder
-    using ChromaDB for storage.
-
-    This version includes a manual loop to process one file at a time,
-    which is a robust workaround for a bug in the library's multi-file handling.
-    """
-    knowledge_dir = "./knowledge"
-    db_dir = "./lance_db"
-    table_name = "local_pdf_knowledge" 
-
-    if not os.path.exists(knowledge_dir) or not os.listdir(knowledge_dir):
-        logger.warning(f"The '{knowledge_dir}' directory is empty or does not exist. No local knowledge base will be loaded.")
-        return None
-
-    # Get a list of all PDF files to process
-    pdf_files = [f for f in os.listdir(knowledge_dir) if f.lower().endswith(".pdf")]
-    if not pdf_files:
-        logger.warning(f"No PDF files found in the '{knowledge_dir}' directory.")
-        return None
-
-    logger.info("Loading Knowledge Base using LanceDb with manual file processing...")
-    
-    try:
-        knowledge_base = PDFKnowledgeBase(
-                path=knowledge_dir, 
-                vector_db=LanceDb(
-                    table_name=table_name,
-                    uri=db_dir,
-                    embedder=GeminiEmbedder(model="models/text-embedding-004"),
-                    search_type=SearchType.keyword,
-                ),
-            reader=PDFReader(chunk=True)  
-            )
-        
-        logger.info("All files processed successfully.")
-
-        return knowledge_base
-
-    except Exception as e:
-        logger.error(f"An unexpected error occurred during manual file loading: {e}")
+from phi.vectordb.lancedb import LanceDb
+from phi.knowledge.pdf import PDFKnowledgeBase, PDFReader
+from phi.embedder.google import GeminiEmbedder
+from phi.vectordb.search import SearchType
+from phi.utils.log import logger
+import os
+
+def load_knowledge_base():
+    """
+    Loads or creates a knowledge base from PDF documents in the 'knowledge' folder
+    using LanceDB for storage.
+
+    This version includes a manual loop to process one file at a time,
+    which is a robust workaround for a bug in the library's multi-file handling.
+    """
+    knowledge_dir = "./knowledge"
+    db_dir = "./lance_db"
+    table_name = "local_pdf_knowledge" 
+
+    if not os.path.exists(knowledge_dir) or not os.listdir(knowledge_dir):
+        logger.warning(f"The '{knowledge_dir}' directory is empty or does not exist. No local knowledge base will be loaded.")
+        return None
+
+    # Get a list of all PDF files to process
+    pdf_files = [f for f in os.listdir(knowledge_dir) if f.lower().endswith(".pdf")]
+    if not pdf_files:
+        logger.warning(f"No PDF files found in the '{knowledge_dir}' directory.")
+        return None
+
+    logger.info("Loading Knowledge Base using LanceDb with manual file processing...")
+    
+    try:
+        knowledge_base = PDFKnowledgeBase(
+                path=knowledge_dir, 
+                vector_db=LanceDb(
+                    table_name=table_name,
+                    uri=db_dir,
+                    embedder=GeminiEmbedder(model="models/text-embedding-004"),
+                    search_type=SearchType.keyword,
+                ),
+            reader=PDFReader(chunk=True)  
+            )
+        
+        logger.info("All files processed successfully.")
+
+        return knowledge_base
+
+    except Exception as e:
+        logger.error(f"An unexpected error occurred during manual file loading: {e}")
         return None