"""Prerequisites:    pip install google-genai pymilvusSet environment variable:    export GOOGLE_API_KEY="your-api-key""""import osimport structimport numpy as npfrom google import genaifrom google.genai import typesfrom pymilvus import MilvusClient, DataType# ── Config ───────────────────────────────────────────────────────────────COLLECTION_NAME = "gemini_multimodal_demo"MILVUS_URI = "http://localhost:19530"  # Change to your Milvus addressDIM = 3072  # gemini-embedding-2-preview output dimensionGEMINI_MODEL = "gemini-embedding-2-preview"# ── Initialize clients ──────────────────────────────────────────────────gemini_client = genai.Client()  # Uses GOOGLE_API_KEY env varmilvus_client = MilvusClient(MILVUS_URI)# ── Helper: generate embedding ──────────────────────────────────────────def embed_texts(texts: list[str], task_type: str = "SEMANTIC_SIMILARITY") -> list[list[float]]:    """Embed a list of text strings."""    result = gemini_client.models.embed_content(        model=GEMINI_MODEL,        contents=texts,        config=types.EmbedContentConfig(task_type=task_type),    )    return [e.values for e in result.embeddings]def embed_image(image_path: str) -> list[float]:    """Embed an image file."""    with open(image_path, "rb") as f:        image_bytes = f.read()    mime = "image/png" if image_path.endswith(".png") else "image/jpeg"    result = gemini_client.models.embed_content(        model=GEMINI_MODEL,        contents=types.Part.from_bytes(data=image_bytes, mime_type=mime),    )    return result.embeddings[0].valuesdef embed_audio(audio_path: str) -> list[float]:    """Embed an audio file."""    with open(audio_path, "rb") as f:        audio_bytes = f.read()    mime_map = {".mp3": "audio/mpeg", ".wav": "audio/wav", ".flac": "audio/flac"}    ext = os.path.splitext(audio_path)[1].lower()    mime = mime_map.get(ext, "audio/mpeg")    result = gemini_client.models.embed_content(        model=GEMINI_MODEL,        contents=types.Part.from_bytes(data=audio_bytes, mime_type=mime),    )    return result.embeddings[0].values# ── 1. Create Milvus collection ─────────────────────────────────────────print("=== Creating collection ===")if milvus_client.has_collection(COLLECTION_NAME):    milvus_client.drop_collection(COLLECTION_NAME)schema = milvus_client.create_schema()schema.add_field("id", DataType.INT64, is_primary=True, auto_id=True)schema.add_field("content", DataType.VARCHAR, max_length=2000)   # description of the contentschema.add_field("modality", DataType.VARCHAR, max_length=20)    # "text", "image", "audio"schema.add_field("vector", DataType.FLOAT_VECTOR, dim=DIM)index_params = milvus_client.prepare_index_params()index_params.add_index(    field_name="vector",    index_type="AUTOINDEX",    metric_type="COSINE",)milvus_client.create_collection(    COLLECTION_NAME,    schema=schema,    index_params=index_params,    consistency_level="Strong",)print(f"Collection '{COLLECTION_NAME}' created (dim={DIM}, metric=COSINE)")# ── 2. Insert text embeddings ───────────────────────────────────────────print("\n=== Inserting text embeddings ===")documents = [    "Artificial intelligence was founded as an academic discipline in 1956.",    "The Mona Lisa is a half-length portrait painting by Leonardo da Vinci.",    "Beethoven's Symphony No. 9 premiered in Vienna on May 7, 1824.",    "The Great Wall of China stretches over 13,000 miles across northern China.",    "Jazz music originated in the African-American communities of New Orleans.",    "The Hubble Space Telescope was launched into orbit on April 24, 1990.",    "Vincent van Gogh painted The Starry Night while in an asylum in Saint-Rémy.",    "Machine learning is a subset of AI focused on learning from data.",]text_vectors = embed_texts(documents)text_rows = [    {"content": doc, "modality": "text", "vector": vec}    for doc, vec in zip(documents, text_vectors)]milvus_client.insert(COLLECTION_NAME, text_rows)print(f"Inserted {len(text_rows)} text documents")# ── 3. (Optional) Insert image embeddings ───────────────────────────────# Uncomment and provide real image paths to test multimodal search## image_files = [#     ("photo of the Mona Lisa painting", "mona_lisa.jpg"),#     ("satellite photo of the Great Wall of China", "great_wall.png"),# ]# for desc, path in image_files:#     if os.path.exists(path):#         vec = embed_image(path)#         milvus_client.insert(COLLECTION_NAME, [#             {"content": desc, "modality": "image", "vector": vec}#         ])#         print(f"Inserted image: {desc}")# ── 4. (Optional) Insert audio embeddings ───────────────────────────────# Uncomment and provide real audio paths to test multimodal search## audio_files = [#     ("Beethoven Symphony No.9 excerpt", "beethoven_9.mp3"),#     ("jazz piano improvisation", "jazz_piano.mp3"),# ]# for desc, path in audio_files:#     if os.path.exists(path):#         vec = embed_audio(path)#         milvus_client.insert(COLLECTION_NAME, [#             {"content": desc, "modality": "audio", "vector": vec}#         ])#         print(f"Inserted audio: {desc}")# ── 5. Search ────────────────────────────────────────────────────────────print("\n=== Searching ===")queries = [    "history of artificial intelligence",    "famous Renaissance paintings",    "classical music concerts",]query_vectors = embed_texts(queries, task_type="SEMANTIC_SIMILARITY")for query_text, query_vec in zip(queries, query_vectors):    results = milvus_client.search(        COLLECTION_NAME,        data=[query_vec],        limit=3,        output_fields=["content", "modality"],        search_params={"metric_type": "COSINE"},    )    print(f"\nQuery: '{query_text}'")    for hits in results:        for rank, hit in enumerate(hits, 1):            print(f"  [{rank}] (score={hit['distance']:.4f}, modality={hit['entity']['modality']}) "                  f"{hit['entity']['content'][:80]}")# ── 6. Cross-modal search example (image query -> text results) ─────────# Uncomment to search text collection using an image as query## print("\n=== Cross-modal search: image -> text ===")# query_image_vec = embed_image("query_image.jpg")# results = milvus_client.search(#     COLLECTION_NAME,#     data=[query_image_vec],#     limit=3,#     output_fields=["content", "modality"],#     search_params={"metric_type": "COSINE"},# )# for hits in results:#     for rank, hit in enumerate(hits, 1):#         print(f"  [{rank}] (score={hit['distance']:.4f}) {hit['entity']['content'][:80]}")# ── Cleanup ──────────────────────────────────────────────────────────────# milvus_client.drop_collection(COLLECTION_NAME)# print(f"\nCollection '{COLLECTION_NAME}' dropped")print("\nDone!")