""" Copyright 2024-2026 ChatterMate Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ from typing import Any, Dict, Iterator, List, Optional, Union import time from datetime import datetime from concurrent.futures import ThreadPoolExecutor, as_completed from queue import Queue, Empty import threading from agno.document import Document from agno.knowledge.agent import AgentKnowledge from agno.embedder import Embedder from agno.document.reader.website_reader import WebsiteReader from app.core.logger import get_logger from app.core.config import settings from app.models.knowledge_queue import ProcessingStage, QueueStatus from pydantic import model_validator from app.knowledge.crawl_scope import DEFAULT_CRAWL_SCOPE from app.knowledge.deadline import Deadline, SourceTimedOut from app.knowledge.enhanced_website_reader import ( EnhancedWebsiteReader, BotProtectionError, EmptyCrawlError, FAILURE_ADVICE, ) from app.knowledge.failure_advice import FailureCause # Initialize logger for this module logger = get_logger(__name__) class EnhancedWebsiteKnowledgeBase(AgentKnowledge): """Enhanced knowledge base for websites with robust content extraction using EnhancedWebsiteReader""" urls: List[str] = [] reader: Optional[WebsiteReader] = None # Reader parameters - using settings from config max_depth: int = settings.KB_MAX_DEPTH max_links: int = settings.KB_MAX_LINKS min_content_length: int = settings.KB_MIN_CONTENT_LENGTH timeout: int = settings.KB_TIMEOUT max_retries: int = settings.KB_MAX_RETRIES max_workers: int = settings.KB_MAX_WORKERS batch_size: int = settings.KB_BATCH_SIZE crawl_scope: str = DEFAULT_CRAWL_SCOPE # These are not part of the model schema, but added as instance attributes _queue_item = None _queue_repo = None _embedding_queue = None _embedding_executor = None _embedding_results = None _embedding_lock = None _deadline = None @property def queue_item(self): return self._queue_item @queue_item.setter def queue_item(self, value): self._queue_item = value @property def queue_repo(self): return self._queue_repo @queue_repo.setter def queue_repo(self, value): self._queue_repo = value @model_validator(mode="after") def set_reader(self) -> "EnhancedWebsiteKnowledgeBase": """Set the reader if not provided""" if self.reader is None: logger.info(f"Initializing default EnhancedWebsiteReader with max_depth={self.max_depth}, max_links={self.max_links}, max_workers={self.max_workers}") self.reader = EnhancedWebsiteReader( max_depth=self.max_depth, max_links=self.max_links, min_content_length=self.min_content_length, timeout=self.timeout, max_retries=self.max_retries, max_workers=self.max_workers, crawl_scope=self.crawl_scope, verify_ssl=False # Disable SSL verification to handle self-signed/invalid certificates ) else: logger.info(f"Using custom reader: {type(self.reader).__name__}") return self def _init_embedding_system(self): """Initialize the embedding queue and thread pool system""" if (settings.ENABLE_IMMEDIATE_EMBEDDING and self.vector_db and hasattr(self.vector_db, 'embedder') and self.vector_db.embedder is not None): self._embedding_queue = Queue() self._embedding_results = [] self._embedding_lock = threading.Lock() self._embedding_executor = ThreadPoolExecutor( max_workers=settings.EMBEDDING_MAX_WORKERS, thread_name_prefix="EmbeddingWorker" ) def _shutdown_embedding_system(self): """Shutdown the embedding system and wait for completion""" if not self._embedding_executor: return try: # Signal workers to stop by adding sentinel values if self._embedding_queue: for i in range(settings.EMBEDDING_MAX_WORKERS): self._embedding_queue.put(None) # Shutdown executor with timeout to prevent hanging self._embedding_executor.shutdown(wait=True) # Clean up resources self._embedding_queue = None self._embedding_executor = None self._embedding_results = None self._embedding_lock = None except Exception as e: logger.error(f"Error during embedding system shutdown: {type(e).__name__}: {str(e)}") # Force shutdown if graceful shutdown fails if self._embedding_executor: self._embedding_executor.shutdown(wait=False) def _embedding_worker(self): """Worker function that processes documents from the embedding queue""" worker_id = threading.current_thread().name while True: document = None try: document = self._embedding_queue.get(timeout=2.0) if document is None: # Sentinel value to stop worker self._embedding_queue.task_done() break # Out of time: drain the queue without the work; the run fails. if self._out_of_time: self._embedding_queue.task_done() continue # Process the document if hasattr(document, 'embedding') and document.embedding is None: document.embed(embedder=self.vector_db.embedder) with self._embedding_lock: self._embedding_results.append(document.id) self._embedding_queue.task_done() except Empty: # Queue is empty - this is expected when waiting for work # Continue looping without logging continue except Exception as e: # Log actual errors with full details logger.error(f"Error in embedding worker {worker_id}: {type(e).__name__}: {str(e)}") if document and hasattr(document, 'id'): logger.error(f"Failed to embed document: {document.id}") # Mark task as done even if failed to prevent queue blocking if document is not None: try: self._embedding_queue.task_done() except ValueError: # task_done called too many times - ignore pass def _process_url(self, url: str) -> List[Document]: """Process a single URL and return its documents with immediate embedding""" try: url_start_time = time.time() logger.info(f"Crawling URL: {url}") # Create a callback for URL crawling progress tracking def on_url_crawled_callback(crawled_url: str): """Callback to track each URL as it's crawled""" try: if self.queue_item and self.queue_repo: logger.debug(f"🌐 Tracking crawled URL: {crawled_url}") # Only add the crawled URL, don't update the processing stage success = self.queue_repo.update_progress( self.queue_item.id, crawled_url=crawled_url ) if success: logger.debug(f"✅ Successfully tracked crawled URL: {crawled_url}") else: logger.warning(f"❌ Failed to track crawled URL: {crawled_url}") except Exception as e: logger.error(f"Error tracking crawled URL {crawled_url}: {str(e)}") # Create a callback for immediate document processing def on_document_callback(document: Document): """Callback to queue documents for embedding""" try: # Queue document for embedding if immediate embedding is enabled if (settings.ENABLE_IMMEDIATE_EMBEDDING and self._embedding_queue is not None and document.embedding is None): self._embedding_queue.put(document) except Exception as e: logger.error(f"Error in document callback for {document.id}: {str(e)}") # Get documents from the reader with both callbacks documents = self.reader.read( url=url, vector_db_callback=on_document_callback, url_crawled_callback=on_url_crawled_callback ) url_end_time = time.time() url_duration = url_end_time - url_start_time # Log document details page_count = len(set(doc.meta_data.get('url', '') for doc in documents)) embedded_count = sum(1 for doc in documents if doc.embedding is not None) logger.info(f"Completed {url} - Extracted {len(documents)} documents ({embedded_count} embedded) from {page_count} pages ({url_duration:.2f}s)") return documents except Exception as e: logger.error(f"Error reading URL {url}: {str(e)}") return [] def _embed_document(self, document: Document, embedder: Embedder) -> None: """Embed a single document (for parallel processing)""" try: if document.embedding is None: document.embed(embedder=embedder) except Exception as e: logger.error(f"Error embedding document '{document.id}': {str(e)}") raise def _process_document_batch(self, documents: List[Document], filters: Optional[Dict[str, Any]] = None) -> int: """Insert a batch into the vector database; return how many were stored. The count is what the caller checks before reporting success. Swallowing the failure and returning nothing is how a run that stored zero vectors still finished as COMPLETED. """ if not documents or not self.vector_db: return 0 try: self.vector_db.upsert(documents=documents, filters=filters) return len(documents) except Exception as e: logger.error(f"Error processing document batch: {str(e)}") return 0 @property def document_lists(self) -> Iterator[List[Document]]: """ Iterate over urls and yield lists of documents with queue-based parallel embedding. Each object yielded by the iterator is a list of documents that are embedded via a dedicated embedding thread pool for optimal performance. Returns: Iterator[List[Document]]: Iterator yielding list of documents """ if self.reader is not None: total_start_time = time.time() logger.info(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Starting to process {len(self.urls)} URLs with queue-based parallel embedding") # Process URLs in parallel using ThreadPoolExecutor with ThreadPoolExecutor(max_workers=self.max_workers) as executor: # Submit all URLs for processing future_to_url = {executor.submit(self._process_url, url): url for url in self.urls} # Process completed futures as they finish for future in as_completed(future_to_url): url = future_to_url[future] try: documents = future.result() if documents: # Documents should already be embedded from immediate processing embedded_count = sum(1 for doc in documents if doc.embedding is not None) logger.debug(f"Yielding {len(documents)} documents ({embedded_count} embedded) from {url}") yield documents except Exception as e: logger.error(f"Error processing URL {url}: {str(e)}") yield [] total_end_time = time.time() total_duration = total_end_time - total_start_time logger.info(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Completed processing all {len(self.urls)} URLs with immediate embedding (Total time: {total_duration:.2f}s)") def _raise_if_nothing_stored(self, stored_documents: int, skipped_existing: int = 0) -> None: """Fail the run when it stored nothing. A zero-page run used to be marked COMPLETED, which made a broken crawl indistinguishable from a healthy source that happens to have no sub-pages — the dashboard showed success while the agent had nothing to answer from. Raising here is enough: process_queue_item already turns any exception into a FAILED queue item with the message shown to the user. The count is of documents *stored*, not crawled: a page that crawls fine and then fails to embed leaves the vector table empty, and counting it as progress reported success for a source with nothing behind it. """ if stored_documents > 0: return # Bot protection first: it is the specific, actionable cause. if getattr(self.reader, '_challenge_blocked', 0) > 0: # The shared advice, so every place that reports a failure says the same raise BotProtectionError( FAILURE_ADVICE[FailureCause.BOT_BLOCKED].format(url=self._source_label()) ) # Every URL was already in the vector store, so there was nothing new to # crawl. That is a no-op, not a failed read: the content is already there # and the caller links this agent to it afterwards. Reporting the crawl # errors here blamed the site ("may require JavaScript") for what was # really a re-add of a source we had already indexed. if skipped_existing > 0 and skipped_existing >= len(self.urls): logger.info( f"Nothing to crawl: all {skipped_existing} URL(s) are already in the " "knowledge base; keeping the existing content." ) return raise EmptyCrawlError(self._failure_message()) @property def _out_of_time(self) -> bool: return self._deadline is not None and self._deadline.expired def _raise_if_out_of_time(self) -> None: """Fail a source that ran past its time limit, before anything of it is stored: a partial copy would make adding it again skip the crawl ("already loaded").""" if self._out_of_time: raise SourceTimedOut( FAILURE_ADVICE[FailureCause.TOO_LONG].format(url=self._source_label()) ) def _source_label(self) -> str: return self.urls[0] if self.urls else 'this source' def _failure_message(self) -> str: """What to tell the user, based on what the crawl actually hit. The reader records a cause per failed page; the most common one is what the source as a whole ran into. Only when it recorded nothing do we fall back to the old catch-all. """ url = self._source_label() cause = None if self.reader is not None and hasattr(self.reader, 'dominant_failure_cause'): cause = self.reader.dominant_failure_cause() advice = FAILURE_ADVICE.get(cause) if advice: return advice.format(url=url) return ( f"No content could be read from {url}. The page may be empty, require " "JavaScript, or be unreachable — check the URL is correct and publicly " "accessible." ) def load( self, recreate: bool = False, upsert: bool = True, skip_existing: bool = True, filters: Optional[Dict[str, Any]] = None, ) -> None: """ Load the website contents to the vector db with parallel processing and batch insertion Args: recreate (bool, optional): Whether to recreate the collection. Defaults to False. upsert (bool, optional): Whether to upsert documents. Defaults to True. skip_existing (bool, optional): Whether to skip existing documents. Defaults to True. filters (Optional[Dict[str, Any]], optional): Filters to apply to the documents. Defaults to None. """ total_start_time = time.time() if self.vector_db is None: logger.warning("No vector db provided") return if self.reader is None: logger.warning("No reader provided") return if recreate: self.vector_db.drop() self.vector_db.create() elif not self.vector_db.exists(): self.vector_db.create() logger.info(f"Starting knowledge base loading - {len(self.urls)} URLs") self._deadline = Deadline(settings.KB_SOURCE_TIME_LIMIT_SECONDS) if hasattr(self.reader, 'deadline'): self.reader.deadline = self._deadline # Initialize the embedding system self._init_embedding_system() # Start embedding workers after initialization if self._embedding_executor and self._embedding_queue is not None: for i in range(settings.EMBEDDING_MAX_WORKERS): self._embedding_executor.submit(self._embedding_worker) # Check if URLs exist in vector db urls_to_read = self.urls.copy() skipped_existing = 0 if not recreate and skip_existing: try: urls_to_read = [url for url in self.urls if not self.vector_db.name_exists(name=url)] skipped_existing = len(self.urls) - len(urls_to_read) if skipped_existing > 0: logger.info(f"Skipping {skipped_existing} already loaded URLs") except Exception as e: logger.error(f"Error checking existing URLs: {str(e)}") urls_to_read = self.urls.copy() skipped_existing = 0 # Process URLs in parallel with batched vector DB insertion total_documents = 0 # Crawled and stored are different numbers, and only the second one means # the source works. A document that fails to embed or upsert still counts # toward total_documents. stored_documents = 0 completed_urls = 0 total_urls = len(urls_to_read) # Initialize progress tracking if we have queue context if self.queue_item and self.queue_repo: self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.CRAWLING, progress_percentage=10.0, total_items=total_urls, processed_items=0, force_stage_update=True ) # Keep track of crawling vs upsert timing crawl_start_time = time.time() with ThreadPoolExecutor(max_workers=self.max_workers) as executor: # Submit all URLs for processing future_to_url = {executor.submit(self._process_url, url): url for url in urls_to_read} all_documents = [] # Process completed futures as they finish for future in as_completed(future_to_url): url = future_to_url[future] try: documents = future.result() total_documents += len(documents) completed_urls += 1 # Update progress after each URL completes if self.queue_item and self.queue_repo and total_urls > 0: # Calculate crawling progress (10% to 60% of total) crawling_progress = 10.0 + (completed_urls / total_urls) * 50.0 self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.CRAWLING, progress_percentage=min(crawling_progress, 60.0), total_items=total_urls, processed_items=completed_urls, force_stage_update=True ) logger.debug(f"Crawling progress: {completed_urls}/{total_urls} URLs ({crawling_progress:.1f}%)") # Collect documents for batch processing all_documents.extend(documents) except Exception as e: logger.error(f"Error processing URL {url}: {str(e)}") completed_urls += 1 # Count failed URLs too crawl_end_time = time.time() crawl_duration = crawl_end_time - crawl_start_time logger.info(f"Crawling completed - {len(urls_to_read)} URLs, {total_documents} documents, {crawl_duration:.2f}s") # Update progress to embedding stage if self.queue_item and self.queue_repo: self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.EMBEDDING, progress_percentage=60.0, force_stage_update=True ) # Get embedding results count before shutdown embedded_count = 0 if self._embedding_results: with self._embedding_lock: embedded_count = len(self._embedding_results) # Shutdown the embedding system (includes waiting for completion) self._shutdown_embedding_system() logger.info(f"Embedding completed - {embedded_count} documents embedded") self._raise_if_out_of_time() # Update embedding progress if self.queue_item and self.queue_repo: self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.EMBEDDING, total_items=total_documents, processed_items=embedded_count ) # Check if any documents missed embedding (should be rare with queue-based embedding) unembedded_count = sum(1 for doc in all_documents if doc.embedding is None) if unembedded_count > 0: logger.warning(f"Found {unembedded_count} unembedded documents, embedding them now") if self.vector_db and hasattr(self.vector_db, 'embedder'): # Embed remaining documents immediately without parallel processing # to avoid multiple model loading for doc in all_documents: if self._out_of_time: break if doc.embedding is None: try: doc.embed(embedder=self.vector_db.embedder) except Exception as e: logger.error(f"Error embedding document {doc.id}: {str(e)}") else: logger.error(f"Cannot embed {unembedded_count} documents - no embedder available") self._raise_if_out_of_time() # Process all documents in optimally sized batches if all_documents and upsert: logger.info(f"Inserting {len(all_documents)} documents into vector database") # Update progress for DB operations if self.queue_item and self.queue_repo: self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.EMBEDDING, total_items=total_documents, processed_items=min(total_documents, embedded_count) ) # Process in larger batches for better DB performance total_batches = (len(all_documents) + self.batch_size - 1) // self.batch_size for i, batch_start in enumerate(range(0, len(all_documents), self.batch_size)): batch = all_documents[batch_start:batch_start + self.batch_size] stored_documents += self._process_document_batch(batch, filters) # Update progress during DB operations if self.queue_item and self.queue_repo and total_batches > 1: processed_docs = min(total_documents, (i + 1) * self.batch_size) self.queue_repo.update_progress( self.queue_item.id, ProcessingStage.EMBEDDING, total_items=total_documents, processed_items=processed_docs ) # No approximate (HNSW/IVFFlat) index on purpose: every agent search # filters by agent_id, and an approximate index applies that filter # after picking candidates, so it can return fewer results than asked. # Exact search over a knowledge table (a few thousand rows) is fast. # (agno's optimize() also failed outright: it binds index parameters, # which Postgres refuses in DDL under psycopg 3.) total_end_time = time.time() total_duration = total_end_time - total_start_time db_duration = total_duration - crawl_duration # Calculate embedding statistics final_embedded_count = sum(1 for doc in all_documents if doc.embedding is not None) embedding_rate = (final_embedded_count / len(all_documents) * 100) if all_documents else 0 # A run that stored nothing is a failure. Raised before the COMPLETED # write below so the queue item cannot report success for an empty crawl. # Retrievable means embedded *and* upserted. An unembedded document is # skipped by the vector store, so counting the upsert alone reported # success for a source whose table stayed empty (chatring.ai); counting # the embedding alone would miss a failed insert. retrievable_documents = ( min(stored_documents, final_embedded_count) if upsert else final_embedded_count ) self._raise_if_nothing_stored(retrievable_documents, skipped_existing) # Mark as completed after all processing is done if self.queue_item and self.queue_repo: self.queue_repo.update_status( self.queue_item.id, QueueStatus.COMPLETED ) logger.info(f"✓ Marked queue item {self.queue_item.id} as completed") logger.info(f"✓ Completed loading {total_documents} documents from {len(urls_to_read)} URLs") logger.info(f"✓ Embedding approach: {'Queue-based parallel' if settings.ENABLE_IMMEDIATE_EMBEDDING else 'Sequential batch'}") logger.info(f"✓ Embedding success rate: {final_embedded_count}/{len(all_documents)} ({embedding_rate:.1f}%)") logger.info(f"✓ Total processing time: {total_duration:.2f}s (Crawling+Embedding: {crawl_duration:.2f}s, DB operations: {db_duration:.2f}s)")