Spaces:
Running on Zero
Running on Zero
File size: 10,715 Bytes
0828c2c | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 | """Main document loader for guaranteed context availability."""
import hashlib
import logging
import time
from pathlib import Path
from .config_loader import get_config
from .document_processor import DocumentProcessor
logger = logging.getLogger(__name__)
class MainDocumentLoader:
"""Load and manage the main profile document."""
def __init__(self):
"""Initialize main document loader with configuration."""
self.config = get_config()
self.doc_processor = DocumentProcessor()
# Cache management
self._cached_content: str | None = None
self._cached_file_hash: str | None = None
self._last_check_time: float = 0
# Config
self.enabled = self.config.get("main_document.enabled", False)
path_str = self.config.get_env(
"MAIN_DOCUMENT_PATH", self.config.get("main_document.path", "")
)
self.path = Path(path_str) if path_str else None
self.max_tokens = self.config.get("main_document.max_tokens", 10000)
self.summarize_if_exceeds = self.config.get("main_document.summarize_if_exceeds", True)
self.cache_enabled = self.config.get("main_document.cache_enabled", True)
self.cache_check_interval = self.config.get("main_document.cache_check_interval", 60)
def load_main_document(self) -> str:
"""Load main document with caching and error handling.
Returns:
Main document content as string (empty if disabled/error)
"""
if not self.enabled:
logger.debug("Main document feature is disabled")
return ""
if not self.path or not self.path.exists():
if self.config.get("main_document.fail_silently", True):
logger.warning(f"Main document not found: {self.path}")
return ""
else:
msg = f"Main document not found: {self.path}"
raise FileNotFoundError(msg)
# Check cache validity
if self._should_use_cache():
logger.debug("Using cached main document")
return self._cached_content or ""
# Load fresh content
try:
content = self._load_document_by_format()
# Token management
token_count = self.count_tokens(content)
logger.info(f"Main document loaded: {token_count} tokens from {self.path.name}")
if token_count > self.max_tokens:
logger.warning(
f"Main document exceeds limit ({token_count} > {self.max_tokens} tokens)"
)
if self.summarize_if_exceeds:
content = self._summarize_content(content)
logger.info(f"Summarized to {self.count_tokens(content)} tokens")
else:
content = self.truncate_to_tokens(content, self.max_tokens)
logger.info(f"Truncated to {self.max_tokens} tokens")
# Update cache
if self.cache_enabled:
self._update_cache(content)
return content
except Exception as e:
logger.error(f"Error loading main document: {e}")
if self.config.get("main_document.fail_silently", True):
return ""
raise
def _load_document_by_format(self) -> str:
"""Auto-detect format and load document.
Returns:
Document content as string
"""
if not self.path:
return ""
extension = self.path.suffix.lower()
logger.debug(f"Auto-detected format: {extension}")
# Use existing DocumentProcessor for format handling
try:
# Load as LangChain documents
if extension == ".pdf":
docs = self.doc_processor.load_pdf(self.path)
elif extension in [".docx", ".doc"]:
docs = self.doc_processor.load_docx(self.path)
elif extension in [".html", ".htm"]:
docs = self.doc_processor.load_html(self.path)
elif extension in [".md", ".txt"]:
docs = self.doc_processor.load_text(self.path)
else:
# Fallback: try as text
logger.warning(f"Unsupported format {extension}, treating as text")
with open(self.path, encoding="utf-8") as f:
return f.read()
# Combine all document chunks (main doc shouldn't be chunked)
content = "\n\n".join(doc.page_content for doc in docs)
return content
except Exception as e:
logger.error(f"Error loading {extension} file: {e}")
raise
def count_tokens(self, text: str) -> int:
"""Count tokens in text using tiktoken.
Args:
text: Text to count tokens for
Returns:
Number of tokens
"""
try:
import tiktoken
# Use cl100k_base encoding (GPT-4, llama compatible)
encoding = tiktoken.get_encoding("cl100k_base")
return len(encoding.encode(text))
except ImportError:
# Fallback: rough estimation (1 token ≈ 4 characters)
logger.warning("tiktoken not available, using rough estimation (1 token ≈ 4 chars)")
return len(text) // 4
def truncate_to_tokens(self, text: str, max_tokens: int) -> str:
"""Truncate text to maximum tokens.
Args:
text: Text to truncate
max_tokens: Maximum tokens allowed
Returns:
Truncated text
"""
try:
import tiktoken
encoding = tiktoken.get_encoding("cl100k_base")
tokens = encoding.encode(text)
if len(tokens) <= max_tokens:
return text
truncated_tokens = tokens[:max_tokens]
truncated_text = encoding.decode(truncated_tokens)
return truncated_text + "\n\n[... content truncated ...]"
except ImportError:
# Fallback: character-based truncation
char_limit = max_tokens * 4
if len(text) <= char_limit:
return text
return text[:char_limit] + "\n\n[... content truncated ...]"
def _summarize_content(self, content: str) -> str:
"""Summarize content using LLM if exceeds token limit.
Args:
content: Original content
Returns:
Summarized content
"""
from .llm_handler import get_llm_handler
logger.info("Summarizing main document using LLM...")
try:
llm_handler = get_llm_handler()
llm = llm_handler.get_llm()
summarization_prompt = self.config.get(
"main_document.summarization_prompt",
"Summarize the following professional profile, preserving all critical information:",
)
target_tokens = self.config.get(
"main_document.summarization_target_tokens", int(self.max_tokens * 0.8)
)
# Truncate input if way too large (to fit in LLM context)
max_input_tokens = 30000 # Conservative limit
if self.count_tokens(content) > max_input_tokens:
content = self.truncate_to_tokens(content, max_input_tokens)
logger.warning(
f"Input content too large, truncated to {max_input_tokens} tokens for summarization"
)
full_prompt = f"""{summarization_prompt}
---
{content}
---
Provide a comprehensive summary (target: ~{target_tokens} tokens):"""
summary = llm.invoke(full_prompt)
# Verify summary is actually shorter
summary_tokens = self.count_tokens(summary)
if summary_tokens > self.max_tokens:
logger.warning(f"Summary still exceeds limit ({summary_tokens} tokens), truncating")
summary = self.truncate_to_tokens(summary, self.max_tokens)
return summary
except Exception as e:
logger.error(f"Error during summarization: {e}")
logger.info("Falling back to truncation")
return self.truncate_to_tokens(content, self.max_tokens)
def _should_use_cache(self) -> bool:
"""Check if cached content is still valid.
Returns:
True if cache should be used, False otherwise
"""
if not self.cache_enabled or self._cached_content is None:
return False
# Check if enough time has passed
current_time = time.time()
if current_time - self._last_check_time < self.cache_check_interval:
return True
# Check if file has changed
current_hash = self._calculate_file_hash()
if current_hash == self._cached_file_hash:
self._last_check_time = current_time
return True
return False
def _calculate_file_hash(self) -> str:
"""Calculate MD5 hash of file for change detection.
Returns:
MD5 hash string
"""
if not self.path or not self.path.exists():
return ""
hash_md5 = hashlib.md5()
with open(self.path, "rb") as f:
for chunk in iter(lambda: f.read(4096), b""):
hash_md5.update(chunk)
return hash_md5.hexdigest()
def _update_cache(self, content: str) -> None:
"""Update cache with new content.
Args:
content: Content to cache
"""
self._cached_content = content
self._cached_file_hash = self._calculate_file_hash()
self._last_check_time = time.time()
logger.debug("Main document cache updated")
def invalidate_cache(self) -> None:
"""Manually invalidate cache (useful for testing/reloading)."""
self._cached_content = None
self._cached_file_hash = None
self._last_check_time = 0
logger.info("Main document cache invalidated")
# Singleton instance
_loader_instance: MainDocumentLoader | None = None
def get_main_document_loader() -> MainDocumentLoader:
"""Get singleton main document loader instance.
Returns:
MainDocumentLoader instance
"""
global _loader_instance
if _loader_instance is None:
_loader_instance = MainDocumentLoader()
return _loader_instance
def reload_main_document_loader() -> MainDocumentLoader:
"""Reload main document loader (useful for testing).
Returns:
New MainDocumentLoader instance
"""
global _loader_instance
_loader_instance = MainDocumentLoader()
return _loader_instance
|