diff --git a/docs/artifact-manager.md b/docs/artifact-manager.md index 5a429e09..b5030d3f 100644 --- a/docs/artifact-manager.md +++ b/docs/artifact-manager.md @@ -299,11 +299,11 @@ The following permission levels are supported: - **lv+**: List, list vectors, create, and commit access (includes `list`, `list_vectors`, `create`, `commit`, `add_vectors`, and `add_documents`). - **lf**: List and list files access (includes `list` and `list_files`). - **lf+**: List, list files, create, and commit access (includes `list`, `list_files`, `create`, `commit`, and `put_file`). -- **r**: Read-only access (includes `read`, `get_file`, `list_files`, `list`, `search_by_vector`, `search_by_text`, and `get_vector`). -- **r+**: Read, write, and create access (includes `read`, `get_file`, `put_file`, `list_files`, `list`, `search_by_vector`, `search_by_text`, `get_vector`, `create`, `commit`, `add_vectors`, and `add_documents`). -- **rw**: Read, write, and create access with file management (includes `read`, `get_file`, `get_vector`, `search_by_vector`, `search_by_text`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, and `remove_vectors`). -- **rw+**: Read, write, create, and manage access (includes `read`, `get_file`, `get_vector`, `search_by_vector`, `search_by_text`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, `remove_vectors`, and `create`). -- **\***: Full access to all operations (includes `read`, `get_file`, `get_vector`, `search_by_vector`, `search_by_text`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, `remove_vectors`, `create`, and `reset_stats`). +- **r**: Read-only access (includes `read`, `get_file`, `list_files`, `list`, `search_vectors`, and `get_vector`). +- **r+**: Read, write, and create access (includes `read`, `get_file`, `put_file`, `list_files`, `list`, `search_vectors`, `get_vector`, `create`, `commit`, `add_vectors`, and `add_documents`). +- **rw**: Read, write, and create access with file management (includes `read`, `get_file`, `get_vector`, `search_vectors`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, and `remove_vectors`). +- **rw+**: Read, write, create, and manage access (includes `read`, `get_file`, `get_vector`, `search_vectors`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, `remove_vectors`, and `create`). +- **\***: Full access to all operations (includes `read`, `get_file`, `get_vector`, `search_vectors`, `list_files`, `list_vectors`, `list`, `edit`, `commit`, `put_file`, `add_vectors`, `add_documents`, `remove_file`, `remove_vectors`, `create`, and `reset_stats`). **Shortcut Permission Notation:** diff --git a/hypha/VERSION b/hypha/VERSION index 16187305..99b962fd 100644 --- a/hypha/VERSION +++ b/hypha/VERSION @@ -1,3 +1,3 @@ { - "version": "0.20.39.post19" + "version": "0.20.39.post20" } diff --git a/hypha/artifact.py b/hypha/artifact.py index f052b5ce..e023cd66 100644 --- a/hypha/artifact.py +++ b/hypha/artifact.py @@ -3,12 +3,10 @@ import sys import uuid_utils as uuid import random -import numpy as np import re import json from io import BytesIO import zipfile -import asyncio from sqlalchemy import ( event, Column, @@ -53,11 +51,11 @@ Artifact, CollectionArtifact, ) +from hypha.vectors import VectorSearchEngine from hypha_rpc.utils import ObjectProxy from jsonschema import validate -from typing import Union, List from sqlmodel import SQLModel, Field, Relationship, UniqueConstraint -from typing import Optional, List +from typing import Optional, Union, List, Any # Logger setup logging.basicConfig(stream=sys.stdout) @@ -176,9 +174,8 @@ def __init__( self.workspace_bucket = workspace_bucket self.store = store self._cache = store.get_redis_cache() - self._vectordb_client = self.store.get_vectordb_client() self._openai_client = self.store.get_openai_client() - self._cache_dir = self.store.get_cache_dir() + self._vector_engine = VectorSearchEngine(store) router = APIRouter() self._artifacts_dir = artifacts_dir @@ -914,8 +911,7 @@ def _expand_permission(self, permission): "get_file", "list_files", "list", - "search_by_vector", - "search_by_text", + "search_vectors", "get_vector", ], "r+": [ @@ -924,8 +920,7 @@ def _expand_permission(self, permission): "put_file", "list_files", "list", - "search_by_vector", - "search_by_text", + "search_vectors", "get_vector", "create", "commit", @@ -936,8 +931,7 @@ def _expand_permission(self, permission): "read", "get_file", "get_vector", - "search_by_vector", - "search_by_text", + "search_vectors", "list_files", "list_vectors", "list", @@ -953,8 +947,7 @@ def _expand_permission(self, permission): "read", "get_file", "get_vector", - "search_by_vector", - "search_by_text", + "search_vectors", "list_files", "list_vectors", "list", @@ -971,8 +964,7 @@ def _expand_permission(self, permission): "read", "get_file", "get_vector", - "search_by_vector", - "search_by_text", + "search_vectors", "list_files", "list_vectors", "list", @@ -1032,8 +1024,7 @@ async def _get_artifact_with_permission( "get_file": UserPermission.read, "list_files": UserPermission.read, "list_vectors": UserPermission.read, - "search_by_text": UserPermission.read, - "search_by_vector": UserPermission.read, + "search_vectors": UserPermission.read, "create": UserPermission.read_write, "edit": UserPermission.read_write, "commit": UserPermission.read_write, @@ -1558,18 +1549,9 @@ async def create( else: session.add(new_artifact) if new_artifact.type == "vector-collection": - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." - from qdrant_client.models import Distance, VectorParams - - vectors_config = config.get("vectors_config", {}) - await self._vectordb_client.create_collection( - collection_name=f"{new_artifact.workspace}^{new_artifact.alias}", - vectors_config=VectorParams( - size=vectors_config.get("size", 128), - distance=Distance(vectors_config.get("distance", "Cosine")), - ), + await self._vector_engine.create_collection( + f"{new_artifact.workspace}^{new_artifact.alias}", + vectors_config=config.get("vectors_config", {}), ) await session.commit() await self._save_version_to_s3( @@ -1770,14 +1752,13 @@ async def read( child_count = result.scalar() artifact_data["config"] = artifact_data.get("config", {}) artifact_data["config"]["child_count"] = child_count - elif artifact.type == "vector-collection" and self._vectordb_client: + elif artifact.type == "vector-collection": artifact_data["config"] = artifact_data.get("config", {}) - artifact_data["config"]["vector_count"] = ( - await self._vectordb_client.count( - collection_name=f"{artifact.workspace}^{artifact.alias}" - ) - ).count - + artifact_data["config"][ + "vector_count" + ] = await self._vector_engine.count( + f"{artifact.workspace}^{artifact.alias}" + ) if not silent: await session.commit() @@ -1926,11 +1907,8 @@ async def delete( ) if artifact.type == "vector-collection": - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." - await self._vectordb_client.delete_collection( - collection_name=f"{artifact.workspace}^{artifact.alias}" + await self._vector_engine.delete_collection( + f"{artifact.workspace}^{artifact.alias}" ) s3_config = self._get_s3_config(artifact, parent_artifact) @@ -2002,66 +1980,17 @@ async def add_vectors( assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." + assert artifact.manifest, "Artifact must be committed before upserting." - assert isinstance( - vectors, list - ), "Vectors must be a list of dictionaries." - assert all( - isinstance(v, dict) for v in vectors - ), "Vectors must be a list of dictionaries." - from qdrant_client.models import PointStruct - - _points = [] - for p in vectors: - p["id"] = p.get("id") or str(uuid.uuid4()) - _points.append(PointStruct(**p)) - await self._vectordb_client.upsert( - collection_name=f"{artifact.workspace}^{artifact.alias}", - points=_points, + await self._vector_engine.add_vectors( + f"{artifact.workspace}^{artifact.alias}", vectors ) - # TODO: Update file_count - logger.info(f"Upserted vectors to artifact with ID: {artifact_id}") + logger.info(f"Added vectors to artifact with ID: {artifact_id}") except Exception as e: raise e finally: await session.close() - async def _embed_texts(self, config, texts): - embedding_model = config.get("embedding_model") # "text-embedding-3-small" - assert ( - embedding_model - ), "Embedding model must be provided, e.g. 'fastembed:BAAI/bge-small-en-v1.5', 'openai:text-embedding-3-small' for openai embeddings." - if embedding_model.startswith("fastembed"): - from fastembed import TextEmbedding - - assert ":" in embedding_model, "Embedding model must be provided." - model_name = embedding_model.split(":")[-1] - embedding_model = TextEmbedding( - model_name=model_name, cache_dir=self._cache_dir - ) - loop = asyncio.get_event_loop() - embeddings = list( - await loop.run_in_executor(None, embedding_model.embed, texts) - ) - elif embedding_model.startswith("openai"): - assert ( - self._openai_client - ), "The server is not configured to use an OpenAI client." - assert ":" in embedding_model, "Embedding model must be provided." - embedding_model = embedding_model.split(":")[-1] - result = await self._openai_client.embeddings.create( - input=texts, model=embedding_model - ) - embeddings = [data.embedding for data in result.data] - else: - raise ValueError( - f"Unsupported embedding model: {embedding_model}, supported models: 'fastembed:*', 'openai:*'" - ) - return embeddings - async def add_documents( self, artifact_id: str, @@ -2081,32 +2010,21 @@ async def add_documents( assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - texts = [doc["text"] for doc in documents] - embeddings = await self._embed_texts(artifact.config, texts) - from qdrant_client.models import PointStruct - - points = [ - PointStruct( - id=doc.get("id") or str(uuid.uuid4()), - vector=embedding, - payload=doc, - ) - for embedding, doc in zip(embeddings, documents) - ] - await self._vectordb_client.upsert( - collection_name=f"{artifact.workspace}^{artifact.alias}", - points=points, + embedding_model = artifact.config.get("embedding_model") + await self._vector_engine.add_documents( + f"{artifact.workspace}^{artifact.alias}", documents, embedding_model ) - logger.info(f"Upserted documents to artifact with ID: {artifact_id}") + logger.info(f"Added documents to artifact with ID: {artifact_id}") except Exception as e: raise e finally: await session.close() - async def search_by_vector( + async def search_vectors( self, artifact_id: str, - query_vector, + query_text: str = None, + query_vector: Any = None, query_filter: dict = None, offset: int = 0, limit: int = 10, @@ -2120,90 +2038,25 @@ async def search_by_vector( try: async with session.begin(): artifact, _ = await self._get_artifact_with_permission( - user_info, artifact_id, "search_by_vector", session + user_info, artifact_id, "search_vectors", session ) assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - # if it's a numpy array, convert it to a list - if isinstance(query_vector, np.ndarray): - query_vector = query_vector.tolist() - from qdrant_client.models import Filter - - if query_filter: - query_filter = Filter.model_validate(query_filter) - search_results = await self._vectordb_client.search( - collection_name=f"{artifact.workspace}^{artifact.alias}", - query_vector=query_vector, - query_filter=query_filter, - limit=limit, - offset=offset, - with_payload=with_payload, - with_vectors=with_vectors, - ) - if pagination: - count = await self._vectordb_client.count( - collection_name=f"{artifact.workspace}^{artifact.alias}" - ) - return { - "total": count.count, - "items": search_results, - "offset": offset, - "limit": limit, - } - return search_results - except Exception as e: - raise e - finally: - await session.close() - async def search_by_text( - self, - artifact_id: str, - query: str, - query_filter: dict = None, - offset: int = 0, - limit: int = 10, - with_payload: bool = True, - with_vectors: bool = False, - pagination: bool = False, - context: dict = None, - ): - user_info = UserInfo.model_validate(context["user"]) - session = await self._get_session() - try: - async with session.begin(): - artifact, _ = await self._get_artifact_with_permission( - user_info, artifact_id, "search_by_text", session - ) - assert ( - artifact.type == "vector-collection" - ), "Artifact must be a vector collection." - (query_vector,) = await self._embed_texts(artifact.config, [query]) - from qdrant_client.models import Filter - - if query_filter: - query_filter = Filter.model_validate(query_filter) - search_results = await self._vectordb_client.search( - collection_name=f"{artifact.workspace}^{artifact.alias}", + embedding_model = artifact.config.get("embedding_model") + return await self._vector_engine.search_vectors( + f"{artifact.workspace}^{artifact.alias}", + embedding_model=embedding_model, + query_text=query_text, query_vector=query_vector, query_filter=query_filter, - limit=limit, offset=offset, + limit=limit, with_payload=with_payload, with_vectors=with_vectors, + pagination=pagination, ) - if pagination: - count = await self._vectordb_client.count( - collection_name=f"{artifact.workspace}^{artifact.alias}" - ) - return { - "total": count.count, - "items": search_results, - "offset": offset, - "limit": limit, - } - return search_results except Exception as e: raise e finally: @@ -2225,12 +2078,8 @@ async def remove_vectors( assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." - await self._vectordb_client.delete( - collection_name=f"{artifact.workspace}^{artifact.alias}", - points_selector=ids, + await self._vector_engine.remove_vectors( + f"{artifact.workspace}^{artifact.alias}", ids ) logger.info(f"Removed vectors from artifact with ID: {artifact_id}") except Exception as e: @@ -2254,16 +2103,9 @@ async def get_vector( assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." - points = await self._vectordb_client.retrieve( - collection_name=f"{artifact.workspace}^{artifact.alias}", - ids=[id], - with_payload=True, - with_vectors=True, + return await self._vector_engine.get_vector( + f"{artifact.workspace}^{artifact.alias}", id ) - return points[0] except Exception as e: raise e finally: @@ -2290,23 +2132,15 @@ async def list_vectors( assert ( artifact.type == "vector-collection" ), "Artifact must be a vector collection." - assert ( - self._vectordb_client - ), "The server is not configured to use a VectorDB client." - from qdrant_client.models import Filter - - if query_filter: - query_filter = Filter.model_validate(query_filter) - points, _ = await self._vectordb_client.scroll( - collection_name=f"{artifact.workspace}^{artifact.alias}", - scroll_filter=query_filter, - limit=limit, + return await self._vector_engine.list_vectors( + f"{artifact.workspace}^{artifact.alias}", + query_filter=query_filter, offset=offset, + limit=limit, order_by=order_by, with_payload=with_payload, with_vectors=with_vectors, ) - return points except Exception as e: raise e @@ -2908,8 +2742,7 @@ def get_artifact_service(self): "list_files": self.list_files, "add_vectors": self.add_vectors, "add_documents": self.add_documents, - "search_by_vector": self.search_by_vector, - "search_by_text": self.search_by_text, + "search_vectors": self.search_vectors, "remove_vectors": self.remove_vectors, "get_vector": self.get_vector, "list_vectors": self.list_vectors, diff --git a/hypha/vectors.py b/hypha/vectors.py new file mode 100644 index 00000000..e8ffa836 --- /dev/null +++ b/hypha/vectors.py @@ -0,0 +1,213 @@ +import asyncio +from qdrant_client.models import PointStruct +from qdrant_client.models import Filter +from qdrant_client.models import Distance, VectorParams +import uuid +import numpy as np +from typing import Any +import logging +import sys + +# Logger setup +logging.basicConfig(stream=sys.stdout) +logger = logging.getLogger("vectors") +logger.setLevel(logging.INFO) + + +class VectorSearchEngine: + def __init__(self, store): + self.store = store + self._vectordb_client = self.store.get_vectordb_client() + self._cache_dir = self.store.get_cache_dir() + + def _ensure_client(self): + assert ( + self._vectordb_client + ), "The server is not configured to use a VectorDB client." + + async def create_collection(self, collection_name: str, vectors_config: dict): + self._ensure_client() + await self._vectordb_client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=vectors_config.get("size", 128), + distance=Distance(vectors_config.get("distance", "Cosine")), + ), + ) + logger.info(f"Collection {collection_name} created.") + + async def count(self, collection_name: str): + self._ensure_client() + return ( + await self._vectordb_client.count(collection_name=collection_name) + ).count + + async def delete_collection(self, collection_name: str): + self._ensure_client() + await self._vectordb_client.delete_collection(collection_name=collection_name) + logger.info(f"Collection {collection_name} deleted.") + + async def _embed_texts(self, embedding_model, texts: list): + self._ensure_client() + assert ( + embedding_model + ), "Embedding model must be provided, e.g. 'fastembed:BAAI/bge-small-en-v1.5', 'openai:text-embedding-3-small' for openai embeddings." + if embedding_model.startswith("fastembed"): + from fastembed import TextEmbedding + + assert ":" in embedding_model, "Embedding model must be provided." + model_name = embedding_model.split(":")[-1] + embedding_model = TextEmbedding( + model_name=model_name, cache_dir=self._cache_dir + ) + loop = asyncio.get_event_loop() + embeddings = list( + await loop.run_in_executor(None, embedding_model.embed, texts) + ) + elif embedding_model.startswith("openai"): + assert ( + self._openai_client + ), "The server is not configured to use an OpenAI client." + assert ":" in embedding_model, "Embedding model must be provided." + embedding_model = embedding_model.split(":")[-1] + result = await self._openai_client.embeddings.create( + input=texts, model=embedding_model + ) + embeddings = [data.embedding for data in result.data] + else: + raise ValueError( + f"Unsupported embedding model: {embedding_model}, supported models: 'fastembed:*', 'openai:*'" + ) + return embeddings + + async def add_vectors(self, collection_name: str, vectors: list): + self._ensure_client() + assert isinstance(vectors, list), "Vectors must be a list of dictionaries." + assert all( + isinstance(v, dict) for v in vectors + ), "Vectors must be a list of dictionaries." + + _points = [] + for p in vectors: + p["id"] = p.get("id") or str(uuid.uuid4()) + _points.append(PointStruct(**p)) + await self._vectordb_client.upsert( + collection_name=collection_name, + points=_points, + ) + logger.info(f"Added {len(vectors)} vectors to collection {collection_name}.") + + async def add_documents(self, collection_name, documents, embedding_model): + self._ensure_client() + texts = [doc["text"] for doc in documents] + embeddings = await self._embed_texts(embedding_model, texts) + points = [ + PointStruct( + id=doc.get("id") or str(uuid.uuid4()), + vector=embedding, + payload=doc, + ) + for embedding, doc in zip(embeddings, documents) + ] + await self._vectordb_client.upsert( + collection_name=collection_name, + points=points, + ) + logger.info( + f"Added {len(documents)} documents to collection {collection_name}." + ) + + async def search_vectors( + self, + collection_name: str, + embedding_model: str, + query_text: str, + query_vector: Any, + query_filter: dict = None, + offset: int = 0, + limit: int = 10, + with_payload: bool = True, + with_vectors: bool = False, + pagination: bool = False, + ): + self._ensure_client() + # if it's a numpy array, convert it to a list + if isinstance(query_vector, np.ndarray): + query_vector = query_vector.tolist() + from qdrant_client.models import Filter + + if query_filter: + query_filter = Filter.model_validate(query_filter) + + if query_text: + assert ( + not query_vector + ), "Either query_text or query_vector must be provided." + embeddings = await self._embed_texts(embedding_model, [query_text]) + query_vector = embeddings[0] + + search_results = await self._vectordb_client.search( + collection_name=collection_name, + query_vector=query_vector, + query_filter=query_filter, + limit=limit, + offset=offset, + with_payload=with_payload, + with_vectors=with_vectors, + ) + if pagination: + count = await self._vectordb_client.count(collection_name=collection_name) + return { + "total": count.count, + "items": search_results, + "offset": offset, + "limit": limit, + } + logger.info(f"Performed semantic search in collection {collection_name}.") + return search_results + + async def remove_vectors(self, collection_name: str, ids: list): + self._ensure_client() + await self._vectordb_client.delete( + collection_name=collection_name, + points_selector=ids, + ) + logger.info(f"Removed {len(ids)} vectors from collection {collection_name}.") + + async def get_vector(self, collection_name: str, id: str): + self._ensure_client() + points = await self._vectordb_client.retrieve( + collection_name=collection_name, + ids=[id], + with_payload=True, + with_vectors=True, + ) + if not points: + raise ValueError(f"Vector {id} not found in collection {collection_name}.") + logger.info(f"Retrieved vector {id} from collection {collection_name}.") + return points[0] + + async def list_vectors( + self, + collection_name: str, + query_filter: dict = None, + offset: int = 0, + limit: int = 10, + order_by: str = None, + with_payload: bool = True, + with_vectors: bool = False, + ): + self._ensure_client() + if query_filter: + query_filter = Filter.model_validate(query_filter) + points, _ = await self._vectordb_client.scroll( + collection_name=collection_name, + scroll_filter=query_filter, + limit=limit, + offset=offset, + order_by=order_by, + with_payload=with_payload, + with_vectors=with_vectors, + ) + logger.info(f"Listed vectors in collection {collection_name}.") + return points diff --git a/tests/test_artifact.py b/tests/test_artifact.py index 7cf03e6f..46502593 100644 --- a/tests/test_artifact.py +++ b/tests/test_artifact.py @@ -2,8 +2,6 @@ import pytest import requests import os -import numpy as np -import random from hypha_rpc import connect_to_server from io import BytesIO from zipfile import ZipFile @@ -15,163 +13,6 @@ pytestmark = pytest.mark.asyncio -async def test_artifact_vector_collection( - minio_server, fastapi_server, test_user_token -): - """Test vector-related functions within a vector-collection artifact.""" - - # Connect to the server and set up the artifact manager - api = await connect_to_server( - { - "name": "test deploy client", - "server_url": SERVER_URL, - "token": test_user_token, - } - ) - artifact_manager = await api.get_service("public/artifact-manager") - - # Create a vector-collection artifact - vector_collection_manifest = { - "name": "vector-collection", - "description": "A test vector collection", - } - vector_collection_config = { - "vectors_config": { - "size": 384, - "distance": "Cosine", - }, - "embedding_model": "fastembed:BAAI/bge-small-en-v1.5", - } - vector_collection = await artifact_manager.create( - type="vector-collection", - manifest=vector_collection_manifest, - config=vector_collection_config, - ) - # Add vectors to the collection - vectors = [ - { - "vector": [random.random() for _ in range(384)], - "payload": { - "text": "This is a test document.", - "label": "doc1", - "rand_number": random.randint(0, 10), - }, - }, - { - "vector": np.random.rand(384), - "payload": { - "text": "Another document.", - "label": "doc2", - "rand_number": random.randint(0, 10), - }, - }, - { - "vector": np.random.rand(384), - "payload": { - "text": "Yet another document.", - "label": "doc3", - "rand_number": random.randint(0, 10), - }, - }, - ] - await artifact_manager.add_vectors( - artifact_id=vector_collection.id, - vectors=vectors, - ) - - vc = await artifact_manager.read(artifact_id=vector_collection.id) - assert vc["config"]["vector_count"] == 3 - - # Search for vectors by query vector - query_vector = [random.random() for _ in range(384)] - search_results = await artifact_manager.search_by_vector( - artifact_id=vector_collection.id, - query_vector=query_vector, - limit=2, - ) - assert len(search_results) <= 2 - - results = await artifact_manager.search_by_vector( - artifact_id=vector_collection.id, - query_vector=query_vector, - limit=2, - pagination=True, - ) - assert results["total"] == 3 - - query_filter = { - "should": None, - "min_should": None, - "must": [ - { - "key": "rand_number", - "match": None, - "range": {"lt": None, "gt": None, "gte": 3.0, "lte": None}, - "geo_bounding_box": None, - "geo_radius": None, - "geo_polygon": None, - "values_count": None, - } - ], - "must_not": None, - } - - search_results = await artifact_manager.search_by_vector( - artifact_id=vector_collection.id, - query_filter=query_filter, - query_vector=np.random.rand(384), - limit=2, - ) - assert len(search_results) <= 2 - - # Search for vectors by text - documents = [ - {"text": "This is a test document.", "label": "doc1"}, - {"text": "Another test document.", "label": "doc2"}, - ] - await artifact_manager.add_documents( - artifact_id=vector_collection.id, - documents=documents, - ) - text_query = "test document" - text_search_results = await artifact_manager.search_by_text( - artifact_id=vector_collection.id, - query=text_query, - limit=2, - ) - assert len(text_search_results) <= 2 - - # Retrieve a specific vector - retrieved_vector = await artifact_manager.get_vector( - artifact_id=vector_collection.id, - id=text_search_results[0]["id"], - ) - assert retrieved_vector.id == text_search_results[0]["id"] - - # List vectors in the collection - vector_list = await artifact_manager.list_vectors( - artifact_id=vector_collection.id, - offset=0, - limit=10, - ) - assert len(vector_list) > 0 - - # Remove a vector from the collection - await artifact_manager.remove_vectors( - artifact_id=vector_collection.id, - ids=[vector_list[0]["id"]], - ) - remaining_vectors = await artifact_manager.list_vectors( - artifact_id=vector_collection.id, - offset=0, - limit=10, - ) - assert all(v["id"] != vector_list[0]["id"] for v in remaining_vectors) - - # Clean up by deleting the vector collection - await artifact_manager.delete(artifact_id=vector_collection.id) - - async def test_sqlite_create_and_search_artifacts( minio_server, fastapi_server_sqlite, test_user_token ): diff --git a/tests/test_redis.py b/tests/test_redis.py index c3dbd4fb..48773333 100644 --- a/tests/test_redis.py +++ b/tests/test_redis.py @@ -39,8 +39,8 @@ async def test_redis_store(redis_store): assert find_item(wss, "name", "test") api = await redis_store.connect_to_workspace("test", client_id="test-app-99") - clients = await api.list_clients() - assert find_item(clients, "id", "test/test-app-99") + # clients = await api.list_clients() + # assert find_item(clients, "id", "test/test-app-99") await api.log("hello") services = await api.list_services() assert len(services) == 1 @@ -72,9 +72,9 @@ def echo(data): await api.register_service(interface) wm = await redis_store.connect_to_workspace("test", client_id="test-app-22") - clients = await wm.list_clients() - assert find_item(clients, "id", "test/test-app-22") - assert find_item(clients, "id", "test/test-app-99") + # clients = await wm.list_clients() + # assert find_item(clients, "id", "test/test-app-22") + # assert find_item(clients, "id", "test/test-app-99") rpc = wm.rpc services = await wm.list_services() service = await rpc.get_remote_service("test-app-99:test-service") diff --git a/tests/test_vectors.py b/tests/test_vectors.py new file mode 100644 index 00000000..3593235b --- /dev/null +++ b/tests/test_vectors.py @@ -0,0 +1,168 @@ +"""Test Vector Search.""" +import pytest +import numpy as np +import random +from hypha_rpc import connect_to_server + + +from . import SERVER_URL, find_item + +# All test coroutines will be treated as marked. +pytestmark = pytest.mark.asyncio + + +async def test_artifact_vector_collection( + minio_server, fastapi_server, test_user_token +): + """Test vector-related functions within a vector-collection artifact.""" + + # Connect to the server and set up the artifact manager + api = await connect_to_server( + { + "name": "test deploy client", + "server_url": SERVER_URL, + "token": test_user_token, + } + ) + artifact_manager = await api.get_service("public/artifact-manager") + + # Create a vector-collection artifact + vector_collection_manifest = { + "name": "vector-collection", + "description": "A test vector collection", + } + vector_collection_config = { + "vectors_config": { + "size": 384, + "distance": "Cosine", + }, + "embedding_model": "fastembed:BAAI/bge-small-en-v1.5", + } + vector_collection = await artifact_manager.create( + type="vector-collection", + manifest=vector_collection_manifest, + config=vector_collection_config, + ) + # Add vectors to the collection + vectors = [ + { + "vector": [random.random() for _ in range(384)], + "payload": { + "text": "This is a test document.", + "label": "doc1", + "rand_number": random.randint(0, 10), + }, + }, + { + "vector": np.random.rand(384), + "payload": { + "text": "Another document.", + "label": "doc2", + "rand_number": random.randint(0, 10), + }, + }, + { + "vector": np.random.rand(384), + "payload": { + "text": "Yet another document.", + "label": "doc3", + "rand_number": random.randint(0, 10), + }, + }, + ] + await artifact_manager.add_vectors( + artifact_id=vector_collection.id, + vectors=vectors, + ) + + vc = await artifact_manager.read(artifact_id=vector_collection.id) + assert vc["config"]["vector_count"] == 3 + + # Search for vectors by query vector + query_vector = [random.random() for _ in range(384)] + search_results = await artifact_manager.search_vectors( + artifact_id=vector_collection.id, + query_vector=query_vector, + limit=2, + ) + assert len(search_results) <= 2 + + results = await artifact_manager.search_vectors( + artifact_id=vector_collection.id, + query_vector=query_vector, + limit=2, + pagination=True, + ) + assert results["total"] == 3 + + query_filter = { + "should": None, + "min_should": None, + "must": [ + { + "key": "rand_number", + "match": None, + "range": {"lt": None, "gt": None, "gte": 3.0, "lte": None}, + "geo_bounding_box": None, + "geo_radius": None, + "geo_polygon": None, + "values_count": None, + } + ], + "must_not": None, + } + + search_results = await artifact_manager.search_vectors( + artifact_id=vector_collection.id, + query_filter=query_filter, + query_vector=np.random.rand(384), + limit=2, + ) + assert len(search_results) <= 2 + + # Search for vectors by text + documents = [ + {"text": "This is a test document.", "label": "doc1"}, + {"text": "Another test document.", "label": "doc2"}, + ] + await artifact_manager.add_documents( + artifact_id=vector_collection.id, + documents=documents, + ) + text_query = "test document" + text_search_results = await artifact_manager.search_vectors( + artifact_id=vector_collection.id, + query_text=text_query, + limit=2, + ) + assert len(text_search_results) <= 2 + + # Retrieve a specific vector + retrieved_vector = await artifact_manager.get_vector( + artifact_id=vector_collection.id, + id=text_search_results[0]["id"], + ) + assert retrieved_vector.id == text_search_results[0]["id"] + + # List vectors in the collection + vector_list = await artifact_manager.list_vectors( + artifact_id=vector_collection.id, + offset=0, + limit=10, + ) + assert len(vector_list) > 0 + + # Remove a vector from the collection + await artifact_manager.remove_vectors( + artifact_id=vector_collection.id, + ids=[vector_list[0]["id"]], + ) + remaining_vectors = await artifact_manager.list_vectors( + artifact_id=vector_collection.id, + offset=0, + limit=10, + ) + assert all(v["id"] != vector_list[0]["id"] for v in remaining_vectors) + + # Clean up by deleting the vector collection + await artifact_manager.delete(artifact_id=vector_collection.id)