import os from pathlib import Path import re from typing import Dict, List, Tuple, Optional, Any import json from tqdm import tqdm import time import psutil import numpy as np from rank_bm25 import BM25Okapi from nltk.tokenize import word_tokenize from nltk.corpus import stopwords from nltk.stem import WordNetLemmatizer from litellm import completion, batch_completion from .async_logger import AsyncLogger import litellm import pickle import hashlib # <--- ADDED for file-hash from fnmatch import fnmatch import glob litellm.set_verbose = False def _compute_file_hash(file_path: Path) -> str: """Compute MD5 hash for the file's entire content.""" hash_md5 = hashlib.md5() with file_path.open("rb") as f: for chunk in iter(lambda: f.read(4096), b""): hash_md5.update(chunk) return hash_md5.hexdigest() class AsyncLLMTextManager: def __init__( self, docs_dir: Path, logger: Optional[AsyncLogger] = None, max_concurrent_calls: int = 5, batch_size: int = 3 ) -> None: self.docs_dir = docs_dir self.logger = logger self.max_concurrent_calls = max_concurrent_calls self.batch_size = batch_size self.bm25_index = None self.document_map: Dict[str, Any] = {} self.tokenized_facts: List[str] = [] self.bm25_index_file = self.docs_dir / "bm25_index.pkl" async def _process_document_batch(self, doc_batch: List[Path]) -> None: """Process a batch of documents in parallel""" contents = [] for file_path in doc_batch: try: with open(file_path, 'r', encoding='utf-8') as f: contents.append(f.read()) except Exception as e: self.logger.error(f"Error reading {file_path}: {str(e)}") contents.append("") # Add empty content to maintain batch alignment prompt = """Given a documentation file, generate a list of atomic facts where each fact: 1. Represents a single piece of knowledge 2. Contains variations in terminology for the same concept 3. References relevant code patterns if they exist 4. Is written in a way that would match natural language queries Each fact should follow this format: : | | Example Facts: browser_config: Configure headless mode and browser type for AsyncWebCrawler | headless, browser_type, chromium, firefox | BrowserConfig(browser_type="chromium", headless=True) redis_connection: Redis client connection requires host and port configuration | redis setup, redis client, connection params | Redis(host='localhost', port=6379, db=0) pandas_filtering: Filter DataFrame rows using boolean conditions | dataframe filter, query, boolean indexing | df[df['column'] > 5] Wrap your response in ... tags. """ # Prepare messages for batch processing messages_list = [ [ {"role": "user", "content": f"{prompt}\n\nGenerate index for this documentation:\n\n{content}"} ] for content in contents if content ] try: responses = batch_completion( model="anthropic/claude-3-5-sonnet-latest", messages=messages_list, logger_fn=None ) # Process responses and save index files for response, file_path in zip(responses, doc_batch): try: index_content_match = re.search( r'(.*?)', response.choices[0].message.content, re.DOTALL ) if not index_content_match: self.logger.warning(f"No ... content found for {file_path}") continue index_content = re.sub( r"\n\s*\n", "\n", index_content_match.group(1) ).strip() if index_content: index_file = file_path.with_suffix('.q.md') with open(index_file, 'w', encoding='utf-8') as f: f.write(index_content) self.logger.info(f"Created index file: {index_file}") else: self.logger.warning(f"No index content found in response for {file_path}") except Exception as e: self.logger.error(f"Error processing response for {file_path}: {str(e)}") except Exception as e: self.logger.error(f"Error in batch completion: {str(e)}") def _validate_fact_line(self, line: str) -> Tuple[bool, Optional[str]]: if "|" not in line: return False, "Missing separator '|'" parts = [p.strip() for p in line.split("|")] if len(parts) != 3: return False, f"Expected 3 parts, got {len(parts)}" concept_part = parts[0] if ":" not in concept_part: return False, "Missing ':' in concept definition" return True, None def _load_or_create_token_cache(self, fact_file: Path) -> Dict: """ Load token cache from .q.tokens if present and matching file hash. Otherwise return a new structure with updated file-hash. """ cache_file = fact_file.with_suffix(".q.tokens") current_hash = _compute_file_hash(fact_file) if cache_file.exists(): try: with open(cache_file, "r") as f: cache = json.load(f) # If the hash matches, return it directly if cache.get("content_hash") == current_hash: return cache # Otherwise, we signal that it's changed self.logger.info(f"Hash changed for {fact_file}, reindex needed.") except json.JSONDecodeError: self.logger.warning(f"Corrupt token cache for {fact_file}, rebuilding.") except Exception as e: self.logger.warning(f"Error reading cache for {fact_file}: {str(e)}") # Return a fresh cache return {"facts": {}, "content_hash": current_hash} def _save_token_cache(self, fact_file: Path, cache: Dict) -> None: cache_file = fact_file.with_suffix(".q.tokens") # Always ensure we're saving the correct file-hash cache["content_hash"] = _compute_file_hash(fact_file) with open(cache_file, "w") as f: json.dump(cache, f) def preprocess_text(self, text: str) -> List[str]: parts = [x.strip() for x in text.split("|")] if "|" in text else [text] # Remove : after the first word of parts[0] parts[0] = re.sub(r"^(.*?):", r"\1", parts[0]) lemmatizer = WordNetLemmatizer() stop_words = set(stopwords.words("english")) - { "how", "what", "when", "where", "why", "which", } tokens = [] for part in parts: if "(" in part and ")" in part: code_tokens = re.findall( r'[\w_]+(?=\()|[\w_]+(?==[\'"]{1}[\w_]+[\'"]{1})', part ) tokens.extend(code_tokens) words = word_tokenize(part.lower()) tokens.extend( [ lemmatizer.lemmatize(token) for token in words if token not in stop_words ] ) return tokens def maybe_load_bm25_index(self, clear_cache=False) -> bool: """ Load existing BM25 index from disk, if present and clear_cache=False. """ if not clear_cache and os.path.exists(self.bm25_index_file): self.logger.info("Loading existing BM25 index from disk.") with open(self.bm25_index_file, "rb") as f: data = pickle.load(f) self.tokenized_facts = data["tokenized_facts"] self.bm25_index = data["bm25_index"] return True return False def build_search_index(self, clear_cache=False) -> None: """ Checks for new or modified .q.md files by comparing file-hash. If none need reindexing and clear_cache is False, loads existing index if available. Otherwise, reindexes only changed/new files and merges or creates a new index. """ # If clear_cache is True, we skip partial logic: rebuild everything from scratch if clear_cache: self.logger.info("Clearing cache and rebuilding full search index.") if self.bm25_index_file.exists(): self.bm25_index_file.unlink() process = psutil.Process() self.logger.info("Checking which .q.md files need (re)indexing...") # Gather all .q.md files q_files = [self.docs_dir / f for f in os.listdir(self.docs_dir) if f.endswith(".q.md")] # We'll store known (unchanged) facts in these lists existing_facts: List[str] = [] existing_tokens: List[List[str]] = [] # Keep track of invalid lines for logging invalid_lines = [] needSet = [] # files that must be (re)indexed for qf in q_files: token_cache_file = qf.with_suffix(".q.tokens") # If no .q.tokens or clear_cache is True → definitely reindex if clear_cache or not token_cache_file.exists(): needSet.append(qf) continue # Otherwise, load the existing cache and compare hash cache = self._load_or_create_token_cache(qf) # If the .q.tokens was out of date (i.e. changed hash), we reindex if len(cache["facts"]) == 0 or cache.get("content_hash") != _compute_file_hash(qf): needSet.append(qf) else: # File is unchanged → retrieve cached token data for line, cache_data in cache["facts"].items(): existing_facts.append(line) existing_tokens.append(cache_data["tokens"]) self.document_map[line] = qf # track the doc for that fact if not needSet and not clear_cache: # If no file needs reindexing, try loading existing index if self.maybe_load_bm25_index(clear_cache=False): self.logger.info("No new/changed .q.md files found. Using existing BM25 index.") return else: # If there's no existing index, we must build a fresh index from the old caches self.logger.info("No existing BM25 index found. Building from cached facts.") if existing_facts: self.logger.info(f"Building BM25 index with {len(existing_facts)} cached facts.") self.bm25_index = BM25Okapi(existing_tokens) self.tokenized_facts = existing_facts with open(self.bm25_index_file, "wb") as f: pickle.dump({ "bm25_index": self.bm25_index, "tokenized_facts": self.tokenized_facts }, f) else: self.logger.warning("No facts found at all. Index remains empty.") return # ----------------------------------------------------- /Users/unclecode/.crawl4ai/docs/14_proxy_security.q.q.tokens '/Users/unclecode/.crawl4ai/docs/14_proxy_security.q.md' # If we reach here, we have new or changed .q.md files # We'll parse them, reindex them, and then combine with existing_facts # ----------------------------------------------------- self.logger.info(f"{len(needSet)} file(s) need reindexing. Parsing now...") # 1) Parse the new or changed .q.md files new_facts = [] new_tokens = [] with tqdm(total=len(needSet), desc="Indexing changed files") as file_pbar: for file in needSet: # We'll build up a fresh cache fresh_cache = {"facts": {}, "content_hash": _compute_file_hash(file)} try: with open(file, "r", encoding="utf-8") as f_obj: content = f_obj.read().strip() lines = [l.strip() for l in content.split("\n") if l.strip()] for line in lines: is_valid, error = self._validate_fact_line(line) if not is_valid: invalid_lines.append((file, line, error)) continue tokens = self.preprocess_text(line) fresh_cache["facts"][line] = { "tokens": tokens, "added": time.time(), } new_facts.append(line) new_tokens.append(tokens) self.document_map[line] = file # Save the new .q.tokens with updated hash self._save_token_cache(file, fresh_cache) mem_usage = process.memory_info().rss / 1024 / 1024 self.logger.debug(f"Memory usage after {file.name}: {mem_usage:.2f}MB") except Exception as e: self.logger.error(f"Error processing {file}: {str(e)}") file_pbar.update(1) if invalid_lines: self.logger.warning(f"Found {len(invalid_lines)} invalid fact lines:") for file, line, error in invalid_lines: self.logger.warning(f"{file}: {error} in line: {line[:50]}...") # 2) Merge newly tokenized facts with the existing ones all_facts = existing_facts + new_facts all_tokens = existing_tokens + new_tokens # 3) Build BM25 index from combined facts self.logger.info(f"Building BM25 index with {len(all_facts)} total facts (old + new).") self.bm25_index = BM25Okapi(all_tokens) self.tokenized_facts = all_facts # 4) Save the updated BM25 index to disk with open(self.bm25_index_file, "wb") as f: pickle.dump({ "bm25_index": self.bm25_index, "tokenized_facts": self.tokenized_facts }, f) final_mem = process.memory_info().rss / 1024 / 1024 self.logger.info(f"Search index updated. Final memory usage: {final_mem:.2f}MB") async def generate_index_files(self, force_generate_facts: bool = False, clear_bm25_cache: bool = False) -> None: """ Generate index files for all documents in parallel batches Args: force_generate_facts (bool): If True, regenerate indexes even if they exist clear_bm25_cache (bool): If True, clear existing BM25 index cache """ self.logger.info("Starting index generation for documentation files.") md_files = [ self.docs_dir / f for f in os.listdir(self.docs_dir) if f.endswith('.md') and not any(f.endswith(x) for x in ['.q.md', '.xs.md']) ] # Filter out files that already have .q files unless force=True if not force_generate_facts: md_files = [ f for f in md_files if not (self.docs_dir / f.name.replace('.md', '.q.md')).exists() ] if not md_files: self.logger.info("All index files exist. Use force=True to regenerate.") else: # Process documents in batches for i in range(0, len(md_files), self.batch_size): batch = md_files[i:i + self.batch_size] self.logger.info(f"Processing batch {i//self.batch_size + 1}/{(len(md_files)//self.batch_size) + 1}") await self._process_document_batch(batch) self.logger.info("Index generation complete, building/updating search index.") self.build_search_index(clear_cache=clear_bm25_cache) def generate(self, sections: List[str], mode: str = "extended") -> str: # Get all markdown files all_files = glob.glob(str(self.docs_dir / "[0-9]*.md")) + \ glob.glob(str(self.docs_dir / "[0-9]*.xs.md")) # Extract base names without extensions base_docs = {Path(f).name.split('.')[0] for f in all_files if not Path(f).name.endswith('.q.md')} # Filter by sections if provided if sections: base_docs = {doc for doc in base_docs if any(section.lower() in doc.lower() for section in sections)} # Get file paths based on mode files = [] for doc in sorted(base_docs, key=lambda x: int(x.split('_')[0]) if x.split('_')[0].isdigit() else 999999): if mode == "condensed": xs_file = self.docs_dir / f"{doc}.xs.md" regular_file = self.docs_dir / f"{doc}.md" files.append(str(xs_file if xs_file.exists() else regular_file)) else: files.append(str(self.docs_dir / f"{doc}.md")) # Read and format content content = [] for file in files: try: with open(file, 'r', encoding='utf-8') as f: fname = Path(file).name content.append(f"{'#'*20}\n# {fname}\n{'#'*20}\n\n{f.read()}") except Exception as e: self.logger.error(f"Error reading {file}: {str(e)}") return "\n\n---\n\n".join(content) if content else "" def search(self, query: str, top_k: int = 5) -> str: if not self.bm25_index: return "No search index available. Call build_search_index() first." query_tokens = self.preprocess_text(query) doc_scores = self.bm25_index.get_scores(query_tokens) mean_score = np.mean(doc_scores) std_score = np.std(doc_scores) score_threshold = mean_score + (0.25 * std_score) file_data = self._aggregate_search_scores( doc_scores=doc_scores, score_threshold=score_threshold, query_tokens=query_tokens, ) ranked_files = sorted( file_data.items(), key=lambda x: ( x[1]["code_match_score"] * 2.0 + x[1]["match_count"] * 1.5 + x[1]["total_score"] ), reverse=True, )[:top_k] results = [] for file, _ in ranked_files: main_doc = str(file).replace(".q.md", ".md") if os.path.exists(self.docs_dir / main_doc): with open(self.docs_dir / main_doc, "r", encoding='utf-8') as f: only_file_name = main_doc.split("/")[-1] content = [ "#" * 20, f"# {only_file_name}", "#" * 20, "", f.read() ] results.append("\n".join(content)) return "\n\n---\n\n".join(results) def _aggregate_search_scores( self, doc_scores: List[float], score_threshold: float, query_tokens: List[str] ) -> Dict: file_data = {} for idx, score in enumerate(doc_scores): if score <= score_threshold: continue fact = self.tokenized_facts[idx] file_path = self.document_map[fact] if file_path not in file_data: file_data[file_path] = { "total_score": 0, "match_count": 0, "code_match_score": 0, "matched_facts": [], } components = fact.split("|") if "|" in fact else [fact] code_match_score = 0 if len(components) == 3: code_ref = components[2].strip() code_tokens = self.preprocess_text(code_ref) code_match_score = len(set(query_tokens) & set(code_tokens)) / len(query_tokens) file_data[file_path]["total_score"] += score file_data[file_path]["match_count"] += 1 file_data[file_path]["code_match_score"] = max( file_data[file_path]["code_match_score"], code_match_score ) file_data[file_path]["matched_facts"].append(fact) return file_data def refresh_index(self) -> None: """Convenience method for a full rebuild.""" self.build_search_index(clear_cache=True)