Julian Carlson
140 lines
4.2 KiB
Python
140 lines
4.2 KiB
Python
import hashlib
|
|
from pathlib import Path
|
|
from datetime import datetime, timezone
|
|
from .config import CONFIG
|
|
from .parsers import parse_file, PARSERS
|
|
from .chunk import chunk_text, count_tokens
|
|
from .embed import embed_batch
|
|
from . import db
|
|
|
|
|
|
def file_hash(path):
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
for chunk in iter(lambda: f.read(8192), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def category_from_path(path):
|
|
try:
|
|
rel = path.relative_to(CONFIG.corpus_root)
|
|
except ValueError:
|
|
return "unknown"
|
|
parts = rel.parts[:-1]
|
|
return "/".join(parts) if parts else "uncategorized"
|
|
|
|
|
|
def build_chunk_prefix(title, category, last_modified, doc_metadata):
|
|
"""Construct a metadata breadcrumb to prepend to each chunk."""
|
|
parts = []
|
|
if last_modified:
|
|
parts.append(last_modified.strftime("%Y-%m-%d"))
|
|
if category:
|
|
parts.append(f"Category: {category}")
|
|
if title:
|
|
parts.append(f"Title: {title}")
|
|
fmt = doc_metadata.get("format")
|
|
if fmt:
|
|
parts.append(f"Format: {fmt}")
|
|
return f"[{' | '.join(parts)}]\n\n" if parts else ""
|
|
|
|
|
|
def ingest_file(path):
|
|
path = Path(path)
|
|
if path.suffix.lower() not in PARSERS:
|
|
return {"path": str(path), "status": "skipped", "reason": "unsupported type"}
|
|
|
|
fhash = file_hash(path)
|
|
category = category_from_path(path)
|
|
|
|
try:
|
|
title, text, metadata = parse_file(path)
|
|
except Exception as e:
|
|
return {"path": str(path), "status": "error", "reason": f"parse failed: {e}"}
|
|
|
|
last_modified = datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
|
|
doc_id, changed = db.upsert_document(
|
|
source_path=str(path),
|
|
file_hash=fhash,
|
|
file_type=path.suffix.lower(),
|
|
category=category,
|
|
title=title,
|
|
last_modified=last_modified,
|
|
metadata=metadata,
|
|
)
|
|
|
|
if not changed:
|
|
return {"path": str(path), "status": "unchanged", "doc_id": doc_id}
|
|
|
|
chunks = chunk_text(text)
|
|
if not chunks:
|
|
return {"path": str(path), "status": "empty", "doc_id": doc_id}
|
|
|
|
# Build per-chunk topic labels using the LLM
|
|
from .topic_extractor import extract_topic
|
|
|
|
base_prefix = build_chunk_prefix(title, category, last_modified, metadata)
|
|
|
|
chunk_records = []
|
|
embeddings_to_compute = []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
# Generate per-chunk topic label
|
|
topic = extract_topic(chunk, doc_title=title, doc_category=category)
|
|
|
|
if topic:
|
|
chunk_with_context = f"{base_prefix}Topic: {topic}\n\n{chunk}"
|
|
else:
|
|
chunk_with_context = base_prefix + chunk
|
|
|
|
embeddings_to_compute.append(chunk_with_context)
|
|
chunk_records.append({
|
|
"chunk_index": i,
|
|
"content_with_context": chunk_with_context,
|
|
"raw_chunk": chunk,
|
|
"topic": topic,
|
|
})
|
|
|
|
# Batch embed
|
|
embeddings = []
|
|
for i in range(0, len(embeddings_to_compute), CONFIG.embed_batch_size):
|
|
batch = embeddings_to_compute[i:i + CONFIG.embed_batch_size]
|
|
embeddings.extend(embed_batch(batch))
|
|
|
|
# Build final records for DB
|
|
final_records = [
|
|
{
|
|
"chunk_index": rec["chunk_index"],
|
|
"content": rec["content_with_context"],
|
|
"embedding": emb,
|
|
"token_count": count_tokens(rec["content_with_context"]),
|
|
"metadata": {
|
|
"raw_chunk": rec["raw_chunk"],
|
|
"topic": rec["topic"],
|
|
},
|
|
}
|
|
for rec, emb in zip(chunk_records, embeddings)
|
|
]
|
|
db.insert_chunks(doc_id, final_records)
|
|
|
|
return {
|
|
"path": str(path),
|
|
"status": "indexed",
|
|
"doc_id": doc_id,
|
|
"chunks": len(chunks),
|
|
"category": category,
|
|
}
|
|
|
|
|
|
def ingest_directory(root=None):
|
|
root = Path(root) if root else Path(CONFIG.corpus_root)
|
|
results = []
|
|
for path in root.rglob("*"):
|
|
if path.is_file() and path.suffix.lower() in PARSERS:
|
|
print(f"Processing: {path}", flush=True)
|
|
result = ingest_file(path)
|
|
results.append(result)
|
|
print(f" → {result['status']}", flush=True)
|
|
return results
|