#這個文件將負責所有的數據準備工作。在運行 app.py 之前，你需要先執行這個文件一次。 import re import os import sqlite3 import datetime import chromadb from pypdf import PdfReader from langchain.text_splitter import RecursiveCharacterTextSplitter from together import Together from langchain_together import TogetherEmbeddings from dotenv import load_dotenv from chromadb.utils import embedding_functions from langchain_core.documents import Document # Load environment variables from .env file load_dotenv() def extract_text_from_pdf(pdf_path): """ Extract all text content from PDF files. """ text = "" try: reader = PdfReader(pdf_path) for page in reader.pages: text += page.extract_text() + "\n" # 提取每頁文字並換行 except Exception as e: print(f"Error reading PDF {pdf_path}: {e}") return text def clean_text_content(text): """ Remove irrelevant content, such as copyright information, legal notices, and long blank spaces. """ text = re.sub(r'\n\s*\n', '\n\n', text) text = re.sub(r'^\s*\d+\s*$', '', text, flags=re.MULTILINE) text = re.sub(r'\s+\d+\s*$', '', text, flags=re.MULTILINE) text = re.sub(r'Copyright © \d{4} [^\n]*\. All Rights Reserved\.', '', text, flags=re.IGNORECASE) text = re.sub(r'Confidential and Proprietary Information[^\n]*', '', text, flags=re.IGNORECASE) text = re.sub(r'Disclaimer:[^.]*\.', '', text, flags=re.IGNORECASE) text = re.sub(r'\s+', ' ', text).strip() return text def standardize_text_format(text): """ Try to organize everything into clear paragraphs. """ text = re.sub(r'([a-zA-Z])-(\n)([a-zA-Z])', r'\1\3', text) text = re.sub(r'([.?!])\s*([A-Z])', r'\1 \2', text) text = re.sub(r'\s+', ' ', text).strip() text = re.sub(r'([^\n])\n([^\n])', r'\1 \2', text) text = re.sub(r'\n{2,}', '\n\n', text) return text def create_database_table(db_path): """ Create a SQLite database and define the table structure. """ conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS documents ( id INTEGER PRIMARY KEY AUTOINCREMENT, original_filename TEXT NOT NULL UNIQUE, -- Add UNIQUE constraint source_type TEXT, processed_text TEXT NOT NULL, processed_date TEXT ) ''') conn.commit() conn.close() print(f"SQLite table 'documents' ensured in {db_path}") def create_chunks_table(db_path): """ Create a table to store the segmented text blocks. """ conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS chunks ( id INTEGER PRIMARY KEY AUTOINCREMENT, document_id INTEGER NOT NULL, chunk_index INTEGER NOT NULL, chunk_content TEXT NOT NULL, chunk_length INTEGER, created_at TEXT, UNIQUE(document_id, chunk_index), -- Ensure chunks are unique per document FOREIGN KEY (document_id) REFERENCES documents(id) ) ''') conn.commit() conn.close() print(f"SQLite table 'chunks' ensured in {db_path}") def insert_document_data(db_path, original_filename, source_type, processed_text, processed_date): """ Insert the processed file data into the database. """ conn = sqlite3.connect(db_path) cursor = conn.cursor() # 檢查是否已存在相同的原始文件 cursor.execute("SELECT id FROM documents WHERE original_filename = ?", (original_filename,)) existing_doc = cursor.fetchone() if existing_doc: print(f"Document '{original_filename}' already exists in DB (ID: {existing_doc[0]}). Skipping insertion.") conn.close() return existing_doc[0] else: cursor.execute(''' INSERT INTO documents (original_filename, source_type, processed_text, processed_date) VALUES (?, ?, ?, ?) ''', (original_filename, source_type, processed_text, processed_date)) conn.commit() doc_id = cursor.lastrowid conn.close() print(f"Document '{original_filename}' saved to database: {db_path} with ID: {doc_id}") return doc_id def get_document_text_from_db(db_path, document_id=None, limit=None): """ Reads the processed text contents of one or more files from a SQLite database. """ conn = None try: conn = sqlite3.connect(db_path) cursor = conn.cursor() query = "SELECT id, processed_text FROM documents" params = [] if document_id: query += " WHERE id = ?" params.append(document_id) if limit: query += " LIMIT ?" params.append(limit) cursor.execute(query, params) if document_id: row = cursor.fetchone() if row: return {'id': row[0], 'processed_text': row[1]} return None else: return [{'id': row[0], 'processed_text': row[1]} for row in cursor.fetchall()] except sqlite3.Error as e: print(f"Error reading from SQLite database: {e}") return [] finally: if conn: conn.close() def chunk_text(text, chunk_size=500, chunk_overlap=100): # Adjusted chunk size and overlap for potentially smaller chunks """ Splits the text using RecursiveCharacterTextSplitter. """ text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, add_start_index=True ) chunks = text_splitter.create_documents([text]) return [chunk.page_content for chunk in chunks] def insert_chunks_to_db(db_path, document_id, chunks): """ Inserts all text blocks from the specified file into the database. Checks if chunks for this document_id already exist and skips if found. """ conn = sqlite3.connect(db_path) cursor = conn.cursor() current_time = datetime.datetime.now().isoformat() # Check if chunks for this document_id already exist cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (document_id,)) if cursor.fetchone()[0] > 0: print(f"Chunks for document ID {document_id} already exist. Skipping chunk insertion.") conn.close() return data_to_insert = [] for i, chunk in enumerate(chunks): data_to_insert.append((document_id, i, chunk, len(chunk), current_time)) cursor.executemany(''' INSERT INTO chunks (document_id, chunk_index, chunk_content, chunk_length, created_at) VALUES (?, ?, ?, ?, ?) ''', data_to_insert) conn.commit() conn.close() print(f"Saved {len(chunks)} chunks into the database for file ID {document_id}.") def generate_embeddings(texts, model_name="togethercomputer/m2-bert-80M-32k-retrieval"): """ Generates text embeddings using Together AI's embedding model. """ try: embeddings_model = TogetherEmbeddings( model=model_name, api_key=os.getenv("TOGETHER_API_KEY") ) vectors = embeddings_model.embed_documents(texts) print(f"Successfully generated embeddings for {len(texts)} text chunks using {model_name}.") return vectors except Exception as e: print(f"Error generating embeddings: {e}") return [] def get_chunks_from_db_for_embedding(db_path): """ Reads all split text chunks from the SQLite database, including their IDs. """ conn = None try: conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute("SELECT id, document_id, chunk_index, chunk_content FROM chunks ORDER BY document_id, chunk_index") rows = cursor.fetchall() return [{'id': row[0], 'source_document_id': row[1], 'chunk_index': row[2], 'text': row[3]} for row in rows] except sqlite3.Error as e: print(f"Error reading chunks from SQLite database: {e}") return [] finally: if conn: conn.close() def load_chunks_to_vector_db(chunks_data, db_path="data/vector_db_chroma", collection_name="document_chunks", embeddings_model_name="togethercomputer/m2-bert-80M-32k-retrieval"): """ Loads text chunks and their embeddings into a ChromaDB vector database. This function will now ADD chunks if they are new (based on their IDs). """ try: client = chromadb.PersistentClient(path=db_path) collection = client.get_or_create_collection(name=collection_name) # Check existing IDs in ChromaDB to avoid adding duplicates existing_chroma_ids = set() if collection.count() > 0: # Fetching all existing IDs can be slow for very large collections. # A more efficient approach for very large databases might involve # querying a batch of IDs or checking after insertion. # For now, this is simpler for demonstration. try: all_ids_in_chroma = collection.get(ids=collection.get()['ids'])['ids'] existing_chroma_ids = set(all_ids_in_chroma) except Exception as e: print(f"Warning: Could not retrieve all existing IDs from ChromaDB. May attempt to add duplicates. Error: {e}") ids_to_add = [] documents_to_add = [] embeddings_to_add = [] metadatas_to_add = [] for item in chunks_data: chunk_id_str = str(item['id']) if chunk_id_str not in existing_chroma_ids: ids_to_add.append(chunk_id_str) documents_to_add.append(item['text']) embeddings_to_add.append(item['embedding']) metadatas_to_add.append({"source_document_id": item.get('source_document_id', 'unknown'), "chunk_index": item.get('chunk_index', None)}) else: print(f"Chunk ID {chunk_id_str} already exists in ChromaDB. Skipping.") if ids_to_add: collection.add( embeddings=embeddings_to_add, documents=documents_to_add, metadatas=metadatas_to_add, ids=ids_to_add ) print(f"Successfully added {len(ids_to_add)} new text chunks into ChromaDB collection '{collection_name}'.") else: print(f"No new chunks to add to ChromaDB collection '{collection_name}'.") except Exception as e: print(f"Error loading data into ChromaDB: {e}") if __name__ == "__main__": # Define the directory containing your PDF files current_dir = os.path.dirname(os.path.abspath(__file__)) pdf_input_directory = os.path.join(current_dir, "input") #pdf_input_directory = "input" # Ensure this directory exists and contains your PDFs db_directory = "database" db_path = os.path.join(db_directory, "processed_documents.db") vector_db_dir = "data/vector_db_chroma" collection_name = "my_instrument_manual_chunks" embeddings_model_name = "togethercomputer/m2-bert-80M-32k-retrieval" # 確保與 main.py 中使用的一致 # Create necessary directories os.makedirs(db_directory, exist_ok=True) os.makedirs(vector_db_dir, exist_ok=True) print("--- Starting knowledge base setup ---") # 1. 確保 SQLite 表格存在 create_database_table(db_path) create_chunks_table(db_path) # 2. 遍歷指定目錄下的所有 PDF 文件 pdf_files = [f for f in os.listdir(pdf_input_directory) if f.lower().endswith('.pdf')] if not pdf_files: print(f"No PDF files found in '{pdf_input_directory}'. Please place your PDF documents there.") for pdf_filename in pdf_files: pdf_full_path = os.path.join(pdf_input_directory, pdf_filename) print(f"\n--- Processing PDF: {pdf_full_path} ---") # 檢查文件是否已在 documents 表中處理過 conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute("SELECT id FROM documents WHERE original_filename = ?", (pdf_filename,)) existing_doc_id = cursor.fetchone() conn.close() if existing_doc_id: doc_id = existing_doc_id[0] print(f"Document '{pdf_filename}' already in 'documents' table with ID: {doc_id}. Skipping PDF extraction and text processing.") # Still process chunks if they aren't in 'chunks' table or ChromaDB else: pdf_content = extract_text_from_pdf(pdf_full_path) if pdf_content: cleaned_text = clean_text_content(pdf_content) final_processed_text = standardize_text_format(cleaned_text) current_date = datetime.date.today().isoformat() doc_id = insert_document_data(db_path, pdf_filename, "PDF", final_processed_text, current_date) else: print(f"No content extracted from {pdf_full_path}. Skipping document insertion for this PDF.") continue # Move to the next PDF if no content # 如果 doc_id 存在（無論是新插入還是已存在的），則進行分塊和嵌入處理 if doc_id is not None: # 從 SQLite documents 表讀取並分塊，存儲到 SQLite chunks 表 document_from_db = get_document_text_from_db(db_path, document_id=doc_id) if document_from_db: full_text = document_from_db['processed_text'] print(f"--- Chunking document ID {doc_id} ('{pdf_filename}') ---") chunks = chunk_text(full_text, chunk_size=500, chunk_overlap=100) # Re-evaluate chunk_size if chunks: insert_chunks_to_db(db_path, doc_id, chunks) else: print(f"Document ID {doc_id} unable to split any chunks.") else: print(f"Could not retrieve processed text for document ID {doc_id}.") # 3. 從 SQLite chunks 表讀取所有塊並生成嵌入，載入到 ChromaDB # 我們需要重新從 DB 獲取所有 chunks，因為可能有多個 PDF 的 chunks chunks_from_db_for_embedding = get_chunks_from_db_for_embedding(db_path) if chunks_from_db_for_embedding: # Filter out chunks that already have an embedding in ChromaDB if checking was efficient # For simplicity, we'll try to generate embeddings for all chunks retrieved from SQLite # and let load_chunks_to_vector_db handle duplicates in ChromaDB. texts_to_embed = [item['text'] for item in chunks_from_db_for_embedding] print(f"\n--- Generating embeddings for {len(texts_to_embed)} total chunks ---") text_embeddings = generate_embeddings(texts_to_embed, model_name=embeddings_model_name) if text_embeddings: data_for_vector_db = [] for i, chunk_data in enumerate(chunks_from_db_for_embedding): if i < len(text_embeddings): chunk_data['embedding'] = text_embeddings[i] data_for_vector_db.append(chunk_data) else: print(f"Warning: Missing embedding for chunk {chunk_data['id']}. Skipping this chunk for ChromaDB.") print(f"--- Loading/Updating {len(data_for_vector_db)} chunks into ChromaDB ---") load_chunks_to_vector_db(data_for_vector_db, db_path=vector_db_dir, collection_name=collection_name, embeddings_model_name=embeddings_model_name) else: print("No embeddings generated. Skipping ChromaDB loading.") else: print("No text chunks read from the database for embedding and loading into ChromaDB.") print("\n--- Knowledge base setup complete ---")