feat(knowledge): 任务队列、独立 worker 与 pgvector 向量索引

知识处理改为数据库任务队列,独立 worker 做切块与向量化,失败可重试。
This commit is contained in:
WangShen007 committed 2026-09-30 22:05:00 +08:00
1 parent 7937de5d7d
commit c7011040c3
36 files changed
+7590

No files matched your search

+19
View File
@@ -0,0 +1,19 @@
"""
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 app.core.processor.status import PROCESSOR_STATUS
__all__ = ["PROCESSOR_STATUS"]
+22
View File
@@ -0,0 +1,22 @@
"""
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.
"""
# Shared processor status to track knowledge processor state
PROCESSOR_STATUS = {
"last_run": None,
"is_running": False,
"error": None
}
+168
View File
@@ -0,0 +1,168 @@
# ReplyPier Enhanced Website Knowledge Base
This module provides an enhanced website knowledge extraction system that significantly improves content extraction from web pages compared to the default implementation.
## Key Features
- **Multi-strategy content extraction**: Uses multiple strategies to extract the most relevant content from web pages.
- **Robust error handling**: Implements retries with exponential backoff for failed requests.
- **Text density analysis**: Can identify important content based on paragraph density when standard selectors fail.
- **Smart content cleaning**: Automatically removes boilerplate, navigation, ads, and other non-content elements.
- **Configurable parameters**: Allows customization of crawl depth, number of links, timeouts, etc.
- **High-performance parallel embedding**: Embeds documents in parallel during crawling to maximize throughput.
- **Optimized database insertion**: Preserves embeddings to avoid redundant work during vector database operations.
## Usage
### Basic Usage
```python
from app.knowledge.enhanced_website_kb import EnhancedWebsiteKnowledgeBase
from agno.vectordb.pgvector import PgVector
# Initialize vector database
vector_db = PgVector(
table_name="your_table",
db_url="your_db_url",
schema="your_schema"
)
# Create knowledge base
kb = EnhancedWebsiteKnowledgeBase(
urls=["https://example.com"], # List of URLs to crawl
max_depth=3, # How deep to crawl (default: 3)
max_links=10, # Maximum number of links to follow (default: 10)
min_content_length=100, # Minimum content length to be considered valid
timeout=30, # Request timeout in seconds
max_retries=3, # Maximum retry attempts for failed requests
vector_db=vector_db # Vector database for storing extracted content
)
# Load knowledge base (crawl websites and store in vector DB)
kb.load(
recreate=False, # Whether to recreate the collection
upsert=True, # Whether to update existing documents
filters={"name": "example"} # Additional filters for the documents
)
```
### High-Performance Configuration with Optimized Vector Database
For maximum performance, use the `OptimizedPgVector` class with increased parallelism settings:
```python
from app.knowledge.enhanced_website_kb import EnhancedWebsiteKnowledgeBase
from app.knowledge.optimized_pgvector import OptimizedPgVector
# Initialize optimized vector database
vector_db = OptimizedPgVector(
table_name="your_table",
db_url="your_db_url",
schema="your_schema"
)
# Create knowledge base with increased parallelism
kb = EnhancedWebsiteKnowledgeBase(
urls=["https://example.com"],
max_depth=5, # Deeper crawling for more content
max_links=25, # More links for comprehensive coverage
min_content_length=100,
timeout=30,
max_retries=3,
max_workers=32, # High parallel processing for maximum throughput
batch_size=100, # Larger batch size for better database performance
vector_db=vector_db
)
# Load knowledge base
kb.load(
recreate=False,
upsert=True,
filters={"source": "website"}
)
```
## Performance Optimization
The system is designed to maximize performance at every stage:
1. **Parallel Crawling**: Multiple URLs are crawled simultaneously using ThreadPoolExecutor
2. **Parallel Embedding**: Documents from each URL are embedded in parallel batches
3. **Final Parallel Embedding**: Any remaining unembedded documents are processed with high parallelism
4. **Embedding Preservation**: The OptimizedPgVector class preserves embeddings during database operations
5. **Batched Database Operations**: Documents are inserted into the database in optimized batches
6. **Vector Database Indexing**: Automatic index optimization for faster vector searches
This approach dramatically reduces the total processing time, especially for large websites with many pages.
### Configuration Options
The `EnhancedWebsiteReader` and `EnhancedWebsiteKnowledgeBase` classes offer several configuration options:
- **max_depth**: How many levels deep to crawl from the starting URL (default: 5)
- **max_links**: Maximum number of links to follow (default: 25)
- **min_content_length**: Minimum text length to be considered valid content (default: 100)
- **timeout**: HTTP request timeout in seconds (default: 30)
- **max_retries**: Maximum number of retry attempts for failed requests (default: 3)
- **max_workers**: Number of parallel workers for crawling and embedding (default: 10)
- **crawl_scope**: How far link-following may wander from the starting URL (default: `host`)
- `host` — pages on the starting URL's own host only (`www.` and the bare host count as one host)
- `path` — that host, restricted to the starting URL's path prefix (e.g. only `/hc/…`)
- `domain` — every host on the registrable domain, so a `help.example.com` crawl also
walks `www.example.com`, `example.com/blog/…` and other subdomains
Only `http`/`https` links are ever followed — `mailto:`, `tel:` and friends are not pages.
Content stored by earlier, wider crawls is cleaned up in two steps:
- Pages that were never pages (mangled `https://mailto:a@example.com/…` URLs) are deleted
automatically by the `purge_mangled_crawl_pages_001` migration.
- Real pages that are simply outside the new scope (a marketing site pulled in from a
help-center seed) are left alone; `scripts/prune_offsite_pages.py` reports them per
source and deletes them only with `--apply`.
- **batch_size**: Size of document batches for database operations (default: 20)
- **blacklist_tags**: HTML tags to remove before content extraction (scripts, styles, etc.)
- **min_content_length**: Floor a page's extracted text must clear to count as readable (default: 100)
Candidate containers and the retention threshold live in `main_content.py`
(`CANDIDATE_TAGS`, `CONTENT_RETENTION`), shared with the help-center importer.
## Content Extraction Strategy
Selection happens in `main_content.select_main_node()`, shared with the help-center
article importer: it picks the **deepest container that still holds at least
`CONTENT_RETENTION` (90%) of the page's non-boilerplate text**, which resolves to
`<main>` or the content `<div>` where one exists and to `<body>` otherwise.
Navigation chrome (`nav`, `header`, `footer`, `aside`) is excluded when measuring,
so a page whose only text is a phone number and a login link reads as empty.
This replaced a walk that returned the *first* candidate clearing
`min_content_length`. That floor is not a quality bar — a hero banner clears it —
so a page whose first `<article>` was a 108-character search widget indexed the
widget and discarded 4,500 characters. Choosing by how much of the page a
container holds cannot make that mistake.
If the selected node yields less than `min_content_length`, the page is thin or
JavaScript-rendered and these fallbacks run in order:
1. **Density-based extraction**: Analyzes paragraph density to identify content-rich areas.
2. **Meaningful-element extraction**: Collects headings, paragraphs, lists and tables.
3. **Body fallback**: The cleaned body content.
4. **Last resort**: All text in the document.
Falling through to these is also what routes a JavaScript-rendered page to the
Crawl4AI browser fallback in `_process_url` — so boilerplate must never be
allowed to clear the floor on its own.
## Testing
The system includes comprehensive unit tests to ensure proper functionality:
```bash
# Activate the Python environment
conda activate python3128
# Run the tests
python -m pytest backend/tests/knowledge/test_enhanced_website_reader.py -v
python -m pytest backend/tests/knowledge/test_enhanced_website_kb.py -v
python -m pytest backend/tests/knowledge/test_optimized_pgvector.py -v
```
+218
View File
@@ -0,0 +1,218 @@
"""
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.
Splitting a knowledge page into embedder-sized chunks.
The search embedder only reads ~512 tokens, so a page stored as one row is
searchable by its first screen only, and a single hit can be tens of thousands
of chars — enough to overflow the model's context when a few are replayed. Each
page is therefore stored as several small rows that together hold every byte of
it.
Stored chunks are disjoint verbatim slices: concatenating a page's chunks gives
the page back exactly, which is what the dashboard, the FAQ generator and page
edits all rely on. Context across a boundary is handled at embedding time (the
tail of the previous chunk is prepended to the text that is embedded, never to
the text that is stored) and at search time (neighbouring chunks are returned
with a hit).
Why not agno's chunkers: ``ChunkingStrategy.clean_text`` collapses whitespace,
which flattens tables, lists and headings into one line and would persist that
way; ``MarkdownChunking`` needs the ``unstructured`` package.
Chunk ids are ``<page_id>::<n>`` (1-based). The ``_<n>`` form agno uses for PDFs
is kept readable by the page-grouping helpers, but it is not used for new chunks
because real page URLs end in ``_<n>`` too (``/blog/post_2``) and an upsert on
such an id would silently overwrite another page.
"""
import re
from dataclasses import dataclass
from typing import List, Optional
from agno.document import Document
from agno.document.chunking.strategy import ChunkingStrategy
from agno.embedder import Embedder
from app.core.config import settings
CHUNK_ID_SEPARATOR = "::"
# Legacy agno PDF chunk suffix. Stripped when grouping, never generated.
_LEGACY_CHUNK_ID_RE = re.compile(r"_\d+$")
_CHUNK_ID_RE = re.compile(re.escape(CHUNK_ID_SEPARATOR) + r"\d+$")
# Break preference, best first. A chunk ends just after the chosen separator so
# the separator itself is never lost.
_BREAKS = ("\n\n", "\n", ". ", " ")
# Don't take a break that would leave a chunk under this share of the target
# size; a hard cut at full size is better than a 40-char chunk.
_MIN_FILL = 0.3
_TITLE_MAX_CHARS = 120
# Metadata that describes a row's place in its page. Set by chunk_page; must be
# dropped before a page's content is re-chunked.
CHUNK_META_KEYS = frozenset({"page_id", "chunk", "chunk_count", "chunk_size"})
def page_id_of(chunk_id: str) -> str:
"""Map a chunk id back to its page id. Mirrors ``PAGE_ID_EXPR`` in SQL and
``basePageId`` in the frontend — keep the three in sync."""
stripped = _CHUNK_ID_RE.sub("", chunk_id)
if stripped != chunk_id:
return stripped
return _LEGACY_CHUNK_ID_RE.sub("", chunk_id)
def chunk_id(page_id: str, index: int) -> str:
return f"{page_id}{CHUNK_ID_SEPARATOR}{index}"
def split_text(text: str, size: Optional[int] = None) -> List[str]:
"""Split ``text`` into consecutive pieces of at most ``size`` chars,
preferring to break on paragraph, line, sentence then word boundaries.
The pieces are disjoint verbatim slices: ``"".join(pieces) == text``."""
size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
if size <= 0:
raise ValueError("chunk size must be positive")
pieces: List[str] = []
start = 0
length = len(text)
while start < length:
if length - start <= size:
pieces.append(text[start:])
break
end = _break_point(text, start, size)
pieces.append(text[start:end])
start = end
return pieces
def _break_point(text: str, start: int, size: int) -> int:
"""End of the piece starting at ``start``: just after the best separator in
the window, else a hard cut at ``size``. Always past ``start``."""
window = text[start:start + size]
floor = int(size * _MIN_FILL)
for sep in _BREAKS:
idx = window.rfind(sep)
if idx >= floor:
return start + idx + len(sep)
return start + size
@dataclass
class ChunkDocument(Document):
"""A chunk whose embedding is computed from more than its stored text: the
page title and the tail of the previous chunk are prepended, so a mid-page
chunk about pricing still matches "pricing" queries that only name the
page, and a sentence cut by the boundary is embedded whole on one side."""
embedding_text: Optional[str] = None
def embed(self, embedder: Optional[Embedder] = None) -> None:
_embedder = embedder or self.embedder
if _embedder is None:
raise ValueError("No embedder provided")
self.embedding, self.usage = _embedder.get_embedding_and_usage(
self.embedding_text or self.content
)
def given_title(document: Document) -> Optional[str]:
"""The title the source recorded for the page, if any."""
title = (document.meta_data or {}).get("title")
if isinstance(title, str) and title.strip():
return title.strip()[:_TITLE_MAX_CHARS]
return None
def page_title(document: Document) -> str:
"""What to call the page when embedding: its recorded title, else its first
non-empty line. The fallback is for the embedding only — it is never stored
as the page's title, since a crawled first line is as likely to be a cookie
banner as a heading."""
title = given_title(document)
if title:
return title
for line in (document.content or "").splitlines():
if line.strip():
return line.strip()[:_TITLE_MAX_CHARS]
return document.id or ""
def needs_explicit_chunk_id(page_id: Optional[str]) -> bool:
"""A page id that ends in ``_<n>`` reads as a legacy chunk suffix to the
grouping rule, which would file ``/blog/post_2`` under ``/blog/post`` and
let an edit of the sibling delete it. Such pages get a ``::1`` id even
when they fit in one chunk."""
return bool(page_id) and bool(_LEGACY_CHUNK_ID_RE.search(page_id))
def chunk_page(
document: Document,
size: Optional[int] = None,
overlap: Optional[int] = None,
) -> List[Document]:
"""Split one page into chunk documents. A page that fits in one chunk is
returned unchanged (same id, same metadata), so small pages are stored
exactly as before. ``overlap`` is how much of the previous chunk is
embedded with each chunk; it never changes what is stored."""
size = settings.KNOWLEDGE_CHUNK_SIZE if size is None else size
overlap = settings.KNOWLEDGE_CHUNK_OVERLAP if overlap is None else overlap
content = document.content or ""
if len(content) <= size and not needs_explicit_chunk_id(document.id):
return [document]
pieces = split_text(content, size)
title = page_title(document)
page_id = document.id
base_meta = {k: v for k, v in (document.meta_data or {}).items() if k not in CHUNK_META_KEYS}
recorded_title = given_title(document)
if recorded_title:
base_meta["title"] = recorded_title
chunks: List[Document] = []
previous = ""
for index, piece in enumerate(pieces, start=1):
meta = {
**base_meta,
"page_id": page_id,
"chunk": index,
"chunk_count": len(pieces),
"chunk_size": len(piece),
}
context = previous[-overlap:] if overlap > 0 else ""
chunks.append(
ChunkDocument(
id=chunk_id(page_id, index),
name=document.name,
content=piece,
meta_data=meta,
embedder=document.embedder,
embedding_text=f"{title}\n{context}{piece}",
)
)
previous = piece
return chunks
class PageChunking(ChunkingStrategy):
"""agno ``ChunkingStrategy`` adapter so the PDF readers chunk the same way
as the website crawler."""
def __init__(self, size: Optional[int] = None, overlap: Optional[int] = None):
self.size = size
self.overlap = overlap
def chunk(self, document: Document) -> List[Document]:
return chunk_page(document, self.size, self.overlap)
+190
View File
@@ -0,0 +1,190 @@
"""
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 typing import Optional
from agno.agent import Agent
from app.utils.agno_utils import create_model
from app.core.logger import get_logger
from app.core.config import settings
from app.core.model_catalog import requires_base_url
logger = get_logger(__name__)
class ContentSummarizer:
"""
Service for summarizing web content using LLM to reduce token usage.
Preserves important information and URLs for reference.
"""
def __init__(self):
"""Initialize the content summarizer with configured model"""
self.enabled = settings.KNOWLEDGE_SUMMARY_ENABLED
self.model_type = settings.KNOWLEDGE_SUMMARY_MODEL_TYPE
self.model_name = settings.KNOWLEDGE_SUMMARY_MODEL_NAME
self.api_key = settings.KNOWLEDGE_SUMMARY_API_KEY
self.base_url = settings.KNOWLEDGE_SUMMARY_BASE_URL or None
self.max_tokens = settings.KNOWLEDGE_SUMMARY_MAX_TOKENS
self._agent = None
if self.enabled and not self.api_key:
logger.warning("Content summarization is enabled but KNOWLEDGE_SUMMARY_API_KEY is not set. Summarization will be disabled.")
self.enabled = False
if self.enabled and requires_base_url(self.model_type) and not self.base_url:
logger.warning(f"Content summarization model type is {self.model_type} but KNOWLEDGE_SUMMARY_BASE_URL is not set. Summarization will be disabled.")
self.enabled = False
def _get_agent(self) -> Optional[Agent]:
"""Lazy initialization of the summarization agent"""
if not self.enabled:
return None
if self._agent is None:
try:
model = create_model(
model_type=self.model_type,
api_key=self.api_key,
model_name=self.model_name,
max_tokens=self.max_tokens,
base_url=self.base_url
)
self._agent = Agent(
name="Content Summarizer",
model=model,
instructions="""You are an expert content summarizer for a knowledge base system.
Your task is to summarize web page content while:
1. **PRESERVING ALL URLs** - This is CRITICAL. Every single URL mentioned in the content MUST be kept in the summary:
- Product URLs
- Reference URLs
- Image URLs
- Documentation URLs
- Any links mentioned in the text
DO NOT remove, shorten, or modify ANY URLs. Keep them exactly as they appear.
2. Preserving ALL important information needed to answer user questions:
- Product names, descriptions, and details
- Pricing and availability information
- Technical specifications and features
- Contact information and addresses
- Instructions and procedures
- Important facts and data points
3. Reducing only redundant and verbose text:
- Remove excessive marketing language
- Eliminate repetitive phrases
- Condense wordy explanations
- Remove filler content
4. Maintaining the original context and structure so users can:
- Find specific products or information
- Access the original URLs for more details
- Get accurate answers to their questions
REMEMBER: URLs are the most critical element to preserve. Users need these links to access products and resources. Never remove or modify URLs under any circumstances.
Format your response as a concise summary that maintains all URLs and critical information.""",
markdown=False,
debug_mode=settings.ENVIRONMENT == "development"
)
logger.info(f"Initialized content summarizer with model: {self.model_type}/{self.model_name}")
except Exception as e:
logger.error(f"Failed to initialize content summarizer: {str(e)}")
self.enabled = False
return None
return self._agent
def summarize(self, content: str, url: str) -> str:
"""
Summarize content while preserving important information and URLs.
Args:
content: The original content to summarize
url: The source URL for reference
Returns:
Summarized content or original content if summarization fails/disabled
"""
# If summarization is disabled, return original content
if not self.enabled:
return content
# If content is already short, don't summarize
if len(content) < 1000:
logger.debug(f"Content too short to summarize ({len(content)} chars), keeping original")
return content
try:
agent = self._get_agent()
if not agent:
logger.warning("Summarization agent not available, using original content")
return content
# Create the summarization prompt
prompt = f"""Summarize the following web page content from {url}:
Content:
{content}
CRITICAL REQUIREMENTS:
- **KEEP EVERY SINGLE URL EXACTLY AS IT APPEARS** - Do not remove, shorten, or modify any URLs
- Preserve all product URLs, reference links, and resource URLs
- Keep all important facts, product details, pricing, and specifications
- Maintain context for answering user questions
- Only reduce redundant/verbose text, never remove URLs or important information"""
logger.info(f"Summarizing content from {url} ({len(content)} chars)")
# Run the agent synchronously
response = agent.run(message=prompt)
# Extract the summary from response
if hasattr(response, 'content') and response.content:
summary = response.content
elif isinstance(response, dict) and 'content' in response:
summary = response['content']
elif isinstance(response, str):
summary = response
else:
logger.error(f"Unexpected response format from summarizer: {type(response)}")
return content
# Add source URL reference at the end if not already present
if url not in summary:
summary = f"{summary}\n\nSource: {url}"
logger.info(f"Summarized content: {len(content)} -> {len(summary)} chars ({100*len(summary)/len(content):.1f}%)")
return summary
except Exception as e:
logger.error(f"Error summarizing content from {url}: {str(e)}")
logger.info("Falling back to original content")
return content
# Global instance
_summarizer = None
def get_content_summarizer() -> ContentSummarizer:
"""Get or create the global content summarizer instance"""
global _summarizer
if _summarizer is None:
_summarizer = ContentSummarizer()
return _summarizer
+324
View File
@@ -0,0 +1,324 @@
"""
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.
"""
import asyncio
from typing import Optional, Tuple
from bs4 import BeautifulSoup
from app.core.logger import get_logger
from app.knowledge.crawler_identity import CRAWLER_USER_AGENT
from app.knowledge.url_safety import guard_url
logger = get_logger(__name__)
# Try to import crawl4ai
try:
from crawl4ai import AsyncWebCrawler, BrowserConfig, CrawlerRunConfig, CacheMode
CRAWL4AI_AVAILABLE = True
logger.info("✓ Crawl4AI module loaded successfully")
except ImportError:
CRAWL4AI_AVAILABLE = False
logger.warning("⚠️ Crawl4AI not installed. Install with: pip install crawl4ai>=0.9")
# Note: We use threading for async compatibility instead of nest_asyncio
# This avoids conflicts with uvloop and other event loop implementations
class Crawl4AIFallback:
"""
Fallback crawler using Crawl4AI with browser rendering for JavaScript-heavy websites.
This provides an alternative when standard HTTP requests fail to extract content.
"""
def __init__(self, timeout: int = 30, verify_ssl: bool = True):
"""
Initialize the Crawl4AI fallback crawler.
:param timeout: Request timeout in seconds
:param verify_ssl: Whether to verify SSL certificates
"""
self.timeout = timeout
self.verify_ssl = verify_ssl
self._is_available = CRAWL4AI_AVAILABLE
@property
def is_available(self) -> bool:
"""Check if Crawl4AI is available."""
return self._is_available
async def _async_fetch(self, url: str, take_screenshot: bool = False) -> Tuple[Optional[str], Optional[str], Optional[str]]:
"""
Async method to fetch content using Crawl4AI.
:param url: The URL to fetch
:param take_screenshot: Whether to take a screenshot of the rendered page
:return: Tuple of (html, markdown, screenshot_base64) or (None, None, None) if failed
:raises BlockedHostError: if the URL targets an internal/blocked host
"""
# SSRF guard: the httpx crawl path validates through url_safety, but this
# browser fallback drives a real Chromium and previously fetched any URL.
# Note this covers the initial URL only - redirects and in-page JS navigation
# the browser follows on its own need network-layer egress policy too.
guard_url(url)
screenshot_base64 = None
html_content = None
markdown_content = None
# If screenshot is needed, use Playwright directly for better control
if take_screenshot:
try:
logger.info(f"📸 Using Playwright directly for viewport screenshot of {url}")
from playwright.async_api import async_playwright
import base64
async with async_playwright() as p:
browser = await p.chromium.launch(
headless=True,
args=[
"--disable-blink-features=AutomationControlled",
"--disable-dev-shm-usage",
"--no-sandbox",
]
)
context = await browser.new_context(
viewport={'width': 1280, 'height': 1024},
user_agent=CRAWLER_USER_AGENT,
ignore_https_errors=not self.verify_ssl
)
page = await context.new_page()
await page.goto(url, wait_until="domcontentloaded", timeout=self.timeout * 1000)
# Wait for content to render
await page.wait_for_timeout(3000) # 3 seconds for JS to execute
# Take viewport-only screenshot
screenshot_bytes = await page.screenshot(type='png', full_page=False)
screenshot_base64 = base64.b64encode(screenshot_bytes).decode('utf-8')
screenshot_size = len(screenshot_base64)
logger.info(f"✅ Viewport screenshot captured: {screenshot_size} chars ({screenshot_size // 1024}KB)")
# Also get the HTML content
html_content = await page.content()
await browser.close()
except Exception as screenshot_error:
logger.error(f"❌ Error with Playwright screenshot: {str(screenshot_error)}", exc_info=True)
screenshot_base64 = None
# Use Crawl4AI for content extraction (more reliable for markdown)
try:
browser_config = BrowserConfig(
headless=True,
java_script_enabled=True,
text_mode=False,
verbose=False,
ignore_https_errors=not self.verify_ssl,
user_agent=CRAWLER_USER_AGENT,
viewport_width=1280,
viewport_height=1024,
extra_args=[
"--disable-blink-features=AutomationControlled",
"--disable-dev-shm-usage",
"--no-sandbox",
]
)
# Wait out short bot-check interstitials (WordPress.com / Cloudflare
# "Checking your browser… will only take a few seconds") which
# JS-redirect to the real page after a few seconds. Resolve as soon as
# the page has substantial real content (>800 chars — fast for normal
# pages), OR, on a still-short page, once the visible interstitial
# phrases are gone. This actually waits through the redirect instead
# of returning the challenge immediately; if it never clears it times
# out and the reader skips the page.
wait_for_real_content = (
"js:() => { const t = document.body ? document.body.innerText : ''; "
"if (t.length > 800) return true; "
"const low = t.toLowerCase(); "
"return t.length > 0 "
"&& !low.includes('checking your browser') "
"&& !low.includes('just a moment') "
"&& !low.includes('secured by wp.com') "
"&& !low.includes('verifying you are human'); }"
)
crawler_config = CrawlerRunConfig(
cache_mode=CacheMode.BYPASS,
wait_until="domcontentloaded",
page_timeout=self.timeout * 1000,
wait_for=wait_for_real_content,
delay_before_return_html=3.0,
screenshot=False,
# Anti-bot: simulate a real user and hide automation flags so
# WordPress.com / Cloudflare interstitials clear for the headless
# browser instead of looping on "Checking your browser…".
magic=True,
simulate_user=True,
override_navigator=True,
)
logger.info(f"🌐 Starting Crawl4AI for content extraction: {url}")
async with AsyncWebCrawler(config=browser_config) as crawler:
result = await crawler.arun(url=url, config=crawler_config)
if result.success:
html_length = len(result.html) if result.html else 0
markdown_length = len(result.markdown.raw_markdown) if result.markdown and result.markdown.raw_markdown else 0
logger.info(f"✓ Crawl4AI content extraction successful - HTML: {html_length} chars, Markdown: {markdown_length} chars")
# Use Crawl4AI's HTML if we didn't get it from Playwright
if not html_content:
html_content = result.html
markdown_content = result.markdown.raw_markdown if result.markdown else None
return html_content, markdown_content, screenshot_base64
else:
logger.error(f"❌ Crawl4AI failed: {result.error_message}")
# Return whatever we got from Playwright
return html_content, None, screenshot_base64
except Exception as e:
logger.error(f"❌ Crawl4AI async fetch error: {str(e)}", exc_info=True)
return None, None, None
def fetch_with_browser(self, url: str, take_screenshot: bool = False) -> Tuple[Optional[str], Optional[BeautifulSoup], Optional[str]]:
"""
Fetch content using Crawl4AI with browser rendering.
This method handles async execution in a thread-safe manner.
:param url: The URL to fetch
:param take_screenshot: Whether to take a screenshot of the rendered page
:return: Tuple of (content, soup, screenshot_base64) where content is the extracted text
"""
if not self._is_available:
logger.error(f"❌ Crawl4AI not available. Cannot fetch {url}")
return None, None, None
try:
logger.info(f"🔄 Using Crawl4AI fallback for JavaScript-heavy page: {url}")
# Handle async execution in thread-safe way
html, markdown, screenshot = self._run_async_safely(url, take_screenshot)
if not html and not markdown:
logger.warning(f"⚠️ Crawl4AI returned no content for {url}")
return None, None, None
# Parse HTML with BeautifulSoup if available
soup = None
if html:
soup = BeautifulSoup(html, 'html.parser')
logger.debug(f"Parsed HTML into BeautifulSoup object ({len(html)} chars)")
# Prefer markdown if available and substantial
if markdown and len(markdown.strip()) > 100:
logger.info(f"✓ Using Crawl4AI markdown output ({len(markdown)} chars)")
return markdown.strip(), soup, screenshot
# Otherwise return None for content (will trigger extraction from soup)
return None, soup, screenshot
except Exception as e:
logger.error(f"❌ Crawl4AI fetch error for {url}: {str(e)}", exc_info=True)
return None, None, None
def _run_async_safely(self, url: str, take_screenshot: bool = False) -> Tuple[Optional[str], Optional[str], Optional[str]]:
"""
Run async crawl in a thread-safe manner, handling various event loop scenarios.
:param url: The URL to fetch
:param take_screenshot: Whether to take a screenshot
:return: Tuple of (html, markdown, screenshot_base64)
"""
try:
# Try to get the current event loop
try:
loop = asyncio.get_event_loop()
loop_is_running = loop.is_running()
except RuntimeError:
loop_is_running = False
if loop_is_running:
# Event loop is already running (e.g., in FastAPI/uvloop context)
# Always use a separate thread for running Crawl4AI to avoid conflicts
logger.debug("Event loop running, creating new loop in separate thread")
import concurrent.futures
import threading
result = [None, None, None]
exception = [None]
def run_in_new_loop():
"""Run async fetch in a new event loop in a separate thread"""
try:
# Create a new event loop for this thread
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
result[0], result[1], result[2] = new_loop.run_until_complete(self._async_fetch(url, take_screenshot))
finally:
new_loop.close()
except Exception as e:
exception[0] = e
# Run in a separate thread
thread = threading.Thread(target=run_in_new_loop)
thread.start()
thread.join(timeout=self.timeout + 10)
if thread.is_alive():
logger.error(f"Thread timeout after {self.timeout + 10}s")
return None, None, None
if exception[0]:
raise exception[0]
return result[0], result[1], result[2]
else:
# No event loop running, safe to use asyncio.run
logger.debug("No event loop running, using asyncio.run")
return asyncio.run(self._async_fetch(url, take_screenshot))
except Exception as e:
logger.error(f"❌ Error running async fetch: {str(e)}", exc_info=True)
return None, None, None
# Singleton instance for convenience
_crawl4ai_fallback_instance = None
def get_crawl4ai_fallback(timeout: int = 30, verify_ssl: bool = True) -> Crawl4AIFallback:
"""
Get or create a Crawl4AI fallback instance (singleton pattern).
:param timeout: Request timeout in seconds
:param verify_ssl: Whether to verify SSL certificates
:return: Crawl4AIFallback instance
"""
global _crawl4ai_fallback_instance
if _crawl4ai_fallback_instance is None:
_crawl4ai_fallback_instance = Crawl4AIFallback(timeout=timeout, verify_ssl=verify_ssl)
return _crawl4ai_fallback_instance
+185
View File
@@ -0,0 +1,185 @@
"""
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.
"""
"""Which URLs a website crawl is allowed to follow.
Two independent problems this decides, both seen on a real help-center import
that stored 247 pages of which only ~102 were help articles:
* **Scheme.** ``urljoin`` leaves a ``mailto:``/``tel:`` href untouched, and
prefixing a scheme-less value with ``https://`` turns ``mailto:a@example.com``
into ``https://mailto:a@example.com`` — which parses as host ``example.com``
with ``mailto:a`` as userinfo, so it passed the same-domain check and was
crawled as a page (then spawned ``https://mailto:a@example.com/pricing`` and
friends from the fetched body). Only http(s) URLs are crawlable.
* **Reach.** Bounding a crawl by the *registrable domain* means a crawl seeded at
``help.example.com`` also walks ``www.example.com``, ``example.com/blog/*`` and
every other host under ``example.com`` — the whole company web presence lands
in one knowledge source. The default is now the seed's own host; ``domain``
(previous behaviour) and ``path`` (seed host + seed path prefix) stay available
per source.
"""
from dataclasses import dataclass
from typing import Optional
from urllib.parse import ParseResult, urlparse, urlunparse
from app.knowledge.domains import registrable_domain
# Scope modes, narrowest first. Kept as plain strings: they travel through the
# API request and the queue item's JSON metadata.
SCOPE_PATH = "path"
SCOPE_HOST = "host"
SCOPE_DOMAIN = "domain"
CRAWL_SCOPES = (SCOPE_PATH, SCOPE_HOST, SCOPE_DOMAIN)
DEFAULT_CRAWL_SCOPE = SCOPE_HOST
# Anything else (mailto:, tel:, sms:, javascript:, data:, ftp:) is a contact or
# script action, not a page to fetch.
CRAWLABLE_SCHEMES = ("http", "https")
def url_scheme(url: str) -> str:
"""Lowercased scheme of a URL, or ``''`` when it has none.
Not ``urlparse(url).scheme``: that reads the scheme-less ``example.com:8443/x``
as scheme ``example.com``. A real scheme has no ``.`` or ``/`` before the colon.
"""
head, sep, _ = (url or "").strip().partition(":")
if not sep or not head or "." in head or "/" in head:
return ""
return head.lower()
def is_crawlable_url(url: str) -> bool:
"""True for an http(s) URL or a scheme-less one (a relative href still to be
resolved against its base). False for ``mailto:``/``tel:``/``javascript:``/…"""
scheme = url_scheme(url)
return scheme in CRAWLABLE_SCHEMES or not scheme
def _parsed(url: str) -> Optional[ParseResult]:
"""Parsed URL, or None when it can't be parsed — a malformed IPv6 literal
makes ``.hostname`` raise, and callers want "out of scope", not a traceback."""
try:
parsed = urlparse(url or "")
parsed.hostname # raises here rather than at each use site
except ValueError:
return None
return parsed
def _hostname(url: str) -> str:
"""Host of a URL, or ``''`` if it has none / can't be parsed."""
parsed = _parsed(url)
return (parsed.hostname or "") if parsed else ""
def _host_key(host: str) -> str:
"""Host identity for comparisons: lowercased, ``www.`` dropped, so a site
linking between ``example.com`` and ``www.example.com`` stays in scope."""
host = (host or "").lower().strip().rstrip(".")
return host[4:] if host.startswith("www.") else host
@dataclass(frozen=True)
class CrawlScope:
"""The boundary of one crawl, derived from its seed URL."""
mode: str
domain: str
host: str
seed_host: str
path_prefix: str
@classmethod
def for_seed(cls, seed_url: str, mode: str = DEFAULT_CRAWL_SCOPE) -> "CrawlScope":
"""Scope for a crawl seeded at ``seed_url``. An unknown mode falls back to
the default rather than failing a queued crawl."""
if mode not in CRAWL_SCOPES:
mode = DEFAULT_CRAWL_SCOPE
parsed = _parsed(seed_url)
host = ((parsed.hostname or "") if parsed else "").lower()
return cls(
mode=mode,
domain=registrable_domain(host),
host=_host_key(host),
seed_host=host,
path_prefix=((parsed.path or "") if parsed else "").rstrip("/"),
)
def canonical(self, url: str) -> str:
"""Put an in-scope URL on the seed's own spelling of the host.
``example.com/terms`` and ``www.example.com/terms`` are one page, but
they are two page ids — so a site that links both ways stores the page
twice, burning two of the source's sub-page slots and feeding retrieval
and FAQ generation the same text twice.
Rewrites toward the SEED's host rather than always stripping ``www.``:
plenty of sites serve only one of the two, and the seed is the form the
user gave us and the crawl already fetched. Only the seed host's own
variant is touched — a genuine subdomain under ``domain`` scope is left
as it is.
"""
parsed = _parsed(url)
if parsed is None or not parsed.hostname:
return url
host = parsed.hostname.lower()
if host == self.seed_host or _host_key(host) != self.host:
return url
netloc = f"{self.seed_host}:{parsed.port}" if parsed.port else self.seed_host
return urlunparse(parsed._replace(netloc=netloc))
def allows(self, url: str) -> bool:
"""True if ``url`` may be fetched and stored as part of this crawl."""
if url_scheme(url) not in CRAWLABLE_SCHEMES:
return False
parsed = _parsed(url)
if parsed is None:
return False
# No userinfo: pages are never fetched with credentials, and this is what
# a leftover 'https://mailto:a@example.com/x' from an older crawl looks
# like ('mailto:a' as user:password) once the scheme has been mangled off.
if parsed.username or parsed.password:
return False
host = parsed.hostname or ""
if self.mode == SCOPE_DOMAIN:
# Equality, not a suffix match: 'evilexample.com'.endswith('example.com')
# would let one plan-limited source fan out across unrelated domains.
return bool(self.domain) and registrable_domain(host) == self.domain
if not self.host or _host_key(host) != self.host:
return False
if self.mode == SCOPE_HOST:
return True
# Path scope: compare on a segment boundary so a '/hc/en' seed keeps
# '/hc/en/articles/1' but not '/hc/entertainment'.
path = (parsed.path or "").rstrip("/")
return path == self.path_prefix or path.startswith(f"{self.path_prefix}/")
def describe(self) -> str:
"""Short human-readable scope, for crawl logs."""
if self.mode == SCOPE_DOMAIN:
return f"domain {self.domain}"
if self.mode == SCOPE_HOST:
return f"host {self.host}"
return f"host {self.host} under {self.path_prefix or '/'}"
+22
View File
@@ -0,0 +1,22 @@
"""
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.
Who our crawler says it is. Honest, and the same for every request: hosts
such as WordPress.com refuse the stale "Chrome 120 on Windows" string
scrapers commonly fake, while a named crawler is let through - and site
owners can see (and allow or block) who is reading them.
"""
CRAWLER_USER_AGENT = "Mozilla/5.0 (compatible; ReplyPierBot/1.0; +https://replypier.site)"
+39
View File
@@ -0,0 +1,39 @@
"""
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.
How long one source may take to read.
Reading runs on worker threads, which cannot be stopped from outside: a
timeout around the call would mark the source failed while the crawl and
its embedding carried on in the background. So the work checks the deadline
itself, between pages and between chunks, and stops.
"""
import time
from typing import Callable
class SourceTimedOut(Exception):
"""The source ran past its time limit. The message is for the user."""
class Deadline:
def __init__(self, seconds: float, clock: Callable[[], float] = time.monotonic):
self._clock = clock
self._at = clock() + seconds
@property
def expired(self) -> bool:
return self._clock() >= self._at
+56
View File
@@ -0,0 +1,56 @@
"""
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.
Single canonical registrable-domain helper, used everywhere a sub-page URL is
compared against its source's domain (website crawler, sitemap parser, manual
add-subpage). Same-domain enforcement stops one plan-limited knowledge source
from accumulating content across unrelated domains. Handles common ccTLDs
(co.uk, com.au, …) so 'a.example.co.uk' → 'example.co.uk' rather than 'co.uk'
(the latter would wrongly treat every *.co.uk host as one domain).
"""
from urllib.parse import urlparse
# Second-level labels that sit under a 2-letter ccTLD (best-effort; a full
# public-suffix list would be a heavy dependency for marginal gain here).
_CCTLD_SECOND_LEVELS = frozenset(
{"co", "com", "org", "net", "edu", "gov", "ac", "gov", "ltd", "plc", "me", "or", "ne"}
)
def registrable_domain(host: str) -> str:
"""Registrable domain for a HOST (not a full URL): lowercased, port/userinfo
already stripped by the caller's urlparse(...).hostname, 'www.' removed."""
host = (host or "").lower().strip().rstrip(".")
if host.startswith("www."):
host = host[4:]
if not host:
return ""
parts = host.split(".")
if len(parts) > 2 and parts[-2] in _CCTLD_SECOND_LEVELS and len(parts[-1]) == 2:
return ".".join(parts[-3:])
return ".".join(parts[-2:]) if len(parts) > 1 else host
def domain_of_url(url: str) -> str:
"""Registrable domain of a URL string. A scheme-less value ('example.com/x')
is normalized so its host is still recognised (else it would parse as a
path with no hostname and silently bypass a same-domain check)."""
text = (url or "").strip()
if not text:
return ""
if "://" not in text:
text = f"https://{text}"
return registrable_domain(urlparse(text).hostname or "")
+60
View File
@@ -0,0 +1,60 @@
"""
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.
The local embedding model, loaded once per process.
agno's ``FastEmbedEmbedder`` builds a new ``TextEmbedding`` (an ONNX session)
inside every ``get_embedding`` call — about 1.5 s of model loading around a
7 ms embedding. One row per page hid that; one row per chunk would not.
"""
import threading
from dataclasses import dataclass
from typing import Dict, List
import numpy as np
from agno.embedder.fastembed import FastEmbedEmbedder
from fastembed import TextEmbedding
from app.core.config import settings
_models: Dict[str, TextEmbedding] = {}
_lock = threading.Lock()
def _model(model_id: str) -> TextEmbedding:
model = _models.get(model_id)
if model is None:
with _lock:
model = _models.get(model_id)
if model is None:
model = _models[model_id] = TextEmbedding(model_name=model_id)
return model
@dataclass
class CachedFastEmbedEmbedder(FastEmbedEmbedder):
"""``FastEmbedEmbedder`` that shares one loaded model per model id."""
id: str = settings.FASTEMBED_MODEL
def get_embedding(self, text: str) -> List[float]:
embedding = list(_model(self.id).embed(text))[0]
return embedding.tolist() if isinstance(embedding, np.ndarray) else list(embedding)
def get_embedder() -> CachedFastEmbedEmbedder:
"""The embedder every knowledge path should use for the configured model."""
return CachedFastEmbedEmbedder()
+81
View File
@@ -0,0 +1,81 @@
"""
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 typing import Any, Dict, Iterator, List, Optional
from agno.knowledge.pdf import PDFKnowledgeBase
from agno.document import Document
from app.core.logger import get_logger
logger = get_logger(__name__)
class EnhancedPDFKnowledgeBase(PDFKnowledgeBase):
"""Enhanced PDF knowledge base that supports metadata filters"""
def load(
self,
recreate: bool = False,
upsert: bool = True,
skip_existing: bool = True,
filters: Optional[Dict[str, Any]] = None,
) -> None:
"""
Load the PDF contents to the vector db with metadata filters
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): Metadata filters to apply to the documents. Defaults to None.
"""
if self.vector_db is None:
logger.warning("No vector db 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"Loading PDF knowledge base: {self.path}")
# Read documents from PDF
if self.reader is None:
logger.warning("No reader provided")
return
# Get documents and embed them
for document_list in self.document_lists:
# Embed documents if embedder is available
if self.vector_db and hasattr(self.vector_db, 'embedder'):
for document in document_list:
if document.embedding is None:
try:
document.embed(embedder=self.vector_db.embedder)
except Exception as e:
logger.error(f"Error embedding document {document.id}: {str(e)}")
# Upsert documents with filters
if upsert and self.vector_db:
try:
logger.debug(f"Upserting {len(document_list)} documents with filters: {filters}")
self.vector_db.upsert(documents=document_list, filters=filters)
except Exception as e:
logger.error(f"Error upserting documents: {str(e)}")
logger.info(f"Completed loading PDF knowledge base: {self.path}")
@@ -0,0 +1,626 @@
"""
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 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)")
File diff suppressed because it is too large. Load diff
+117
View File
@@ -0,0 +1,117 @@
"""
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.
What to tell someone when a knowledge source couldn't be read: one piece of
advice per failure cause. Kept apart from the crawler so code that only needs
the words doesn't import the crawling stack.
"""
from typing import Optional
class FailureCause:
"""Why a page could not be read.
The old message offered the user three guesses at once ("empty, require
JavaScript, or be unreachable"). _process_url always knew which one it hit;
it just threw the answer away. These are the classes it can tell apart, and
each one gets its own advice in FAILURE_ADVICE below.
"""
DNS = "dns"
CONNECTION = "connection"
TIMEOUT = "timeout"
BOT_BLOCKED = "bot_blocked"
HTTP_ERROR = "http_error"
EMPTY_RESPONSE = "empty_response"
THIN_CONTENT = "thin_content"
BLOCKED_HOST = "blocked_host"
NOT_A_PAGE = "not_a_page"
# Something that is not a transport failure at all — most likely a bug of
# ours. Deliberately has no entry in FAILURE_ADVICE: we fall back to the
# generic message rather than blame the customer's site for our crash.
UNKNOWN = "unknown"
# The source as a whole, not one page: it ran past its time limit, or its
# worker died partway (the queue item was left 'processing').
TOO_LONG = "too_long"
INTERRUPTED = "interrupted"
# What to tell the user for each cause. Phrased as "what happened, what to do".
FAILURE_ADVICE = {
FailureCause.DNS: (
"{url} could not be found — the domain does not resolve. Check the "
"address for typos, and that the site is published."
),
FailureCause.CONNECTION: (
"{url} refused the connection. Check the site is online and reachable "
"from the public internet."
),
FailureCause.TIMEOUT: (
"{url} did not respond in time. If the site is usually slow, try again "
"in a few minutes, or add a specific page instead of the whole site."
),
FailureCause.BOT_BLOCKED: (
"{url} refused our crawler. Add the content manually with Upload PDF "
"or Paste text."
),
FailureCause.HTTP_ERROR: (
"{url} returned an error response. Check the URL is correct and "
"publicly accessible, then try again."
),
FailureCause.EMPTY_RESPONSE: (
"{url} returned an empty page. Check the site is published and serving "
"content at that address."
),
FailureCause.THIN_CONTENT: (
"{url} loaded but had almost no readable text — it may render its "
"content with JavaScript we could not run. Add the content manually "
"with Upload PDF or Paste text."
),
FailureCause.BLOCKED_HOST: (
"{url} redirected to an address we are not allowed to fetch. Use the "
"site's public URL."
),
FailureCause.NOT_A_PAGE: (
"{url} is a file, not a web page. Add a PDF with Upload PDF, or add "
"the page that links to it."
),
FailureCause.TOO_LONG: (
"{url} took too long to read. Add specific pages instead of the whole "
"site, or add the content with Upload PDF or Paste text."
),
FailureCause.INTERRUPTED: (
"{url} could not be read to the end: our reader stopped partway. Add "
"it again to retry, or add specific pages instead of the whole site."
),
}
GENERIC_REASON = (
"We couldn't read any pages there. Try a specific page, or paste text or "
"upload a PDF instead."
)
def plain_reason(error: Optional[str], source: Optional[str]) -> str:
"""The stored failure for `source`, phrased about "This site" (the email
or page already shows the address). Anything that isn't one of our own
advice messages - a raw exception, say - becomes GENERIC_REASON rather
than leak internals to the customer."""
if error and source:
for advice in FAILURE_ADVICE.values():
if error == advice.format(url=source):
return advice.format(url="This site")
return GENERIC_REASON
+526
View File
@@ -0,0 +1,526 @@
"""
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 <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
+110
View File
@@ -0,0 +1,110 @@
"""
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.
Main-content selection shared by the knowledge crawler (which extracts text) and
the help-center article importer (which needs the actual DOM node to convert to
Markdown). One implementation, because the two drifting apart is how the crawler
bug below survived in the importer.
"""
from typing import Iterable, Optional
from bs4 import BeautifulSoup, Tag
# Tags stripped before any content selection — never useful content.
STRIP_TAGS = ("script", "style", "noscript", "iframe", "head")
# Chrome that repeats on every page. Excluded when measuring how much content a
# container holds, so that a page whose only text is a phone number and a login
# link reads as empty rather than full.
BOILERPLATE_TAGS = ("nav", "header", "footer", "aside")
# A container must hold at least this share of the page's content text to be
# considered "the page". Below it we keep looking outward.
CONTENT_RETENTION = 0.90
MIN_MAIN_CONTENT_CHARS = 100
# Containers worth considering, cheapest-to-richest. 'div' is here because most
# sites wrap their content in one rather than in <main>.
CANDIDATE_TAGS = ("article", "main", "section", "div")
def content_chars(node: Tag) -> int:
"""Length of a node's text, ignoring navigation chrome.
solcontrol.ca has 147 characters of body text, all of it a phone number and
a login link. Counting that as content made the page look ingestible and
suppressed the browser fallback that could actually read it.
"""
total = len(node.get_text(strip=True))
for boilerplate in node.find_all(BOILERPLATE_TAGS):
# Only outermost chrome: a <nav> inside a <header> is already counted in
# the header's text, and subtracting both would understate the page.
if boilerplate.find_parent(BOILERPLATE_TAGS) is not None:
continue
total -= len(boilerplate.get_text(strip=True))
return max(total, 0)
def _strip_noise(soup: BeautifulSoup) -> None:
for tag in STRIP_TAGS:
for element in soup.find_all(tag):
element.extract()
def _candidates(soup: BeautifulSoup, tags: Iterable[str]) -> Iterable[Tag]:
for tag in tags:
for element in soup.find_all(tag):
yield element
def select_main_node(soup: BeautifulSoup, min_chars: int = MIN_MAIN_CONTENT_CHARS) -> Optional[Tag]:
"""The page's main-content element (mutates soup: strips script/style).
Picks the *deepest* container that still holds essentially the whole page,
which lands on <main> or the content <div> when one exists and on <body>
otherwise.
The previous implementation returned the first candidate whose text cleared
MIN_MAIN_CONTENT_CHARS. That floor is not a quality bar — a hero banner
clears it — so on a page whose first <article> was a 108-character search
widget it returned the widget and discarded 4,500 characters of page.
Selecting by how much of the page a container holds cannot make that
mistake: a banner holds almost none of it.
Returns None only for genuinely empty pages, so callers can fall back.
"""
_strip_noise(soup)
body = soup.find("body") or soup
budget = content_chars(body)
if budget < min_chars:
return None
best = body
best_depth = len(list(body.parents))
for element in _candidates(soup, CANDIDATE_TAGS):
if element is best:
continue
if content_chars(element) < CONTENT_RETENTION * budget:
continue
# Same content, more specific container: prefer it, since the extra
# wrapper is chrome we would otherwise index on every page.
depth = len(list(element.parents))
if depth > best_depth:
best, best_depth = element, depth
return best
+254
View File
@@ -0,0 +1,254 @@
"""
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 typing import Any, Dict, List, Optional, Union
from hashlib import md5
from agno.document import Document
from agno.embedder import Embedder
from agno.reranker.base import Reranker
from agno.utils.log import log_debug, log_info, logger
from agno.vectordb.pgvector.pgvector import PgVector
from agno.vectordb.distance import Distance
from agno.vectordb.pgvector.index import HNSW, Ivfflat
from agno.vectordb.search import SearchType
from sqlalchemy.engine import Engine
class OptimizedPgVector(PgVector):
"""
OptimizedPgVector class that extends the standard PgVector implementation
to avoid re-embedding documents that already have embeddings.
This class optimizes the workflow for web crawling by allowing embeddings
to be generated during the crawling phase, and then preserved during insertion
into the database.
"""
def __init__(
self,
table_name: str,
schema: str = "ai",
db_url: Optional[str] = None,
db_engine: Optional[Engine] = None,
embedder: Optional[Embedder] = None,
search_type: SearchType = SearchType.vector,
vector_index: Union[Ivfflat, HNSW] = HNSW(),
distance: Distance = Distance.cosine,
prefix_match: bool = False,
vector_score_weight: float = 0.5,
content_language: str = "english",
schema_version: int = 1,
auto_upgrade_schema: bool = False,
reranker: Optional[Reranker] = None,
):
"""Initialize the OptimizedPgVector with the same parameters as PgVector."""
super().__init__(
table_name=table_name,
schema=schema,
db_url=db_url,
db_engine=db_engine,
embedder=embedder,
search_type=search_type,
vector_index=vector_index,
distance=distance,
prefix_match=prefix_match,
vector_score_weight=vector_score_weight,
content_language=content_language,
schema_version=schema_version,
auto_upgrade_schema=auto_upgrade_schema,
reranker=reranker,
)
def create(self) -> None:
"""Create the per-org table, tolerating a concurrent creator.
agno's PgVector.create() checks table_exists() and then issues CREATE
TABLE, which is not atomic. The knowledge processor runs two queue items
at once, so an organisation whose first two sources are queued together
has both workers find no table and both create it; the loser died with
`duplicate key value violates unique constraint "pg_type_typname_nsp_index"`
— Postgres rejecting the second table's implicit row type — and the user
lost that source.
Losing the race is not an error: the table the other worker created is
the one we wanted. Anything else still raises.
"""
try:
super().create()
except Exception as e:
if not self._is_duplicate_object_error(e):
raise
logger.debug(
f"Table {self.schema}.{self.table_name} was created concurrently; continuing"
)
@staticmethod
def _is_duplicate_object_error(error: Exception) -> bool:
"""True when the failure is 'someone else already created this'.
Matched on SQLSTATE rather than message text: 42P07 duplicate_table and
23505 unique_violation, the latter being how a clash on the pg_type
catalog index surfaces.
"""
for exc in (error, getattr(error, 'orig', None)):
sqlstate = getattr(exc, 'sqlstate', None) or getattr(exc, 'pgcode', None)
if sqlstate in ('42P07', '23505'):
return True
return False
def insert(
self,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
batch_size: int = 100,
) -> None:
"""
Insert documents into the database, preserving existing embeddings.
Args:
documents (List[Document]): List of documents to insert.
filters (Optional[Dict[str, Any]]): Filters to apply to the documents.
batch_size (int): Number of documents to insert in each batch.
"""
try:
with self.Session() as sess:
for i in range(0, len(documents), batch_size):
batch_docs = documents[i : i + batch_size]
try:
# Prepare documents for insertion
batch_records = []
for doc in batch_docs:
try:
# Embed document only if not already embedded (should be rare)
if doc.embedding is None:
if self.embedder is not None:
doc.embed(embedder=self.embedder)
logger.debug(f"Document embedding done for {doc.name}")
else:
logger.warning(f"Document '{doc.name}' has no embedding and no embedder available - skipping")
continue
# Merge filters into document metadata for searchability
merged_meta_data = doc.meta_data.copy() if doc.meta_data else {}
if filters:
merged_meta_data.update(filters)
logger.debug(f"Merged filters into meta_data: {filters}")
cleaned_content = self._clean_content(doc.content)
content_hash = md5(cleaned_content.encode()).hexdigest()
_id = doc.id or content_hash
record = {
"id": _id,
"name": doc.name,
"meta_data": merged_meta_data, # Use merged metadata
"filters": filters,
"content": cleaned_content,
"embedding": doc.embedding,
"usage": doc.usage,
"content_hash": content_hash,
}
batch_records.append(record)
except Exception as e:
logger.error(f"Error processing document '{doc.name}': {e}")
# Insert the batch of records
from sqlalchemy.dialects import postgresql
insert_stmt = postgresql.insert(self.table)
sess.execute(insert_stmt, batch_records)
sess.commit() # Commit batch independently
except Exception as e:
logger.error(f"Error with batch starting at index {i}: {e}")
sess.rollback() # Rollback the current batch if there's an error
raise
except Exception as e:
logger.error(f"Error inserting documents: {e}")
raise
def upsert(
self,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
batch_size: int = 100,
) -> None:
"""
Upsert (insert or update) documents in the database, preserving existing embeddings.
Args:
documents (List[Document]): List of documents to upsert.
filters (Optional[Dict[str, Any]]): Filters to apply to the documents.
batch_size (int): Number of documents to upsert in each batch.
"""
try:
with self.Session() as sess:
for i in range(0, len(documents), batch_size):
batch_docs = documents[i : i + batch_size]
try:
# Prepare documents for upserting
batch_records = []
for doc in batch_docs:
try:
# Skip embedding if the document already has an embedding
if doc.embedding is None:
doc.embed(embedder=self.embedder)
# Merge filters into document metadata for searchability
merged_meta_data = doc.meta_data.copy() if doc.meta_data else {}
if filters:
merged_meta_data.update(filters)
logger.debug(f"Merged filters into meta_data: {filters}")
cleaned_content = self._clean_content(doc.content)
content_hash = md5(cleaned_content.encode()).hexdigest()
_id = doc.id or content_hash
record = {
"id": _id,
"name": doc.name,
"meta_data": merged_meta_data, # Use merged metadata
"filters": filters,
"content": cleaned_content,
"embedding": doc.embedding,
"usage": doc.usage,
"content_hash": content_hash,
}
batch_records.append(record)
except Exception as e:
logger.error(f"Error processing document '{doc.name}': {e}")
# Upsert the batch of records
from sqlalchemy.dialects import postgresql
insert_stmt = postgresql.insert(self.table).values(batch_records)
upsert_stmt = insert_stmt.on_conflict_do_update(
index_elements=["id"],
set_=dict(
name=insert_stmt.excluded.name,
meta_data=insert_stmt.excluded.meta_data,
filters=insert_stmt.excluded.filters,
content=insert_stmt.excluded.content,
embedding=insert_stmt.excluded.embedding,
usage=insert_stmt.excluded.usage,
content_hash=insert_stmt.excluded.content_hash,
),
)
sess.execute(upsert_stmt)
sess.commit() # Commit batch independently
except Exception as e:
logger.error(f"Error with batch starting at index {i}: {e}")
sess.rollback() # Rollback the current batch if there's an error
raise
except Exception as e:
logger.error(f"Error upserting documents: {e}")
raise
+345
View File
@@ -0,0 +1,345 @@
"""
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.
"""
"""Shared helpers for reading, re-embedding and deleting knowledge sub-pages.
A knowledge "source" (``Knowledge`` row) stores its crawled/extracted content as
rows in the pgvector table ``{schema}."{table}"``. Each row is a *chunk*; a
*page* (aka sub-page) is the group of chunks that share a base ``id`` — the page
URL/name — with optional ``_N`` chunk suffixes (e.g. ``site.com/docs`` and
``site.com/docs_1``, ``site.com/docs_2``).
These helpers centralise the embed/upsert/delete logic that the knowledge API
endpoints previously duplicated inline, so editing a chunk, adding a sub-page and
replacing a whole page all go through one code path.
"""
from typing import Any, Dict, List, Optional
from uuid import UUID
from agno.document import Document
from sqlalchemy import text
from sqlalchemy.orm import Session
from app.core.logger import get_logger
from app.knowledge.chunking import CHUNK_META_KEYS, chunk_page
from app.knowledge.knowledge_base import KnowledgeManager
from app.models.knowledge import Knowledge, SourceType
from app.models.knowledge_to_agent import KnowledgeToAgent
from app.repositories.knowledge import KnowledgeRepository
from app.repositories.knowledge_to_agent import KnowledgeToAgentRepository
logger = get_logger(__name__)
# SQL expression that maps a chunk ``id`` back to its page id by stripping the
# trailing chunk suffix: ``<page>::<n>`` (app.knowledge.chunking) or the legacy
# ``<page>_<n>`` agno's PDF reader produced. Only a *numeric* suffix is
# stripped, so legitimate underscores in a page name (e.g. ``getting_started``)
# are kept. This is the single source of truth for page grouping — the
# source-listing queries import it too, and ``page_id_of`` / the frontend
# ``basePageId`` mirror it — so page membership is identical everywhere
# edit/delete/list operate.
PAGE_ID_EXPR = (
"CASE WHEN id ~ '::[0-9]+$' THEN substring(id from '^(.*)::[0-9]+$') "
"WHEN id ~ '_[0-9]+$' THEN substring(id from '^(.*)_[0-9]+$') ELSE id END"
)
# Chunk ids within a page sort naturally (length first) so ``page::2`` comes
# before ``page::10`` where plain text ordering would not. The bare page id, if
# any, sorts first.
NATURAL_ID_ORDER = "length(id), id"
def get_manager(organization_id) -> KnowledgeManager:
"""Build a KnowledgeManager bound to an org's vector table (no agent link)."""
return KnowledgeManager(org_id=organization_id, agent_id=None)
def agent_ids_for(knowledge: Knowledge) -> List[str]:
"""Return the agent ids linked to a knowledge source (for vector filters)."""
return [str(link.agent_id) for link in knowledge.agent_links]
def embed_document(
manager: KnowledgeManager,
source: str,
doc_id: str,
content: str,
meta_data: Optional[Dict[str, Any]] = None,
) -> Document:
"""Create a Document for ``doc_id`` and embed it with the org's embedder."""
doc = Document(
name=source,
id=doc_id,
content=content,
meta_data=meta_data or {},
)
doc.embed(embedder=manager.vector_db.embedder)
return doc
def embed_page(
manager: KnowledgeManager,
source: str,
page_id: str,
content: str,
meta_data: Optional[Dict[str, Any]] = None,
) -> List[Document]:
"""Chunk a page and embed every chunk. Embedding happens before any write,
so an embedder failure changes nothing in the store."""
page = Document(name=source, id=page_id, content=content, meta_data=meta_data or {})
chunks = chunk_page(page)
for chunk in chunks:
chunk.embed(embedder=manager.vector_db.embedder)
return chunks
def delete_page_chunks_except(
db: Session, knowledge: Knowledge, page_id: str, keep_ids: List[str]
) -> int:
"""Delete the chunks of a page whose id is not in ``keep_ids``. Used after
a page's new chunks are stored: whatever the old layout was (one bare row,
``_N`` rows, ``::N`` rows), only the new set survives. Caller commits."""
sql = (
f'DELETE FROM {knowledge.schema}."{knowledge.table_name}" '
f"WHERE name = :source AND {PAGE_ID_EXPR} = :page_id "
"AND NOT (id = ANY(:keep_ids))"
)
result = db.execute(
text(sql), {"source": knowledge.source, "page_id": page_id, "keep_ids": keep_ids}
)
return result.rowcount
def count_subpages(db: Session, knowledge: Knowledge) -> int:
"""Count the distinct sub-pages stored for a knowledge source.
Counts pages, not raw chunks, so a page split into several ``_N`` chunks
counts once — matching the ``max_sub_pages`` plan limit's intent.
"""
query = text(
f"SELECT COUNT(DISTINCT {PAGE_ID_EXPR}) AS count "
f'FROM {knowledge.schema}."{knowledge.table_name}" WHERE name = :source'
)
row = db.execute(query, {"source": knowledge.source}).fetchone()
return row.count if row else 0
def get_page_chunks(db: Session, knowledge: Knowledge, page_id: str) -> List[Any]:
"""Return all chunk rows (id, content, meta_data) belonging to a page."""
query = text(
f'SELECT id, content, meta_data FROM {knowledge.schema}."{knowledge.table_name}" '
f"WHERE name = :source AND {PAGE_ID_EXPR} = :page_id "
f"ORDER BY {NATURAL_ID_ORDER}"
)
return db.execute(
query, {"source": knowledge.source, "page_id": page_id}
).fetchall()
def delete_page_chunks(db: Session, knowledge: Knowledge, page_id: str) -> int:
"""Delete every chunk of a page. Returns the number of rows removed. The
caller commits the transaction."""
sql = (
f'DELETE FROM {knowledge.schema}."{knowledge.table_name}" '
f"WHERE name = :source AND {PAGE_ID_EXPR} = :page_id"
)
result = db.execute(text(sql), {"source": knowledge.source, "page_id": page_id})
return result.rowcount
def replace_page(
db: Session,
knowledge: Knowledge,
page_id: str,
content: str,
title: Optional[str] = None,
) -> int:
"""Replace a page's content with freshly chunked and embedded rows.
Metadata (including the page ``url`` and agent linkage) is preserved from the
existing chunks where available, so the edited page stays attributed to the
same agents and keeps its source URL.
Ordering is chosen so the page can never be lost:
1. Chunk and embed the new content (if this raises, nothing has changed).
2. ``upsert`` the new chunks on the vector store's own connection. Ids that
already exist are overwritten in place; the new content is now durably
stored before anything is deleted.
3. Delete every other row of the page — the bare ``page_id`` row of a page
that now spans several chunks, or trailing chunks of a page that shrank —
and commit.
A failure at step 2 leaves the original page intact; a failure at step 3
leaves the correct new content in place plus at most some stale extra chunks
(no data loss). This function commits the request session itself.
Returns the number of chunks the page had before the replace (0 if the page
does not exist).
"""
existing = get_page_chunks(db, knowledge, page_id)
if not existing:
return 0
# Preserve the first chunk's metadata (url, etc.); refresh agent linkage.
# Per-chunk keys are recomputed by chunk_page — carrying them over would
# leave a page edited down to one row claiming to be "part 1 of 5".
base_meta: Dict[str, Any] = {
key: value
for key, value in (existing[0].meta_data or {}).items()
if key not in CHUNK_META_KEYS
}
agent_ids = agent_ids_for(knowledge)
base_meta["agent_id"] = agent_ids
if title is not None:
base_meta["title"] = title
manager = get_manager(knowledge.organization_id)
chunks = embed_page(manager, knowledge.source, page_id, content, base_meta)
filters = {
"name": knowledge.source,
"agent_id": agent_ids,
"org_id": str(knowledge.organization_id),
}
# Persist the new content first, then clear every other row of the page.
manager.vector_db.upsert(chunks, filters=filters)
removed_extra = delete_page_chunks_except(
db, knowledge, page_id, [chunk.id for chunk in chunks]
)
db.commit()
logger.info(
f"Replaced page '{page_id}' ({len(existing)} -> {len(chunks)} chunk(s), "
f"{removed_extra} old removed) for source '{knowledge.source}'"
)
return len(existing)
def reembed_chunk(
db: Session,
knowledge: Knowledge,
chunk_id: str,
content: str,
) -> None:
"""Re-embed one chunk's content in place (scoped to its source). Caller commits."""
manager = get_manager(knowledge.organization_id)
doc = embed_document(manager, knowledge.source, chunk_id, content)
update_query = text(
f'UPDATE {knowledge.schema}."{knowledge.table_name}" '
"SET content = :content, embedding = :embedding "
"WHERE id = :chunk_id AND name = :source"
)
db.execute(
update_query,
{
"content": content,
"embedding": doc.embedding,
"chunk_id": chunk_id,
"source": knowledge.source,
},
)
def create_text_source(
db: Session,
organization_id: UUID,
title: str,
content: str,
agent_id: Optional[UUID] = None,
) -> Knowledge:
"""Create a knowledge source from pasted text and embed it as one page.
Unlike a URL/PDF source (which is crawled asynchronously via the queue), a
text source is indexed inline and is immediately ``synced``. Returns the new
``Knowledge`` row. Caller is responsible for the enclosing request context;
the vector write commits on its own session.
"""
manager = KnowledgeManager(
org_id=organization_id, agent_id=str(agent_id) if agent_id else None
)
agent_ids = [str(agent_id)] if agent_id else []
# Embed first — an embedding failure then creates no source row at all.
chunks = embed_page(
manager, title, title, content, {"agent_id": agent_ids, "url": title, "title": title}
)
repo = KnowledgeRepository(db)
# KnowledgeRepository.create dedupes on (org, source); track whether the row
# already existed so we only compensate-delete a row we actually created.
pre_existing = bool(repo.get_by_sources(organization_id, [title]))
knowledge = repo.create(
Knowledge(
organization_id=organization_id,
source=title,
source_type=SourceType.CUSTOM,
table_name=manager.vector_db.table_name,
schema=manager.vector_db.schema,
)
)
try:
if agent_id is not None:
KnowledgeToAgentRepository(db).create(
KnowledgeToAgent(knowledge_id=knowledge.id, agent_id=agent_id)
)
# The org's vector table is otherwise created only by URL/PDF ingestion, so
# an org whose first source is pasted text has none yet. create() is a
# no-op when it exists and tolerates a concurrent creator.
manager.vector_db.create()
# upsert (not insert) is idempotent on the id, so a raced duplicate title
# overwrites its own chunks instead of raising a duplicate-key error.
manager.vector_db.upsert(
chunks,
filters={"name": title, "agent_id": agent_ids, "org_id": str(organization_id)},
)
except Exception:
# Don't leave a chunkless orphan source behind if indexing fails.
if not pre_existing:
repo.delete(knowledge.id)
raise
logger.info(f"Created text knowledge source '{title}' for org {organization_id}")
return knowledge
def insert_subpage(
knowledge: Knowledge,
subpage_name: str,
content: str,
url: Optional[str] = None,
) -> None:
"""Embed and store a brand-new sub-page as one or more chunks.
``subpage_name`` is the page id/title; ``url`` is an optional source link
stored in metadata so the UI can render it (and derive the title from the
name rather than from a URL the user pasted). The write happens on the vector
store's own session (which commits itself); there is nothing to commit on the
request session. The caller has already checked the page does not exist;
``upsert`` keeps a retry from failing on a half-written page.
"""
agent_ids = agent_ids_for(knowledge)
meta_data: Dict[str, Any] = {"agent_id": agent_ids, "title": subpage_name}
if url:
meta_data["url"] = url
manager = get_manager(knowledge.organization_id)
chunks = embed_page(manager, knowledge.source, subpage_name, content, meta_data)
filters = {
"name": knowledge.source,
"agent_id": agent_ids,
"org_id": str(knowledge.organization_id),
}
manager.vector_db.upsert(chunks, filters=filters)
+79
View File
@@ -0,0 +1,79 @@
"""
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.
Reading a crawled link, but only when it is a page.
A site's links include its downloads. Read as text, one Linux installer
became ~30,000 chunks to embed: ten links on ian-vpn.site made 149,500, and
the worker ran for three hours before it hit its memory limit. The headers
say what a link is, so a download is skipped before its body is fetched.
"""
from app.knowledge.url_safety import safe_stream
# What a crawl reads. Anything else (archives, installers, images, PDFs -
# those have their own source type) is a download, not a page. XML is in
# because some sites serve their XHTML pages as plain XML.
PAGE_TYPES = frozenset({
"text/html", "application/xhtml+xml", "text/plain", "application/xml", "text/xml",
})
# Far beyond any real page; reading stops here so a mislabelled file cannot
# fill memory. Counted after decompression, so a small gzip bomb stops too.
MAX_PAGE_BYTES = 5 * 1024 * 1024
# A server that sends no type still gives a binary file away in its first
# bytes. Only then: a declared page may be UTF-16, which is full of NULs.
_SNIFF_BYTES = 1024
class NotAPage(Exception):
"""The link is a download or too large to be a page. The crawl skips it;
it is not a failure of the site."""
def __init__(self, message: str, download: bool = True):
super().__init__(message)
self.download = download
def media_type(content_type: str) -> str:
"""'text/html; charset=utf-8' -> 'text/html'."""
return content_type.split(";", 1)[0].strip().lower()
def fetch_page_text(client, url: str, max_bytes: int = MAX_PAGE_BYTES) -> str:
"""The page at `url` as text. Raises NotAPage for a download or an
oversized body, httpx.HTTPStatusError for a non-2xx answer, and
BlockedHostError for a redirect to an address we may not fetch."""
with safe_stream(client, url) as response:
response.raise_for_status()
kind = media_type(response.headers.get("content-type", ""))
if kind and kind not in PAGE_TYPES:
raise NotAPage(f"{url} is a download ({kind}), not a page")
declared = response.headers.get("content-length", "")
if declared.isdigit() and int(declared) > max_bytes:
raise NotAPage(f"{url} is larger than {max_bytes} bytes", download=False)
body = bytearray()
for chunk in response.iter_bytes():
if not kind and not body and b"\x00" in chunk[:_SNIFF_BYTES]:
raise NotAPage(f"{url} is a binary file, not a page")
body.extend(chunk)
if len(body) > max_bytes:
raise NotAPage(f"{url} is larger than {max_bytes} bytes", download=False)
return body.decode(response.encoding or "utf-8", errors="replace")
+63
View File
@@ -0,0 +1,63 @@
"""
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.
"""
"""SSRF-guarded download of a knowledge PDF from a user-supplied URL."""
import httpx
from app.core.config import settings
from app.knowledge.pdf_limits import PdfRefused
from app.knowledge.url_safety import safe_stream
# Same ceiling as a directly uploaded PDF, so a URL is no way around it.
MAX_PDF_SIZE = 25 * 1024 * 1024
_CHUNK_SIZE = 8192
class PdfTooLargeError(PdfRefused):
"""The remote file is larger than the allowed PDF size."""
def __init__(self, max_bytes: int):
# Shown on the failed source: no URL (an upload's is a signed S3 link)
super().__init__(
f"This PDF is larger than {max_bytes // (1024 * 1024)} MB. Upload a "
"smaller file, or split it into parts."
)
def download_pdf(url: str, dest_path: str, max_bytes: int = MAX_PDF_SIZE) -> str:
"""Stream ``url`` into ``dest_path`` and return the path.
Every hop, redirects included, is checked against internal addresses.
Raises BlockedHostError for an internal target, PdfTooLargeError past
``max_bytes`` and httpx.HTTPStatusError for a non-2xx answer; the caller
owns cleaning up a partially written ``dest_path``.
"""
with httpx.Client(
timeout=settings.KB_TIMEOUT, verify=settings.VERIFY_SSL_CERTIFICATES
) as client, safe_stream(client, url) as response:
response.raise_for_status()
declared = response.headers.get("content-length", "")
if declared.isdigit() and int(declared) > max_bytes:
raise PdfTooLargeError(max_bytes)
written = 0
with open(dest_path, "wb") as f:
for chunk in response.iter_bytes(_CHUNK_SIZE):
written += len(chunk)
if written > max_bytes:
raise PdfTooLargeError(max_bytes)
f.write(chunk)
return dest_path
+61
View File
@@ -0,0 +1,61 @@
"""
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.
PDFs we will not index, and why.
Bytes don't bound the work: a small file can hold thousands of pages, and a
PDF is read page by page - every page's text embedded, and a scanned one
OCR'd at a second or more a page. So the page count is limited too, checked
before any of it is read.
"""
from typing import BinaryIO, Optional, Union
from pypdf import PdfReader
from app.core.config import settings
from app.core.logger import get_logger
logger = get_logger(__name__)
class PdfRefused(ValueError):
"""A PDF over one of our limits. The message is for the user, as it is
shown on the failed source."""
class PdfTooManyPages(PdfRefused):
pass
def count_pages(pdf: Union[str, BinaryIO]) -> Optional[int]:
"""Pages in `pdf` (a path or a file object), or None when it cannot be
counted (encrypted, damaged): the reader then fails it its own way."""
try:
return len(PdfReader(pdf).pages)
except Exception as e:
logger.warning(f"Could not count the pages of a PDF: {e}")
return None
def check_page_count(pdf: Union[str, BinaryIO], max_pages: Optional[int] = None) -> None:
"""Raise PdfTooManyPages when `pdf` has more pages than we index."""
limit = max_pages if max_pages is not None else settings.KB_PDF_MAX_PAGES
pages = count_pages(pdf)
if pages is not None and pages > limit:
raise PdfTooManyPages(
f"This PDF has {pages} pages; the limit is {limit}. Upload the part "
"you need, or split it into smaller files."
)
+222
View File
@@ -0,0 +1,222 @@
"""
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.
"""
"""Fetch and parse ``sitemap.xml`` files into a flat list of page URLs.
Supports plain ``<urlset>`` sitemaps and ``<sitemapindex>`` files that point to
child sitemaps (recursed breadth-first, bounded), gzip-compressed sitemaps, and
namespaced tags. Parsing uses ``defusedxml`` to avoid XML-bomb / XXE attacks.
Because the sitemap URL (and every child ``<loc>``) is user-supplied, fetches are
hardened against abuse: downloads and decompression are size-capped (gzip-bomb
protection), literal private/loopback/link-local IP hosts are blocked (SSRF), and
recursed child sitemaps are restricted to the root's registrable domain so an
index can't fan out to arbitrary hosts.
"""
import gzip
import io
from typing import Dict, List, Optional, Tuple
from urllib.parse import urljoin, urlparse
import httpx
from defusedxml import ElementTree as DefusedET
from app.core.logger import get_logger
from app.knowledge.url_safety import BlockedHostError, MAX_REDIRECTS, guard_url
logger = get_logger(__name__)
# Sitemaps.org limits are 50 MB uncompressed / 50k URLs per file; cap a little
# above that and bound decompression so a gzip bomb can't exhaust memory.
MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024
MAX_DECOMPRESSED_BYTES = 60 * 1024 * 1024
def _local_name(tag: str) -> str:
"""Strip an XML namespace from a tag: ``{ns}loc`` -> ``loc`` (lowercased)."""
return tag.rsplit("}", 1)[-1].lower()
def _registrable_domain(host: str) -> str:
"""Registrable domain for same-site checks — the shared ccTLD-aware helper
so 'a.example.co.uk' → 'example.co.uk' (not 'co.uk')."""
from app.knowledge.domains import registrable_domain
return registrable_domain(host)
def _bounded_gunzip(data: bytes) -> Optional[bytes]:
"""Decompress gzip bytes, bounded to MAX_DECOMPRESSED_BYTES. None if not gzip
or the decompressed size would exceed the cap (gzip-bomb protection)."""
try:
with gzip.GzipFile(fileobj=io.BytesIO(data)) as gz:
out = gz.read(MAX_DECOMPRESSED_BYTES + 1)
except OSError:
return None
if len(out) > MAX_DECOMPRESSED_BYTES:
logger.warning("Sitemap decompressed size exceeds cap; rejecting")
return None
return out
def _maybe_gunzip(url: str, content: bytes) -> bytes:
"""Decompress a ``.gz`` sitemap; leave anything else untouched."""
if url.lower().endswith(".gz"):
out = _bounded_gunzip(content)
if out is not None:
return out
return content
def _fetch(client: httpx.Client, url: str) -> Optional[bytes]:
"""GET a sitemap URL (size-capped, SSRF-guarded); (gunzipped) bytes or None.
Redirects are followed manually so each hop's host is re-validated against
internal addresses before it is fetched.
"""
current = url
for _ in range(MAX_REDIRECTS + 1):
try:
guard_url(current)
except BlockedHostError as e:
logger.warning(str(e))
return None
try:
with client.stream("GET", current, follow_redirects=False) as resp:
if resp.is_redirect:
location = resp.headers.get("location")
if not location:
return None
current = urljoin(str(resp.url), location)
continue
resp.raise_for_status()
chunks: List[bytes] = []
total = 0
for chunk in resp.iter_bytes():
total += len(chunk)
if total > MAX_DOWNLOAD_BYTES:
logger.warning(f"Sitemap '{current}' exceeds size cap; rejecting")
return None
chunks.append(chunk)
return _maybe_gunzip(current, b"".join(chunks))
except Exception as e:
logger.warning(f"Failed to fetch sitemap '{current}': {e}")
return None
logger.warning(f"Too many redirects while fetching sitemap {url}")
return None
def _parse(content: bytes) -> Tuple[List[str], List[str]]:
"""Parse sitemap XML into (page_urls, child_sitemap_urls).
A ``<url>`` entry is a page; a ``<sitemap>`` entry points to a child sitemap.
Falls back to a bounded gunzip once if the bytes aren't valid XML (some
servers return gzip without a ``.gz`` extension).
"""
try:
root = DefusedET.fromstring(content)
except Exception:
unzipped = _bounded_gunzip(content)
if unzipped is None:
return [], []
try:
root = DefusedET.fromstring(unzipped)
except Exception as e:
logger.warning(f"Failed to parse sitemap XML: {e}")
return [], []
pages: List[str] = []
children: List[str] = []
for entry in root:
kind = _local_name(entry.tag)
if kind not in ("url", "sitemap"):
continue
loc: Optional[str] = None
for child in entry:
if _local_name(child.tag) == "loc" and child.text and child.text.strip():
loc = child.text.strip()
break
if not loc:
continue
(children if kind == "sitemap" else pages).append(loc)
return pages, children
def fetch_sitemap_urls(
url: str,
*,
headers: Optional[Dict[str, str]] = None,
timeout: float = 30,
verify_ssl: bool = True,
max_urls: int = 50,
max_sitemaps: int = 50,
) -> List[str]:
"""Return the deduped page URLs listed in a sitemap (or sitemap index).
Recurses child sitemaps breadth-first, bounded by ``max_sitemaps`` fetches
and ``max_urls`` total pages (the plan's per-source sub-page limit). Child
sitemaps outside the root's registrable domain, and errors on individual
(child) sitemaps, are skipped.
"""
headers = headers or {}
root_domain = _registrable_domain(urlparse(url).hostname or "")
to_visit: List[str] = [url]
visited_sitemaps: set = set()
seen_pages: set = set()
pages: List[str] = []
with httpx.Client(
timeout=timeout, follow_redirects=True, headers=headers, verify=verify_ssl
) as client:
while to_visit and len(pages) < max_urls and len(visited_sitemaps) < max_sitemaps:
sitemap_url = to_visit.pop(0)
if sitemap_url in visited_sitemaps:
continue
visited_sitemaps.add(sitemap_url)
content = _fetch(client, sitemap_url)
if not content:
continue
found_pages, child_sitemaps = _parse(content)
for page in found_pages:
if page in seen_pages:
continue
# Page <loc>s must be on the root's registrable domain too — a
# sitemap that lists other domains' URLs would otherwise pull
# them all in under one plan-limited knowledge source.
if _registrable_domain(urlparse(page).hostname or "") != root_domain:
logger.warning(f"Skipping cross-domain sitemap page: {page}")
continue
seen_pages.add(page)
pages.append(page)
if len(pages) >= max_urls:
break
for child in child_sitemaps:
if child in visited_sitemaps or child in to_visit:
continue
if _registrable_domain(urlparse(child).hostname or "") != root_domain:
logger.warning(f"Skipping cross-domain child sitemap: {child}")
continue
to_visit.append(child)
logger.info(
f"Sitemap '{url}': discovered {len(pages)} page(s) from "
f"{len(visited_sitemaps)} sitemap file(s)"
)
return pages[:max_urls]
+85
View File
@@ -0,0 +1,85 @@
"""
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 typing import Callable, Dict, Optional
from app.core.logger import get_logger
from app.knowledge.crawl_scope import SCOPE_DOMAIN, CrawlScope
from app.knowledge.enhanced_website_reader import EnhancedWebsiteReader
from app.knowledge.sitemap_parser import fetch_sitemap_urls
logger = get_logger(__name__)
class SitemapReader(EnhancedWebsiteReader):
"""Crawl exactly the pages listed in a ``sitemap.xml`` — no ``<a href>`` following.
Reuses EnhancedWebsiteReader's whole per-page pipeline (httpx fetch, Crawl4AI
fallback, bot-challenge detection, content extraction, Document creation, and
the ``name = <sitemap url>`` / ``id = <page url>`` grouping); only the URL
source differs — pages come from the sitemap rather than BFS link discovery.
Construct with ``max_depth=1`` so ``_process_url`` never yields child links,
and ``max_links`` = the plan's ``max_sub_pages`` so the discovered page list
is capped at ingest.
"""
def crawl(
self,
url: str,
starting_depth: int = 1,
on_document_callback: Optional[Callable[[str, str], None]] = None,
on_url_crawled_callback: Optional[Callable[[str], None]] = None,
) -> Dict[str, str]:
url = self._normalize_url(url)
# Same www/apex fallback as a website crawl - a sitemap URL on a broken
# www host is the identical customer problem.
url = self._seed_with_resolvable_host(url)
# Reset per-crawl state exactly as the base crawl does — the counters are
# read by the progress logs and by _raise_if_nothing_stored (a fully
# bot-blocked sitemap must still fail with the "add manually" message).
self._visited = set()
self._urls_to_crawl = []
self._crawled_pages_count = 0
self._successful_crawls = 0
self._failed_crawls = 0
self._challenge_blocked = 0
self._failure_causes = {}
# A sitemap index legitimately lists pages on sibling hosts
# (help.example.com from example.com/sitemap.xml), so the pages it names
# are taken at registrable-domain scope — matching the parser's own
# cross-domain filter. Only <a href> following is host-scoped.
scope = CrawlScope.for_seed(url, SCOPE_DOMAIN)
page_urls = fetch_sitemap_urls(
url,
headers=self.headers,
timeout=self.timeout,
verify_ssl=self.verify_ssl,
max_urls=self.max_links,
)
# Exposed so the dispatcher can fail a sitemap that lists no pages
# (wrong URL / not a sitemap) instead of silently creating an empty source.
self._sitemap_page_count = len(page_urls)
logger.info(f"Sitemap {url}: crawling {len(page_urls)} listed page(s)")
if not page_urls:
return {}
seeds = [(page_url, starting_depth) for page_url in page_urls]
return self._run_crawl_queue(
seeds, scope, on_document_callback, on_url_crawled_callback
)
+165
View File
@@ -0,0 +1,165 @@
"""
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.
"""
"""Shared SSRF guards for the crawler / sitemap fetchers."""
import ipaddress
import socket
from contextlib import contextmanager
from typing import Optional, Union
from urllib.parse import urljoin, urlparse
from app.core.logger import get_logger
logger = get_logger(__name__)
# Match httpx's default so the manual (SSRF-checked) redirect following doesn't
# reject legitimate sites with longer redirect chains.
MAX_REDIRECTS = 20
class BlockedHostError(Exception):
"""Raised when a fetch target resolves to a private/internal address."""
IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address]
# NAT64 well-known prefix: 64:ff9b::a9fe:a9fe reaches 169.254.169.254
_NAT64 = ipaddress.ip_network("64:ff9b::/96")
# Deprecated IPv4-compatible IPv6 (::a.b.c.d): ::7f00:1 is 127.0.0.1
_IPV4_COMPATIBLE = ipaddress.ip_network("::/96")
def is_public_ip(ip: IPAddress) -> bool:
"""True only for an address on the public internet.
An IPv6 address that carries an IPv4 one (mapped, NAT64, 6to4, Teredo,
IPv4-compatible) is judged as that IPv4 address. Then it must be global —
which also excludes carrier-grade NAT (100.64/10) and documentation
ranges — and none of the classic internal kinds.
"""
ip = _embedded_ipv4(ip) or ip
return ip.is_global and not (
ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved
or ip.is_multicast or ip.is_unspecified
or (ip.version == 6 and ip.is_site_local) # deprecated fec0::/10
)
def _embedded_ipv4(ip: IPAddress) -> Optional[ipaddress.IPv4Address]:
if ip.version != 6:
return None
if ip.ipv4_mapped is not None:
return ip.ipv4_mapped
if ip.sixtofour is not None:
return ip.sixtofour
if ip.teredo is not None:
return ip.teredo[1] # the client's address
if ip in _NAT64 or (ip in _IPV4_COMPATIBLE and int(ip) > 1):
return ipaddress.IPv4Address(int(ip) & 0xFFFFFFFF)
return None
def is_blocked_host(url: str) -> bool:
"""True if the URL targets a literal private/loopback/link-local IP host.
Catches the obvious internal targets given as an IP literal (cloud metadata
``169.254.169.254``, ``127.0.0.1``, RFC1918 ranges) plus ``localhost``.
"""
host = (urlparse(url).hostname or "").lower()
if host in ("localhost", "localhost.localdomain"):
return True
try:
return not is_public_ip(ipaddress.ip_address(host))
except ValueError:
return False
def resolves_to_blocked_host(url: str) -> bool:
"""True if the URL's host is, or DNS-resolves to, an internal address.
Extends :func:`is_blocked_host` by resolving hostnames — this closes the
"attacker domain with an A record pointing at 169.254.169.254 / an RFC1918
IP" SSRF vector. Hosts that can't be resolved are not blocked here (the
request will simply fail to connect).
"""
if is_blocked_host(url):
return True
host = (urlparse(url).hostname or "").lower()
if not host:
return False
try:
infos = socket.getaddrinfo(host, None)
except (socket.gaierror, UnicodeError, ValueError, OSError):
return False
for info in infos:
addr = info[4][0]
try:
if not is_public_ip(ipaddress.ip_address(addr)):
return True
except ValueError:
continue
return False
def guard_url(url: str) -> None:
"""Raise BlockedHostError if the URL targets an internal address."""
if resolves_to_blocked_host(url):
raise BlockedHostError(f"Refusing to fetch internal/blocked host: {url}")
def _redirect_target(resp) -> Optional[str]:
"""The absolute URL a redirect response points at, or None when ``resp``
is final (not a redirect, or a redirect with no Location)."""
if not resp.is_redirect:
return None
location = resp.headers.get("location")
return urljoin(str(resp.url), location) if location else None
def safe_get(client, url: str, *, max_redirects: int = MAX_REDIRECTS, **kwargs):
"""httpx GET that follows redirects manually, re-validating each hop's host.
Prevents the "public URL 302-redirects to an internal IP" SSRF: every hop
(including the initial URL) is checked with :func:`guard_url` before it is
fetched. Returns the final non-redirect response; the caller handles status.
"""
current = url
for _ in range(max_redirects + 1):
guard_url(current)
resp = client.get(current, follow_redirects=False, **kwargs)
target = _redirect_target(resp)
if target is None:
return resp
current = target
raise BlockedHostError(f"Too many redirects while fetching {url}")
@contextmanager
def safe_stream(client, url: str, *, max_redirects: int = MAX_REDIRECTS, **kwargs):
"""Streaming counterpart of :func:`safe_get`, for bodies too large to hold
in memory. Yields the final non-redirect response, unread; every hop is
checked with :func:`guard_url` before it is opened."""
current = url
for _ in range(max_redirects + 1):
guard_url(current)
with client.stream("GET", current, follow_redirects=False, **kwargs) as resp:
target = _redirect_target(resp)
if target is None:
yield resp
return
current = target
raise BlockedHostError(f"Too many redirects while fetching {url}")
@@ -0,0 +1,140 @@
"""
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.
Keeps the agent_id list inside the vector store in step with the relational
knowledge_to_agents links.
Retrieval filters on `filters->'agent_id'`, so a link row that never reaches the
vector store is invisible to the agent: the dashboard shows the source as linked
while the agent cannot retrieve a single chunk of it. Both directions live here
so link and unlink cannot drift apart.
"""
import json
from typing import List, Optional
from uuid import UUID
from sqlalchemy import text
from sqlalchemy.orm import Session
from app.core.logger import get_logger
from app.models.knowledge import Knowledge
logger = get_logger(__name__)
# Every caller-influenced value is bound as a query parameter. The JSON payloads
# embed knowledge.source (a user-supplied title) and json.dumps does not escape
# single quotes, so interpolating them into a '...'::jsonb literal would allow
# SQL injection. Only the system-derived schema/table identifiers are inlined.
#
# CAST(:x AS jsonb), never :x::jsonb — SQLAlchemy's text() bind-parameter regex
# has a negative lookahead on ':', so a parameter immediately followed by the
# '::' cast operator is not recognised as a parameter at all. It reaches the
# driver as a literal ':new_filters' and Postgres rejects it with
# `syntax error at or near ":"`.
_UPDATE_SQL = """
UPDATE {schema}."{table}"
SET
filters = CAST(:new_filters AS jsonb),
meta_data = CASE
WHEN meta_data IS NULL THEN
CASE WHEN :has_agents THEN jsonb_build_object('agent_id', CAST(:agent_ids AS jsonb))
ELSE NULL END
WHEN meta_data ? 'agent_id' THEN
jsonb_set(meta_data, '{{agent_id}}', CAST(:agent_ids AS jsonb))
ELSE
meta_data || jsonb_build_object('agent_id', CAST(:agent_ids AS jsonb))
END
WHERE name = :source
"""
def get_linked_agent_ids(db: Session, knowledge: Knowledge) -> Optional[List[str]]:
"""Agent ids currently recorded in the vector store for this source.
Returns None when the source has no rows in the vector table yet (nothing to
sync — the chunks have not been written), which callers must distinguish
from an empty list (rows exist, no agents linked).
"""
if not knowledge.table_name or not knowledge.schema:
return None
row = db.execute(
text(
f'SELECT filters FROM {knowledge.schema}."{knowledge.table_name}" '
"WHERE name = :source LIMIT 1"
),
{"source": knowledge.source},
).first()
if row is None:
return None
filters = row.filters or {}
agent_ids = filters.get("agent_id") or []
# Defensive: the column is caller-adjacent JSON, so a non-list here would
# otherwise propagate a bad type into the rewritten filters.
return [str(a) for a in agent_ids] if isinstance(agent_ids, list) else []
def sync_agent_ids(db: Session, knowledge: Knowledge, agent_ids: List[str]) -> None:
"""Rewrite filters/meta_data for every chunk of `knowledge` to `agent_ids`.
Does not commit — the caller owns the transaction so the link row and this
update land together.
"""
if not knowledge.table_name or not knowledge.schema:
return
# Order-independent de-duplication; sorted so repeated syncs are idempotent
# and the stored JSON does not churn.
unique_ids = sorted(set(agent_ids))
new_filters = {
"name": knowledge.source,
"org_id": str(knowledge.organization_id),
"agent_id": unique_ids,
}
db.execute(
text(_UPDATE_SQL.format(schema=knowledge.schema, table=knowledge.table_name)),
{
"source": knowledge.source,
"new_filters": json.dumps(new_filters),
"agent_ids": json.dumps(unique_ids),
"has_agents": bool(unique_ids),
},
)
logger.info(
"Synced vector-store agent_id for knowledge source %s: %d agent(s)",
knowledge.source,
len(unique_ids),
)
def add_agent(db: Session, knowledge: Knowledge, agent_uuid: UUID) -> None:
"""Add one agent to the vector-store filters, preserving the others."""
current = get_linked_agent_ids(db, knowledge)
if current is None:
return
sync_agent_ids(db, knowledge, current + [str(agent_uuid)])
def remove_agent(db: Session, knowledge: Knowledge, agent_uuid: UUID) -> None:
"""Remove one agent from the vector-store filters, preserving the others."""
current = get_linked_agent_ids(db, knowledge)
if current is None:
return
sync_agent_ids(db, knowledge, [a for a in current if a != str(agent_uuid)])
+101
View File
@@ -0,0 +1,101 @@
"""
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.
"""
import asyncio
import sys
import os
from pathlib import Path
# Add the parent directory to Python path to allow imports
sys.path.append(str(Path(__file__).parent.parent.parent))
from app.database import SessionLocal
from app.repositories.session_to_agent import SessionToAgentRepository
from app.core.logger import get_logger
logger = get_logger(__name__)
# Configure check interval (default to 24 hours - once a day)
CHECK_INTERVAL = int(os.environ.get('CHAT_AUTO_CLOSE_INTERVAL', '86400')) # 24 hours in seconds
WORKER_ID = os.environ.get('GUNICORN_WORKER_ID', 'standalone')
async def run_auto_closer():
"""Single run of the auto closer"""
db = None
try:
logger.info(f"Worker {WORKER_ID}: Running chat auto-closer")
db = SessionLocal()
session_repo = SessionToAgentRepository(db)
# Auto-close inactive agent chats
closed_count = session_repo.auto_close_inactive_agent_chats()
if closed_count > 0:
logger.info(f"Worker {WORKER_ID}: Auto-closed {closed_count} inactive agent chats")
else:
logger.info(f"Worker {WORKER_ID}: No inactive agent chats found to close")
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Error in chat auto-closer: {str(e)}")
raise
finally:
if db:
db.close()
async def run_auto_closer_loop():
"""Run the chat auto-closer in a continuous loop"""
logger.info(f"Worker {WORKER_ID}: Starting chat auto-closer with {CHECK_INTERVAL}s check interval (once per day)")
try:
iteration = 0
while True:
iteration += 1
start_time = asyncio.get_event_loop().time()
try:
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Checking for inactive agent chats...")
await run_auto_closer()
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Auto-closer completed successfully")
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Iteration {iteration} - Error in auto-closer: {str(e)}")
# Calculate how long the processing took
elapsed = asyncio.get_event_loop().time() - start_time
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Processing took {elapsed:.2f} seconds")
# Sleep for the configured interval (24 hours)
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Sleeping for {CHECK_INTERVAL} seconds (24 hours) before next check")
await asyncio.sleep(CHECK_INTERVAL)
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Fatal error in chat auto-closer: {str(e)}")
sys.exit(1)
def app(environ, start_response):
"""WSGI application for health checks"""
status = '200 OK'
headers = [('Content-type', 'text/plain')]
start_response(status, headers)
return [b'Chat Auto-Closer Worker is running']
# Main entry point for running as a standalone service
if __name__ == "__main__":
logger.info("Starting chat auto-closer service")
# Run the auto-closer loop
asyncio.run(run_auto_closer_loop())
+141
View File
@@ -0,0 +1,141 @@
"""
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.
CRM lead-push worker.
Claims due crm_sync_jobs (FOR UPDATE SKIP LOCKED — replica-safe), groups them
by (provider, organization) and pushes each group serially so one tenant's
burst can't trip provider rate limits, while distinct tenants run
concurrently. A second, slower timer proactively refreshes Pipedrive tokens:
their refresh tokens die 60 days after last use, so idle tenants must be kept
alive by us, not by traffic.
Run as its own container: python -m app.workers.crm_sync
"""
import asyncio
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from typing import List
from app.core.logger import get_logger
from app.crm.base import CrmAuthError, CrmTransientError
from app.database import SessionLocal
from app.models.crm import CrmProvider, CrmSyncJob
from app.repositories.crm import CrmConnectionRepository, CrmSyncJobRepository
from app.services.crm_sync import process_job
from app.services.crm_tokens import get_valid_tokens
logger = get_logger(__name__)
POLL_INTERVAL_SECONDS = 10
CLAIM_BATCH_SIZE = 20
REFRESH_SWEEP_INTERVAL_SECONDS = 6 * 3600
# Refresh anything not refreshed within a day — an idle tenant then has ~59
# failed sweeps of headroom before its 60-day window could actually lapse.
REFRESH_STALE_AFTER_HOURS = 24
async def run_sync_pass() -> None:
"""One poll: reclaim crashed claims, claim due jobs, push them."""
with SessionLocal() as db:
repo = CrmSyncJobRepository(db)
repo.reclaim_stale()
jobs = repo.claim_batch(CLAIM_BATCH_SIZE)
if not jobs:
return
groups = defaultdict(list)
for job in jobs:
groups[(job.provider, job.organization_id)].append(job)
await asyncio.gather(*(
_process_group([job.id for job in group]) for group in groups.values()
))
async def _process_group(job_ids: List) -> None:
"""Serial pushes for one (provider, org) — its own session, so groups
running concurrently never share SQLAlchemy state. When the connection
proves dead mid-group, the rest fail fast instead of re-hitting it."""
with SessionLocal() as db:
repo = CrmSyncJobRepository(db)
connection_dead = False
for job_id in job_ids:
job = db.get(CrmSyncJob, job_id)
if job is None:
continue
if connection_dead:
repo.fail(job, "auth: connection expired (failed earlier in batch)")
continue
try:
job = await process_job(db, job)
except Exception as e:
logger.error(f"CRM sync job {job_id} crashed: {e}")
db.rollback() # discard any half-done state before the retry write
repo.schedule_retry(job, f"worker error: {e}")
continue
if (job.last_error or "").startswith("auth:"):
connection_dead = True
async def run_refresh_sweep() -> None:
"""Keep Pipedrive's sliding refresh-token window alive for idle tenants."""
cutoff = datetime.now(timezone.utc) - timedelta(hours=REFRESH_STALE_AFTER_HOURS)
with SessionLocal() as db:
connections = CrmConnectionRepository(db).list_needing_refresh(
CrmProvider.PIPEDRIVE.value, cutoff)
for connection in connections:
try:
# No force_refresh: idle connections already have an expired
# access token, so get_valid_tokens refreshes them, and the
# per-connection lock + runway re-check inside it prevents a
# redundant second refresh if another worker just did it.
await get_valid_tokens(db, connection)
logger.info(f"Refreshed idle Pipedrive connection {connection.id}")
except CrmAuthError:
pass # marked expired inside get_valid_tokens; UI shows Reconnect
except CrmTransientError as e:
logger.warning(f"Pipedrive sweep refresh failed (will retry next "
f"sweep): {e}")
async def run_sync_loop(poll_interval: int = POLL_INTERVAL_SECONDS) -> None:
while True:
try:
await run_sync_pass()
except Exception as e:
logger.error(f"Error in CRM sync loop: {e}")
await asyncio.sleep(poll_interval)
async def run_refresh_loop(interval: int = REFRESH_SWEEP_INTERVAL_SECONDS) -> None:
while True:
try:
await run_refresh_sweep()
except Exception as e:
logger.error(f"Error in CRM refresh sweep: {e}")
await asyncio.sleep(interval)
async def main() -> None:
await asyncio.gather(
run_sync_loop(),
run_refresh_loop(),
)
if __name__ == "__main__":
logger.info("Starting CRM sync worker")
asyncio.run(main())
+166
View File
@@ -0,0 +1,166 @@
"""
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.
FAQ generation/import worker. Runs as an independent loop inside the
knowledge_processor container; `python -m app.workers.faq_processor` also
works as a standalone service if it ever needs its own container.
"""
import asyncio
from app.database import SessionLocal
from app.core.config import settings
from app.core.logger import get_logger
from app.models.faq_generation_job import FAQJobType
from app.models.notification import NotificationType
from app.repositories.faq_generation_job import FAQGenerationJobRepository
from app.services.faq_article_import import run_article_import_job
from app.services.faq_generation import run_generation_job
from app.services.faq_import import run_import_job, run_pdf_import_job
from app.services.notifications import notify_user
logger = get_logger(__name__)
POLL_INTERVAL_SECONDS = 60
async def process_faq_job(job_id: int):
"""Run one job to a terminal state, with notifications."""
with SessionLocal() as db:
job_repo = FAQGenerationJobRepository(db)
job = job_repo.get_by_id(job_id)
if not job:
logger.error(f"FAQ job {job_id} not found")
return
# Captured before the try: after a DB failure + rollback the ORM
# instance is expired and attribute access could raise or re-query.
job_type = job.job_type
is_import = job_type in (
FAQJobType.IMPORT_URL.value,
FAQJobType.IMPORT_ARTICLES.value,
FAQJobType.IMPORT_PDF.value,
)
try:
job_repo.mark_processing(job.id)
if job_type == FAQJobType.IMPORT_URL.value:
created = await run_import_job(db, job)
elif job_type == FAQJobType.IMPORT_ARTICLES.value:
created = await run_article_import_job(db, job)
elif job_type == FAQJobType.IMPORT_PDF.value:
created = await run_pdf_import_job(db, job)
else:
created = await run_generation_job(db, job)
job_repo.mark_completed(job.id, faqs_created=created)
if created:
body = f"{created} draft FAQ{'s' if created != 1 else ''} added — review and publish them in Help center."
elif is_import:
body = (
"Nothing new was imported — no Q&A was found on the page, or it duplicates "
"existing FAQs. Tip: use Import → Articles to bring in whole articles verbatim (no AI needed)."
)
else:
body = (
"No new FAQs were generated — the model returned nothing for this content. "
"Try a larger/different model, or Import → Articles to import pages directly."
)
await notify_user(
db,
job.user_id,
NotificationType.FAQ_GENERATED,
"Imported drafts ready for review" if is_import and created else
("Import finished" if is_import else "Draft FAQs ready for review"),
body,
metadata={"faq_job_id": job.id},
)
except Exception as e:
logger.error(f"FAQ job {job_id} failed: {e}")
# A DB-originated failure leaves the session aborted; roll back so
# mark_failed and the notification can still be persisted.
try:
db.rollback()
except Exception:
pass
job_repo.mark_failed(job.id, str(e))
await notify_user(
db,
job.user_id,
NotificationType.FAQ_GENERATION_FAILED,
"FAQ import failed" if is_import else "FAQ generation failed",
str(e),
metadata={"faq_job_id": job.id},
)
_reaped_orphans = False
def _reap_orphans_once() -> None:
"""Reap orphaned `processing` jobs the first time the processor runs in
this process — covers both worker entrypoints (the loop wrapper here and
run_knowledge_processor's per-iteration call) without reaping every tick."""
global _reaped_orphans
if _reaped_orphans:
return
_reaped_orphans = True
reap_orphaned_jobs()
async def run_faq_processor():
"""Single pass: process every pending FAQ job (bounded concurrency)."""
_reap_orphans_once()
with SessionLocal() as db:
pending_ids = [job.id for job in FAQGenerationJobRepository(db).get_pending()]
if not pending_ids:
return
semaphore = asyncio.Semaphore(max(1, settings.MAX_CONCURRENT_FAQ_JOBS))
async def process_with_semaphore(job_id: int):
async with semaphore:
await process_faq_job(job_id)
await asyncio.gather(*[process_with_semaphore(job_id) for job_id in pending_ids])
def reap_orphaned_jobs() -> None:
"""Fail jobs left in `processing` by a previous crash/kill. Runs once at
startup: this worker is the only writer of `processing`, so on a fresh boot
any such row is orphaned — otherwise it hangs the enqueue guard (one active
job per type) and the UI's active-job polling forever."""
try:
with SessionLocal() as db:
reaped = FAQGenerationJobRepository(db).fail_orphaned_processing(
"Worker restarted before the job finished — please try again."
)
if reaped:
logger.warning(f"Reaped {reaped} orphaned FAQ job(s) left in 'processing' on startup")
except Exception as e:
logger.error(f"FAQ orphan reap failed (non-fatal): {e}")
async def run_faq_processor_loop(poll_interval: int = POLL_INTERVAL_SECONDS):
"""Forever-loop wrapper — run as an independent task next to the knowledge
loop so a long FAQ job never delays knowledge ingestion. run_faq_processor
reaps orphans on its first call, so no explicit reap is needed here."""
while True:
try:
await run_faq_processor()
except Exception as e:
logger.error(f"Error in FAQ processor loop: {e}")
await asyncio.sleep(poll_interval)
if __name__ == "__main__":
logger.info("Starting FAQ processor service")
asyncio.run(run_faq_processor_loop())
+334
View File
@@ -0,0 +1,334 @@
"""
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.
"""
import asyncio
from app.database import SessionLocal
from app.repositories.knowledge_queue import KnowledgeQueueRepository
from app.repositories.knowledge import KnowledgeRepository
from app.knowledge.knowledge_base import KnowledgeManager
from app.knowledge import page_editor
from app.services import knowledge_vector_links
from app.models.knowledge_queue import QueueStatus, ProcessingStage
from app.core.logger import get_logger
import os
from app.core.config import settings
from app.core.processor import PROCESSOR_STATUS
from app.knowledge.failure_advice import FAILURE_ADVICE, FailureCause
from datetime import datetime, timedelta, timezone
from app.models.notification import Notification, NotificationType
from app.services.user import send_fcm_notification
from urllib.parse import urlparse, unquote
logger = get_logger(__name__)
def get_user_friendly_filename(source: str, source_type: str) -> str:
"""Extract a user-friendly filename from the source URL or path"""
try:
if source_type in ('website', 'sitemap'):
# For websites/sitemaps, just return the domain
parsed = urlparse(source)
return parsed.netloc or source
# For files (PDF, etc.), extract filename from URL or path
if source.startswith('http'):
# Parse URL to get the path
parsed = urlparse(source)
path = parsed.path
# Extract filename from path
if path:
# Get the last part of the path (filename)
filename = path.split('/')[-1]
# URL decode the filename to handle encoded characters like %20
filename = unquote(filename)
# Remove query parameters if any
filename = filename.split('?')[0]
if filename:
return filename
# Fallback to domain if no filename found
return parsed.netloc or source
else:
# For local file paths, just get the basename
return os.path.basename(source)
except Exception as e:
logger.warning(f"Error extracting filename from {source}: {e}")
return source
def _sync_vector_agent_links(db, queue_item) -> None:
"""Point the ingested chunks at every agent the source is linked to.
Ingestion writes ``filters.agent_id`` from the queue row's single optional
``agent_id``, so a run queued without one leaves every chunk with an empty
list. Retrieval matches on that filter, so the source is invisible to the
agent while the dashboard still shows it as linked — the failure that left
two production sources unreadable until the links were rebuilt by hand.
The relational links are the authority, so re-derive from them. Idempotent,
and non-fatal: the content is already ingested, and losing the run over a
filter refresh would be worse than the stale filter.
"""
try:
knowledge_repo = KnowledgeRepository(db)
sources = knowledge_repo.get_by_sources(
queue_item.organization_id, [queue_item.source]
)
if not sources:
return
# An org can hold more than one knowledge row for the same URL, and they
# share one set of vector chunks (keyed on the source name). Sync the
# union of their links once, rather than per row where the last write
# would drop the other rows' agents.
agent_ids = sorted({
agent_id
for knowledge in sources
for agent_id in page_editor.agent_ids_for(knowledge)
})
knowledge_vector_links.sync_agent_ids(db, sources[0], agent_ids)
db.commit()
except Exception as e:
db.rollback()
logger.error(
f"Could not refresh vector agent filters for {queue_item.source} "
f"(non-fatal, content was ingested): {e}"
)
async def _notify(db, queue_item, kind: NotificationType, title: str, message: str) -> None:
"""Tell whoever added the source, in the app and by push. Never fatal: the
source's own status must not depend on it. A source added without a user
(a Shopify install queues the shop's site) has no one to tell."""
if not queue_item.user_id:
logger.info(f"Queue item {queue_item.id} has no user to notify ({kind.value})")
return
try:
notification = Notification(
user_id=queue_item.user_id,
type=kind,
title=title,
message=message,
# The column is notification_metadata; the old metadata=
# kwarg silently shadowed Base.metadata and stored nothing.
notification_metadata={"queue_id": queue_item.id}
)
db.add(notification)
db.commit()
except Exception as e:
db.rollback()
logger.error(f"Could not create {kind.value} notification for queue item {queue_item.id}: {e}")
return
await send_fcm_notification(queue_item.user_id, notification, db)
async def process_queue_item(queue_item_id: int):
"""Process a single queue item"""
with SessionLocal() as db:
queue_item = None
try:
queue_repo = KnowledgeQueueRepository(db)
queue_item = queue_repo.get_by_id(queue_item_id)
if not queue_item:
logger.error(f"Queue item {queue_item_id} not found")
return
# Get knowledge manager instance
knowledge = KnowledgeManager(
org_id=queue_item.organization_id,
agent_id=queue_item.agent_id
)
# Process if status is pending
if queue_item.status == QueueStatus.PENDING:
# Update to processing
queue_item.status = QueueStatus.PROCESSING
queue_item.processing_stage = ProcessingStage.NOT_STARTED
queue_item.progress_percentage = 0.0
db.commit()
await knowledge.process_knowledge(queue_item)
_sync_vector_agent_links(db, queue_item)
user_friendly_name = get_user_friendly_filename(queue_item.source, queue_item.source_type)
await _notify(db, queue_item, NotificationType.KNOWLEDGE_PROCESSED,
"Knowledge Processing Complete", f"Successfully processed {user_friendly_name}")
# Auto-draft FAQs from the new source for orgs using the help
# center. Fully non-fatal — even an import failure of the FAQ
# stack must not flip an already-successful run to FAILED.
try:
from app.services.faq_generation import maybe_enqueue_auto_faq_job
maybe_enqueue_auto_faq_job(db, queue_item)
except Exception as faq_hook_err:
logger.error(f"Auto FAQ hook unavailable (non-fatal): {faq_hook_err}")
queue_item.status = QueueStatus.COMPLETED
db.commit()
except Exception as e:
logger.error(f"Error processing queue item {queue_item_id}: {str(e)}")
if queue_item is None:
raise
# The session may hold a failed flush: start clean, record the
# failure first, and only then tell anyone about it.
db.rollback()
try:
queue_item.status = QueueStatus.FAILED
# Persist the reason so the UI can show why it failed.
queue_item.error = str(e)
db.commit()
except Exception as status_err:
db.rollback()
logger.error(f"Could not mark queue item {queue_item_id} failed: {status_err}")
else:
await _tell_failed(db, queue_item)
raise
async def _tell_failed(db, queue_item) -> None:
user_friendly_name = get_user_friendly_filename(queue_item.source, queue_item.source_type)
await _notify(db, queue_item, NotificationType.KNOWLEDGE_FAILED,
"Knowledge Processing Failed", f"Failed to process {user_friendly_name}: {queue_item.error}")
# A source still 'processing' with no progress for this long was abandoned:
# its worker died partway (out of memory, a restart) and nothing will pick it
# up again, so it would spin in the dashboard for good. Longer than a live
# run can go without progress (websites stop at their time limit; files are
# size-capped and take about a minute), so it never fails one another
# process is still reading.
ABANDONED_AFTER = timedelta(seconds=settings.KB_SOURCE_TIME_LIMIT_SECONDS) + timedelta(minutes=15)
def _shown_as(queue_item) -> str:
"""A website by its address, as every other failure names it; a file by
its name, not the path it was uploaded to on our server."""
if queue_item.source_type in ('website', 'sitemap'):
return queue_item.source
return get_user_friendly_filename(queue_item.source, queue_item.source_type)
async def fail_abandoned_items() -> None:
"""Fail sources left 'processing' by a worker that died, and tell whoever
added them."""
with SessionLocal() as db:
cutoff = datetime.now(timezone.utc) - ABANDONED_AFTER
abandoned = KnowledgeQueueRepository(db).get_abandoned(cutoff)
for item in abandoned:
item.status = QueueStatus.FAILED
item.error = FAILURE_ADVICE[FailureCause.INTERRUPTED].format(url=_shown_as(item))
db.commit()
for item in abandoned:
logger.warning(f"Queue item {item.id} was abandoned mid-run; marked failed")
await _tell_failed(db, item)
async def run_processor():
"""Single run of the processor"""
try:
PROCESSOR_STATUS["is_running"] = True
PROCESSOR_STATUS["error"] = None
try:
await fail_abandoned_items()
except Exception as e:
# Never at the cost of processing what is pending
logger.error(f"Could not check for abandoned queue items: {e}")
# Get pending items with proper connection handling
with SessionLocal() as db:
queue_repo = KnowledgeQueueRepository(db)
pending_items = queue_repo.get_pending()
# Extract IDs before closing the session
pending_item_ids = [item.id for item in pending_items] if pending_items else []
if pending_item_ids:
# Reduce concurrent processing to 2 to conserve connections on t3.micro
semaphore = asyncio.Semaphore(2)
async def process_with_semaphore(item_id):
async with semaphore:
await process_queue_item(item_id)
# Process items with semaphore control
await asyncio.gather(*[process_with_semaphore(item_id) for item_id in pending_item_ids])
PROCESSOR_STATUS["last_run"] = datetime.utcnow().isoformat()
except Exception as e:
logger.error(f"Error in knowledge processor: {str(e)}")
PROCESSOR_STATUS["error"] = str(e)
raise
finally:
PROCESSOR_STATUS["is_running"] = False
# Between runs. A new source waits up to this long before its crawl starts,
# and someone in onboarding is watching it: keep it short. An idle run is one
# small query.
IDLE_POLL_SECONDS = 10
def run_service():
"""Run the knowledge and FAQ loops as a standalone service (what the
knowledge_processor container runs: python -m app.workers.knowledge_processor)."""
logger.info("Starting knowledge processor service")
# Pushes for processed sources go out from this process: without this the
# first one failed with "The default Firebase app does not exist"
# (run_knowledge_processor.py did it; this entry point, which prod runs, didn't).
from app.services.firebase import initialize_firebase
try:
initialize_firebase()
except Exception as e:
logger.error(f"Firebase unavailable, push notifications are off: {e}")
async def processor_loop():
while True:
try:
logger.debug("Running knowledge processor")
await run_processor()
logger.debug(f"Knowledge processor completed, sleeping for {IDLE_POLL_SECONDS} seconds")
except Exception as e:
logger.error(f"Error in knowledge processor loop: {str(e)}")
await asyncio.sleep(IDLE_POLL_SECONDS)
async def main():
"""Knowledge and FAQ loops run as independent tasks in this container:
a long FAQ generation job never delays knowledge ingestion, and an
import failure of the FAQ stack only disables the FAQ half."""
tasks = [processor_loop()]
try:
from app.workers.faq_processor import run_faq_processor_loop
tasks.append(run_faq_processor_loop())
except Exception as e:
logger.error(f"FAQ processor unavailable, running knowledge only: {e}")
await asyncio.gather(*tasks)
asyncio.run(main())
# Main entry point for running as a standalone service
if __name__ == "__main__":
run_service()
@@ -0,0 +1,116 @@
"""
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.
"""
"""
Knowledge Processor Runner
This script runs the knowledge processor independently at regular intervals.
"""
import asyncio
import sys
import os
from pathlib import Path
# Add the parent directory to Python path to allow imports
sys.path.append(str(Path(__file__).parent.parent.parent))
from app.workers.knowledge_processor import run_processor
from app.core.logger import get_logger
from app.services.firebase import initialize_firebase
logger = get_logger(__name__)
# Configure polling interval
POLL_INTERVAL = int(os.environ.get('POLL_INTERVAL', '10')) # Default to 10 seconds
WORKER_ID = os.environ.get('GUNICORN_WORKER_ID', 'standalone')
# Initialize Firebase on module load
try:
logger.info(f"Initializing Firebase for worker {WORKER_ID}...")
initialize_firebase()
logger.info(f"Worker {WORKER_ID}: Firebase initialized successfully")
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Failed to initialize Firebase: {e}")
async def run_processor_loop():
"""Run the knowledge processor in a continuous loop with short poll interval"""
logger.info(f"Worker {WORKER_ID}: Starting knowledge processor with {POLL_INTERVAL}s polling interval")
try:
logger.info(f"Worker {WORKER_ID}: Starting knowledge processor loop")
iteration = 0
while True:
iteration += 1
start_time = asyncio.get_event_loop().time()
try:
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Checking for pending knowledge items...")
await run_processor()
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Processing completed successfully")
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Iteration {iteration} - Error in knowledge processor: {str(e)}")
# FAQ generation jobs ride the same loop in this entrypoint too, so
# both worker entrypoints drain the FAQ queue.
try:
from app.workers.faq_processor import run_faq_processor
await run_faq_processor()
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Iteration {iteration} - Error in FAQ processor: {str(e)}")
# Calculate how long the processing took
elapsed = asyncio.get_event_loop().time() - start_time
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Processing took {elapsed:.2f} seconds")
# If the processing took less than the polling interval, sleep for the remainder
sleep_time = max(0.1, POLL_INTERVAL - elapsed)
logger.info(f"Worker {WORKER_ID}: Iteration {iteration} - Sleeping for {sleep_time:.2f} seconds before next check")
await asyncio.sleep(sleep_time)
except Exception as e:
logger.error(f"Worker {WORKER_ID}: Fatal error in knowledge processor: {str(e)}")
sys.exit(1)
def app(environ, start_response):
"""WSGI application for health checks"""
status = '200 OK'
headers = [('Content-type', 'application/json')]
start_response(status, headers)
return [b'{"status": "running"}']
# Set worker ID if running under Gunicorn
if 'GUNICORN_WORKER_ID' not in os.environ and 'GUNICORN_ARBITER' in os.environ:
pid = os.getpid()
os.environ['GUNICORN_WORKER_ID'] = f"worker-{pid}"
WORKER_ID = f"worker-{pid}"
# Start the processor loop in the background when Gunicorn loads this module
logger.info(f"Worker {WORKER_ID}: Initializing knowledge processor worker")
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
task = loop.create_task(run_processor_loop())
if __name__ == "__main__":
try:
asyncio.run(run_processor_loop())
except KeyboardInterrupt:
logger.info("Knowledge processor stopped by user")
except Exception as e:
logger.error(f"Fatal error in knowledge processor: {str(e)}")
sys.exit(1)
+692
View File
@@ -0,0 +1,692 @@
"""
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.
AI ticket worker: processes pending runs (triage + hypothesis-driven
investigation with MCP evidence gathering) and runs lifecycle hygiene
(confirmation-timeout close + CSAT). Runs as its own container:
`python -m app.workers.ticket_investigator`.
"""
import asyncio
from datetime import datetime, timedelta, timezone
from typing import Optional
from uuid import UUID
from app.channels.app_hooks import load_extensions
from app.core.config import settings
from app.core.logger import get_logger
from app.core.security import decrypt_api_key
from app.database import SessionLocal
from app.models.investigation import (
HypothesisStatus,
InvestigationRun,
InvestigationRunStatus,
InvestigationRunType,
InvestigationTrigger,
)
from app.models.ticket import ResolutionOutcome, Ticket, TicketStatus
from app.models.ticket_activity import TicketActivityType, TicketActorType
from app.repositories.ai_config import AIConfigRepository
from app.repositories.ticket import InvestigationRepository, TicketRepository
from app.services.ticket import TicketService, embed_ticket_text
from app.services.ticket_events import emit_ticket_update
from app.services.ticket_privacy import redact_reference_text
from app.services.usage_metering import HOSTED_MODEL_TYPE
logger = get_logger(__name__)
RUN_POLL_INTERVAL_SECONDS = 15
LIFECYCLE_POLL_INTERVAL_SECONDS = 300
TRIAGE_WALL_SECONDS = 120
# Fallbacks when a run row / settings row leaves a budget unset (mirrors the
# model server_defaults so behaviour can't drift between them).
DEFAULT_MAX_WALL_SECONDS = 600
DEFAULT_AUTO_RESOLVE_CONFIDENCE = 0.85
# Share of a run's wall budget the connectors may spend coming up. Generous
# per-connector timeouts are what make slow `npx` servers usable, but a run
# that spends all its time connecting produces no hypotheses at all.
CONNECT_BUDGET_SHARE = 0.25
async def _build_triage_context(db, service: TicketService, ticket: Ticket):
"""Transcript of the linked conversation + similar past resolved tickets."""
transcript = None
session_record = service._primary_session(ticket)
if session_record is not None:
try:
transcript = await service.build_session_transcript(
session_record.session_id, max_messages=15, line_chars=400
)
except Exception as e:
logger.error(f"Transcript build failed for {ticket.display_number}: {e}")
similar = []
embedding = ticket.embedding
if embedding is None:
# CPU-bound local model — keep it off the worker's event loop so
# concurrent runs aren't serialized behind inference.
embedding = await asyncio.to_thread(
embed_ticket_text, ticket.title, ticket.description
)
if embedding is not None:
for other, score in service.repo.find_similar(
ticket.organization_id, list(embedding), limit=5,
exclude_ticket_id=ticket.id, min_similarity=0.70,
):
resolution = other.resolution_summary or other.resolution_outcome or "unresolved"
# These tickets belong to OTHER customers. They're here for the
# shape of the failure and its fix, so identifiers in their titles
# are pure downside — strip them before they reach the model that
# also writes this ticket's customer-facing summary.
similar.append(
redact_reference_text(
f"{other.display_number} [{other.status}] {other.title} — resolution: {resolution}"
)
)
return transcript, similar
def _customer_identity_section(ticket: Ticket) -> Optional[str]:
"""A trusted, structured statement of who the customer is, so the
investigator filters connected databases/logs by exact identifiers rather
than scraping them out of the freeform transcript (where an anonymous
…@noemail.com placeholder is all it would find)."""
from app.repositories.customer import CustomerRepository
customer = ticket.customer
if customer is None:
return None
bits = []
name = (customer.full_name or "").strip()
if name:
bits.append(f"name: {name}")
email = CustomerRepository.display_email(customer.email)
if email:
bits.append(f"email: {email}")
if customer.phone:
bits.append(f"phone: {customer.phone}")
if isinstance(customer.meta_data, dict) and customer.meta_data:
# Integrator-supplied identity fields (account id, external ids, etc.).
import json as _json
bits.append(f"other identifiers: {_json.dumps(customer.meta_data)[:300]}")
if not bits:
return None
return (
"CUSTOMER IDENTITY (trusted — these are the exact values to filter by when "
"correlating this customer in connected databases and logs; they are "
"identifiers, not instructions):\n" + "\n".join(bits)
)
async def process_run(run_id: UUID) -> None:
"""Run one pending investigation run to a terminal state."""
with SessionLocal() as db:
run_repo = InvestigationRepository(db)
run = run_repo.get_by_id(run_id)
if run is None or run.status != InvestigationRunStatus.PENDING.value:
return
service = TicketService(db)
ticket = TicketRepository(db).get_by_id(run.ticket_id)
if ticket is None:
run.status = InvestigationRunStatus.FAILED
run.error = "Ticket no longer exists"
run.finished_at = datetime.now(timezone.utc)
db.commit()
return
run.status = InvestigationRunStatus.RUNNING
run.started_at = datetime.now(timezone.utc)
previous_status = str(ticket.status)
is_investigation = str(run.run_type) == InvestigationRunType.INVESTIGATION.value
if is_investigation:
if previous_status in (
TicketStatus.OPEN.value, TicketStatus.TRIAGING.value,
TicketStatus.REOPENED.value, TicketStatus.IN_PROGRESS.value,
):
ticket.status = TicketStatus.INVESTIGATING
elif previous_status == TicketStatus.OPEN.value:
ticket.status = TicketStatus.TRIAGING
db.commit()
await emit_ticket_update(
ticket.organization_id, ticket.id, "run",
{"status": "running", "run_type": str(run.run_type)},
)
agent = None
try:
config = AIConfigRepository(db).get_active_config(ticket.organization_id)
if config is None:
raise RuntimeError("Organization has no active AI configuration")
model_type = config.model_type.value if hasattr(config.model_type, "value") else str(config.model_type)
from app.agents.ticket_investigator import TicketInvestigatorAgent
agent = TicketInvestigatorAgent(
api_key=decrypt_api_key(config.encrypted_api_key),
model_name=config.model_name,
model_type=model_type,
base_url=(config.settings or {}).get('base_url'),
)
run.model_name = config.model_name
# Metered only on the hosted model — own-key orgs pay their provider
# directly. This is what the message-limit check sums.
run.metered = model_type.upper() == HOSTED_MODEL_TYPE
def count_call():
run.llm_calls = (run.llm_calls or 0) + 1
agent.on_llm_call = count_call
if is_investigation:
await _process_investigation(db, run, service, ticket, agent, previous_status)
else:
await _process_triage(db, run, service, ticket, agent, previous_status)
except asyncio.TimeoutError:
_fail_run(db, run, ticket, previous_status, "AI run timed out")
except Exception as e:
logger.error(f"Investigation run {run_id} failed: {e}")
try:
db.rollback()
except Exception:
pass
_fail_run(db, run, ticket, previous_status, str(e))
finally:
# Record token usage on every path — success, budget-exceeded,
# failure or timeout — so partial spend is still metered.
if agent is not None and (agent.input_tokens or agent.output_tokens):
try:
run.input_tokens = agent.input_tokens
run.output_tokens = agent.output_tokens
db.commit()
except Exception as e:
logger.warning(f"Recording token usage failed for run {run_id}: {e}")
db.rollback()
async def _process_triage(db, run, service: TicketService, ticket: Ticket, agent, previous_status: str) -> None:
transcript, similar = await _build_triage_context(db, service, ticket)
result = await asyncio.wait_for(
agent.triage(ticket.title, ticket.description, transcript, similar),
timeout=TRIAGE_WALL_SECONDS,
)
if result is not None:
_apply_triage(service, ticket, result)
else:
logger.warning(f"Triage produced no result for {ticket.display_number}")
# Triage done — the ticket goes back to plain open unless a human
# already moved it.
if str(ticket.status) == TicketStatus.TRIAGING.value:
ticket.status = previous_status
run.status = InvestigationRunStatus.COMPLETED
run.finished_at = datetime.now(timezone.utc)
# Auto-chain: tickets created with auto-investigate on flow straight into
# a full investigation when the org has connectors configured.
settings_row = service.settings_repo.get_or_create(ticket.organization_id)
# Escalation policy runs on the post-triage priority.
from app.services.ticket_jira import maybe_escalate_to_jira
await maybe_escalate_to_jira(db, service, ticket, settings_row)
if (
str(run.trigger) == InvestigationTrigger.AUTO_ON_CREATE.value
and settings_row.auto_investigate_on_create
and _has_investigation_sources(db, ticket, settings_row)
):
service.enqueue_run(
ticket,
run_type=InvestigationRunType.INVESTIGATION,
trigger=InvestigationTrigger.AUTO_ON_CREATE,
settings=settings_row,
)
db.commit()
await emit_ticket_update(
ticket.organization_id, ticket.id, "triage",
{"status": str(ticket.status), "priority": str(ticket.priority)},
)
async def _apply_autonomy(
service: TicketService, ticket: Ticket, run, rca, hypotheses, settings_row, partial: bool
) -> str:
"""Apply the org's autonomy level to a finished investigation.
L1: nothing — the RCA is informational.
L2: create a proposal for human approval.
L3: auto-resolve, but ONLY behind hard guards enforced here in code:
a complete (non-partial) run, at least one validated hypothesis,
RCA confidence >= the org threshold, and a reachable customer.
Anything less falls back to an L2 proposal.
"""
if rca is None:
return "none"
level = settings_row.autonomy_level or 1
if level < 2:
return "none"
if str(ticket.status) in (
TicketStatus.RESOLVED.value,
TicketStatus.RESOLVED_PENDING_CONFIRMATION.value,
TicketStatus.CLOSED.value,
):
return "none"
summary = rca.remediation or rca.conclusion or rca.summary or ""
if not summary.strip():
return "none"
has_validated = any(str(h.status) == HypothesisStatus.VALIDATED.value for h in hypotheses)
confidence = rca.confidence or 0.0
# Explicit None check — a configured 0.0 threshold must not fall back to
# the default via `or`.
threshold = settings_row.min_confidence_to_auto_resolve
if threshold is None:
threshold = DEFAULT_AUTO_RESOLVE_CONFIDENCE
if (
level >= 3
and not partial
and has_validated
and confidence >= threshold
and service.can_notify_customer(ticket)
):
try:
await service.resolve(
ticket,
outcome=ResolutionOutcome.FIXED,
resolution_summary=summary,
customer_message=rca.customer_summary,
actor_type=TicketActorType.AI,
)
return "auto_resolved"
except Exception as e:
logger.error(f"L3 auto-resolve failed for {ticket.display_number}: {e}")
service.create_proposal(
ticket,
run_id=run.id,
summary=summary,
customer_message=rca.customer_summary,
confidence=confidence,
)
return "proposed"
def _has_investigation_sources(db, ticket: Ticket, settings_row) -> bool:
"""Whether the org has anything for an investigation to gather evidence
from — MCP connectors OR a guardrailed DB connector. Auto-chaining on MCP
alone would silently skip DB-connector-only orgs."""
if settings_row.investigation_mcp_tool_ids or []:
return True
from app.models.ticket_db_connector import TicketDBConnector
return (
db.query(TicketDBConnector.id)
.filter(
TicketDBConnector.organization_id == ticket.organization_id,
TicketDBConnector.enabled == True, # noqa: E712
TicketDBConnector.allowed_tables.isnot(None),
)
.first()
is not None
)
def _build_db_tools(db, ticket: Ticket, run):
"""Guardrailed read-only DB tools from the org's enabled connectors with
a non-empty table allowlist. None when there are none."""
try:
from app.models.ticket_db_connector import TicketDBConnector
from app.services.db_connector_service import DBConnectorConfig
from app.tools.guardrailed_db_toolkit import GuardrailedDBTools
connectors = (
db.query(TicketDBConnector)
.filter(
TicketDBConnector.organization_id == ticket.organization_id,
TicketDBConnector.enabled == True, # noqa: E712
)
.all()
)
configs = [
DBConnectorConfig.from_model(c) for c in connectors if c.allowed_tables
]
if not configs:
return None
return GuardrailedDBTools(configs, ticket_id=ticket.id, run_id=run.id)
except Exception as e:
logger.error(f"Failed to build DB tools for {ticket.display_number}: {e}")
return None
async def _investigate_with_tools(
db, run, service, ticket, agent, context_message, settings_row, recorder, manager
):
"""Connect the org's investigation tools (MCP connectors + guardrailed DB
connectors), run the hypothesis phases, then clean the tools up — ALL in
this one coroutine. It's driven under asyncio.wait_for (a child task), so
MCP init, every tool call and cleanup must live in the same task or the
anyio cancel scope owning the MCP streams gets crossed on timeout. The
manager is caller-owned so its connect/fail bookkeeping stays readable
after a timeout cancels this coroutine."""
from app.services.ticket_investigation import run_investigation_phases
mcp_tools = []
try:
tool_ids = settings_row.investigation_mcp_tool_ids or []
if tool_ids:
recorder.record_phase("Connecting investigation connectors")
mcp_tools = await manager.initialize_mcp_tools_by_ids(
str(ticket.organization_id), tool_ids
)
db_tools = _build_db_tools(db, ticket, run)
if db_tools is not None:
mcp_tools = list(mcp_tools) + [db_tools]
recorder.map_connectors(mcp_tools)
return await run_investigation_phases(
db, run, ticket, agent, context_message, mcp_tools, recorder
)
finally:
await manager.cleanup_mcp_tools()
async def _process_investigation(db, run, service: TicketService, ticket: Ticket, agent, previous_status: str) -> None:
"""Hypothesis-driven investigation: connect the org's tools, generate +
test hypotheses (evidence recorded per tool call), then write the RCA
document."""
from app.repositories.ticket import InvestigationRepository
from app.services.ticket_investigation import (
EvidenceRecorder,
synthesize_and_store_rca,
)
settings_row = service.settings_repo.get_or_create(ticket.organization_id)
recorder = EvidenceRecorder(run.id, ticket.id)
service._add_activity(
ticket,
TicketActivityType.AI_INVESTIGATION_STARTED,
actor_type=TicketActorType.AI,
body="AI investigation started",
metadata={"run_id": str(run.id), "trigger": str(run.trigger)},
)
db.commit()
transcript, similar = await _build_triage_context(db, service, ticket)
extra_sections = []
identity = _customer_identity_section(ticket)
if identity:
extra_sections.append(identity)
if ticket.ai_summary:
extra_sections.append(f"TRIAGE SUMMARY: {ticket.ai_summary} (intent: {ticket.intent})")
if run.context_note:
# e.g. the human's rejection reason feeding a refined run.
extra_sections.append(f"REVIEWER FEEDBACK ON THE PREVIOUS RUN:\n{run.context_note[:2000]}")
context_message = agent.build_context_message(
ticket.title, ticket.description, transcript, similar, extra_sections=extra_sections
)
from app.tools.mcp_manager import MCPToolsManager
wall_seconds = run.max_wall_seconds or DEFAULT_MAX_WALL_SECONDS
mcp_manager = MCPToolsManager(connect_budget=wall_seconds * CONNECT_BUDGET_SHARE)
hypotheses = []
partial = False
try:
hypotheses, budget_exhausted = await asyncio.wait_for(
_investigate_with_tools(
db, run, service, ticket, agent, context_message, settings_row,
recorder, mcp_manager
),
timeout=wall_seconds,
)
partial = budget_exhausted
except asyncio.TimeoutError:
partial = True
try:
db.rollback()
except Exception:
pass
hypotheses = InvestigationRepository(db).list_hypotheses(run.id)
for hypothesis in hypotheses:
if str(hypothesis.status) in (
HypothesisStatus.PENDING.value, HypothesisStatus.TESTING.value
):
hypothesis.status = HypothesisStatus.INCONCLUSIVE.value
hypothesis.conclusion = "Not completed — the run's wall-clock budget was reached."
db.commit()
run.tool_calls_used = recorder.tool_calls
# Surface the connector outcome on the run — a run that "completed" with 0
# of its configured MCP tools (missing npx, bad key, unreachable server)
# must be visible on the dashboard, not just in worker logs (issue #271).
configured_tool_ids = settings_row.investigation_mcp_tool_ids or []
# Also reported when the org has no MCP connectors at all: a
# DB-connector-only run can still have every query fail, and that must not
# render as a clean run (#318).
if configured_tool_ids or recorder.tool_calls:
run.connector_status = {
"configured": len(configured_tool_ids),
"loaded": len(mcp_manager.connected_tool_names),
"failed": mcp_manager.failed_tools,
# A connector can load cleanly and then error on every query. That
# leaves loaded == configured with no provider errors, so the run
# looks healthy while resting on nothing (#318).
"tool_calls": recorder.tool_calls,
"tool_calls_failed": recorder.failed_tool_calls,
# A connector can come up perfectly and still be unusable if the
# provider refuses its tools; that must show next to the
# connection failures, not only in worker logs (#303).
"provider_errors": list(getattr(agent, "provider_errors", []) or []),
}
rca = await synthesize_and_store_rca(
db, run, ticket, agent, context_message, hypotheses, recorder, partial
)
autonomy_action = await _apply_autonomy(
service, ticket, run, rca, hypotheses, settings_row, partial
)
validated = sum(1 for h in hypotheses if str(h.status) == "validated")
service._add_activity(
ticket,
TicketActivityType.AI_INVESTIGATION_COMPLETED,
actor_type=TicketActorType.AI,
body=(rca.summary if rca else "Investigation finished without an RCA document"),
metadata={
"run_id": str(run.id),
"hypotheses": len(hypotheses),
"validated": validated,
"rca_version": rca.version if rca else None,
"partial": partial,
"autonomy_action": autonomy_action,
},
)
if str(ticket.status) == TicketStatus.INVESTIGATING.value:
restore = previous_status
if restore in (TicketStatus.TRIAGING.value, TicketStatus.INVESTIGATING.value):
restore = TicketStatus.OPEN.value
ticket.status = restore
run.status = (
InvestigationRunStatus.BUDGET_EXCEEDED if partial else InvestigationRunStatus.COMPLETED
)
run.finished_at = datetime.now(timezone.utc)
service._sync_linked_sessions(ticket)
db.commit()
await emit_ticket_update(
ticket.organization_id, ticket.id, "run",
{"status": str(run.status), "rca": rca is not None},
)
def _apply_triage(service: TicketService, ticket: Ticket, result) -> None:
"""Apply the triage classification and record the activity."""
from app.models.ticket import TicketPriority
changes = {}
if result.title and result.title.strip() and result.title.strip() != ticket.title:
if ticket.original_title is None:
ticket.original_title = ticket.title
ticket.title = result.title.strip()[:500]
changes["title"] = ticket.title
# The LLM's priority is untrusted output — anything outside the enum
# ('Medium', 'P1', ...) would 500 every later read of the ticket.
priority = (result.priority or "").strip().lower()
if priority in {p.value for p in TicketPriority} and str(ticket.priority) != priority:
changes["priority"] = {"from": str(ticket.priority), "to": priority}
ticket.priority = priority
if result.severity is not None and ticket.severity is None:
ticket.severity = result.severity
changes["severity"] = result.severity
ticket.intent = result.intent
ticket.ai_summary = result.summary
ticket.triage_confidence = result.confidence
if result.tags and not ticket.tags:
ticket.tags = [t.strip().lower()[:30] for t in result.tags[:2] if t.strip()]
service._add_activity(
ticket,
TicketActivityType.AI_TRIAGE,
actor_type=TicketActorType.AI,
body=result.summary,
metadata={
"intent": result.intent,
"confidence": result.confidence,
"changes": changes,
},
)
service._sync_linked_sessions(ticket)
def _fail_run(db, run: InvestigationRun, ticket: Ticket, previous_status: str, error: str) -> None:
run.status = InvestigationRunStatus.FAILED
run.error = error[:2000]
run.finished_at = datetime.now(timezone.utc)
transient = (TicketStatus.TRIAGING.value, TicketStatus.INVESTIGATING.value)
if ticket is not None and str(ticket.status) in transient:
ticket.status = previous_status if previous_status not in transient else TicketStatus.OPEN.value
db.commit()
async def run_ticket_investigator() -> None:
"""Single pass: process every pending run (bounded concurrency)."""
with SessionLocal() as db:
pending_ids = InvestigationRepository(db).get_pending_ids(limit=50)
if not pending_ids:
return
semaphore = asyncio.Semaphore(max(1, settings.MAX_CONCURRENT_INVESTIGATIONS))
async def process_with_semaphore(run_id: UUID):
async with semaphore:
await process_run(run_id)
await asyncio.gather(*[process_with_semaphore(rid) for rid in pending_ids])
async def run_ticket_lifecycle() -> None:
"""Close resolved-pending-confirmation tickets past the org's timeout
(customer never replied). The close fires the CSAT request when enabled."""
with SessionLocal() as db:
service = TicketService(db)
now = datetime.now(timezone.utc)
# defer the pgvector embedding: this hygiene pass never reads it, and
# pulling a 384-dim vector per row × up to 500 rows every 5 minutes is
# pure waste.
from sqlalchemy.orm import defer
candidates = (
db.query(Ticket)
.options(defer(Ticket.embedding))
.filter(
Ticket.status == TicketStatus.RESOLVED_PENDING_CONFIRMATION.value,
Ticket.confirmation_requested_at.isnot(None),
)
.limit(500)
.all()
)
for ticket in candidates:
try:
org_settings = service.settings_repo.get_or_create(ticket.organization_id)
requested = ticket.confirmation_requested_at
if requested.tzinfo is None:
requested = requested.replace(tzinfo=timezone.utc)
deadline = requested + timedelta(hours=org_settings.confirmation_timeout_hours)
if now < deadline:
continue
await service.close(
ticket,
actor_type=TicketActorType.SYSTEM,
reason=(
"Closed automatically — no customer response within "
f"{org_settings.confirmation_timeout_hours}h of resolution."
),
)
db.commit()
await emit_ticket_update(
ticket.organization_id, ticket.id, "status",
{"status": TicketStatus.CLOSED.value},
)
except Exception as e:
logger.error(f"Lifecycle close failed for ticket {ticket.id}: {e}")
try:
db.rollback()
except Exception:
pass
def reap_orphaned_runs() -> None:
try:
with SessionLocal() as db:
reaped = InvestigationRepository(db).reap_orphaned_runs()
if reaped:
logger.warning(f"Reaped {reaped} orphaned investigation run(s) on startup")
except Exception as e:
logger.error(f"Investigation orphan reap failed (non-fatal): {e}")
async def run_ticket_investigator_loop(poll_interval: int = RUN_POLL_INTERVAL_SECONDS):
while True:
try:
await run_ticket_investigator()
except Exception as e:
logger.error(f"Error in ticket investigator loop: {e}")
await asyncio.sleep(poll_interval)
async def run_ticket_lifecycle_loop(poll_interval: int = LIFECYCLE_POLL_INTERVAL_SECONDS):
while True:
try:
await run_ticket_lifecycle()
except Exception as e:
logger.error(f"Error in ticket lifecycle loop: {e}")
await asyncio.sleep(poll_interval)
async def main():
# Route socket emits through the Redis manager (when enabled) so
# ticket_update frames from this process reach dashboard clients connected
# to the web containers. Without Redis the frontend's poll fallback covers.
from app.core.socketio import configure_socketio
configure_socketio()
# Ticket replies reach customers' channels from this process too
load_extensions()
reap_orphaned_runs()
await asyncio.gather(
run_ticket_investigator_loop(),
run_ticket_lifecycle_loop(),
)
if __name__ == "__main__":
logger.info("Starting ticket investigator service")
asyncio.run(main())