Function ensure_collection

tests/test_hybrid.py:722–766

Create or recreate the collection with proper configuration

(name: str)

    return results

def ensure_collection(name: str):
    """Create or recreate the collection with proper configuration"""
    try:
        print(f"Creating collection {name}...")
        collection = client.create_collection(
            name=name,
            dimension=DIMENSION,
            tf_idf_options={"enabled": True},
        )
        print("Collection created successfully")

        print("Creating dense HNSW index...")
        collection.create_index(distance_metric="cosine")

        print("Creating TF-IDF index...")
        collection.create_tf_idf_index(name=name, k1=K1, b=B)

        # Add hybrid search methods to collection via monkey patching
        collection.search.hybrid = hybrid.__get__(collection.search)
        collection.search.batch_hybrid = batch_hybrid.__get__(collection.search)

        return collection

    except Exception as e:
        print(
            f"Collection {name} may already exist, attempting to delete and recreate: {e}"
        )
        try:
            existing = client.get_collection(name)
            existing.delete()
            time.sleep(2)
        except Exception as delete_error:
            print(f"Error deleting collection: {delete_error}")

        collection = client.create_collection(
            name=name, dimension=DIMENSION, tf_idf_options={"enabled": True}
        )
        collection.create_index(distance_metric="cosine")
        collection.create_tf_idf_index(name=name, k1=K1, b=B)

        # Add hybrid search methods to collection via monkey patching
        collection.search.hybrid = hybrid.__get__(collection.search)
        collection.search.batch_hybrid = batch_hybrid.__get__(collection.search)

        return collection

def upsert_hybrid_vectors(collection, corpus_ids, corpus_embeddings, corpus):
    """Upsert both dense embeddings and raw text for hybrid search"""