Fix Milvus collection initialization and Towhee connection

app.py

@@ -243,20 +243,74 @@ else:
     print(f"Connecting to Milvus Host: {host_milvus}")
     connections.connect(host=host_milvus, port='19530')
 
+# Check if collection exists, if not create and populate it
+if not utility.has_collection(COLLECTION_NAME):
+    print(f"Collection {COLLECTION_NAME} not found. Creating and populating...")
+    
+    # 1. Define Schema
+    fields = [
+        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
+        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768) # DPR uses 768 dims
+    ]
+    schema = CollectionSchema(fields, "Medical Chatbot QA")
+    collection = Collection(COLLECTION_NAME, schema)
+    
+    # 2. Generate Embeddings
+    print("Generating embeddings for initial data...")
+    embedding_pipe = (
+        pipe.input('question')
+        .map('question', 'vec', lambda x: x[:500])
+        .map('vec', 'vec', ops.text_embedding.dpr(model_name='facebook/dpr-ctx_encoder-single-nq-base'))
+        .map('vec', 'vec', lambda x: x / np.linalg.norm(x, axis=0))
+        .output('vec')
+    )
+    
+    vectors = []
+    # Process in batches to be safe
+    for q in df['question']:
+        res = embedding_pipe(q)
+        vectors.append(res.get()[0])
+        
+    # 3. Insert Data
+    print("Inserting data into Zilliz...")
+    collection.insert([df['id'].tolist(), vectors])
+    
+    # 4. Create Index
+    print("Creating index...")
+    index_params = {
+        "metric_type": "IP",
+        "index_type": "AUTOINDEX",
+        "params": {}
+    }
+    collection.create_index(field_name="embedding", index_params=index_params)
+    print("Collection setup complete.")
 
 collection = Collection(COLLECTION_NAME)      
-collection.load(replica_number=1)
+collection.load()
 utility.load_state(COLLECTION_NAME)
 utility.loading_progress(COLLECTION_NAME)
 
 max_input_length = 500  # Maximum length allowed by the model
+
+# Configure Towhee Milvus Client arguments based on connection type
+milvus_args = {
+    "collection_name": COLLECTION_NAME,
+    "limit": 1
+}
+if milvus_uri and milvus_token:
+    milvus_args["uri"] = milvus_uri
+    milvus_args["token"] = milvus_token
+else:
+    milvus_args["host"] = host_milvus
+    milvus_args["port"] = '19530'
+
 # Create the combined pipe for question encoding and answer retrieval
 combined_pipe = (
     pipe.input('question')
         .map('question', 'vec', lambda x: x[:max_input_length])  # Truncate the question if longer than 512 tokens
         .map('vec', 'vec', ops.text_embedding.dpr(model_name='facebook/dpr-ctx_encoder-single-nq-base'))
         .map('vec', 'vec', lambda x: x / np.linalg.norm(x, axis=0))
-        .map('vec', 'res', ops.ann_search.milvus_client(host=host_milvus, port='19530', collection_name=COLLECTION_NAME, limit=1))
+        .map('vec', 'res', ops.ann_search.milvus_client(**milvus_args))
         .map('res', 'answer', lambda x: [id_answer[int(i[0])] for i in x])
         .output('question', 'answer')
 )