from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection

# Milvusに接続
connections.connect("default", host="localhost", port="19530")

# フィールドの定義
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536)  # 例として1536次元の埋め込みベクトル
]

# コレクションスキーマの定義
schema = CollectionSchema(fields, "Text collection")

# コレクションの作成
collection = Collection("text_collection", schema)

# テキストデータの準備と挿入
texts = ["This is a sample text.", "Another example text.", "More text data for embedding."]
embeddings = [get_embedding(text) for text in texts]

collection.insert([texts, embeddings])

# インデックスの作成
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128}
}
collection.create_index("embedding", index_params)
collection.load()

# 検索クエリテキストの埋め込みを取得
query_text = "Sample search query."
query_embedding = get_embedding(query_text)

# 検索の実行
search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
results = collection.search([query_embedding], "embedding", search_params, limit=5, output_fields=["text"])

# 検索結果の表示
for result in results[0]:
    print(f"ID: {result.id}, Distance: {result.distance}, Text: {result.entity.get('text')}")
"""
ID: 450921762817052061, Distance: 0.3035614490509033, Text: This is a sample text.
ID: 450921762817052062, Distance: 0.37800854444503784, Text: Another example text.
ID: 450921762817052063, Distance: 0.4625099301338196, Text: More text data for embedding.
"""

import openai
import os
from pymilvus import connections, Collection

# Milvusに接続
connections.connect("default", host="localhost", port="19530")

# 既存のコレクションに接続
collection = Collection("text_collection")

# 新しいテキストデータの準備と挿入
new_texts = ["New sample text.", "Additional text for testing."]
new_embeddings = [get_embedding(text) for text in new_texts]

# 新しいデータの挿入
collection.insert([new_texts, new_embeddings])