ragmill 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ragmill/__init__.py +3 -0
- ragmill/embeddings.py +95 -0
- ragmill/engine.py +152 -0
- ragmill/parsers.py +33 -0
- ragmill/sync.py +74 -0
- ragmill/vector_store.py +163 -0
- ragmill-0.2.0.dist-info/METADATA +135 -0
- ragmill-0.2.0.dist-info/RECORD +10 -0
- ragmill-0.2.0.dist-info/WHEEL +4 -0
- ragmill-0.2.0.dist-info/licenses/LICENSE +21 -0
ragmill/__init__.py
ADDED
ragmill/embeddings.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Local embedding generation via ONNX Runtime.
|
|
3
|
+
|
|
4
|
+
Downloads a small quantized sentence-embedding model on first use and
|
|
5
|
+
caches it under ~/.cache/ragmill/models — every call after that runs
|
|
6
|
+
fully offline. Imports for onnxruntime/tokenizers are lazy so importing
|
|
7
|
+
this module doesn't force the 'embeddings' extra on core users.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import urllib.request
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
from typing import List, Optional, Union
|
|
13
|
+
|
|
14
|
+
import numpy as np
|
|
15
|
+
|
|
16
|
+
DEFAULT_MODEL = "Xenova/all-MiniLM-L6-v2"
|
|
17
|
+
DEFAULT_CACHE_DIR = Path.home() / ".cache" / "ragmill" / "models"
|
|
18
|
+
EMBEDDING_DIM = 384
|
|
19
|
+
|
|
20
|
+
_MODEL_FILES = {
|
|
21
|
+
"model.onnx": "onnx/model_quantized.onnx",
|
|
22
|
+
"tokenizer.json": "tokenizer.json",
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _model_dir(model_name: str, cache_dir: Path) -> Path:
|
|
27
|
+
return cache_dir / model_name.replace("/", "__")
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _download(model_name: str, cache_dir: Path) -> Path:
|
|
31
|
+
target_dir = _model_dir(model_name, cache_dir)
|
|
32
|
+
target_dir.mkdir(parents=True, exist_ok=True)
|
|
33
|
+
|
|
34
|
+
for local_name, remote_path in _MODEL_FILES.items():
|
|
35
|
+
local_path = target_dir / local_name
|
|
36
|
+
if local_path.exists():
|
|
37
|
+
continue
|
|
38
|
+
url = f"https://huggingface.co/{model_name}/resolve/main/{remote_path}"
|
|
39
|
+
urllib.request.urlretrieve(url, local_path)
|
|
40
|
+
|
|
41
|
+
return target_dir
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class EmbeddingModel:
|
|
45
|
+
"""Wraps an ONNX sentence-embedding model + tokenizer for local inference."""
|
|
46
|
+
|
|
47
|
+
def __init__(self, model_name: str = DEFAULT_MODEL, cache_dir: Optional[Union[str, Path]] = None):
|
|
48
|
+
try:
|
|
49
|
+
import onnxruntime as ort
|
|
50
|
+
except ImportError as exc:
|
|
51
|
+
raise ImportError(
|
|
52
|
+
"Embeddings require the 'embeddings' extra. Install it with: pip install ragmill[embeddings]"
|
|
53
|
+
) from exc
|
|
54
|
+
try:
|
|
55
|
+
from tokenizers import Tokenizer
|
|
56
|
+
except ImportError as exc:
|
|
57
|
+
raise ImportError(
|
|
58
|
+
"Embeddings require the 'embeddings' extra. Install it with: pip install ragmill[embeddings]"
|
|
59
|
+
) from exc
|
|
60
|
+
|
|
61
|
+
resolved_cache_dir = Path(cache_dir) if cache_dir else DEFAULT_CACHE_DIR
|
|
62
|
+
model_dir = _download(model_name, resolved_cache_dir)
|
|
63
|
+
|
|
64
|
+
self.tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
|
|
65
|
+
self.tokenizer.enable_padding()
|
|
66
|
+
self.tokenizer.enable_truncation(max_length=256)
|
|
67
|
+
self.session = ort.InferenceSession(str(model_dir / "model.onnx"))
|
|
68
|
+
|
|
69
|
+
def embed(self, texts: List[str]) -> np.ndarray:
|
|
70
|
+
"""Encodes a batch of strings into L2-normalized dense vectors."""
|
|
71
|
+
if not texts:
|
|
72
|
+
return np.zeros((0, EMBEDDING_DIM), dtype=np.float32)
|
|
73
|
+
|
|
74
|
+
encodings = self.tokenizer.encode_batch(texts)
|
|
75
|
+
input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
|
|
76
|
+
attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
|
|
77
|
+
token_type_ids = np.zeros_like(input_ids)
|
|
78
|
+
|
|
79
|
+
outputs = self.session.run(
|
|
80
|
+
None,
|
|
81
|
+
{
|
|
82
|
+
"input_ids": input_ids,
|
|
83
|
+
"attention_mask": attention_mask,
|
|
84
|
+
"token_type_ids": token_type_ids,
|
|
85
|
+
},
|
|
86
|
+
)
|
|
87
|
+
token_embeddings = outputs[0] # (batch, seq_len, dim)
|
|
88
|
+
|
|
89
|
+
mask = attention_mask[:, :, np.newaxis].astype(np.float32)
|
|
90
|
+
summed = (token_embeddings * mask).sum(axis=1)
|
|
91
|
+
counts = np.clip(mask.sum(axis=1), 1e-9, None)
|
|
92
|
+
pooled = summed / counts
|
|
93
|
+
|
|
94
|
+
norms = np.clip(np.linalg.norm(pooled, axis=1, keepdims=True), 1e-9, None)
|
|
95
|
+
return (pooled / norms).astype(np.float32)
|
ragmill/engine.py
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
"""
|
|
2
|
+
RAGMill Core Engine
|
|
3
|
+
Optimized for high-performance directory crawling, semantic splitting,
|
|
4
|
+
and clean data structural alignment.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import os
|
|
8
|
+
import re
|
|
9
|
+
from typing import List, Dict, Any, Generator
|
|
10
|
+
|
|
11
|
+
from ragmill.parsers import extract_pdf_text, extract_docx_text
|
|
12
|
+
|
|
13
|
+
PLAIN_TEXT_EXTENSIONS = ('.txt', '.md', '.log', '.rst')
|
|
14
|
+
PDF_EXTENSIONS = ('.pdf',)
|
|
15
|
+
DOCX_EXTENSIONS = ('.docx',)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class RAGEngine:
|
|
19
|
+
def __init__(self, chunk_size: int = 500, overlap: int = 50):
|
|
20
|
+
"""
|
|
21
|
+
Initializes the AI-Native Data Pipeline Engine.
|
|
22
|
+
|
|
23
|
+
:param chunk_size: Maximum structural character size permitted per single block.
|
|
24
|
+
:param overlap: Token/character historical window size to carry context forward.
|
|
25
|
+
"""
|
|
26
|
+
self.chunk_size = chunk_size
|
|
27
|
+
self.overlap = max(0, overlap)
|
|
28
|
+
|
|
29
|
+
if self.overlap >= self.chunk_size:
|
|
30
|
+
raise ValueError("Overlap threshold cannot be greater than or equal to total chunk size.")
|
|
31
|
+
|
|
32
|
+
def stream_directory(self, directory_path: str) -> Generator[Dict[str, Any], None, None]:
|
|
33
|
+
"""
|
|
34
|
+
Performs high-efficiency traversal over a target local directory,
|
|
35
|
+
streaming extracted text payloads while preserving system memory.
|
|
36
|
+
"""
|
|
37
|
+
if not os.path.exists(directory_path):
|
|
38
|
+
raise FileNotFoundError(f"Target path tracking validation failed for: '{directory_path}'")
|
|
39
|
+
|
|
40
|
+
for root, _, files in os.walk(directory_path):
|
|
41
|
+
for file in files:
|
|
42
|
+
extension = os.path.splitext(file)[1].lower()
|
|
43
|
+
if extension not in PLAIN_TEXT_EXTENSIONS + PDF_EXTENSIONS + DOCX_EXTENSIONS:
|
|
44
|
+
continue
|
|
45
|
+
|
|
46
|
+
full_path = os.path.join(root, file)
|
|
47
|
+
try:
|
|
48
|
+
content = self._extract_content(full_path, extension)
|
|
49
|
+
|
|
50
|
+
yield {
|
|
51
|
+
"source_path": os.path.abspath(full_path),
|
|
52
|
+
"filename": file,
|
|
53
|
+
"raw_content": content.strip(),
|
|
54
|
+
"modified_at": os.path.getmtime(full_path)
|
|
55
|
+
}
|
|
56
|
+
except Exception as e:
|
|
57
|
+
# Log error safely to console without terminating pipeline execution
|
|
58
|
+
print(f"[⚠️ Pipeline Warning] Unable to parse file {full_path}: {str(e)}")
|
|
59
|
+
|
|
60
|
+
def _extract_content(self, full_path: str, extension: str) -> str:
|
|
61
|
+
if extension in PLAIN_TEXT_EXTENSIONS:
|
|
62
|
+
with open(full_path, 'r', encoding='utf-8', errors='ignore') as f:
|
|
63
|
+
return f.read()
|
|
64
|
+
if extension in PDF_EXTENSIONS:
|
|
65
|
+
return extract_pdf_text(full_path)
|
|
66
|
+
if extension in DOCX_EXTENSIONS:
|
|
67
|
+
return extract_docx_text(full_path)
|
|
68
|
+
raise ValueError(f"Unsupported extension: {extension}")
|
|
69
|
+
|
|
70
|
+
def semantic_chunking(self, text: str) -> List[str]:
|
|
71
|
+
"""
|
|
72
|
+
Recursively splits text payloads based on logical paragraph, structural,
|
|
73
|
+
and grammatical sentence boundaries to protect semantic context integrity.
|
|
74
|
+
"""
|
|
75
|
+
if not text:
|
|
76
|
+
return []
|
|
77
|
+
|
|
78
|
+
# Split along logical structural breaks (paragraphs, list blocks, markdown breaks)
|
|
79
|
+
paragraphs = re.split(r'\n\s*\n', text)
|
|
80
|
+
chunks: List[str] = []
|
|
81
|
+
current_buffer = ""
|
|
82
|
+
|
|
83
|
+
for paragraph in paragraphs:
|
|
84
|
+
paragraph = paragraph.strip()
|
|
85
|
+
if not paragraph:
|
|
86
|
+
continue
|
|
87
|
+
|
|
88
|
+
# Handle edge cases where single paragraphs wildly exceed standard target size limits
|
|
89
|
+
if len(paragraph) > self.chunk_size:
|
|
90
|
+
# If buffer already holds content, clear it to start clean sentence processing
|
|
91
|
+
if current_buffer:
|
|
92
|
+
chunks.append(current_buffer.strip())
|
|
93
|
+
current_buffer = ""
|
|
94
|
+
|
|
95
|
+
# Split down into sentence tokens
|
|
96
|
+
sentences = re.split(r'(?<=[.!?])\s+', paragraph)
|
|
97
|
+
for sentence in sentences:
|
|
98
|
+
sentence = sentence.strip()
|
|
99
|
+
if not sentence:
|
|
100
|
+
continue
|
|
101
|
+
|
|
102
|
+
if len(current_buffer) + len(sentence) + 1 <= self.chunk_size:
|
|
103
|
+
current_buffer = f"{current_buffer} {sentence}".strip()
|
|
104
|
+
else:
|
|
105
|
+
if current_buffer:
|
|
106
|
+
chunks.append(current_buffer)
|
|
107
|
+
|
|
108
|
+
# Handle long sentence edge case: verify slice safety before copying historical context
|
|
109
|
+
overlap_prefix = current_buffer[-self.overlap:] if len(current_buffer) >= self.overlap else current_buffer
|
|
110
|
+
current_buffer = f"{overlap_prefix} {sentence}".strip() if self.overlap > 0 else sentence
|
|
111
|
+
else:
|
|
112
|
+
# Standard appending logic for typical sized semantic paragraphs
|
|
113
|
+
spacing = "\n\n" if current_buffer else ""
|
|
114
|
+
if len(current_buffer) + len(spacing) + len(paragraph) <= self.chunk_size:
|
|
115
|
+
current_buffer = f"{current_buffer}{spacing}{paragraph}"
|
|
116
|
+
else:
|
|
117
|
+
if current_buffer:
|
|
118
|
+
chunks.append(current_buffer.strip())
|
|
119
|
+
|
|
120
|
+
# Establish overlap baseline from preceding content block
|
|
121
|
+
overlap_prefix = current_buffer[-self.overlap:] if len(current_buffer) >= self.overlap else current_buffer
|
|
122
|
+
current_buffer = f"{overlap_prefix}\n\n{paragraph}".strip() if self.overlap > 0 else paragraph
|
|
123
|
+
|
|
124
|
+
if current_buffer:
|
|
125
|
+
chunks.append(current_buffer.strip())
|
|
126
|
+
|
|
127
|
+
return chunks
|
|
128
|
+
|
|
129
|
+
def execute_pipeline(self, directory_path: str) -> List[Dict[str, Any]]:
|
|
130
|
+
"""
|
|
131
|
+
Compiles the full ingestion and data structure cycle across a target directory.
|
|
132
|
+
|
|
133
|
+
:return: Array containing distinct dictionary models containing contextual metadata maps.
|
|
134
|
+
"""
|
|
135
|
+
pipeline_payloads: List[Dict[str, Any]] = []
|
|
136
|
+
|
|
137
|
+
for file_manifest in self.stream_directory(directory_path):
|
|
138
|
+
text_chunks = self.semantic_chunking(file_manifest["raw_content"])
|
|
139
|
+
|
|
140
|
+
for index, chunk in enumerate(text_chunks):
|
|
141
|
+
pipeline_payloads.append({
|
|
142
|
+
"metadata": {
|
|
143
|
+
"source_file": file_manifest["source_path"],
|
|
144
|
+
"filename": file_manifest["filename"],
|
|
145
|
+
"chunk_index": index,
|
|
146
|
+
"character_length": len(chunk),
|
|
147
|
+
"modified_at": file_manifest["modified_at"]
|
|
148
|
+
},
|
|
149
|
+
"content": chunk
|
|
150
|
+
})
|
|
151
|
+
|
|
152
|
+
return pipeline_payloads
|
ragmill/parsers.py
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Format-specific text extractors.
|
|
3
|
+
|
|
4
|
+
Each extractor takes a file path and returns its plain-text content.
|
|
5
|
+
PDF and DOCX support are optional extras — the imports are lazy so the
|
|
6
|
+
core package stays installable with zero dependencies.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def extract_pdf_text(path: str) -> str:
|
|
11
|
+
try:
|
|
12
|
+
from pypdf import PdfReader
|
|
13
|
+
except ImportError as exc:
|
|
14
|
+
raise ImportError(
|
|
15
|
+
"PDF support requires the 'pdf' extra. Install it with: pip install ragmill[pdf]"
|
|
16
|
+
) from exc
|
|
17
|
+
|
|
18
|
+
reader = PdfReader(path)
|
|
19
|
+
pages = [page.extract_text() or "" for page in reader.pages]
|
|
20
|
+
return "\n\n".join(page.strip() for page in pages if page.strip())
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def extract_docx_text(path: str) -> str:
|
|
24
|
+
try:
|
|
25
|
+
import docx
|
|
26
|
+
except ImportError as exc:
|
|
27
|
+
raise ImportError(
|
|
28
|
+
"DOCX support requires the 'docx' extra. Install it with: pip install ragmill[docx]"
|
|
29
|
+
) from exc
|
|
30
|
+
|
|
31
|
+
document = docx.Document(path)
|
|
32
|
+
paragraphs = [p.text.strip() for p in document.paragraphs]
|
|
33
|
+
return "\n\n".join(p for p in paragraphs if p)
|
ragmill/sync.py
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Incremental sync between a directory and a VectorStore.
|
|
3
|
+
|
|
4
|
+
Re-running execute_pipeline() on every call re-embeds and re-inserts every
|
|
5
|
+
file, every time. sync_directory() instead tracks a content hash per file
|
|
6
|
+
(in VectorStore.file_state) so unchanged files are skipped entirely, changed
|
|
7
|
+
files have their old chunks replaced, and files removed from disk have their
|
|
8
|
+
chunks removed from the store too.
|
|
9
|
+
|
|
10
|
+
This module has no hard dependencies of its own — it only calls methods on
|
|
11
|
+
the engine/model/store objects passed in, so it stays zero-dependency at
|
|
12
|
+
import time just like engine.py.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
import hashlib
|
|
16
|
+
from typing import Any, Dict
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def _hash_content(text: str) -> str:
|
|
20
|
+
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def sync_directory(directory_path: str, engine, model, store) -> Dict[str, int]:
|
|
24
|
+
"""
|
|
25
|
+
Walks directory_path with `engine`, embeds changed/new files with `model`,
|
|
26
|
+
and reconciles the result into `store`.
|
|
27
|
+
|
|
28
|
+
:return: counts of {"added", "updated", "skipped", "deleted"} files.
|
|
29
|
+
"""
|
|
30
|
+
seen_sources = set()
|
|
31
|
+
added = updated = skipped = 0
|
|
32
|
+
|
|
33
|
+
for file_manifest in engine.stream_directory(directory_path):
|
|
34
|
+
source_file = file_manifest["source_path"]
|
|
35
|
+
seen_sources.add(source_file)
|
|
36
|
+
|
|
37
|
+
content_hash = _hash_content(file_manifest["raw_content"])
|
|
38
|
+
existing_state = store.get_file_state(source_file)
|
|
39
|
+
|
|
40
|
+
if existing_state is not None and existing_state["content_hash"] == content_hash:
|
|
41
|
+
skipped += 1
|
|
42
|
+
continue
|
|
43
|
+
|
|
44
|
+
text_chunks = engine.semantic_chunking(file_manifest["raw_content"])
|
|
45
|
+
payloads = [
|
|
46
|
+
{
|
|
47
|
+
"metadata": {
|
|
48
|
+
"source_file": source_file,
|
|
49
|
+
"filename": file_manifest["filename"],
|
|
50
|
+
"chunk_index": index,
|
|
51
|
+
"character_length": len(chunk),
|
|
52
|
+
"modified_at": file_manifest["modified_at"],
|
|
53
|
+
},
|
|
54
|
+
"content": chunk,
|
|
55
|
+
}
|
|
56
|
+
for index, chunk in enumerate(text_chunks)
|
|
57
|
+
]
|
|
58
|
+
|
|
59
|
+
# Replace any previous chunks for this file before inserting the new ones,
|
|
60
|
+
# so an updated file doesn't leave stale chunks from its old content behind.
|
|
61
|
+
store.delete_by_source(source_file)
|
|
62
|
+
if payloads:
|
|
63
|
+
vectors = model.embed([p["content"] for p in payloads])
|
|
64
|
+
store.add(payloads, vectors)
|
|
65
|
+
store.upsert_file_state(source_file, content_hash, file_manifest["modified_at"])
|
|
66
|
+
|
|
67
|
+
if existing_state is None:
|
|
68
|
+
added += 1
|
|
69
|
+
else:
|
|
70
|
+
updated += 1
|
|
71
|
+
|
|
72
|
+
deleted = store.delete_missing_sources(seen_sources)
|
|
73
|
+
|
|
74
|
+
return {"added": added, "updated": updated, "skipped": skipped, "deleted": deleted}
|
ragmill/vector_store.py
ADDED
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
"""
|
|
2
|
+
A minimal local vector store backed by SQLite.
|
|
3
|
+
|
|
4
|
+
Stores chunk payloads alongside their embedding vectors and performs
|
|
5
|
+
brute-force cosine similarity search in-memory. No external vector index
|
|
6
|
+
needed at the scale this library targets (a local folder's worth of
|
|
7
|
+
documents) — every payload is loaded once per search and scored with a
|
|
8
|
+
single matrix multiply.
|
|
9
|
+
|
|
10
|
+
Also tracks per-file content hashes (the `file_state` table) so callers
|
|
11
|
+
can detect which files changed since the last run without re-embedding
|
|
12
|
+
everything — see sync.py for the orchestration that uses this.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
import sqlite3
|
|
16
|
+
from pathlib import Path
|
|
17
|
+
from typing import Any, Dict, List, Optional, Union
|
|
18
|
+
|
|
19
|
+
import numpy as np
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class VectorStore:
|
|
23
|
+
def __init__(self, db_path: Union[str, Path] = ":memory:"):
|
|
24
|
+
self.connection = sqlite3.connect(str(db_path))
|
|
25
|
+
self.connection.execute(
|
|
26
|
+
"""
|
|
27
|
+
CREATE TABLE IF NOT EXISTS chunks (
|
|
28
|
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
29
|
+
source_file TEXT,
|
|
30
|
+
filename TEXT,
|
|
31
|
+
chunk_index INTEGER,
|
|
32
|
+
content TEXT NOT NULL,
|
|
33
|
+
modified_at REAL,
|
|
34
|
+
embedding BLOB NOT NULL
|
|
35
|
+
)
|
|
36
|
+
"""
|
|
37
|
+
)
|
|
38
|
+
self.connection.execute(
|
|
39
|
+
"""
|
|
40
|
+
CREATE TABLE IF NOT EXISTS file_state (
|
|
41
|
+
source_file TEXT PRIMARY KEY,
|
|
42
|
+
content_hash TEXT NOT NULL,
|
|
43
|
+
modified_at REAL
|
|
44
|
+
)
|
|
45
|
+
"""
|
|
46
|
+
)
|
|
47
|
+
self.connection.commit()
|
|
48
|
+
|
|
49
|
+
def add(self, payloads: List[Dict[str, Any]], embeddings: np.ndarray) -> None:
|
|
50
|
+
if len(payloads) != len(embeddings):
|
|
51
|
+
raise ValueError("payloads and embeddings must be the same length")
|
|
52
|
+
|
|
53
|
+
rows = [
|
|
54
|
+
(
|
|
55
|
+
payload["metadata"]["source_file"],
|
|
56
|
+
payload["metadata"]["filename"],
|
|
57
|
+
payload["metadata"]["chunk_index"],
|
|
58
|
+
payload["content"],
|
|
59
|
+
payload["metadata"].get("modified_at"),
|
|
60
|
+
np.asarray(vector, dtype=np.float32).tobytes(),
|
|
61
|
+
)
|
|
62
|
+
for payload, vector in zip(payloads, embeddings)
|
|
63
|
+
]
|
|
64
|
+
self.connection.executemany(
|
|
65
|
+
"""
|
|
66
|
+
INSERT INTO chunks (source_file, filename, chunk_index, content, modified_at, embedding)
|
|
67
|
+
VALUES (?, ?, ?, ?, ?, ?)
|
|
68
|
+
""",
|
|
69
|
+
rows,
|
|
70
|
+
)
|
|
71
|
+
self.connection.commit()
|
|
72
|
+
|
|
73
|
+
def search(
|
|
74
|
+
self,
|
|
75
|
+
query_embedding: np.ndarray,
|
|
76
|
+
top_k: int = 5,
|
|
77
|
+
filename: Optional[str] = None,
|
|
78
|
+
source_file: Optional[str] = None,
|
|
79
|
+
modified_after: Optional[float] = None,
|
|
80
|
+
modified_before: Optional[float] = None,
|
|
81
|
+
) -> List[Dict[str, Any]]:
|
|
82
|
+
query_vector = np.asarray(query_embedding, dtype=np.float32)
|
|
83
|
+
|
|
84
|
+
clauses = []
|
|
85
|
+
params: List[Any] = []
|
|
86
|
+
if filename is not None:
|
|
87
|
+
clauses.append("filename = ?")
|
|
88
|
+
params.append(filename)
|
|
89
|
+
if source_file is not None:
|
|
90
|
+
clauses.append("source_file = ?")
|
|
91
|
+
params.append(source_file)
|
|
92
|
+
if modified_after is not None:
|
|
93
|
+
clauses.append("modified_at >= ?")
|
|
94
|
+
params.append(modified_after)
|
|
95
|
+
if modified_before is not None:
|
|
96
|
+
clauses.append("modified_at <= ?")
|
|
97
|
+
params.append(modified_before)
|
|
98
|
+
|
|
99
|
+
query = "SELECT source_file, filename, chunk_index, content, embedding FROM chunks"
|
|
100
|
+
if clauses:
|
|
101
|
+
query += " WHERE " + " AND ".join(clauses)
|
|
102
|
+
|
|
103
|
+
rows = self.connection.execute(query, params).fetchall()
|
|
104
|
+
if not rows:
|
|
105
|
+
return []
|
|
106
|
+
|
|
107
|
+
vectors = np.stack([np.frombuffer(row[4], dtype=np.float32) for row in rows])
|
|
108
|
+
# Embeddings are pre-normalized (see EmbeddingModel.embed), so the dot
|
|
109
|
+
# product against a normalized query is equivalent to cosine similarity.
|
|
110
|
+
scores = vectors @ query_vector
|
|
111
|
+
|
|
112
|
+
top_indices = np.argsort(-scores)[:top_k]
|
|
113
|
+
return [
|
|
114
|
+
{
|
|
115
|
+
"score": float(scores[i]),
|
|
116
|
+
"metadata": {
|
|
117
|
+
"source_file": rows[i][0],
|
|
118
|
+
"filename": rows[i][1],
|
|
119
|
+
"chunk_index": rows[i][2],
|
|
120
|
+
},
|
|
121
|
+
"content": rows[i][3],
|
|
122
|
+
}
|
|
123
|
+
for i in top_indices
|
|
124
|
+
]
|
|
125
|
+
|
|
126
|
+
def count(self) -> int:
|
|
127
|
+
return self.connection.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
|
|
128
|
+
|
|
129
|
+
def delete_by_source(self, source_file: str) -> None:
|
|
130
|
+
self.connection.execute("DELETE FROM chunks WHERE source_file = ?", (source_file,))
|
|
131
|
+
self.connection.execute("DELETE FROM file_state WHERE source_file = ?", (source_file,))
|
|
132
|
+
self.connection.commit()
|
|
133
|
+
|
|
134
|
+
def delete_missing_sources(self, known_sources: set) -> int:
|
|
135
|
+
"""Removes chunks/file_state for any source_file not in known_sources. Returns count removed."""
|
|
136
|
+
rows = self.connection.execute("SELECT source_file FROM file_state").fetchall()
|
|
137
|
+
stale = [row[0] for row in rows if row[0] not in known_sources]
|
|
138
|
+
for source_file in stale:
|
|
139
|
+
self.delete_by_source(source_file)
|
|
140
|
+
return len(stale)
|
|
141
|
+
|
|
142
|
+
def get_file_state(self, source_file: str) -> Optional[Dict[str, Any]]:
|
|
143
|
+
row = self.connection.execute(
|
|
144
|
+
"SELECT content_hash, modified_at FROM file_state WHERE source_file = ?",
|
|
145
|
+
(source_file,),
|
|
146
|
+
).fetchone()
|
|
147
|
+
if row is None:
|
|
148
|
+
return None
|
|
149
|
+
return {"content_hash": row[0], "modified_at": row[1]}
|
|
150
|
+
|
|
151
|
+
def upsert_file_state(self, source_file: str, content_hash: str, modified_at: Optional[float]) -> None:
|
|
152
|
+
self.connection.execute(
|
|
153
|
+
"""
|
|
154
|
+
INSERT INTO file_state (source_file, content_hash, modified_at)
|
|
155
|
+
VALUES (?, ?, ?)
|
|
156
|
+
ON CONFLICT(source_file) DO UPDATE SET content_hash = excluded.content_hash, modified_at = excluded.modified_at
|
|
157
|
+
""",
|
|
158
|
+
(source_file, content_hash, modified_at),
|
|
159
|
+
)
|
|
160
|
+
self.connection.commit()
|
|
161
|
+
|
|
162
|
+
def close(self) -> None:
|
|
163
|
+
self.connection.close()
|
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ragmill
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: A lightweight, zero-config local pipeline engine for AI data ingestion, semantic chunking, embeddings, and vector search.
|
|
5
|
+
Project-URL: Homepage, https://github.com/Abdullahbinaqeel/RAGMill
|
|
6
|
+
Project-URL: Repository, https://github.com/Abdullahbinaqeel/RAGMill
|
|
7
|
+
Project-URL: Issues, https://github.com/Abdullahbinaqeel/RAGMill/issues
|
|
8
|
+
Author-email: Abdullah Bin Aqeel <abdulbinaqeel@gmail.com>
|
|
9
|
+
License: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Keywords: data-pipeline,llm-ingestion,onnx,rag,semantic-chunking,vector-embeddings,vector-search
|
|
12
|
+
Classifier: Development Status :: 4 - Beta
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
15
|
+
Classifier: Programming Language :: Python :: 3
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
21
|
+
Classifier: Topic :: Text Processing :: Linguistic
|
|
22
|
+
Requires-Python: >=3.9
|
|
23
|
+
Provides-Extra: all
|
|
24
|
+
Requires-Dist: numpy>=1.22.0; extra == 'all'
|
|
25
|
+
Requires-Dist: onnxruntime>=1.14.0; extra == 'all'
|
|
26
|
+
Requires-Dist: pypdf>=4.0; extra == 'all'
|
|
27
|
+
Requires-Dist: python-docx>=1.0; extra == 'all'
|
|
28
|
+
Requires-Dist: tokenizers>=0.15.0; extra == 'all'
|
|
29
|
+
Provides-Extra: dev
|
|
30
|
+
Requires-Dist: black>=23.0; extra == 'dev'
|
|
31
|
+
Requires-Dist: mypy>=1.0; extra == 'dev'
|
|
32
|
+
Requires-Dist: numpy>=1.22.0; extra == 'dev'
|
|
33
|
+
Requires-Dist: onnxruntime>=1.14.0; extra == 'dev'
|
|
34
|
+
Requires-Dist: pypdf>=4.0; extra == 'dev'
|
|
35
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
36
|
+
Requires-Dist: python-docx>=1.0; extra == 'dev'
|
|
37
|
+
Requires-Dist: reportlab>=4.0; extra == 'dev'
|
|
38
|
+
Requires-Dist: tokenizers>=0.15.0; extra == 'dev'
|
|
39
|
+
Provides-Extra: docx
|
|
40
|
+
Requires-Dist: python-docx>=1.0; extra == 'docx'
|
|
41
|
+
Provides-Extra: embeddings
|
|
42
|
+
Requires-Dist: numpy>=1.22.0; extra == 'embeddings'
|
|
43
|
+
Requires-Dist: onnxruntime>=1.14.0; extra == 'embeddings'
|
|
44
|
+
Requires-Dist: tokenizers>=0.15.0; extra == 'embeddings'
|
|
45
|
+
Provides-Extra: pdf
|
|
46
|
+
Requires-Dist: pypdf>=4.0; extra == 'pdf'
|
|
47
|
+
Description-Content-Type: text/markdown
|
|
48
|
+
|
|
49
|
+
# RAGMill
|
|
50
|
+
|
|
51
|
+
[](https://pypi.org/project/ragmill/)
|
|
52
|
+
[](https://github.com/Abdullahbinaqeel/RAGMill/actions/workflows/ci.yml)
|
|
53
|
+
[](LICENSE)
|
|
54
|
+
[](https://pypi.org/project/ragmill/)
|
|
55
|
+
|
|
56
|
+
A lightweight, zero-config local pipeline engine for AI data ingestion, semantic chunking, embeddings, and vector search.
|
|
57
|
+
|
|
58
|
+
## Install
|
|
59
|
+
|
|
60
|
+
```bash
|
|
61
|
+
pip install ragmill[all] # includes PDF + DOCX + embeddings support
|
|
62
|
+
# or
|
|
63
|
+
pip install ragmill # core only (txt/md), zero dependencies
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
Developing locally instead? Clone the repo and use an editable install:
|
|
67
|
+
|
|
68
|
+
```bash
|
|
69
|
+
pip install -e ".[dev]"
|
|
70
|
+
pytest tests/ -v
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
## Usage
|
|
74
|
+
|
|
75
|
+
### Ingest + chunk
|
|
76
|
+
|
|
77
|
+
```python
|
|
78
|
+
from ragmill import RAGEngine
|
|
79
|
+
|
|
80
|
+
engine = RAGEngine(chunk_size=500, overlap=50)
|
|
81
|
+
chunks = engine.execute_pipeline("./my_documents")
|
|
82
|
+
|
|
83
|
+
for chunk in chunks:
|
|
84
|
+
print(chunk["metadata"]["filename"], chunk["content"][:80])
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
Supports `.txt`, `.md`, `.log`, `.rst`, `.pdf`, and `.docx` out of the box.
|
|
88
|
+
|
|
89
|
+
### Embed + search locally
|
|
90
|
+
|
|
91
|
+
Requires the `embeddings` extra (`pip install -e ".[embeddings]"`). The model
|
|
92
|
+
(a quantized MiniLM ONNX export, ~23MB) downloads once to
|
|
93
|
+
`~/.cache/ragmill/models` and runs fully offline after that.
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
from ragmill import RAGEngine
|
|
97
|
+
from ragmill.embeddings import EmbeddingModel
|
|
98
|
+
from ragmill.vector_store import VectorStore
|
|
99
|
+
|
|
100
|
+
chunks = RAGEngine().execute_pipeline("./my_documents")
|
|
101
|
+
|
|
102
|
+
model = EmbeddingModel()
|
|
103
|
+
vectors = model.embed([c["content"] for c in chunks])
|
|
104
|
+
|
|
105
|
+
store = VectorStore("my_documents.db") # or VectorStore() for in-memory
|
|
106
|
+
store.add(chunks, vectors)
|
|
107
|
+
|
|
108
|
+
query_vector = model.embed(["how does the overlap window work?"])[0]
|
|
109
|
+
for result in store.search(query_vector, top_k=3):
|
|
110
|
+
print(round(result["score"], 3), result["metadata"]["filename"], "->", result["content"][:80])
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
Filter a search to a specific file or a time window:
|
|
114
|
+
|
|
115
|
+
```python
|
|
116
|
+
store.search(query_vector, top_k=3, filename="report.pdf")
|
|
117
|
+
store.search(query_vector, top_k=3, modified_after=1704067200.0) # since 2024-01-01
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
### Keep a store in sync with a folder
|
|
121
|
+
|
|
122
|
+
Re-embedding every file on every run is wasteful once a folder is large.
|
|
123
|
+
`sync_directory` tracks a content hash per file and only touches what
|
|
124
|
+
actually changed:
|
|
125
|
+
|
|
126
|
+
```python
|
|
127
|
+
from ragmill.sync import sync_directory
|
|
128
|
+
|
|
129
|
+
result = sync_directory("./my_documents", engine, model, store)
|
|
130
|
+
print(result) # {"added": 2, "updated": 1, "skipped": 40, "deleted": 1}
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
Unchanged files are skipped without re-embedding. A changed file has its old
|
|
134
|
+
chunks replaced with new ones. A file removed from disk has its chunks
|
|
135
|
+
removed from the store on the next sync.
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
ragmill/__init__.py,sha256=Uunco4CvQroa93Q8DmkxquA1iuERwOSqZKw1Ee94NGU,62
|
|
2
|
+
ragmill/embeddings.py,sha256=Xf0jYRHuSsOrIwiy_2GgZntiVPXlfFCB2scsZoF9Pyg,3502
|
|
3
|
+
ragmill/engine.py,sha256=8gTS-LdnaUBwytS3G3ydtH9blQLhAQ_MdbRPiP3px2Y,6603
|
|
4
|
+
ragmill/parsers.py,sha256=N42QDvA1jCU0eCG4IgX32vASfz_8wFRIlR2rWe-XOko,1059
|
|
5
|
+
ragmill/sync.py,sha256=EV1DnSQbjZCmhWsqnSQnNGPZ41y33zcLDFipIoj7BPY,2724
|
|
6
|
+
ragmill/vector_store.py,sha256=PNAsYBSBfpQ7phAgm9WSZxv7dcTuOtvlzHHbdCFw7-M,6078
|
|
7
|
+
ragmill-0.2.0.dist-info/METADATA,sha256=MEqj_kZ5XbQBTRk2WAY9fnajYaGWXjqh_eqG8EpxeP0,4974
|
|
8
|
+
ragmill-0.2.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
|
|
9
|
+
ragmill-0.2.0.dist-info/licenses/LICENSE,sha256=aE9ynIl05eKXtuUT40jNfGLnjwkP_aJHW2IJEUmQdQI,1075
|
|
10
|
+
ragmill-0.2.0.dist-info/RECORD,,
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Abdullah Bin Aqeel
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|