File size: 2,788 Bytes
477651d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """
Utility helpers: document loading, env validation, stats.
Reconstructed to match the interface app.py imports.
"""
import os
from pathlib import Path
from typing import List
from langchain.schema import Document
from langchain_community.document_loaders import (
PyPDFLoader,
TextLoader,
Docx2txtLoader,
)
REQUIRED_ENV_VARS = [
"PINECONE_API_KEY",
"PINECONE_INDEX_NAME",
"OPENAI_API_KEY",
]
SUPPORTED_EXTENSIONS = {".pdf", ".txt", ".md", ".docx"}
def validate_env_variables() -> bool:
"""Return True only if every required environment variable is set and non-empty."""
return all(os.getenv(var) for var in REQUIRED_ENV_VARS)
def _loader_for(file_path: str):
"""Pick the right LangChain loader for a file extension."""
ext = Path(file_path).suffix.lower()
if ext == ".pdf":
return PyPDFLoader(file_path)
if ext == ".docx":
return Docx2txtLoader(file_path)
if ext in {".txt", ".md"}:
return TextLoader(file_path, encoding="utf-8")
raise ValueError(f"Unsupported file type: {ext}")
def load_document(file_path: str) -> List[Document]:
"""Load a single document into a list of LangChain Document objects."""
path = Path(file_path)
if not path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
docs = _loader_for(str(path)).load()
for doc in docs:
doc.metadata.setdefault("source", path.name)
return docs
def load_documents_from_directory(directory_path: str) -> List[Document]:
"""Load every supported document in a directory (recursively)."""
directory = Path(directory_path)
if not directory.exists():
raise FileNotFoundError(f"Directory not found: {directory_path}")
documents: List[Document] = []
for path in sorted(directory.rglob("*")):
if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS:
try:
documents.extend(load_document(str(path)))
except Exception as exc: # noqa: BLE001
print(f"Skipped {path.name}: {exc}")
return documents
def get_document_stats(documents: List[Document]) -> dict:
"""Summarize a batch of loaded documents."""
total_characters = sum(len(doc.page_content) for doc in documents)
sources = sorted(
{doc.metadata.get("source", "Unknown") for doc in documents}
)
return {
"total_documents": len(documents),
"total_characters": total_characters,
"sources": sources,
}
def format_file_size(num_bytes: int) -> str:
"""Human-readable file size."""
size = float(num_bytes)
for unit in ["B", "KB", "MB", "GB", "TB"]:
if size < 1024.0:
return f"{size:.1f} {unit}"
size /= 1024.0
return f"{size:.1f} PB"
|