feat(knowledge): 任务队列、独立 worker 与 pgvector 向量索引
知识处理改为数据库任务队列，独立 worker 做切块与向量化，失败可重试。
This commit is contained in:
WangShen007 committed
1 parent
53825b0cfb
commit
738e37fbb0
36 files changed
+7590
No files matched your search
@@ -0,0 +1,526 @@
|
||||
"""
|
||||
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 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 <loc> 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 <a href>; 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
|
||||
Reference in new issue
Block a user