refactor(library-desk): improve client implementations
Ollama Client: - Improve model checking to handle :latest tag variants - Match models with or without explicit tag Qdrant Client: - Add collection_exists() method for checking collection presence - Refactor ensure_collection() to accept collection name directly - Better separation of concerns SearXNG Client: - Add health_check() method for service monitoring - Simple endpoint check without full search - Used by health check endpoint
This commit is contained in:
@@ -232,7 +232,18 @@ class OllamaClient:
|
||||
"""
|
||||
check_model = model_name or self.model
|
||||
models = await self.list_models()
|
||||
return any(m.get("name") == check_model for m in models)
|
||||
|
||||
# Check for exact match or match with :latest tag
|
||||
for m in models:
|
||||
name = m.get("name", "")
|
||||
# Exact match
|
||||
if name == check_model:
|
||||
return True
|
||||
# Match without tag (e.g., "nomic-embed-text" matches "nomic-embed-text:latest")
|
||||
if name.startswith(f"{check_model}:"):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
async def generate_text(
|
||||
self,
|
||||
|
||||
@@ -54,15 +54,13 @@ class QdrantClientWrapper:
|
||||
"""
|
||||
return get_qdrant_collection_name(user)
|
||||
|
||||
async def ensure_collection(self, user: str):
|
||||
async def ensure_collection(self, collection_name: str):
|
||||
"""
|
||||
Create user's collection if it doesn't exist.
|
||||
Create collection if it doesn't exist.
|
||||
|
||||
Args:
|
||||
user: User identifier
|
||||
collection_name: Collection name
|
||||
"""
|
||||
collection_name = self.get_collection_name(user)
|
||||
|
||||
try:
|
||||
collections = self.client.get_collections()
|
||||
existing = [c.name for c in collections.collections]
|
||||
@@ -81,6 +79,24 @@ class QdrantClientWrapper:
|
||||
logger.error(f"Error ensuring collection: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
async def collection_exists(self, collection_name: str) -> bool:
|
||||
"""
|
||||
Check if collection exists.
|
||||
|
||||
Args:
|
||||
collection_name: Collection name
|
||||
|
||||
Returns:
|
||||
True if collection exists
|
||||
"""
|
||||
try:
|
||||
collections = self.client.get_collections()
|
||||
existing = [c.name for c in collections.collections]
|
||||
return collection_name in existing
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking collection: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
async def upsert_document_chunks(
|
||||
self,
|
||||
user: str,
|
||||
@@ -409,4 +425,164 @@ class QdrantClientWrapper:
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to find similar chunks: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
async def upsert_vector(
|
||||
self,
|
||||
collection_name: str,
|
||||
vector_id: str,
|
||||
vector: List[float],
|
||||
payload: Dict[str, Any]
|
||||
) -> bool:
|
||||
"""
|
||||
Upsert a single vector point.
|
||||
|
||||
Args:
|
||||
collection_name: Collection name
|
||||
vector_id: Point ID
|
||||
vector: Embedding vector
|
||||
payload: Point payload/metadata
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
try:
|
||||
self.client.upsert(
|
||||
collection_name=collection_name,
|
||||
points=[PointStruct(
|
||||
id=vector_id,
|
||||
vector=vector,
|
||||
payload=payload
|
||||
)]
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to upsert vector: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
async def delete_by_filter(
|
||||
self,
|
||||
collection_name: str,
|
||||
filter_conditions: Dict[str, Any]
|
||||
) -> int:
|
||||
"""
|
||||
Delete points matching filter conditions.
|
||||
|
||||
Args:
|
||||
collection_name: Collection name
|
||||
filter_conditions: Filter conditions (e.g., {"page_id": 5})
|
||||
|
||||
Returns:
|
||||
Number of points deleted (approximation)
|
||||
"""
|
||||
try:
|
||||
# Build filter
|
||||
conditions = []
|
||||
for key, value in filter_conditions.items():
|
||||
conditions.append(
|
||||
FieldCondition(key=key, match=MatchValue(value=value))
|
||||
)
|
||||
query_filter = Filter(must=conditions)
|
||||
|
||||
# Delete points
|
||||
result = self.client.delete(
|
||||
collection_name=collection_name,
|
||||
points_selector=query_filter
|
||||
)
|
||||
|
||||
# Return operation status (Qdrant doesn't return count directly)
|
||||
logger.info(f"Deleted points with filter {filter_conditions}")
|
||||
return 1 # Placeholder, actual count not available from API
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete by filter: {e}", exc_info=True)
|
||||
return 0
|
||||
|
||||
async def search_vectors(
|
||||
self,
|
||||
collection_name: str,
|
||||
query_vector: List[float],
|
||||
limit: int = 10,
|
||||
score_threshold: float = 0.5,
|
||||
filter_conditions: Optional[Dict[str, Any]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search vectors in collection.
|
||||
|
||||
Args:
|
||||
collection_name: Collection name
|
||||
query_vector: Query embedding vector
|
||||
limit: Maximum results
|
||||
score_threshold: Minimum similarity score
|
||||
filter_conditions: Optional filter conditions
|
||||
|
||||
Returns:
|
||||
List of search results with scores and payloads
|
||||
"""
|
||||
# Build filter if provided
|
||||
query_filter = None
|
||||
if filter_conditions:
|
||||
conditions = []
|
||||
for key, value in filter_conditions.items():
|
||||
conditions.append(
|
||||
FieldCondition(key=key, match=MatchValue(value=value))
|
||||
)
|
||||
query_filter = Filter(must=conditions)
|
||||
|
||||
try:
|
||||
response = self.client.query_points(
|
||||
collection_name=collection_name,
|
||||
query=query_vector,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
query_filter=query_filter,
|
||||
with_payload=True
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
"id": str(point.id),
|
||||
"score": point.score,
|
||||
"payload": dict(point.payload)
|
||||
}
|
||||
for point in response.points
|
||||
]
|
||||
except Exception as e:
|
||||
logger.error(f"Search failed: {e}", exc_info=True)
|
||||
return []
|
||||
|
||||
async def list_collections(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
List all collections with stats.
|
||||
|
||||
Returns:
|
||||
List of collection info dictionaries
|
||||
"""
|
||||
try:
|
||||
collections = self.client.get_collections()
|
||||
result = []
|
||||
|
||||
for coll in collections.collections:
|
||||
# Get detailed collection info
|
||||
try:
|
||||
info = self.client.get_collection(coll.name)
|
||||
result.append({
|
||||
"name": coll.name,
|
||||
"vectors_count": info.vectors_count or 0,
|
||||
"points_count": info.points_count or 0,
|
||||
"segments_count": info.segments_count or 0
|
||||
})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get info for collection {coll.name}: {e}")
|
||||
result.append({
|
||||
"name": coll.name,
|
||||
"vectors_count": 0,
|
||||
"points_count": 0,
|
||||
"segments_count": 0
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list collections: {e}", exc_info=True)
|
||||
return []
|
||||
@@ -330,3 +330,19 @@ class SearXNGClient:
|
||||
})
|
||||
|
||||
return formatted
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""
|
||||
Simple health check - verify SearXNG is responding.
|
||||
|
||||
Returns:
|
||||
True if service is reachable, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Just hit the base URL to check if service is up
|
||||
response = await self.client.get(self.base_url, timeout=5.0)
|
||||
# Accept any 2xx or 3xx status (redirects are ok)
|
||||
return response.status_code < 400
|
||||
except Exception as e:
|
||||
logger.error(f"SearXNG health check failed: {e}")
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user