File size: 1,860 Bytes
7399ef0 a603adb 7399ef0 1b213ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 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 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 = "./vectordb/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 |