ryan
b8ccc88b58
Compaction replaces a tool's text with its artifact, and the search artifact only carries UI snippets, so the full chunk the retriever found never reached the model. Search tools now attach the full text as model_context, which compaction prefers and shortens under budget pressure while keeping source ids; the persisted artifact and UI are unchanged. The per-chunk header drops the raw metadata dict (about a fifth of search output tokens) for title, date and source. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
573 lines
19 KiB
Python
573 lines
19 KiB
Python
import datetime
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
from langchain_core.documents import Document
|
|
from langchain_openai import OpenAIEmbeddings
|
|
from langchain_postgres import PGVector
|
|
from langchain_text_splitters import (
|
|
MarkdownHeaderTextSplitter,
|
|
RecursiveCharacterTextSplitter,
|
|
)
|
|
from sqlalchemy import create_engine, text
|
|
|
|
from .fetchers import PaperlessNGXService
|
|
from .retrieval import build_filter, citation_url, document_date, document_dates
|
|
from utils.obsidian_service import ObsidianService
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_embedding_server_url = os.getenv("EMBEDDING_SERVER_URL")
|
|
_embedding_model = os.getenv("EMBEDDING_MODEL_NAME", "text-embedding-3-small")
|
|
|
|
if _embedding_server_url:
|
|
embeddings = OpenAIEmbeddings(
|
|
model=_embedding_model,
|
|
base_url=_embedding_server_url,
|
|
api_key="not-needed",
|
|
check_embedding_ctx_length=False,
|
|
)
|
|
else:
|
|
embeddings = OpenAIEmbeddings(model=_embedding_model)
|
|
|
|
# Convert Tortoise-style postgres:// URL to SQLAlchemy-style postgresql+psycopg://
|
|
_db_url = os.getenv(
|
|
"DATABASE_URL", "postgres://raggr:raggr_dev_password@localhost:5432/raggr"
|
|
)
|
|
_pgvector_url = _db_url.replace("postgres://", "postgresql+psycopg://")
|
|
|
|
# Lazy-initialized vector store (defers DB connection to first use)
|
|
_vector_store = None
|
|
|
|
|
|
def _get_vector_store() -> PGVector:
|
|
global _vector_store
|
|
if _vector_store is None:
|
|
_vector_store = PGVector(
|
|
embeddings=embeddings,
|
|
collection_name="simba_docs",
|
|
connection=_pgvector_url,
|
|
use_jsonb=True,
|
|
create_extension=False, # created by docker init script
|
|
async_mode=True,
|
|
)
|
|
return _vector_store
|
|
|
|
|
|
def _get_engine():
|
|
"""Get a SQLAlchemy engine for direct queries."""
|
|
if not hasattr(_get_engine, "_engine"):
|
|
_get_engine._engine = create_engine(_pgvector_url)
|
|
return _get_engine._engine
|
|
|
|
|
|
# Chunks per embedding request when indexing Paperless documents.
|
|
EMBED_BATCH_SIZE = 32
|
|
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000, # chunk size (characters)
|
|
chunk_overlap=200, # chunk overlap (characters)
|
|
add_start_index=True, # track index in original document
|
|
)
|
|
|
|
md_header_splitter = MarkdownHeaderTextSplitter(
|
|
headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")],
|
|
strip_headers=False,
|
|
)
|
|
|
|
md_chunk_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
add_start_index=True,
|
|
)
|
|
|
|
|
|
def _split_markdown_document(doc: Document) -> list[Document]:
|
|
"""Split a markdown document by headers first, then by size.
|
|
|
|
Prepends the note filename to each chunk so chunks are self-contained.
|
|
"""
|
|
note_name = (
|
|
Path(doc.metadata.get("filepath", "")).stem
|
|
if doc.metadata.get("filepath")
|
|
else ""
|
|
)
|
|
|
|
# Split by markdown headers
|
|
header_splits = md_header_splitter.split_text(doc.page_content)
|
|
|
|
# Carry over original document metadata to each header split
|
|
for split in header_splits:
|
|
split.metadata.update(doc.metadata)
|
|
|
|
# Then apply size-based splitting on large sections
|
|
sized_splits = md_chunk_splitter.split_documents(header_splits)
|
|
|
|
# Prepend note name for self-contained context
|
|
if note_name:
|
|
for split in sized_splits:
|
|
split.page_content = f"[Note: {note_name}]\n{split.page_content}"
|
|
|
|
return sized_splits
|
|
|
|
|
|
def _get_collection_id():
|
|
"""Get the UUID of our collection from the langchain_pg_collection table."""
|
|
engine = _get_engine()
|
|
try:
|
|
with engine.connect() as conn:
|
|
result = conn.execute(
|
|
text("SELECT uuid FROM langchain_pg_collection WHERE name = :name"),
|
|
{"name": "simba_docs"},
|
|
)
|
|
row = result.fetchone()
|
|
return row[0] if row else None
|
|
except Exception:
|
|
# Table doesn't exist yet (first run before any indexing)
|
|
return None
|
|
|
|
|
|
def date_to_epoch(date_str: str) -> float:
|
|
split_date = date_str.split("-")
|
|
date = datetime.datetime(
|
|
int(split_date[0]),
|
|
int(split_date[1]),
|
|
int(split_date[2]),
|
|
0,
|
|
0,
|
|
0,
|
|
tzinfo=datetime.timezone.utc,
|
|
)
|
|
|
|
return date.timestamp()
|
|
|
|
|
|
async def fetch_documents_from_paperless_ngx() -> list[Document]:
|
|
ppngx = PaperlessNGXService()
|
|
data = ppngx.get_data()
|
|
doctypes = ppngx.get_doctypes()
|
|
documents = []
|
|
for doc in data:
|
|
metadata = {
|
|
"source": "paperless",
|
|
"document_id": doc.get("id"),
|
|
"document_date": doc["created_date"][:10],
|
|
"indexed_at": time.time(),
|
|
"created_date": date_to_epoch(doc["created_date"]),
|
|
"filename": doc["original_file_name"],
|
|
"document_type": doctypes.get(doc["document_type"], ""),
|
|
}
|
|
documents.append(Document(page_content=doc["content"], metadata=metadata))
|
|
|
|
return documents
|
|
|
|
|
|
def _make_serializable(value):
|
|
"""Convert a value to a JSON-serializable type."""
|
|
if isinstance(value, (str, int, float, bool, type(None))):
|
|
return value
|
|
if isinstance(value, (list, tuple)):
|
|
return [_make_serializable(v) for v in value]
|
|
if isinstance(value, dict):
|
|
return {k: _make_serializable(v) for k, v in value.items()}
|
|
return str(value)
|
|
|
|
|
|
def _sanitize_text(text_content: str) -> str:
|
|
"""Strip non-printable and invalid characters that break embedding tokenizers."""
|
|
# Remove null bytes and control characters (keep newlines and tabs)
|
|
text_content = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]", "", text_content)
|
|
# Remove Unicode surrogates and other problematic Unicode
|
|
text_content = re.sub(r"[\ud800-\udfff\ufffe\uffff]", "", text_content)
|
|
# Remove replacement character clusters
|
|
text_content = text_content.replace("\ufffd", "")
|
|
# Collapse excessive whitespace
|
|
text_content = re.sub(r" {3,}", " ", text_content)
|
|
return text_content.strip()
|
|
|
|
|
|
def _sanitize_documents(documents: list[Document]) -> list[Document]:
|
|
"""Sanitize page_content of all documents for embedding compatibility."""
|
|
for doc in documents:
|
|
doc.page_content = _sanitize_text(doc.page_content)
|
|
return [doc for doc in documents if doc.page_content]
|
|
|
|
|
|
async def index_documents():
|
|
"""Index Paperless-NGX documents into vector store."""
|
|
documents = await fetch_documents_from_paperless_ngx()
|
|
|
|
splits = text_splitter.split_documents(documents)
|
|
splits = _sanitize_documents(splits)
|
|
logger.info(f"Indexing {len(splits)} chunks from {len(documents)} documents")
|
|
vector_store = _get_vector_store()
|
|
for start in range(0, len(splits), EMBED_BATCH_SIZE):
|
|
batch = splits[start : start + EMBED_BATCH_SIZE]
|
|
try:
|
|
await vector_store.aadd_documents(documents=batch)
|
|
except Exception as batch_error:
|
|
# Embedding fails before anything is inserted; retry the batch one
|
|
# chunk at a time to name the chunk the embedding server rejects.
|
|
logger.warning(
|
|
f"Batch of {len(batch)} chunks failed ({batch_error}); retrying one at a time"
|
|
)
|
|
for i, split in enumerate(batch, start):
|
|
try:
|
|
await vector_store.aadd_documents(documents=[split])
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to embed chunk {i} from {split.metadata.get('filename', 'unknown')}: {e}"
|
|
)
|
|
logger.debug(f"Chunk content preview: {split.page_content[:200]!r}")
|
|
raise
|
|
|
|
|
|
async def fetch_obsidian_documents() -> list[Document]:
|
|
"""Fetch all markdown documents from Obsidian vault.
|
|
|
|
Returns:
|
|
List of LangChain Document objects with source='obsidian' metadata.
|
|
"""
|
|
obsidian_service = ObsidianService()
|
|
documents = []
|
|
|
|
for md_path in obsidian_service.walk_vault():
|
|
try:
|
|
# Read markdown file
|
|
with open(md_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
|
|
# Parse metadata
|
|
parsed = obsidian_service.parse_markdown(content, md_path)
|
|
|
|
# Create LangChain Document with obsidian source
|
|
metadata = {
|
|
"source": "obsidian",
|
|
"filepath": parsed["filepath"],
|
|
"folder": str(Path(parsed["filepath"]).parent)
|
|
if parsed["filepath"]
|
|
else "",
|
|
"tags": parsed["tags"],
|
|
"created_at": parsed["metadata"].get("created_at"),
|
|
"indexed_at": time.time(),
|
|
**{
|
|
k: v
|
|
for k, v in parsed["metadata"].items()
|
|
if k not in ["created_at", "created_by"]
|
|
},
|
|
}
|
|
document = Document(
|
|
page_content=parsed["content"],
|
|
metadata=_make_serializable(metadata),
|
|
)
|
|
documents.append(document)
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Error reading {md_path}: {e}")
|
|
continue
|
|
|
|
return documents
|
|
|
|
|
|
async def index_obsidian_documents():
|
|
"""Index all Obsidian markdown documents into vector store.
|
|
|
|
Deletes existing obsidian source chunks before re-indexing.
|
|
"""
|
|
documents = await fetch_obsidian_documents()
|
|
|
|
if not documents:
|
|
logger.info("No Obsidian documents found to index")
|
|
return {"indexed": 0}
|
|
|
|
# Delete existing obsidian chunks
|
|
delete_documents_by_metadata("source", "obsidian")
|
|
|
|
# Split using markdown-aware chunking, sanitize, and index
|
|
splits = []
|
|
for doc in documents:
|
|
splits.extend(_split_markdown_document(doc))
|
|
splits = _sanitize_documents(splits)
|
|
vector_store = _get_vector_store()
|
|
await vector_store.aadd_documents(documents=splits)
|
|
|
|
return {"indexed": len(documents)}
|
|
|
|
|
|
# In-memory cache of indexed obsidian files: {filepath: indexed_at}
|
|
_obsidian_index_cache: dict[str, float] = {}
|
|
|
|
|
|
def _load_obsidian_index_cache() -> dict[str, float]:
|
|
"""Load indexed obsidian files from DB into cache (cold start only)."""
|
|
collection_id = _get_collection_id()
|
|
if not collection_id:
|
|
return {}
|
|
engine = _get_engine()
|
|
with engine.connect() as conn:
|
|
result = conn.execute(
|
|
text(
|
|
"SELECT DISTINCT cmetadata->>'filepath' AS filepath, "
|
|
"MAX((cmetadata->>'indexed_at')::float) AS indexed_at "
|
|
"FROM langchain_pg_embedding "
|
|
"WHERE collection_id = :cid AND cmetadata->>'source' = 'obsidian' "
|
|
"GROUP BY cmetadata->>'filepath'"
|
|
),
|
|
{"cid": collection_id},
|
|
)
|
|
return {row[0]: row[1] for row in result if row[0] is not None}
|
|
|
|
|
|
async def sync_obsidian_documents() -> dict[str, int]:
|
|
"""Incrementally sync Obsidian documents to pgvector.
|
|
|
|
Compares file mtimes against stored indexed_at timestamps to only
|
|
re-index changed/new files and remove deleted ones.
|
|
|
|
Returns:
|
|
Dict with counts of added, updated, and deleted files.
|
|
"""
|
|
global _obsidian_index_cache
|
|
|
|
obsidian_service = ObsidianService()
|
|
|
|
# Load cache from DB on first run
|
|
if not _obsidian_index_cache:
|
|
_obsidian_index_cache = _load_obsidian_index_cache()
|
|
|
|
# Build map of current vault files -> mtime
|
|
vault_files: dict[str, float] = {}
|
|
for md_path in obsidian_service.walk_vault():
|
|
vault_files[str(md_path)] = md_path.stat().st_mtime
|
|
|
|
added = 0
|
|
updated = 0
|
|
deleted = 0
|
|
|
|
# Find files to add or update
|
|
files_to_index: list[str] = []
|
|
for filepath, mtime in vault_files.items():
|
|
indexed_at = _obsidian_index_cache.get(filepath)
|
|
if indexed_at is None:
|
|
files_to_index.append(filepath)
|
|
added += 1
|
|
elif mtime > indexed_at:
|
|
# Delete old chunks first
|
|
delete_documents_by_metadata("filepath", filepath)
|
|
files_to_index.append(filepath)
|
|
updated += 1
|
|
|
|
# Find deleted files (in cache but not on disk)
|
|
for filepath in list(_obsidian_index_cache):
|
|
if filepath not in vault_files:
|
|
delete_documents_by_metadata("filepath", filepath)
|
|
del _obsidian_index_cache[filepath]
|
|
deleted += 1
|
|
|
|
# Index new/changed files
|
|
if files_to_index:
|
|
now = time.time()
|
|
documents = []
|
|
for filepath in files_to_index:
|
|
try:
|
|
with open(filepath, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
parsed = obsidian_service.parse_markdown(content, filepath)
|
|
metadata = {
|
|
"source": "obsidian",
|
|
"filepath": parsed["filepath"],
|
|
"folder": str(Path(parsed["filepath"]).parent)
|
|
if parsed["filepath"]
|
|
else "",
|
|
"tags": parsed["tags"],
|
|
"created_at": parsed["metadata"].get("created_at"),
|
|
"indexed_at": now,
|
|
**{
|
|
k: v
|
|
for k, v in parsed["metadata"].items()
|
|
if k not in ["created_at", "created_by"]
|
|
},
|
|
}
|
|
document = Document(
|
|
page_content=parsed["content"],
|
|
metadata=_make_serializable(metadata),
|
|
)
|
|
documents.append(document)
|
|
except Exception as e:
|
|
logger.warning(f"Error reading {filepath}: {e}")
|
|
continue
|
|
|
|
if documents:
|
|
splits = []
|
|
for doc in documents:
|
|
splits.extend(_split_markdown_document(doc))
|
|
splits = _sanitize_documents(splits)
|
|
if splits:
|
|
vector_store = _get_vector_store()
|
|
await vector_store.aadd_documents(documents=splits)
|
|
|
|
# Update cache for successfully processed files
|
|
for filepath in files_to_index:
|
|
_obsidian_index_cache[filepath] = now
|
|
|
|
logger.info(
|
|
f"Obsidian sync complete: {added} added, {updated} updated, {deleted} deleted"
|
|
)
|
|
return {"added": added, "updated": updated, "deleted": deleted}
|
|
|
|
|
|
async def query_vector_store(
|
|
query: str,
|
|
source: str | None = None,
|
|
k: int = 8,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
):
|
|
"""Query the vector store with optional source filtering and MMR.
|
|
|
|
Args:
|
|
query: Search query text
|
|
source: Filter by source metadata (e.g., "obsidian", "paperless")
|
|
k: Number of results to return
|
|
"""
|
|
if not isinstance(query, str) or not query.strip():
|
|
raise ValueError("query must not be empty")
|
|
if not isinstance(k, int) or not 1 <= k <= 20:
|
|
raise ValueError("k must be between 1 and 20")
|
|
filter_dict = build_filter(source, start_date, end_date)
|
|
vector_store = _get_vector_store()
|
|
retrieved_docs = await vector_store.amax_marginal_relevance_search(
|
|
query,
|
|
k=k,
|
|
fetch_k=k * 3,
|
|
filter=filter_dict,
|
|
)
|
|
if start_date or end_date:
|
|
retrieved_docs = [
|
|
doc
|
|
for doc in retrieved_docs
|
|
# Same "any date field in range" rule as the database filter.
|
|
if any(
|
|
(not start_date or stamp >= start_date)
|
|
and (not end_date or stamp <= end_date)
|
|
for stamp in document_dates(doc.metadata)
|
|
)
|
|
]
|
|
serialized = "\n\n".join(
|
|
(
|
|
f"{_chunk_header(doc)}\n"
|
|
+ (f"Citation: {citation_url(doc)}\n" if citation_url(doc) else "")
|
|
+ f"Content: {doc.page_content}"
|
|
)
|
|
for doc in retrieved_docs
|
|
)
|
|
return serialized, retrieved_docs
|
|
|
|
|
|
def _chunk_title(doc: Document) -> str:
|
|
from blueprints.conversation.artifacts_util import note_title, vault_relative_path
|
|
|
|
meta = doc.metadata or {}
|
|
if meta.get("source") == "obsidian":
|
|
path = vault_relative_path(meta.get("filepath") or "")
|
|
return note_title(path, meta) if path else "Untitled note"
|
|
filename = meta.get("filename") or ""
|
|
return Path(filename).stem or filename or "Document"
|
|
|
|
|
|
def _chunk_header(doc: Document) -> str:
|
|
"""Compact per-chunk header (title, date, source) replacing the raw metadata dict."""
|
|
meta = doc.metadata or {}
|
|
parts = [f"Title: {_chunk_title(doc)}"]
|
|
if date := document_date(meta):
|
|
parts.append(f"Date: {date}")
|
|
if source := meta.get("source"):
|
|
parts.append(f"Source: {source}")
|
|
return " | ".join(parts)
|
|
|
|
|
|
def delete_all_documents():
|
|
"""Delete all documents from the vector store collection."""
|
|
collection_id = _get_collection_id()
|
|
if not collection_id:
|
|
return
|
|
engine = _get_engine()
|
|
with engine.connect() as conn:
|
|
conn.execute(
|
|
text("DELETE FROM langchain_pg_embedding WHERE collection_id = :cid"),
|
|
{"cid": collection_id},
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def delete_documents_by_metadata(key: str, value: str):
|
|
"""Delete documents matching a metadata key/value pair."""
|
|
collection_id = _get_collection_id()
|
|
if not collection_id:
|
|
return
|
|
engine = _get_engine()
|
|
with engine.connect() as conn:
|
|
conn.execute(
|
|
text(
|
|
"DELETE FROM langchain_pg_embedding "
|
|
"WHERE collection_id = :cid AND cmetadata->>:key = :value"
|
|
),
|
|
{"cid": collection_id, "key": key, "value": value},
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def get_vector_store_stats():
|
|
"""Get statistics about the vector store."""
|
|
collection_id = _get_collection_id()
|
|
count = 0
|
|
if collection_id:
|
|
engine = _get_engine()
|
|
with engine.connect() as conn:
|
|
result = conn.execute(
|
|
text(
|
|
"SELECT COUNT(*) FROM langchain_pg_embedding WHERE collection_id = :cid"
|
|
),
|
|
{"cid": collection_id},
|
|
)
|
|
count = result.scalar()
|
|
return {
|
|
"total_documents": count,
|
|
"collection_name": "simba_docs",
|
|
}
|
|
|
|
|
|
def list_all_documents(limit: int = 10):
|
|
"""List documents in the vector store with their metadata."""
|
|
collection_id = _get_collection_id()
|
|
if not collection_id:
|
|
return []
|
|
|
|
engine = _get_engine()
|
|
with engine.connect() as conn:
|
|
result = conn.execute(
|
|
text(
|
|
"SELECT id, document, cmetadata FROM langchain_pg_embedding "
|
|
"WHERE collection_id = :cid LIMIT :limit"
|
|
),
|
|
{"cid": collection_id, "limit": limit},
|
|
)
|
|
documents = []
|
|
for row in result:
|
|
documents.append(
|
|
{
|
|
"id": str(row[0]),
|
|
"metadata": row[2],
|
|
"content_preview": row[1][:200] if row[1] else None,
|
|
}
|
|
)
|
|
|
|
return documents
|