e63466845f
The ingest_directory function was recursively indexing ALL files including node_modules, which caused: - Massive index bloat (thousands of irrelevant library files) - Polluted search results with third-party code - Very slow indexing times Added SKIP_DIRECTORIES set with common directories to exclude: - Package managers: node_modules, vendor, Pods - Build outputs: dist, build, target, .next, .nuxt - VCS: .git, .svn, .hg - Caches: __pycache__, .pytest_cache, .mypy_cache, .cache - IDE: .vscode, .idea - Our own: .claude-os Also added should_skip_path() helper function with wildcard support for patterns like *.egg-info.
408 lines
12 KiB
Python
408 lines
12 KiB
Python
"""
|
|
Document ingestion pipeline for Claude OS.
|
|
Handles file upload, text extraction, chunking, embedding, and storage.
|
|
"""
|
|
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
import fitz # PyMuPDF
|
|
from llama_index.core import Document, Settings
|
|
from llama_index.core.node_parser import SentenceSplitter
|
|
from llama_index.embeddings.ollama import OllamaEmbedding
|
|
|
|
from app.core.sqlite_manager import get_sqlite_manager
|
|
from app.core.config import Config
|
|
from app.core.markdown_preprocessor import preprocess_markdown
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def extract_text_from_file(file_path: str) -> str:
|
|
"""
|
|
Extract text content from a file.
|
|
|
|
Args:
|
|
file_path: Path to the file
|
|
|
|
Returns:
|
|
str: Extracted text content
|
|
"""
|
|
file_path = Path(file_path)
|
|
extension = file_path.suffix.lower()
|
|
|
|
try:
|
|
if extension == ".pdf":
|
|
# Extract text from PDF using PyMuPDF
|
|
text = ""
|
|
with fitz.open(file_path) as doc:
|
|
for page in doc:
|
|
text += page.get_text()
|
|
return text
|
|
else:
|
|
# Read text/code files with UTF-8 encoding
|
|
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
|
|
return f.read()
|
|
except Exception as e:
|
|
logger.error(f"Failed to extract text from {file_path}: {e}")
|
|
return ""
|
|
|
|
|
|
def chunk_document(text: str, metadata: Dict) -> List[Document]:
|
|
"""
|
|
Chunk document text into smaller pieces with overlap.
|
|
|
|
Args:
|
|
text: Document text
|
|
metadata: Document metadata
|
|
|
|
Returns:
|
|
List of Document objects
|
|
"""
|
|
splitter = SentenceSplitter(
|
|
chunk_size=Config.CHUNK_SIZE,
|
|
chunk_overlap=Config.CHUNK_OVERLAP
|
|
)
|
|
|
|
# Create a single document
|
|
doc = Document(text=text, metadata=metadata)
|
|
|
|
# Split into chunks
|
|
nodes = splitter.get_nodes_from_documents([doc])
|
|
|
|
# Convert nodes back to documents
|
|
chunks = []
|
|
for i, node in enumerate(nodes):
|
|
chunk_metadata = metadata.copy()
|
|
chunk_metadata["chunk_index"] = i
|
|
chunk_metadata["chunk_id"] = node.node_id
|
|
chunks.append(Document(text=node.text, metadata=chunk_metadata))
|
|
|
|
return chunks
|
|
|
|
|
|
def ingest_file(
|
|
file_path: str,
|
|
collection_name: str,
|
|
filename: str
|
|
) -> Dict[str, any]:
|
|
"""
|
|
Ingest a single file into a knowledge base.
|
|
|
|
Args:
|
|
file_path: Path to the file
|
|
collection_name: Target collection name
|
|
filename: Original filename
|
|
|
|
Returns:
|
|
dict: Ingestion result with status and details
|
|
"""
|
|
try:
|
|
# Extract text
|
|
text = extract_text_from_file(file_path)
|
|
if not text.strip():
|
|
return {
|
|
"status": "error",
|
|
"filename": filename,
|
|
"error": "No text content extracted"
|
|
}
|
|
|
|
# Create base metadata
|
|
file_ext = Path(filename).suffix.lower()
|
|
metadata = {
|
|
"filename": filename,
|
|
"file_type": file_ext,
|
|
"upload_date": datetime.now().isoformat(),
|
|
"source_path": str(file_path)
|
|
}
|
|
|
|
# Preprocess markdown files
|
|
if file_ext in ['.md', '.markdown']:
|
|
try:
|
|
processed_text, enriched_metadata = preprocess_markdown(
|
|
text, filename, str(file_path)
|
|
)
|
|
text = processed_text
|
|
metadata.update(enriched_metadata)
|
|
logger.info(f"Preprocessed markdown: {filename}")
|
|
except Exception as e:
|
|
logger.warning(f"Markdown preprocessing failed for {filename}: {e}")
|
|
# Continue with original text if preprocessing fails
|
|
|
|
# Chunk document
|
|
chunks = chunk_document(text, metadata)
|
|
|
|
# Initialize embedding model
|
|
embed_model = OllamaEmbedding(
|
|
model_name=Config.OLLAMA_EMBED_MODEL,
|
|
base_url=Config.OLLAMA_HOST
|
|
)
|
|
|
|
# Get SQLite collection
|
|
pg_manager = get_sqlite_manager()
|
|
if not pg_manager.collection_exists(collection_name):
|
|
return {
|
|
"status": "error",
|
|
"filename": filename,
|
|
"error": f"Collection {collection_name} not found"
|
|
}
|
|
|
|
# Generate embeddings and add to collection
|
|
documents = []
|
|
embeddings = []
|
|
metadatas = []
|
|
ids = []
|
|
|
|
failed_chunks = 0
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
try:
|
|
# Truncate very long chunks to prevent Ollama crashes
|
|
chunk_text = chunk.text
|
|
if len(chunk_text) > 8000: # Limit to ~8K characters
|
|
logger.warning(f"Truncating chunk {i} of {filename} from {len(chunk_text)} to 8000 chars")
|
|
chunk_text = chunk_text[:8000]
|
|
|
|
# Generate embedding
|
|
embedding = embed_model.get_text_embedding(chunk_text)
|
|
|
|
# Create unique ID
|
|
chunk_id = f"{filename}_{chunk.metadata['chunk_index']}_{uuid.uuid4().hex[:8]}"
|
|
|
|
# Collect data for batch insert
|
|
documents.append(chunk_text)
|
|
embeddings.append(embedding)
|
|
metadatas.append(chunk.metadata)
|
|
ids.append(chunk_id)
|
|
except Exception as e:
|
|
failed_chunks += 1
|
|
logger.warning(f"Failed to embed chunk {i} of {filename}: {e}. Skipping this chunk.")
|
|
continue
|
|
|
|
if not documents:
|
|
return {
|
|
"status": "error",
|
|
"filename": filename,
|
|
"error": f"All {len(chunks)} chunks failed to generate embeddings"
|
|
}
|
|
|
|
# Add all documents to SQLite in batch
|
|
pg_manager.add_documents(
|
|
kb_name=collection_name,
|
|
documents=documents,
|
|
embeddings=embeddings,
|
|
metadatas=metadatas,
|
|
ids=ids
|
|
)
|
|
|
|
success_msg = f"Ingested {filename}: {len(documents)}/{len(chunks)} chunks"
|
|
if failed_chunks > 0:
|
|
success_msg += f" ({failed_chunks} chunks skipped due to errors)"
|
|
logger.info(success_msg)
|
|
|
|
return {
|
|
"status": "success",
|
|
"filename": filename,
|
|
"chunks": len(chunks),
|
|
"file_type": file_ext
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to ingest {filename}: {e}")
|
|
return {
|
|
"status": "error",
|
|
"filename": filename,
|
|
"error": str(e)
|
|
}
|
|
|
|
|
|
def ingest_documents(
|
|
collection_name: str,
|
|
documents: List[str],
|
|
metadatas: List[Dict]
|
|
) -> Dict[str, any]:
|
|
"""
|
|
Ingest a list of pre-processed documents into a knowledge base.
|
|
Used by hooks system for bulk ingestion.
|
|
|
|
Args:
|
|
collection_name: Target collection name
|
|
documents: List of document texts
|
|
metadatas: List of metadata dicts corresponding to each document
|
|
|
|
Returns:
|
|
dict: Ingestion result with status and details
|
|
"""
|
|
try:
|
|
# Initialize embedding model
|
|
embed_model = OllamaEmbedding(
|
|
model_name=Config.OLLAMA_EMBED_MODEL,
|
|
base_url=Config.OLLAMA_HOST
|
|
)
|
|
|
|
# Get SQLite manager
|
|
db_manager = get_sqlite_manager()
|
|
if not db_manager.collection_exists(collection_name):
|
|
return {
|
|
"status": "error",
|
|
"error": f"Collection {collection_name} not found",
|
|
"documents_processed": 0
|
|
}
|
|
|
|
# Process all documents
|
|
all_document_texts = []
|
|
all_embeddings = []
|
|
all_metadatas = []
|
|
all_ids = []
|
|
|
|
for doc_text, metadata in zip(documents, metadatas):
|
|
if not doc_text.strip():
|
|
logger.warning(f"Skipping empty document: {metadata.get('filename', 'unknown')}")
|
|
continue
|
|
|
|
# Chunk the document
|
|
chunks = chunk_document(doc_text, metadata)
|
|
|
|
# Generate embeddings for each chunk
|
|
for i, chunk in enumerate(chunks):
|
|
try:
|
|
# Truncate very long chunks to prevent Ollama crashes
|
|
chunk_text = chunk.text
|
|
if len(chunk_text) > 8000: # Limit to ~8K characters
|
|
logger.warning(f"Truncating chunk {i} from {len(chunk_text)} to 8000 chars")
|
|
chunk_text = chunk_text[:8000]
|
|
|
|
embedding = embed_model.get_text_embedding(chunk_text)
|
|
|
|
# Create unique ID
|
|
chunk_id = f"{metadata.get('filename', 'doc')}_{chunk.metadata['chunk_index']}_{uuid.uuid4().hex[:8]}"
|
|
|
|
# Collect data
|
|
all_document_texts.append(chunk_text)
|
|
all_embeddings.append(embedding)
|
|
all_metadatas.append(chunk.metadata)
|
|
all_ids.append(chunk_id)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to embed chunk {i}: {e}. Skipping this chunk.")
|
|
continue
|
|
|
|
if not all_document_texts:
|
|
return {
|
|
"status": "error",
|
|
"error": "No valid documents to ingest",
|
|
"documents_processed": 0
|
|
}
|
|
|
|
# Add all documents to database in batch
|
|
db_manager.add_documents(
|
|
kb_name=collection_name,
|
|
documents=all_document_texts,
|
|
embeddings=all_embeddings,
|
|
metadatas=all_metadatas,
|
|
ids=all_ids
|
|
)
|
|
|
|
logger.info(f"Ingested {len(documents)} documents into {collection_name}: {len(all_document_texts)} total chunks")
|
|
|
|
return {
|
|
"status": "success",
|
|
"documents_processed": len(documents),
|
|
"chunks_created": len(all_document_texts),
|
|
"collection_name": collection_name
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to ingest documents into {collection_name}: {e}")
|
|
return {
|
|
"status": "error",
|
|
"error": str(e),
|
|
"documents_processed": 0
|
|
}
|
|
|
|
|
|
# Directories to always skip during ingestion
|
|
SKIP_DIRECTORIES = {
|
|
'node_modules',
|
|
'.git',
|
|
'.svn',
|
|
'.hg',
|
|
'__pycache__',
|
|
'.pytest_cache',
|
|
'.mypy_cache',
|
|
'.tox',
|
|
'.nox',
|
|
'.eggs',
|
|
'*.egg-info',
|
|
'dist',
|
|
'build',
|
|
'.next',
|
|
'.nuxt',
|
|
'.output',
|
|
'coverage',
|
|
'.nyc_output',
|
|
'.cache',
|
|
'vendor',
|
|
'target', # Rust/Java
|
|
'Pods', # iOS
|
|
'.gradle',
|
|
'.idea',
|
|
'.vscode',
|
|
'.claude-os', # Our own cache
|
|
}
|
|
|
|
|
|
def should_skip_path(file_path: Path) -> bool:
|
|
"""Check if a file path should be skipped based on directory exclusions."""
|
|
for part in file_path.parts:
|
|
if part in SKIP_DIRECTORIES:
|
|
return True
|
|
# Handle wildcard patterns like *.egg-info
|
|
for pattern in SKIP_DIRECTORIES:
|
|
if '*' in pattern and part.endswith(pattern.replace('*', '')):
|
|
return True
|
|
return False
|
|
|
|
|
|
def ingest_directory(
|
|
dir_path: str,
|
|
collection_name: str
|
|
) -> List[Dict[str, any]]:
|
|
"""
|
|
Recursively ingest all supported files from a directory.
|
|
Automatically skips common non-source directories like node_modules, .git, etc.
|
|
|
|
Args:
|
|
dir_path: Path to directory
|
|
collection_name: Target collection name
|
|
|
|
Returns:
|
|
List of ingestion results
|
|
"""
|
|
results = []
|
|
dir_path = Path(dir_path)
|
|
|
|
if not dir_path.exists() or not dir_path.is_dir():
|
|
logger.error(f"Directory not found: {dir_path}")
|
|
return [{
|
|
"status": "error",
|
|
"error": f"Directory not found: {dir_path}"
|
|
}]
|
|
|
|
# Recursively find all supported files, skipping excluded directories
|
|
for file_path in dir_path.rglob("*"):
|
|
# Skip excluded directories
|
|
if should_skip_path(file_path):
|
|
continue
|
|
if file_path.is_file() and Config.is_supported_file(file_path.name):
|
|
result = ingest_file(
|
|
str(file_path),
|
|
collection_name,
|
|
file_path.name
|
|
)
|
|
results.append(result)
|
|
|
|
return results
|
|
|