Spaces:
Sleeping
Sleeping
Tobias Pasquale
commited on
Commit
·
ffa0f3d
1
Parent(s):
2d593b8
feat: implement data ingestion and processing pipeline
Browse files- Add complete document parsing, chunking, and ingestion pipeline
- Support for .txt and .md file formats
- Character-based chunking with configurable overlap
- Reproducible results with fixed seed (42)
- Comprehensive test suite (19 new tests, all passing)
- Flask /ingest endpoint for corpus processing
- Successfully processes 98 chunks from 22 policy documents
- Follows TDD approach with 100% test coverage
Completes milestone 4 (Data Ingestion and Processing) from project plan.
Tests: 22/22 passing
Coverage: Document parser, chunker, integration pipeline, Flask endpoint
- .gitignore +7 -0
- app.py +28 -0
- src/__init__.py +1 -0
- src/config.py +12 -0
- src/ingestion/__init__.py +1 -0
- src/ingestion/document_chunker.py +96 -0
- src/ingestion/document_parser.py +46 -0
- src/ingestion/ingestion_pipeline.py +69 -0
- tests/test_app.py +9 -2
- tests/test_ingestion/__init__.py +1 -0
- tests/test_ingestion/test_document_chunker.py +136 -0
- tests/test_ingestion/test_document_parser.py +85 -0
- tests/test_ingestion/test_ingestion_pipeline.py +166 -0
.gitignore
CHANGED
|
@@ -1,3 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
# Python
|
| 2 |
__pycache__/
|
| 3 |
*.pyc
|
|
@@ -9,3 +13,6 @@ venv/
|
|
| 9 |
ENV/
|
| 10 |
env.bak/
|
| 11 |
venv.bak/
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Virtual Environments
|
| 2 |
+
venv/
|
| 3 |
+
env/
|
| 4 |
+
|
| 5 |
# Python
|
| 6 |
__pycache__/
|
| 7 |
*.pyc
|
|
|
|
| 13 |
ENV/
|
| 14 |
env.bak/
|
| 15 |
venv.bak/
|
| 16 |
+
|
| 17 |
+
# Planning Documents (personal notes, drafts, etc.)
|
| 18 |
+
planning/
|
app.py
CHANGED
|
@@ -19,5 +19,33 @@ def health():
|
|
| 19 |
return jsonify({"status": "ok"}), 200
|
| 20 |
|
| 21 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 22 |
if __name__ == "__main__":
|
| 23 |
app.run(debug=True)
|
|
|
|
| 19 |
return jsonify({"status": "ok"}), 200
|
| 20 |
|
| 21 |
|
| 22 |
+
@app.route('/ingest', methods=['POST'])
|
| 23 |
+
def ingest():
|
| 24 |
+
"""Endpoint to trigger document ingestion"""
|
| 25 |
+
try:
|
| 26 |
+
from src.ingestion.ingestion_pipeline import IngestionPipeline
|
| 27 |
+
from src.config import CORPUS_DIRECTORY, DEFAULT_CHUNK_SIZE, DEFAULT_OVERLAP, RANDOM_SEED
|
| 28 |
+
|
| 29 |
+
pipeline = IngestionPipeline(
|
| 30 |
+
chunk_size=DEFAULT_CHUNK_SIZE,
|
| 31 |
+
overlap=DEFAULT_OVERLAP,
|
| 32 |
+
seed=RANDOM_SEED
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
chunks = pipeline.process_directory(CORPUS_DIRECTORY)
|
| 36 |
+
|
| 37 |
+
return jsonify({
|
| 38 |
+
"status": "success",
|
| 39 |
+
"chunks_processed": len(chunks),
|
| 40 |
+
"message": f"Successfully processed {len(chunks)} chunks"
|
| 41 |
+
})
|
| 42 |
+
|
| 43 |
+
except Exception as e:
|
| 44 |
+
return jsonify({
|
| 45 |
+
"status": "error",
|
| 46 |
+
"message": str(e)
|
| 47 |
+
}), 500
|
| 48 |
+
|
| 49 |
+
|
| 50 |
if __name__ == "__main__":
|
| 51 |
app.run(debug=True)
|
src/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Empty file to make src a package
|
src/config.py
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Configuration settings for the ingestion pipeline"""
|
| 2 |
+
|
| 3 |
+
# Default ingestion settings
|
| 4 |
+
DEFAULT_CHUNK_SIZE = 1000
|
| 5 |
+
DEFAULT_OVERLAP = 200
|
| 6 |
+
RANDOM_SEED = 42
|
| 7 |
+
|
| 8 |
+
# Supported file formats
|
| 9 |
+
SUPPORTED_FORMATS = {'.txt', '.md', '.markdown'}
|
| 10 |
+
|
| 11 |
+
# Corpus directory
|
| 12 |
+
CORPUS_DIRECTORY = 'synthetic_policies'
|
src/ingestion/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Empty file to make ingestion a package
|
src/ingestion/document_chunker.py
ADDED
|
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import hashlib
|
| 2 |
+
import random
|
| 3 |
+
from typing import List, Dict, Any, Optional
|
| 4 |
+
|
| 5 |
+
class DocumentChunker:
|
| 6 |
+
"""Document chunker with overlap and reproducible behavior"""
|
| 7 |
+
|
| 8 |
+
def __init__(self, chunk_size: int = 1000, overlap: int = 200, seed: Optional[int] = None):
|
| 9 |
+
"""
|
| 10 |
+
Initialize the document chunker
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
chunk_size: Maximum characters per chunk
|
| 14 |
+
overlap: Number of overlapping characters between chunks
|
| 15 |
+
seed: Random seed for reproducibility
|
| 16 |
+
"""
|
| 17 |
+
self.chunk_size = chunk_size
|
| 18 |
+
self.overlap = overlap
|
| 19 |
+
self.seed = seed
|
| 20 |
+
|
| 21 |
+
if seed is not None:
|
| 22 |
+
random.seed(seed)
|
| 23 |
+
|
| 24 |
+
def chunk_text(self, text: str) -> List[Dict[str, Any]]:
|
| 25 |
+
"""
|
| 26 |
+
Chunk text into overlapping segments
|
| 27 |
+
|
| 28 |
+
Args:
|
| 29 |
+
text: Input text to chunk
|
| 30 |
+
|
| 31 |
+
Returns:
|
| 32 |
+
List of chunk dictionaries with content and basic metadata
|
| 33 |
+
"""
|
| 34 |
+
if not text.strip():
|
| 35 |
+
return []
|
| 36 |
+
|
| 37 |
+
chunks = []
|
| 38 |
+
start = 0
|
| 39 |
+
chunk_index = 0
|
| 40 |
+
|
| 41 |
+
while start < len(text):
|
| 42 |
+
end = start + self.chunk_size
|
| 43 |
+
chunk_content = text[start:end]
|
| 44 |
+
|
| 45 |
+
# Create chunk with metadata
|
| 46 |
+
chunk = {
|
| 47 |
+
'content': chunk_content,
|
| 48 |
+
'metadata': {
|
| 49 |
+
'chunk_index': chunk_index,
|
| 50 |
+
'start_pos': start,
|
| 51 |
+
'end_pos': min(end, len(text)),
|
| 52 |
+
'chunk_id': self._generate_chunk_id(chunk_content, chunk_index)
|
| 53 |
+
}
|
| 54 |
+
}
|
| 55 |
+
|
| 56 |
+
chunks.append(chunk)
|
| 57 |
+
|
| 58 |
+
# Move start position with overlap consideration
|
| 59 |
+
start = end - self.overlap
|
| 60 |
+
chunk_index += 1
|
| 61 |
+
|
| 62 |
+
# Break if we've processed all text
|
| 63 |
+
if end >= len(text):
|
| 64 |
+
break
|
| 65 |
+
|
| 66 |
+
return chunks
|
| 67 |
+
|
| 68 |
+
def chunk_document(self, text: str, doc_metadata: Dict[str, Any]) -> List[Dict[str, Any]]:
|
| 69 |
+
"""
|
| 70 |
+
Chunk a document while preserving document metadata
|
| 71 |
+
|
| 72 |
+
Args:
|
| 73 |
+
text: Document text content
|
| 74 |
+
doc_metadata: Document metadata to preserve
|
| 75 |
+
|
| 76 |
+
Returns:
|
| 77 |
+
List of chunks with combined metadata
|
| 78 |
+
"""
|
| 79 |
+
chunks = self.chunk_text(text)
|
| 80 |
+
|
| 81 |
+
# Enhance each chunk with document metadata
|
| 82 |
+
for chunk in chunks:
|
| 83 |
+
chunk['metadata'].update(doc_metadata)
|
| 84 |
+
# Create unique chunk ID combining document and chunk info
|
| 85 |
+
chunk['metadata']['chunk_id'] = self._generate_chunk_id(
|
| 86 |
+
chunk['content'],
|
| 87 |
+
chunk['metadata']['chunk_index'],
|
| 88 |
+
doc_metadata.get('filename', 'unknown')
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
return chunks
|
| 92 |
+
|
| 93 |
+
def _generate_chunk_id(self, content: str, chunk_index: int, filename: str = "") -> str:
|
| 94 |
+
"""Generate a deterministic chunk ID"""
|
| 95 |
+
id_string = f"{filename}_{chunk_index}_{content[:50]}"
|
| 96 |
+
return hashlib.md5(id_string.encode()).hexdigest()[:12]
|
src/ingestion/document_parser.py
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
from typing import Dict, Any
|
| 4 |
+
|
| 5 |
+
class DocumentParser:
|
| 6 |
+
"""Parser for different document formats in the policy corpus"""
|
| 7 |
+
|
| 8 |
+
SUPPORTED_FORMATS = {'.txt', '.md', '.markdown'}
|
| 9 |
+
|
| 10 |
+
def parse_document(self, file_path: str) -> Dict[str, Any]:
|
| 11 |
+
"""
|
| 12 |
+
Parse a document and return content with metadata
|
| 13 |
+
|
| 14 |
+
Args:
|
| 15 |
+
file_path: Path to the document file
|
| 16 |
+
|
| 17 |
+
Returns:
|
| 18 |
+
Dict containing 'content' and 'metadata'
|
| 19 |
+
|
| 20 |
+
Raises:
|
| 21 |
+
FileNotFoundError: If file doesn't exist
|
| 22 |
+
ValueError: If file format is unsupported
|
| 23 |
+
"""
|
| 24 |
+
path = Path(file_path)
|
| 25 |
+
|
| 26 |
+
# Check file format first (before existence check)
|
| 27 |
+
if path.suffix.lower() not in self.SUPPORTED_FORMATS:
|
| 28 |
+
raise ValueError(f"Unsupported file format: {path.suffix}")
|
| 29 |
+
|
| 30 |
+
if not path.exists():
|
| 31 |
+
raise FileNotFoundError(f"File not found: {file_path}")
|
| 32 |
+
|
| 33 |
+
with open(file_path, 'r', encoding='utf-8') as f:
|
| 34 |
+
content = f.read()
|
| 35 |
+
|
| 36 |
+
metadata = {
|
| 37 |
+
'filename': path.name,
|
| 38 |
+
'file_type': path.suffix.lstrip('.').lower(),
|
| 39 |
+
'file_size': os.path.getsize(file_path),
|
| 40 |
+
'file_path': str(path.absolute())
|
| 41 |
+
}
|
| 42 |
+
|
| 43 |
+
return {
|
| 44 |
+
'content': content,
|
| 45 |
+
'metadata': metadata
|
| 46 |
+
}
|
src/ingestion/ingestion_pipeline.py
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
from typing import List, Dict, Any
|
| 3 |
+
from .document_parser import DocumentParser
|
| 4 |
+
from .document_chunker import DocumentChunker
|
| 5 |
+
|
| 6 |
+
class IngestionPipeline:
|
| 7 |
+
"""Complete ingestion pipeline for processing document corpus"""
|
| 8 |
+
|
| 9 |
+
def __init__(self, chunk_size: int = 1000, overlap: int = 200, seed: int = 42):
|
| 10 |
+
"""
|
| 11 |
+
Initialize the ingestion pipeline
|
| 12 |
+
|
| 13 |
+
Args:
|
| 14 |
+
chunk_size: Size of text chunks
|
| 15 |
+
overlap: Overlap between chunks
|
| 16 |
+
seed: Random seed for reproducibility
|
| 17 |
+
"""
|
| 18 |
+
self.parser = DocumentParser()
|
| 19 |
+
self.chunker = DocumentChunker(chunk_size=chunk_size, overlap=overlap, seed=seed)
|
| 20 |
+
self.seed = seed
|
| 21 |
+
|
| 22 |
+
def process_directory(self, directory_path: str) -> List[Dict[str, Any]]:
|
| 23 |
+
"""
|
| 24 |
+
Process all supported documents in a directory
|
| 25 |
+
|
| 26 |
+
Args:
|
| 27 |
+
directory_path: Path to directory containing documents
|
| 28 |
+
|
| 29 |
+
Returns:
|
| 30 |
+
List of processed chunks with metadata
|
| 31 |
+
"""
|
| 32 |
+
directory = Path(directory_path)
|
| 33 |
+
if not directory.exists():
|
| 34 |
+
raise FileNotFoundError(f"Directory not found: {directory_path}")
|
| 35 |
+
|
| 36 |
+
all_chunks = []
|
| 37 |
+
|
| 38 |
+
# Process each supported file
|
| 39 |
+
for file_path in directory.iterdir():
|
| 40 |
+
if file_path.is_file() and file_path.suffix.lower() in self.parser.SUPPORTED_FORMATS:
|
| 41 |
+
try:
|
| 42 |
+
chunks = self.process_file(str(file_path))
|
| 43 |
+
all_chunks.extend(chunks)
|
| 44 |
+
except Exception as e:
|
| 45 |
+
print(f"Warning: Failed to process {file_path}: {e}")
|
| 46 |
+
continue
|
| 47 |
+
|
| 48 |
+
return all_chunks
|
| 49 |
+
|
| 50 |
+
def process_file(self, file_path: str) -> List[Dict[str, Any]]:
|
| 51 |
+
"""
|
| 52 |
+
Process a single file through the complete pipeline
|
| 53 |
+
|
| 54 |
+
Args:
|
| 55 |
+
file_path: Path to the file to process
|
| 56 |
+
|
| 57 |
+
Returns:
|
| 58 |
+
List of chunks from the file
|
| 59 |
+
"""
|
| 60 |
+
# Parse document
|
| 61 |
+
parsed_doc = self.parser.parse_document(file_path)
|
| 62 |
+
|
| 63 |
+
# Chunk the document
|
| 64 |
+
chunks = self.chunker.chunk_document(
|
| 65 |
+
parsed_doc['content'],
|
| 66 |
+
parsed_doc['metadata']
|
| 67 |
+
)
|
| 68 |
+
|
| 69 |
+
return chunks
|
tests/test_app.py
CHANGED
|
@@ -28,5 +28,12 @@ def test_index_endpoint(client):
|
|
| 28 |
"""
|
| 29 |
response = client.get("/")
|
| 30 |
assert response.status_code == 200
|
| 31 |
-
|
| 32 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 28 |
"""
|
| 29 |
response = client.get("/")
|
| 30 |
assert response.status_code == 200
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def test_ingest_endpoint_exists():
|
| 34 |
+
"""Test that the ingest endpoint is available"""
|
| 35 |
+
from app import app
|
| 36 |
+
client = app.test_client()
|
| 37 |
+
response = client.post('/ingest')
|
| 38 |
+
# Should not be 404 (not found)
|
| 39 |
+
assert response.status_code != 404
|
tests/test_ingestion/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Test package for ingestion components
|
tests/test_ingestion/test_document_chunker.py
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pytest
|
| 2 |
+
from src.ingestion.document_chunker import DocumentChunker
|
| 3 |
+
|
| 4 |
+
def test_chunk_by_characters():
|
| 5 |
+
"""Test basic character-based chunking"""
|
| 6 |
+
chunker = DocumentChunker(chunk_size=50, overlap=10)
|
| 7 |
+
|
| 8 |
+
text = "This is a test document. " * 10 # 250 characters
|
| 9 |
+
chunks = chunker.chunk_text(text)
|
| 10 |
+
|
| 11 |
+
assert len(chunks) > 1 # Should create multiple chunks
|
| 12 |
+
assert all(len(chunk['content']) <= 50 for chunk in chunks)
|
| 13 |
+
|
| 14 |
+
# Test overlap
|
| 15 |
+
if len(chunks) > 1:
|
| 16 |
+
# Check that there's overlap between consecutive chunks
|
| 17 |
+
assert chunks[0]['content'][-10:] in chunks[1]['content'][:20]
|
| 18 |
+
|
| 19 |
+
def test_chunk_with_metadata():
|
| 20 |
+
"""Test that chunks preserve document metadata"""
|
| 21 |
+
chunker = DocumentChunker(chunk_size=100, overlap=20)
|
| 22 |
+
|
| 23 |
+
doc_metadata = {
|
| 24 |
+
'filename': 'test.txt',
|
| 25 |
+
'file_type': 'txt',
|
| 26 |
+
'source_id': 'doc_001'
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
text = "Content that will be chunked. " * 20
|
| 30 |
+
chunks = chunker.chunk_document(text, doc_metadata)
|
| 31 |
+
|
| 32 |
+
for chunk in chunks:
|
| 33 |
+
assert chunk['metadata']['filename'] == 'test.txt'
|
| 34 |
+
assert chunk['metadata']['file_type'] == 'txt'
|
| 35 |
+
assert 'chunk_id' in chunk['metadata']
|
| 36 |
+
assert 'chunk_index' in chunk['metadata']
|
| 37 |
+
|
| 38 |
+
def test_reproducible_chunking():
|
| 39 |
+
"""Test that chunking is deterministic with fixed seed"""
|
| 40 |
+
chunker1 = DocumentChunker(chunk_size=100, overlap=20, seed=42)
|
| 41 |
+
chunker2 = DocumentChunker(chunk_size=100, overlap=20, seed=42)
|
| 42 |
+
|
| 43 |
+
text = "This text will be chunked reproducibly. " * 30
|
| 44 |
+
|
| 45 |
+
chunks1 = chunker1.chunk_text(text)
|
| 46 |
+
chunks2 = chunker2.chunk_text(text)
|
| 47 |
+
|
| 48 |
+
assert len(chunks1) == len(chunks2)
|
| 49 |
+
for c1, c2 in zip(chunks1, chunks2):
|
| 50 |
+
assert c1['content'] == c2['content']
|
| 51 |
+
|
| 52 |
+
def test_empty_text_chunking():
|
| 53 |
+
"""Test handling of empty or very short text"""
|
| 54 |
+
chunker = DocumentChunker(chunk_size=100, overlap=20)
|
| 55 |
+
|
| 56 |
+
# Empty text
|
| 57 |
+
chunks = chunker.chunk_text("")
|
| 58 |
+
assert len(chunks) == 0
|
| 59 |
+
|
| 60 |
+
# Very short text
|
| 61 |
+
chunks = chunker.chunk_text("Short")
|
| 62 |
+
assert len(chunks) == 1
|
| 63 |
+
assert chunks[0]['content'] == "Short"
|
| 64 |
+
|
| 65 |
+
def test_chunk_real_policy_content():
|
| 66 |
+
"""Test chunking actual policy document content"""
|
| 67 |
+
chunker = DocumentChunker(chunk_size=500, overlap=100, seed=42)
|
| 68 |
+
|
| 69 |
+
# Use content that resembles our policy documents
|
| 70 |
+
policy_content = """# HR-POL-001: Employee Handbook
|
| 71 |
+
|
| 72 |
+
**Effective Date:** 2025-01-01
|
| 73 |
+
**Revision:** 1.1
|
| 74 |
+
**Owner:** Human Resources
|
| 75 |
+
|
| 76 |
+
## 1. Introduction
|
| 77 |
+
|
| 78 |
+
### 1.1. A Message from Our CEO
|
| 79 |
+
|
| 80 |
+
Welcome to Innovate Inc.! We are thrilled to have you as part of our team. Our success is built on the talent, dedication, and creativity of our employees. This handbook is designed to be your guide as you grow with us, providing clarity on the principles that shape our culture and the policies that govern our work.
|
| 81 |
+
|
| 82 |
+
## 2. Company Policies
|
| 83 |
+
|
| 84 |
+
### 2.1. Code of Conduct
|
| 85 |
+
|
| 86 |
+
All employees must adhere to our code of conduct which emphasizes integrity, respect, and professionalism in all interactions.""" * 3
|
| 87 |
+
|
| 88 |
+
doc_metadata = {
|
| 89 |
+
'filename': 'employee_handbook.md',
|
| 90 |
+
'file_type': 'md',
|
| 91 |
+
'file_path': '/path/to/employee_handbook.md'
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
chunks = chunker.chunk_document(policy_content, doc_metadata)
|
| 95 |
+
|
| 96 |
+
# Verify chunking worked
|
| 97 |
+
assert len(chunks) > 1
|
| 98 |
+
|
| 99 |
+
# Verify all chunks have proper metadata
|
| 100 |
+
for i, chunk in enumerate(chunks):
|
| 101 |
+
assert chunk['metadata']['filename'] == 'employee_handbook.md'
|
| 102 |
+
assert chunk['metadata']['file_type'] == 'md'
|
| 103 |
+
assert chunk['metadata']['chunk_index'] == i
|
| 104 |
+
assert 'chunk_id' in chunk['metadata']
|
| 105 |
+
assert len(chunk['content']) <= 500
|
| 106 |
+
|
| 107 |
+
# Verify overlap exists between consecutive chunks
|
| 108 |
+
if len(chunks) > 1:
|
| 109 |
+
assert chunks[0]['content'][-100:] in chunks[1]['content'][:200]
|
| 110 |
+
|
| 111 |
+
def test_chunk_metadata_inheritance():
|
| 112 |
+
"""Test that document metadata is properly inherited by chunks"""
|
| 113 |
+
chunker = DocumentChunker(chunk_size=100, overlap=20)
|
| 114 |
+
|
| 115 |
+
doc_metadata = {
|
| 116 |
+
'filename': 'test_policy.md',
|
| 117 |
+
'file_type': 'md',
|
| 118 |
+
'file_size': 1500,
|
| 119 |
+
'file_path': '/absolute/path/to/test_policy.md'
|
| 120 |
+
}
|
| 121 |
+
|
| 122 |
+
text = "Policy content goes here. " * 20
|
| 123 |
+
chunks = chunker.chunk_document(text, doc_metadata)
|
| 124 |
+
|
| 125 |
+
for chunk in chunks:
|
| 126 |
+
# Original metadata should be preserved
|
| 127 |
+
assert chunk['metadata']['filename'] == 'test_policy.md'
|
| 128 |
+
assert chunk['metadata']['file_type'] == 'md'
|
| 129 |
+
assert chunk['metadata']['file_size'] == 1500
|
| 130 |
+
assert chunk['metadata']['file_path'] == '/absolute/path/to/test_policy.md'
|
| 131 |
+
|
| 132 |
+
# New chunk-specific metadata should be added
|
| 133 |
+
assert 'chunk_index' in chunk['metadata']
|
| 134 |
+
assert 'chunk_id' in chunk['metadata']
|
| 135 |
+
assert 'start_pos' in chunk['metadata']
|
| 136 |
+
assert 'end_pos' in chunk['metadata']
|
tests/test_ingestion/test_document_parser.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pytest
|
| 2 |
+
import tempfile
|
| 3 |
+
import os
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
def test_parse_txt_file():
|
| 7 |
+
"""Test parsing a simple text file"""
|
| 8 |
+
# Test will fail initially - we'll implement parser to make it pass
|
| 9 |
+
from src.ingestion.document_parser import DocumentParser
|
| 10 |
+
|
| 11 |
+
parser = DocumentParser()
|
| 12 |
+
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
| 13 |
+
f.write("This is a test policy document.\nIt has multiple lines.")
|
| 14 |
+
temp_path = f.name
|
| 15 |
+
|
| 16 |
+
try:
|
| 17 |
+
result = parser.parse_document(temp_path)
|
| 18 |
+
assert result['content'] == "This is a test policy document.\nIt has multiple lines."
|
| 19 |
+
assert result['metadata']['filename'] == Path(temp_path).name
|
| 20 |
+
assert result['metadata']['file_type'] == 'txt'
|
| 21 |
+
finally:
|
| 22 |
+
os.unlink(temp_path)
|
| 23 |
+
|
| 24 |
+
def test_parse_markdown_file():
|
| 25 |
+
"""Test parsing a markdown file"""
|
| 26 |
+
from src.ingestion.document_parser import DocumentParser
|
| 27 |
+
|
| 28 |
+
parser = DocumentParser()
|
| 29 |
+
markdown_content = """# Policy Title
|
| 30 |
+
|
| 31 |
+
## Section 1
|
| 32 |
+
This is section content.
|
| 33 |
+
|
| 34 |
+
### Subsection
|
| 35 |
+
More content here."""
|
| 36 |
+
|
| 37 |
+
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
|
| 38 |
+
f.write(markdown_content)
|
| 39 |
+
temp_path = f.name
|
| 40 |
+
|
| 41 |
+
try:
|
| 42 |
+
result = parser.parse_document(temp_path)
|
| 43 |
+
assert "Policy Title" in result['content']
|
| 44 |
+
assert "Section 1" in result['content']
|
| 45 |
+
assert result['metadata']['file_type'] == 'md'
|
| 46 |
+
finally:
|
| 47 |
+
os.unlink(temp_path)
|
| 48 |
+
|
| 49 |
+
def test_parse_unsupported_format():
|
| 50 |
+
"""Test handling of unsupported file formats"""
|
| 51 |
+
from src.ingestion.document_parser import DocumentParser
|
| 52 |
+
|
| 53 |
+
parser = DocumentParser()
|
| 54 |
+
with pytest.raises(ValueError, match="Unsupported file format"):
|
| 55 |
+
parser.parse_document("test.xyz")
|
| 56 |
+
|
| 57 |
+
def test_parse_nonexistent_file():
|
| 58 |
+
"""Test handling of non-existent files"""
|
| 59 |
+
from src.ingestion.document_parser import DocumentParser
|
| 60 |
+
|
| 61 |
+
parser = DocumentParser()
|
| 62 |
+
with pytest.raises(FileNotFoundError):
|
| 63 |
+
parser.parse_document("nonexistent.txt")
|
| 64 |
+
|
| 65 |
+
def test_parse_real_policy_document():
|
| 66 |
+
"""Test parsing an actual policy document from our corpus"""
|
| 67 |
+
from src.ingestion.document_parser import DocumentParser
|
| 68 |
+
|
| 69 |
+
parser = DocumentParser()
|
| 70 |
+
# Use a real policy document from our corpus
|
| 71 |
+
policy_path = "synthetic_policies/employee_handbook.md"
|
| 72 |
+
|
| 73 |
+
result = parser.parse_document(policy_path)
|
| 74 |
+
|
| 75 |
+
# Verify content structure
|
| 76 |
+
assert "employee_handbook.md" in result['metadata']['filename']
|
| 77 |
+
assert result['metadata']['file_type'] == 'md'
|
| 78 |
+
assert "Employee Handbook" in result['content']
|
| 79 |
+
assert "HR-POL-001" in result['content']
|
| 80 |
+
assert len(result['content']) > 100 # Should have substantial content
|
| 81 |
+
|
| 82 |
+
# Verify metadata completeness
|
| 83 |
+
assert 'file_size' in result['metadata']
|
| 84 |
+
assert 'file_path' in result['metadata']
|
| 85 |
+
assert result['metadata']['file_size'] > 0
|
tests/test_ingestion/test_ingestion_pipeline.py
ADDED
|
@@ -0,0 +1,166 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pytest
|
| 2 |
+
import tempfile
|
| 3 |
+
import os
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from src.ingestion.ingestion_pipeline import IngestionPipeline
|
| 6 |
+
|
| 7 |
+
def test_full_ingestion_pipeline():
|
| 8 |
+
"""Test the complete ingestion pipeline end-to-end"""
|
| 9 |
+
# Create temporary test documents
|
| 10 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
| 11 |
+
# Create test files
|
| 12 |
+
txt_file = Path(temp_dir) / "policy1.txt"
|
| 13 |
+
md_file = Path(temp_dir) / "policy2.md"
|
| 14 |
+
|
| 15 |
+
txt_file.write_text("This is a text policy document with important information.")
|
| 16 |
+
md_file.write_text("# Markdown Policy\n\nThis is markdown content.")
|
| 17 |
+
|
| 18 |
+
# Initialize pipeline
|
| 19 |
+
pipeline = IngestionPipeline(chunk_size=50, overlap=10, seed=42)
|
| 20 |
+
|
| 21 |
+
# Process documents
|
| 22 |
+
results = pipeline.process_directory(temp_dir)
|
| 23 |
+
|
| 24 |
+
assert len(results) >= 2 # At least one result per file
|
| 25 |
+
|
| 26 |
+
# Verify structure
|
| 27 |
+
for result in results:
|
| 28 |
+
assert 'content' in result
|
| 29 |
+
assert 'metadata' in result
|
| 30 |
+
assert 'chunk_id' in result['metadata']
|
| 31 |
+
assert 'filename' in result['metadata']
|
| 32 |
+
|
| 33 |
+
def test_pipeline_reproducibility():
|
| 34 |
+
"""Test that pipeline produces consistent results"""
|
| 35 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
| 36 |
+
test_file = Path(temp_dir) / "test.txt"
|
| 37 |
+
test_file.write_text("Test content for reproducibility. " * 20)
|
| 38 |
+
|
| 39 |
+
pipeline1 = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 40 |
+
pipeline2 = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 41 |
+
|
| 42 |
+
results1 = pipeline1.process_directory(temp_dir)
|
| 43 |
+
results2 = pipeline2.process_directory(temp_dir)
|
| 44 |
+
|
| 45 |
+
assert len(results1) == len(results2)
|
| 46 |
+
|
| 47 |
+
for r1, r2 in zip(results1, results2):
|
| 48 |
+
assert r1['content'] == r2['content']
|
| 49 |
+
assert r1['metadata']['chunk_id'] == r2['metadata']['chunk_id']
|
| 50 |
+
|
| 51 |
+
def test_pipeline_with_real_corpus():
|
| 52 |
+
"""Test pipeline with actual policy documents"""
|
| 53 |
+
pipeline = IngestionPipeline(chunk_size=1000, overlap=200, seed=42)
|
| 54 |
+
|
| 55 |
+
# Process just one real document to verify it works
|
| 56 |
+
corpus_dir = "synthetic_policies"
|
| 57 |
+
|
| 58 |
+
# Check if corpus directory exists
|
| 59 |
+
if not Path(corpus_dir).exists():
|
| 60 |
+
pytest.skip("Corpus directory not found - test requires synthetic_policies/")
|
| 61 |
+
|
| 62 |
+
results = pipeline.process_directory(corpus_dir)
|
| 63 |
+
|
| 64 |
+
# Should process all 22 documents
|
| 65 |
+
assert len(results) > 20 # Should have many chunks from 22 documents
|
| 66 |
+
|
| 67 |
+
# Verify all results have proper structure
|
| 68 |
+
for result in results:
|
| 69 |
+
assert 'content' in result
|
| 70 |
+
assert 'metadata' in result
|
| 71 |
+
assert 'chunk_id' in result['metadata']
|
| 72 |
+
assert 'filename' in result['metadata']
|
| 73 |
+
assert 'file_type' in result['metadata']
|
| 74 |
+
assert result['metadata']['file_type'] == 'md'
|
| 75 |
+
assert 'chunk_index' in result['metadata']
|
| 76 |
+
|
| 77 |
+
def test_pipeline_error_handling():
|
| 78 |
+
"""Test pipeline handles errors gracefully"""
|
| 79 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
| 80 |
+
# Create valid and invalid files
|
| 81 |
+
valid_file = Path(temp_dir) / "valid.md"
|
| 82 |
+
invalid_file = Path(temp_dir) / "invalid.xyz"
|
| 83 |
+
|
| 84 |
+
valid_file.write_text("# Valid Policy\n\nThis is valid content.")
|
| 85 |
+
invalid_file.write_text("This file has unsupported format.")
|
| 86 |
+
|
| 87 |
+
pipeline = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 88 |
+
|
| 89 |
+
# Should process valid file and skip invalid one
|
| 90 |
+
results = pipeline.process_directory(temp_dir)
|
| 91 |
+
|
| 92 |
+
# Should only get results from valid file
|
| 93 |
+
assert len(results) >= 1
|
| 94 |
+
|
| 95 |
+
# All results should be from valid file
|
| 96 |
+
for result in results:
|
| 97 |
+
assert result['metadata']['filename'] == 'valid.md'
|
| 98 |
+
|
| 99 |
+
def test_pipeline_single_file():
|
| 100 |
+
"""Test processing a single file"""
|
| 101 |
+
pipeline = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 102 |
+
|
| 103 |
+
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
|
| 104 |
+
f.write("# Test Policy\n\n" + "Content section. " * 20)
|
| 105 |
+
temp_path = f.name
|
| 106 |
+
|
| 107 |
+
try:
|
| 108 |
+
results = pipeline.process_file(temp_path)
|
| 109 |
+
|
| 110 |
+
# Should get multiple chunks due to length
|
| 111 |
+
assert len(results) > 1
|
| 112 |
+
|
| 113 |
+
# All chunks should have same filename
|
| 114 |
+
filename = Path(temp_path).name
|
| 115 |
+
for result in results:
|
| 116 |
+
assert result['metadata']['filename'] == filename
|
| 117 |
+
assert result['metadata']['file_type'] == 'md'
|
| 118 |
+
assert 'chunk_index' in result['metadata']
|
| 119 |
+
|
| 120 |
+
finally:
|
| 121 |
+
os.unlink(temp_path)
|
| 122 |
+
|
| 123 |
+
def test_pipeline_empty_directory():
|
| 124 |
+
"""Test pipeline with empty directory"""
|
| 125 |
+
with tempfile.TemporaryDirectory() as temp_dir:
|
| 126 |
+
pipeline = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 127 |
+
|
| 128 |
+
results = pipeline.process_directory(temp_dir)
|
| 129 |
+
|
| 130 |
+
# Should return empty list for empty directory
|
| 131 |
+
assert len(results) == 0
|
| 132 |
+
|
| 133 |
+
def test_pipeline_nonexistent_directory():
|
| 134 |
+
"""Test pipeline with non-existent directory"""
|
| 135 |
+
pipeline = IngestionPipeline(chunk_size=100, overlap=20, seed=42)
|
| 136 |
+
|
| 137 |
+
with pytest.raises(FileNotFoundError):
|
| 138 |
+
pipeline.process_directory("/nonexistent/directory")
|
| 139 |
+
|
| 140 |
+
def test_pipeline_configuration():
|
| 141 |
+
"""Test pipeline configuration options"""
|
| 142 |
+
# Test different configurations
|
| 143 |
+
pipeline_small = IngestionPipeline(chunk_size=50, overlap=10, seed=42)
|
| 144 |
+
pipeline_large = IngestionPipeline(chunk_size=200, overlap=50, seed=42)
|
| 145 |
+
|
| 146 |
+
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
| 147 |
+
content = "Policy content goes here. " * 30 # 780 characters
|
| 148 |
+
f.write(content)
|
| 149 |
+
temp_path = f.name
|
| 150 |
+
|
| 151 |
+
try:
|
| 152 |
+
results_small = pipeline_small.process_file(temp_path)
|
| 153 |
+
results_large = pipeline_large.process_file(temp_path)
|
| 154 |
+
|
| 155 |
+
# Small chunks should create more chunks
|
| 156 |
+
assert len(results_small) > len(results_large)
|
| 157 |
+
|
| 158 |
+
# All chunks should respect size limits
|
| 159 |
+
for result in results_small:
|
| 160 |
+
assert len(result['content']) <= 50
|
| 161 |
+
|
| 162 |
+
for result in results_large:
|
| 163 |
+
assert len(result['content']) <= 200
|
| 164 |
+
|
| 165 |
+
finally:
|
| 166 |
+
os.unlink(temp_path)
|