Spaces:
Running
Running
File size: 3,976 Bytes
1e68437 599e860 1e68437 | 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 93 94 95 96 97 98 99 100 101 102 103 | import os
import logging
from pathlib import Path
from typing import List
logger = logging.getLogger("DocIngestion")
def extract_pdf_text(pdf_path: Path) -> str:
"""
Safely extracts text from a PDF file using pypdf.
If pypdf is not installed, logs a warning and returns an empty string.
"""
try:
import pypdf
reader = pypdf.PdfReader(str(pdf_path))
text_parts = []
for i, page in enumerate(reader.pages):
page_text = page.extract_text()
if page_text:
text_parts.append(page_text)
return "\n".join(text_parts)
except ImportError:
logger.warning(
f"pypdf is not installed. Cannot extract text from PDF: {pdf_path.name}. "
"Please run 'pip install pypdf' to enable PDF ingestion."
)
return ""
except Exception as e:
logger.error(f"Error reading PDF file {pdf_path.name}: {e}", exc_info=True)
return ""
def compile_knowledge_context(directories: List[str] = ["docs", "knowledge"]) -> str:
"""
Scans the specified directories relative to the project root for PDF, Markdown,
and text files. Combines their contents into a single structured string block.
Target files are related to WorldQuant BRAIN documentation:
- Double Neutralization
- Fast D1 Documentation
- Sentiment1
- Model77
- earnings4
Returns:
A compiled string of all matching documents, formatted with markdown headers.
"""
project_root = Path(__file__).resolve().parent.parent
compiled_parts = []
for dir_name in directories:
target_dir = project_root / dir_name
if not target_dir.exists():
logger.debug(f"Directory {dir_name} does not exist at {target_dir}")
continue
logger.info(f"Ingesting documents from directory: {target_dir}")
# Search recursively or directly
for file_path in target_dir.rglob("*"):
if file_path.is_dir():
continue
suffix = file_path.suffix.lower()
if suffix not in (".md", ".txt", ".pdf"):
continue
logger.info(f"Processing doc: {file_path.name}")
content = ""
try:
if suffix in (".md", ".txt"):
with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
content = f.read()
elif suffix == ".pdf":
content = extract_pdf_text(file_path)
except Exception as e:
logger.error(f"Failed to read file {file_path.name}: {e}")
continue
if content.strip():
# Keep the prompt bounded when consultant PDFs are present.
# Preserve both the opening definitions and closing tests;
# otherwise a full multi-megabyte corpus can crowd out the
# actual generation constraints sent to the model.
content = content.strip()
max_chars = 40000
if len(content) > max_chars:
head = 30000
content = (
content[:head]
+ "\n\n[DOCUMENT CONTENT TRUNCATED FOR PROMPT BUDGET]\n\n"
+ content[-(max_chars - head):]
)
logger.info("Truncated document %s to %s characters for model context.", file_path.name, max_chars)
doc_name = file_path.stem.replace("_", " ").title()
compiled_parts.append(
f"### DOCUMENT: {doc_name} ({file_path.name})\n\n{content.strip()}\n\n"
)
if not compiled_parts:
logger.warning("No knowledge documents were found or successfully ingested.")
return "No local platform documentation available in docs/ or knowledge/ directories."
return "\n".join(compiled_parts)
|