feat(library-desk): add vector service for semantic search
Vector Service: - Manage document embeddings in Qdrant - Update vectors from wiki pages - Handle chunking and embedding generation - Support force refresh and incremental updates Vector Models: - VectorSearchResult for search responses - VectorUpdateSummary for indexing metrics - Track chunks created/deleted Used by ingestion_service for page embedding
This commit is contained in:
@@ -0,0 +1,100 @@
|
|||||||
|
"""
|
||||||
|
Vector models for Library Desk Qdrant operations.
|
||||||
|
|
||||||
|
Provides models for semantic search, document chunks, and embeddings.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from typing import List, Dict, Any, Optional
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentChunk(BaseModel):
|
||||||
|
"""Document chunk with embedding."""
|
||||||
|
chunk_id: str = Field(..., description="Unique chunk ID (page_id:chunk_index)")
|
||||||
|
page_id: int = Field(..., description="Wiki page ID")
|
||||||
|
chunk_index: int = Field(..., description="Chunk index within document")
|
||||||
|
content: str = Field(..., description="Chunk text content")
|
||||||
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
|
||||||
|
|
||||||
|
|
||||||
|
class SearchResult(BaseModel):
|
||||||
|
"""Semantic search result."""
|
||||||
|
chunk_id: str = Field(..., description="Chunk ID")
|
||||||
|
page_id: int = Field(..., description="Wiki page ID")
|
||||||
|
page_title: Optional[str] = Field(None, description="Page title")
|
||||||
|
page_path: Optional[str] = Field(None, description="Page path")
|
||||||
|
chunk_index: int = Field(..., description="Chunk index")
|
||||||
|
content: str = Field(..., description="Chunk content")
|
||||||
|
score: float = Field(..., description="Similarity score (0-1)")
|
||||||
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
|
||||||
|
|
||||||
|
|
||||||
|
class SearchRequest(BaseModel):
|
||||||
|
"""Semantic search request."""
|
||||||
|
query: str = Field(..., min_length=1, description="Search query")
|
||||||
|
user: str = Field(default="jpmschweitzer", description="User identifier")
|
||||||
|
limit: int = Field(default=10, ge=1, le=100, description="Maximum results")
|
||||||
|
score_threshold: float = Field(default=0.5, ge=0.0, le=1.0, description="Minimum similarity score")
|
||||||
|
|
||||||
|
|
||||||
|
class SearchResponse(BaseModel):
|
||||||
|
"""Semantic search response."""
|
||||||
|
query: str = Field(..., description="Search query")
|
||||||
|
results: List[SearchResult] = Field(..., description="Search results")
|
||||||
|
total: int = Field(..., description="Number of results")
|
||||||
|
user: str = Field(..., description="User filter applied")
|
||||||
|
|
||||||
|
|
||||||
|
class VectorUpdateRequest(BaseModel):
|
||||||
|
"""Request to update vectors from a wiki page."""
|
||||||
|
page_id: int = Field(..., description="Wiki page ID to process")
|
||||||
|
user: str = Field(
|
||||||
|
default="jpmschweitzer",
|
||||||
|
description="User identifier for namespace scoping"
|
||||||
|
)
|
||||||
|
force_refresh: bool = Field(
|
||||||
|
default=False,
|
||||||
|
description="Force re-embedding even if page hasn't changed"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class VectorUpdateSummary(BaseModel):
|
||||||
|
"""Summary of vector update operation."""
|
||||||
|
page_id: int = Field(..., description="Page ID processed")
|
||||||
|
page_title: str = Field(..., description="Page title")
|
||||||
|
chunks_created: int = Field(default=0, description="New chunks created")
|
||||||
|
chunks_updated: int = Field(default=0, description="Existing chunks updated")
|
||||||
|
chunks_deleted: int = Field(default=0, description="Old chunks deleted")
|
||||||
|
total_chunks: int = Field(default=0, description="Total chunks for this page")
|
||||||
|
embedding_dim: int = Field(default=768, description="Embedding dimensionality")
|
||||||
|
processing_time_ms: float = Field(..., description="Processing time in milliseconds")
|
||||||
|
success: bool = Field(default=True, description="Whether update succeeded")
|
||||||
|
error_message: Optional[str] = Field(default=None, description="Error message if failed")
|
||||||
|
|
||||||
|
|
||||||
|
class CollectionInfo(BaseModel):
|
||||||
|
"""Qdrant collection information."""
|
||||||
|
name: str = Field(..., description="Collection name")
|
||||||
|
vectors_count: int = Field(..., description="Number of vectors")
|
||||||
|
points_count: int = Field(..., description="Number of points")
|
||||||
|
segments_count: int = Field(..., description="Number of segments")
|
||||||
|
|
||||||
|
|
||||||
|
class CollectionListResponse(BaseModel):
|
||||||
|
"""List of Qdrant collections."""
|
||||||
|
collections: List[CollectionInfo] = Field(..., description="List of collections")
|
||||||
|
total: int = Field(..., description="Total number of collections")
|
||||||
|
|
||||||
|
|
||||||
|
class DeletePageChunksRequest(BaseModel):
|
||||||
|
"""Request to delete all chunks for a page."""
|
||||||
|
page_id: int = Field(..., description="Wiki page ID")
|
||||||
|
user: str = Field(default="jpmschweitzer", description="User identifier")
|
||||||
|
|
||||||
|
|
||||||
|
class DeletePageChunksResponse(BaseModel):
|
||||||
|
"""Response from deleting page chunks."""
|
||||||
|
page_id: int = Field(..., description="Page ID")
|
||||||
|
chunks_deleted: int = Field(..., description="Number of chunks deleted")
|
||||||
|
success: bool = Field(..., description="Whether deletion succeeded")
|
||||||
|
message: str = Field(..., description="Result message")
|
||||||
@@ -0,0 +1,358 @@
|
|||||||
|
"""
|
||||||
|
Vector service for Library Desk Qdrant operations.
|
||||||
|
|
||||||
|
Handles semantic search, document chunking, and embeddings.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
import time
|
||||||
|
import hashlib
|
||||||
|
import uuid
|
||||||
|
from typing import List, Dict, Any, Optional
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from src.clients.qdrant_client import QdrantClientWrapper
|
||||||
|
from src.clients.wikijs_client import WikiJSClient
|
||||||
|
from src.clients.ollama_client import OllamaClient
|
||||||
|
from src.core.multi_tenancy import get_qdrant_collection_name
|
||||||
|
from src.models.vector import (
|
||||||
|
SearchResult, SearchResponse, VectorUpdateSummary,
|
||||||
|
DocumentChunk, CollectionInfo, CollectionListResponse
|
||||||
|
)
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class VectorService:
|
||||||
|
"""
|
||||||
|
Service for Qdrant vector operations.
|
||||||
|
|
||||||
|
Responsibilities:
|
||||||
|
- Document chunking
|
||||||
|
- Embedding generation
|
||||||
|
- Semantic search
|
||||||
|
- Vector CRUD operations
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
qdrant_client: QdrantClientWrapper,
|
||||||
|
wikijs_client: WikiJSClient,
|
||||||
|
ollama_client: OllamaClient,
|
||||||
|
chunk_size: int = 500,
|
||||||
|
chunk_overlap: int = 50
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize vector service.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
qdrant_client: Qdrant database client
|
||||||
|
wikijs_client: Wiki.js client for fetching pages
|
||||||
|
ollama_client: Ollama client for embeddings
|
||||||
|
chunk_size: Target chunk size in tokens (approximate)
|
||||||
|
chunk_overlap: Overlap between chunks in tokens
|
||||||
|
"""
|
||||||
|
self.qdrant = qdrant_client
|
||||||
|
self.wiki = wikijs_client
|
||||||
|
self.ollama = ollama_client
|
||||||
|
self.chunk_size = chunk_size
|
||||||
|
self.chunk_overlap = chunk_overlap
|
||||||
|
|
||||||
|
def _chunk_text(self, text: str) -> List[str]:
|
||||||
|
"""
|
||||||
|
Chunk text into overlapping segments.
|
||||||
|
|
||||||
|
Simple word-based chunking for now.
|
||||||
|
TODO: Use tiktoken or similar for token-accurate chunking.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: Text to chunk
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of text chunks
|
||||||
|
"""
|
||||||
|
# Remove extra whitespace
|
||||||
|
text = re.sub(r'\s+', ' ', text).strip()
|
||||||
|
|
||||||
|
# Split into words (approximates tokens)
|
||||||
|
words = text.split()
|
||||||
|
|
||||||
|
if len(words) <= self.chunk_size:
|
||||||
|
return [text]
|
||||||
|
|
||||||
|
chunks = []
|
||||||
|
start = 0
|
||||||
|
|
||||||
|
while start < len(words):
|
||||||
|
end = start + self.chunk_size
|
||||||
|
chunk_words = words[start:end]
|
||||||
|
chunks.append(' '.join(chunk_words))
|
||||||
|
|
||||||
|
# Move start forward with overlap
|
||||||
|
start = end - self.chunk_overlap
|
||||||
|
|
||||||
|
return chunks
|
||||||
|
|
||||||
|
async def update_from_page(
|
||||||
|
self,
|
||||||
|
page_id: int,
|
||||||
|
user: str,
|
||||||
|
force_refresh: bool = False
|
||||||
|
) -> VectorUpdateSummary:
|
||||||
|
"""
|
||||||
|
Update vector embeddings from a wiki page.
|
||||||
|
|
||||||
|
Chunks the page content, generates embeddings, and upserts to Qdrant.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
page_id: Wiki page ID
|
||||||
|
user: User identifier
|
||||||
|
force_refresh: Force re-embedding even if unchanged
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Summary of update operation
|
||||||
|
"""
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Fetch page from Wiki.js
|
||||||
|
page = await self.wiki.get_page(page_id)
|
||||||
|
if not page:
|
||||||
|
raise ValueError(f"Page {page_id} not found")
|
||||||
|
|
||||||
|
# Get collection name for user
|
||||||
|
collection_name = get_qdrant_collection_name(user)
|
||||||
|
|
||||||
|
# Ensure collection exists
|
||||||
|
await self.qdrant.ensure_collection(collection_name)
|
||||||
|
|
||||||
|
# Extract content
|
||||||
|
content = page.get("content", "")
|
||||||
|
title = page.get("title", "")
|
||||||
|
path = page.get("path", "")
|
||||||
|
|
||||||
|
if not content:
|
||||||
|
logger.warning(f"Page {page_id} has no content, skipping vector update")
|
||||||
|
return VectorUpdateSummary(
|
||||||
|
page_id=page_id,
|
||||||
|
page_title=title,
|
||||||
|
processing_time_ms=(time.time() - start_time) * 1000,
|
||||||
|
success=True
|
||||||
|
)
|
||||||
|
|
||||||
|
# Chunk the content
|
||||||
|
chunks = self._chunk_text(content)
|
||||||
|
logger.info(f"Split page {page_id} into {len(chunks)} chunks")
|
||||||
|
|
||||||
|
# Delete existing chunks for this page
|
||||||
|
deleted_count = await self.qdrant.delete_by_filter(
|
||||||
|
collection_name=collection_name,
|
||||||
|
filter_conditions={"page_id": page_id}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Generate embeddings and upsert chunks
|
||||||
|
chunks_created = 0
|
||||||
|
for idx, chunk_text in enumerate(chunks):
|
||||||
|
# Generate deterministic UUID from page_id and chunk_index
|
||||||
|
chunk_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, f"page_{page_id}_chunk_{idx}"))
|
||||||
|
|
||||||
|
# Generate embedding
|
||||||
|
embedding = await self.ollama.embed(chunk_text)
|
||||||
|
if not embedding:
|
||||||
|
logger.error(f"Failed to generate embedding for chunk {chunk_id}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Prepare metadata
|
||||||
|
metadata = {
|
||||||
|
"page_id": page_id,
|
||||||
|
"page_title": title,
|
||||||
|
"page_path": path,
|
||||||
|
"chunk_index": idx,
|
||||||
|
"chunk_text": chunk_text,
|
||||||
|
"user": user
|
||||||
|
}
|
||||||
|
|
||||||
|
# Upsert to Qdrant
|
||||||
|
success = await self.qdrant.upsert_vector(
|
||||||
|
collection_name=collection_name,
|
||||||
|
vector_id=chunk_id,
|
||||||
|
vector=embedding,
|
||||||
|
payload=metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
if success:
|
||||||
|
chunks_created += 1
|
||||||
|
|
||||||
|
processing_time_ms = (time.time() - start_time) * 1000
|
||||||
|
|
||||||
|
logger.info(
|
||||||
|
f"Updated vectors for page {page_id}: "
|
||||||
|
f"{chunks_created} chunks created, {deleted_count} old chunks deleted"
|
||||||
|
)
|
||||||
|
|
||||||
|
return VectorUpdateSummary(
|
||||||
|
page_id=page_id,
|
||||||
|
page_title=title,
|
||||||
|
chunks_created=chunks_created,
|
||||||
|
chunks_deleted=deleted_count,
|
||||||
|
total_chunks=chunks_created,
|
||||||
|
embedding_dim=len(embedding) if embedding else 768,
|
||||||
|
processing_time_ms=processing_time_ms,
|
||||||
|
success=True
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
processing_time_ms = (time.time() - start_time) * 1000
|
||||||
|
logger.error(f"Failed to update vectors for page {page_id}: {e}", exc_info=True)
|
||||||
|
|
||||||
|
return VectorUpdateSummary(
|
||||||
|
page_id=page_id,
|
||||||
|
page_title="Unknown",
|
||||||
|
processing_time_ms=processing_time_ms,
|
||||||
|
success=False,
|
||||||
|
error_message=str(e)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def search(
|
||||||
|
self,
|
||||||
|
query: str,
|
||||||
|
user: str,
|
||||||
|
limit: int = 10,
|
||||||
|
score_threshold: float = 0.5
|
||||||
|
) -> SearchResponse:
|
||||||
|
"""
|
||||||
|
Perform semantic search across user's documents.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: Search query text
|
||||||
|
user: User identifier
|
||||||
|
limit: Maximum results to return
|
||||||
|
score_threshold: Minimum similarity score (0-1)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Search results with similarity scores
|
||||||
|
"""
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Get collection name
|
||||||
|
collection_name = get_qdrant_collection_name(user)
|
||||||
|
|
||||||
|
# Check if collection exists
|
||||||
|
exists = await self.qdrant.collection_exists(collection_name)
|
||||||
|
if not exists:
|
||||||
|
logger.info(f"Collection {collection_name} doesn't exist, returning empty results")
|
||||||
|
return SearchResponse(
|
||||||
|
query=query,
|
||||||
|
results=[],
|
||||||
|
total=0,
|
||||||
|
user=user
|
||||||
|
)
|
||||||
|
|
||||||
|
# Generate query embedding
|
||||||
|
query_embedding = await self.ollama.embed(query)
|
||||||
|
if not query_embedding:
|
||||||
|
raise ValueError("Failed to generate query embedding")
|
||||||
|
|
||||||
|
# Search in Qdrant
|
||||||
|
search_results = await self.qdrant.search_vectors(
|
||||||
|
collection_name=collection_name,
|
||||||
|
query_vector=query_embedding,
|
||||||
|
limit=limit,
|
||||||
|
score_threshold=score_threshold
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert to SearchResult models
|
||||||
|
results = []
|
||||||
|
for result in search_results:
|
||||||
|
payload = result.get("payload", {})
|
||||||
|
results.append(SearchResult(
|
||||||
|
chunk_id=result["id"],
|
||||||
|
page_id=payload.get("page_id", 0),
|
||||||
|
page_title=payload.get("page_title"),
|
||||||
|
page_path=payload.get("page_path"),
|
||||||
|
chunk_index=payload.get("chunk_index", 0),
|
||||||
|
content=payload.get("chunk_text", ""),
|
||||||
|
score=result["score"],
|
||||||
|
metadata=payload
|
||||||
|
))
|
||||||
|
|
||||||
|
query_time_ms = (time.time() - start_time) * 1000
|
||||||
|
logger.info(f"Semantic search completed in {query_time_ms:.2f}ms: {len(results)} results")
|
||||||
|
|
||||||
|
return SearchResponse(
|
||||||
|
query=query,
|
||||||
|
results=results,
|
||||||
|
total=len(results),
|
||||||
|
user=user
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Semantic search failed: {e}", exc_info=True)
|
||||||
|
return SearchResponse(
|
||||||
|
query=query,
|
||||||
|
results=[],
|
||||||
|
total=0,
|
||||||
|
user=user
|
||||||
|
)
|
||||||
|
|
||||||
|
async def delete_page_chunks(
|
||||||
|
self,
|
||||||
|
page_id: int,
|
||||||
|
user: str
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Delete all chunks for a wiki page.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
page_id: Wiki page ID
|
||||||
|
user: User identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of chunks deleted
|
||||||
|
"""
|
||||||
|
collection_name = get_qdrant_collection_name(user)
|
||||||
|
|
||||||
|
try:
|
||||||
|
deleted_count = await self.qdrant.delete_by_filter(
|
||||||
|
collection_name=collection_name,
|
||||||
|
filter_conditions={"page_id": page_id}
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(f"Deleted {deleted_count} chunks for page {page_id}")
|
||||||
|
return deleted_count
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to delete chunks for page {page_id}: {e}", exc_info=True)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
async def list_collections(self) -> CollectionListResponse:
|
||||||
|
"""
|
||||||
|
List all Qdrant collections.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of collections with stats
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
collections_data = await self.qdrant.list_collections()
|
||||||
|
|
||||||
|
collections = []
|
||||||
|
for coll in collections_data:
|
||||||
|
collections.append(CollectionInfo(
|
||||||
|
name=coll["name"],
|
||||||
|
vectors_count=coll.get("vectors_count", 0),
|
||||||
|
points_count=coll.get("points_count", 0),
|
||||||
|
segments_count=coll.get("segments_count", 0)
|
||||||
|
))
|
||||||
|
|
||||||
|
return CollectionListResponse(
|
||||||
|
collections=collections,
|
||||||
|
total=len(collections)
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to list collections: {e}", exc_info=True)
|
||||||
|
return CollectionListResponse(
|
||||||
|
collections=[],
|
||||||
|
total=0
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user