""" Copyright 2024-2026 ReplyPier 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 agno.knowledge.pdf import PDFKnowledgeBase, PDFImageReader, PDFReader from agno.vectordb.pgvector import PgVector, SearchType from app.knowledge.optimized_pgvector import OptimizedPgVector from app.core.config import settings from app.core.logger import get_logger from app.knowledge.crawl_scope import DEFAULT_CRAWL_SCOPE from app.knowledge.enhanced_website_kb import EnhancedWebsiteKnowledgeBase from app.knowledge.enhanced_website_reader import EnhancedWebsiteReader from app.knowledge.sitemap_reader import SitemapReader from app.knowledge.chunking import PageChunking from app.knowledge.enhanced_pdf_kb import EnhancedPDFKnowledgeBase from app.core.s3 import delete_file_from_s3, owned_s3_key from app.knowledge.pdf_download import download_pdf from app.knowledge.pdf_limits import PdfRefused, check_page_count from app.models.knowledge import Knowledge, SourceType from app.models.knowledge_to_agent import KnowledgeToAgent from app.models.knowledge_queue import ProcessingStage, QueueStatus from app.repositories.ai_config import AIConfigRepository from app.database import SessionLocal from app.repositories.knowledge_to_agent import KnowledgeToAgentRepository from app.repositories.knowledge import KnowledgeRepository from app.repositories.knowledge_queue import KnowledgeQueueRepository from typing import List, Optional, Dict, Union import os import shutil import tempfile import asyncio from urllib.parse import unquote, urlparse from uuid import UUID from app.knowledge.embedder import get_embedder # Try to import enterprise modules try: from app.enterprise.repositories.subscription import SubscriptionRepository HAS_ENTERPRISE = True except ImportError: HAS_ENTERPRISE = False logger = get_logger(__name__) # Downloaded PDFs land under backend/temp rather than the system temp dir so # they stay reachable inside the Docker image. PDF_TEMP_ROOT = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "temp") _MAX_PDF_NAME_CHARS = 200 def _pdf_name_from_url(url: str) -> str: """Knowledge-source name for a PDF URL: the decoded last path segment without its extension, falling back to the host. basename() after decoding strips any separator an encoded "..%2F" smuggled in.""" parsed = urlparse(url) name = os.path.basename(unquote(parsed.path.rstrip("/").split("/")[-1])) name = os.path.splitext(name)[0].strip(". ") return (name or parsed.hostname or "document")[:_MAX_PDF_NAME_CHARS] def _raise_unless(success: bool) -> None: """Fail the source before anything marks it completed.""" if not success: raise Exception("Failed to process knowledge source") class KnowledgeManager: def __init__(self, org_id: UUID, agent_id: Optional[str] = None): self.org_id = org_id self.agent_id = agent_id # Get API key from AI config using context manager with SessionLocal() as db: ai_config_repo = AIConfigRepository(db) ai_config = ai_config_repo.get_active_config(org_id) # Default to SentenceTransformer embedder if no AI config is found embedder = None table_name = f"d_{org_id}" # One loaded model per process; agno's embedder reloads it per call. embedder = get_embedder() # Dimensions will be automatically set by the model # Import the shared engine to avoid creating separate connection pools from app.database import engine as shared_engine self.vector_db = OptimizedPgVector( table_name=table_name, db_engine=shared_engine, schema="ai", search_type=SearchType.vector, embedder=embedder, auto_upgrade_schema=True ) def _add_knowledge_source(self, source: str, source_type: SourceType): """Track knowledge source in database""" with SessionLocal() as db: knowledge_repo = KnowledgeRepository(db) link_repo = KnowledgeToAgentRepository(db) # Create or get knowledge source knowledge = Knowledge( organization_id=self.org_id, source=source, source_type=source_type, table_name=self.vector_db.table_name, schema=self.vector_db.schema ) logger.debug(f"Adding knowledge source: {knowledge}") knowledge = knowledge_repo.create(knowledge) logger.debug(f"Knowledge source added: {knowledge}") # Link to agent if specified if self.agent_id: link = KnowledgeToAgent( knowledge_id=knowledge.id, agent_id=self.agent_id ) link_repo.create(link) return knowledge async def add_pdf_urls(self, urls: List[str]) -> bool: """Add knowledge from PDF URLs: user-supplied links, and the presigned S3 links of PDFs uploaded through the dashboard. Every fetch goes through download_pdf, which refuses internal hosts on each redirect hop and caps the size.""" try: for url in urls: logger.info(f"Adding PDF URL: {url}") filename = _pdf_name_from_url(url) os.makedirs(PDF_TEMP_ROOT, exist_ok=True) # A private directory per download: the readable file name # comes from the URL, and must not collide or escape. temp_dir = tempfile.mkdtemp(dir=PDF_TEMP_ROOT) try: temp_path = os.path.join(temp_dir, f"{filename}.pdf") await asyncio.to_thread(download_pdf, url, temp_path) logger.info(f"Downloaded {url} to {temp_path}, processing as: {filename}") await self.add_pdf_files([temp_path], filename=filename) await self._delete_uploaded_pdf(url) finally: shutil.rmtree(temp_dir, ignore_errors=True) return True except PdfRefused: raise # its message is the reason shown on the source except Exception as e: logger.error(f"Error adding PDF URLs: {str(e)}") return False async def _delete_uploaded_pdf(self, url: str) -> None: """Remove a dashboard-uploaded PDF from S3 once it is indexed. Only this org's own knowledge uploads qualify: the URL is user-supplied, and any other S3 object it names is not ours to delete.""" key = owned_s3_key(url) if not key or not key.startswith(f"knowledge/{self.org_id}/"): return if await delete_file_from_s3(url): logger.info(f"Deleted S3 object: {key}") else: logger.warning(f"Failed to delete S3 file: {key}") async def add_websites(self, urls: List[str], max_links: int = 10) -> bool: """Add knowledge from websites using the enhanced website reader""" try: # Convert agent_id to string if it exists agent_id_filter = [str(self.agent_id)] if self.agent_id else [] for url in urls: logger.debug(f"Adding website: {url}") # Use config values when not in enterprise mode if not HAS_ENTERPRISE: knowledge_base = EnhancedWebsiteKnowledgeBase( urls=[url], # Pass single URL in a list max_links=settings.KB_MAX_LINKS, max_depth=settings.KB_MAX_DEPTH, min_content_length=settings.KB_MIN_CONTENT_LENGTH, timeout=settings.KB_TIMEOUT, max_retries=settings.KB_MAX_RETRIES, max_workers=min(settings.KB_MAX_WORKERS, settings.EMBEDDING_MAX_WORKERS), # Use the smaller of the two vector_db=self.vector_db ) else: # Use enterprise-provided values where available, config defaults otherwise knowledge_base = EnhancedWebsiteKnowledgeBase( urls=[url], # Pass single URL in a list max_links=max_links, # Use enterprise subscription limit max_depth=settings.KB_MAX_DEPTH, # Use config for depth min_content_length=settings.KB_MIN_CONTENT_LENGTH, # Use config timeout=settings.KB_TIMEOUT, # Use config for timeout max_retries=settings.KB_MAX_RETRIES, # Use config for retries max_workers=min(settings.KB_MAX_WORKERS, settings.EMBEDDING_MAX_WORKERS), # Use the smaller of the two vector_db=self.vector_db ) logger.debug(f"Enhanced knowledge base created for: {url}") knowledge_base.load(recreate=False, upsert=True, filters={ "name": url, "agent_id": agent_id_filter, "org_id": str(self.org_id) }) logger.debug(f"Enhanced knowledge base loaded for: {url}") self._add_knowledge_source(url, SourceType.WEBSITE) return True except Exception as e: logger.error(f"Error adding websites: {str(e)}") return False async def add_pdf_files(self, files: List[str], chunk: bool = True, reader: Optional[Union[PDFReader, PDFImageReader]] = None, filename: Optional[str] = None) -> bool: """Add knowledge from PDF files""" try: # Process each file individually since PDFKnowledgeBase expects a single path for file_path in files: temp_file = None path_to_use = file_path # Check if file exists before processing if not os.path.exists(path_to_use): logger.error(f"PDF file not found: {path_to_use}") raise FileNotFoundError(f"PDF file not found: {path_to_use}") logger.info(f"Processing PDF file: {path_to_use} (size: {os.path.getsize(path_to_use)} bytes)") # Before any page is read: past the limit, reading is the cost check_page_count(path_to_use) try: # First try with regular PDFReader try: knowledge_base = EnhancedPDFKnowledgeBase( path=path_to_use, vector_db=self.vector_db, reader=PDFReader(chunk=chunk, chunking_strategy=PageChunking()) if not reader else reader ) if filename is None: filename = os.path.splitext(os.path.basename(file_path))[0] # Convert agent_id to string if it exists agent_id_filter = [str(self.agent_id)] if self.agent_id else [] logger.info(f"Adding PDF knowledge source: {filename}") filters = { "name": filename, "agent_id": agent_id_filter, "org_id": str(self.org_id) } # Run PDF processing in thread pool to avoid blocking other APIs await asyncio.to_thread( knowledge_base.load, recreate=False, upsert=True, filters=filters ) except Exception as e: logger.warning(f"Regular PDFReader failed for {path_to_use}, trying PDFImageReader: {str(e)}") # Fallback to PDFImageReader if PDFReader fails knowledge_base = EnhancedPDFKnowledgeBase( path=path_to_use, vector_db=self.vector_db, reader=PDFImageReader(chunk=chunk, chunking_strategy=PageChunking()) ) if filename is None: filename = os.path.splitext(os.path.basename(file_path))[0] # Convert agent_id to string if it exists agent_id_filter = [str(self.agent_id)] if self.agent_id else [] filters = { "name": filename, "agent_id": agent_id_filter, "org_id": str(self.org_id) } # Run PDF processing in thread pool to avoid blocking other APIs await asyncio.to_thread( knowledge_base.load, recreate=False, upsert=True, filters=filters ) self._add_knowledge_source(filename, SourceType.FILE) finally: # Clean up temp file if we created one if temp_file and os.path.exists(temp_file.name): os.unlink(temp_file.name) return True except PdfRefused: raise # its message is the reason shown on the source except Exception as e: logger.error(f"Error adding PDF files: {str(e)}") return False def get_knowledge_base(self) -> List[Dict]: """Get all knowledge sources for the organization or specific agent""" try: with SessionLocal() as db: knowledge_repo = KnowledgeRepository(db) if self.agent_id: sources = knowledge_repo.get_by_agent(self.agent_id) else: sources = knowledge_repo.get_by_org(self.org_id) return [{ 'id': source.id, 'source': source.source, 'source_type': source.source_type, 'table_name': source.table_name, 'schema': source.schema, 'agent_ids': [link.agent_id for link in source.agent_links] } for source in sources] except Exception as e: logger.error(f"Error getting knowledge base: {str(e)}") return [] async def process_knowledge(self, queue_item): """Process knowledge based on source type""" try: # Get DB session for progress updates with SessionLocal() as db: queue_repo = KnowledgeQueueRepository(db) # Process based on source type if queue_item.source_type == 'pdf_file': # Update to crawling stage (PDF extraction) queue_repo.update_progress( queue_item.id, ProcessingStage.CRAWLING, total_items=1, processed_items=0, force_stage_update=True ) # Process the PDF (already async, but contains blocking operations) try: success = await self.add_pdf_files([queue_item.source]) _raise_unless(success) finally: # The uploaded copy goes either way, refused PDFs too if not queue_item.source.startswith('http') and os.path.exists(queue_item.source): os.remove(queue_item.source) # Update to embedding/training stage queue_repo.update_progress( queue_item.id, ProcessingStage.EMBEDDING, total_items=1, processed_items=1, force_stage_update=True ) # Mark as completed queue_repo.update_status( queue_item.id, QueueStatus.COMPLETED ) elif queue_item.source_type == 'pdf_url': # Update to crawling stage (downloading PDF) queue_repo.update_progress( queue_item.id, ProcessingStage.CRAWLING, total_items=1, processed_items=0, force_stage_update=True ) # Process the PDF URL (already async, but contains blocking operations) success = await self.add_pdf_urls([queue_item.source]) _raise_unless(success) # Update to embedding/training stage queue_repo.update_progress( queue_item.id, ProcessingStage.EMBEDDING, total_items=1, processed_items=1, force_stage_update=True ) # Mark as completed queue_repo.update_status( queue_item.id, QueueStatus.COMPLETED ) elif queue_item.source_type == 'website': # Update to crawling stage (don't add main URL here, let the crawler handle it) queue_repo.update_progress( queue_item.id, ProcessingStage.CRAWLING ) metadata = queue_item.queue_metadata or {} max_links = metadata.get('max_links', settings.KB_MAX_LINKS) # How far link-following may wander from the seed URL; older # queue items carry no scope and get the default (seed host). crawl_scope = metadata.get('crawl_scope') or DEFAULT_CRAWL_SCOPE # Use EnhancedWebsiteReader for all website crawling logger.info(f"Using EnhancedWebsiteReader for queue item: {queue_item.source}") reader = EnhancedWebsiteReader( max_depth=settings.KB_MAX_DEPTH, max_links=max_links, min_content_length=settings.KB_MIN_CONTENT_LENGTH, timeout=settings.KB_TIMEOUT, max_retries=settings.KB_MAX_RETRIES, max_workers=settings.KB_MAX_WORKERS, crawl_scope=crawl_scope, verify_ssl=False # Disable SSL verification to support websites with invalid/self-signed certificates ) # Process the website - pass queue item and repo for URL tracking knowledge_base = EnhancedWebsiteKnowledgeBase( urls=[queue_item.source], max_links=max_links, crawl_scope=crawl_scope, vector_db=self.vector_db, reader=reader ) # Add queue_item and repo for URL tracking knowledge_base.queue_item = queue_item knowledge_base.queue_repo = queue_repo # Process the website in a thread pool to avoid blocking other APIs logger.info(f"Starting website processing in background thread for {queue_item.source}") agent_id_filter = [str(self.agent_id)] if self.agent_id else [] filters = { "name": queue_item.source, "agent_id": agent_id_filter, "org_id": str(self.org_id) } # Run the heavy synchronous operation in a thread pool await asyncio.to_thread( knowledge_base.load, recreate=False, upsert=True, filters=filters ) logger.info(f"Completed website processing for {queue_item.source}") self._add_knowledge_source(queue_item.source, SourceType.WEBSITE) success = True # Completion status is handled internally by EnhancedWebsiteKnowledgeBase elif queue_item.source_type == 'sitemap': # Crawl exactly the pages listed in the sitemap.xml (parsed # from entries, incl. sitemap-index recursion) — no BFS. queue_repo.update_progress(queue_item.id, ProcessingStage.CRAWLING) max_links = queue_item.queue_metadata.get( 'max_links', settings.KB_MAX_LINKS) if queue_item.queue_metadata else settings.KB_MAX_LINKS logger.info(f"Using SitemapReader for queue item: {queue_item.source}") reader = SitemapReader( max_depth=1, # never follow ; pages come from the sitemap max_links=max_links, min_content_length=settings.KB_MIN_CONTENT_LENGTH, timeout=settings.KB_TIMEOUT, max_retries=settings.KB_MAX_RETRIES, max_workers=settings.KB_MAX_WORKERS, verify_ssl=False ) # One source (the sitemap URL); its listed pages are the sub-pages. knowledge_base = EnhancedWebsiteKnowledgeBase( urls=[queue_item.source], max_links=max_links, vector_db=self.vector_db, reader=reader ) knowledge_base.queue_item = queue_item knowledge_base.queue_repo = queue_repo agent_id_filter = [str(self.agent_id)] if self.agent_id else [] filters = { "name": queue_item.source, "agent_id": agent_id_filter, "org_id": str(self.org_id) } await asyncio.to_thread( knowledge_base.load, recreate=False, upsert=True, filters=filters ) # Fail clearly if the sitemap listed no pages (wrong URL / not a # sitemap) rather than creating an empty knowledge source. if getattr(reader, '_sitemap_page_count', 0) == 0: raise ValueError( "No pages found in the sitemap. Make sure the URL points " "to a valid sitemap.xml." ) logger.info(f"Completed sitemap processing for {queue_item.source}") self._add_knowledge_source(queue_item.source, SourceType.WEBSITE) success = True else: raise ValueError(f"Unsupported source type: { queue_item.source_type}") _raise_unless(success) return success except Exception as e: logger.error(f"Error processing knowledge: {str(e)}") raise