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