Files
Project-Thoth/processors/rag_ingestion/ingest.py
T

303 lines
10 KiB
Python

"""Command-line entry point for single-file Markdown ingestion."""
import argparse
import sys
from collections.abc import Sequence
from .chunking import (
DEFAULT_CHUNK_OVERLAP,
DEFAULT_CHUNK_SIZE,
ChunkConfiguration,
chunk_document,
)
from .config import EmbeddingConfig, load_embedding_config
from .embeddings import (
EmbeddingConnectionError,
EmbeddingRequestError,
embed_chunk,
provider_from_config,
)
from .loader import load_markdown_document
from .models import Chunk, Document, Embedding
from .similarity import cosine_similarity, vector_norm
PREVIEW_LENGTH = 120
VECTOR_PREVIEW_LENGTH = 8
CONFIGURATION_TEST_TEXT = "Project Thoth embedding configuration test."
def format_document(document: Document) -> str:
"""Format the mechanical document record for human inspection."""
return "\n".join(
(
f"Document ID: {document.document_id}",
f"Source Path: {document.source_path}",
f"Filename: {document.filename}",
f"Source Format: {document.source_format}",
f"File Size: {document.file_size} bytes",
f"Modified At: {document.modified_at.isoformat()}",
f"Raw Text Length: {len(document.raw_text)} characters",
)
)
def preview(text: str) -> str:
"""Make a chunk boundary visible on one console line."""
return text.replace("\r", "\\r").replace("\n", "\\n")
def format_chunks(
document: Document,
chunks: list[Chunk],
configuration: ChunkConfiguration,
show_chunks: bool = False,
) -> str:
"""Format aggregate statistics and inspectable chunk boundaries."""
sizes = [chunk.size for chunk in chunks]
minimum_size = min(sizes, default=0)
maximum_size = max(sizes, default=0)
average_size = sum(sizes) / len(sizes) if sizes else 0.0
lines = [
"",
"Chunking Summary",
f"Source Size: {len(document.raw_text)} characters",
f"Chunk Size: {configuration.chunk_size} characters",
f"Chunk Overlap: {configuration.chunk_overlap} characters",
f"Chunk Count: {len(chunks)}",
f"Minimum Chunk Size: {minimum_size} characters",
f"Maximum Chunk Size: {maximum_size} characters",
f"Average Chunk Size: {average_size:.2f} characters",
]
for chunk in chunks:
lines.extend(
(
"",
f"Chunk {chunk.chunk_number}",
f" Chunk ID: {chunk.chunk_id}",
f" Chunk Number: {chunk.chunk_number}",
f" Size: {chunk.size} {chunk.size_unit}",
f" Offsets: [{chunk.start_offset}, {chunk.end_offset}) characters",
f" Overlap With Previous: {chunk.overlap_with_previous} characters",
f" Overlap With Next: {chunk.overlap_with_next} characters",
f' Starts With: "{preview(chunk.text[:PREVIEW_LENGTH])}"',
f' Ends With: "{preview(chunk.text[-PREVIEW_LENGTH:])}"',
)
)
if show_chunks:
lines.extend((" Complete Text:", chunk.text))
return "\n".join(lines)
def format_embedding(chunk: Chunk, embedding: Embedding) -> str:
"""Format inspectable properties without dumping the complete vector."""
leading_values = ", ".join(
f"{value:.6f}" for value in embedding.vector[:VECTOR_PREVIEW_LENGTH]
)
return "\n".join(
(
"",
f"Embedding for Chunk {chunk.chunk_number}",
f" Chunk ID: {embedding.chunk_id}",
f" Chunk Number: {chunk.chunk_number}",
f" Provider: {embedding.embedding_provider}",
f" Model: {embedding.embedding_model}",
f" Dimension: {embedding.embedding_dimension}",
f" Vector Norm: {vector_norm(embedding.vector):.8f}",
f" Vector Preview: [{leading_values}, ...] (incomplete)",
f' Text Preview: "{preview(chunk.text[:PREVIEW_LENGTH])}"',
)
)
def selected_chunk(chunks: list[Chunk], chunk_number: int) -> Chunk:
"""Return a zero-based chunk selection or fail with an actionable error."""
if chunk_number < 0 or chunk_number >= len(chunks):
raise ValueError(
f"chunk number {chunk_number} is out of range; "
f"valid range is 0-{len(chunks) - 1}"
)
return chunks[chunk_number]
def check_embedding_config(config: EmbeddingConfig) -> tuple[str, bool]:
"""Make one explicit live request and report configuration diagnostics."""
lines = [
"Oracle Embedding Configuration",
f"Oracle Base URL: {config.oracle_base_url or '(missing)'}",
f"Embedding Model: {config.embedding_model or '(missing)'}",
"Configured Dimension: "
+ (str(config.embedding_dimension) if config.embedding_dimension else "unknown"),
]
try:
provider = provider_from_config(config)
vector = provider.embed(CONFIGURATION_TEST_TEXT)
if (
config.embedding_dimension is not None
and len(vector) != config.embedding_dimension
):
raise ValueError(
f"model {provider.model_name} returned {len(vector)} dimensions; "
f"expected {config.embedding_dimension}"
)
except EmbeddingConnectionError as error:
lines.extend(
(
"Connection: failed",
f"Embedding Request: failed - {error}",
"Returned Dimension: unavailable",
)
)
return "\n".join(lines), False
except (EmbeddingRequestError, ValueError) as error:
lines.extend(
(
"Connection: succeeded",
f"Embedding Request: failed - {error}",
"Returned Dimension: unavailable",
)
)
return "\n".join(lines), False
lines.extend(
(
"Connection: succeeded",
"Embedding Request: succeeded",
f"Returned Dimension: {len(vector)}",
)
)
return "\n".join(lines), True
def build_parser() -> argparse.ArgumentParser:
"""Build the command-line parser."""
parser = argparse.ArgumentParser(
description="Ingest one Markdown Primary Source for inspection."
)
parser.add_argument(
"source", nargs="?", help="Path to one Markdown (.md) source file"
)
parser.add_argument(
"--chunk-size",
type=int,
default=DEFAULT_CHUNK_SIZE,
help=f"Characters per chunk (default: {DEFAULT_CHUNK_SIZE})",
)
parser.add_argument(
"--chunk-overlap",
type=int,
default=DEFAULT_CHUNK_OVERLAP,
help=f"Characters repeated between chunks (default: {DEFAULT_CHUNK_OVERLAP})",
)
parser.add_argument(
"--show-chunks",
action="store_true",
help="Print complete chunk text in addition to boundary previews",
)
parser.add_argument(
"--embed-chunk",
type=int,
action="append",
default=[],
help="Embed one zero-based chunk number; may be repeated",
)
parser.add_argument(
"--compare-chunks",
type=int,
nargs=2,
metavar=("CHUNK_A", "CHUNK_B"),
help="Embed and compare two zero-based chunk numbers",
)
parser.add_argument(
"--check-embedding-config",
action="store_true",
help="Validate configured Oracle embedding inference with one small request",
)
return parser
def main(argv: Sequence[str] | None = None) -> int:
"""Run ingestion from command-line arguments."""
# Vault Markdown may contain characters outside a host console's legacy
# code page. UTF-8 keeps previews and --show-chunks faithful to the source.
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
parser = build_parser()
arguments = parser.parse_args(argv)
try:
embedding_config = load_embedding_config()
except ValueError as error:
parser.error(str(error))
if arguments.check_embedding_config:
if arguments.source:
parser.error("source must be omitted with --check-embedding-config")
report, succeeded = check_embedding_config(embedding_config)
print(report)
return 0 if succeeded else 1
if not arguments.source:
parser.error("source is required unless --check-embedding-config is used")
try:
document = load_markdown_document(arguments.source)
configuration = ChunkConfiguration(
chunk_size=arguments.chunk_size,
chunk_overlap=arguments.chunk_overlap,
)
chunks = chunk_document(document, configuration)
requested_numbers = list(arguments.embed_chunk)
if arguments.compare_chunks:
requested_numbers.extend(arguments.compare_chunks)
embeddings_by_chunk: dict[int, Embedding] = {}
if requested_numbers:
provider = provider_from_config(embedding_config)
for chunk_number in dict.fromkeys(requested_numbers):
chunk = selected_chunk(chunks, chunk_number)
embeddings_by_chunk[chunk_number] = embed_chunk(chunk, provider)
except (FileNotFoundError, RuntimeError, ValueError, OSError, UnicodeError) as error:
parser.error(str(error))
print(format_document(document))
print(format_chunks(document, chunks, configuration, arguments.show_chunks))
for chunk_number in dict.fromkeys(requested_numbers):
print(
format_embedding(
selected_chunk(chunks, chunk_number),
embeddings_by_chunk[chunk_number],
)
)
if arguments.compare_chunks:
chunk_a, chunk_b = arguments.compare_chunks
score = cosine_similarity(
embeddings_by_chunk[chunk_a].vector,
embeddings_by_chunk[chunk_b].vector,
)
print(
"\n".join(
(
"",
"Chunk Comparison",
f" Chunk A: {chunk_a}",
f" Chunk B: {chunk_b}",
f" Cosine Similarity: {score:.8f}",
)
)
)
return 0
if __name__ == "__main__":
raise SystemExit(main())