Replace towhee milvus-client with direct MilvusClient search for Zilliz Serverless

app.py

@@ -317,29 +317,46 @@ 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
-}
+# Initialize MilvusClient for search (compatible with Zilliz Serverless)
+from pymilvus import MilvusClient as SearchClient
 if milvus_uri and milvus_token:
-    milvus_args["uri"] = milvus_uri
-    milvus_args["token"] = milvus_token
+    search_client = SearchClient(uri=milvus_uri, token=milvus_token)
 else:
-    milvus_args["host"] = host_milvus
-    milvus_args["port"] = '19530'
+    search_client = SearchClient(uri=f"http://{host_milvus}:19530")
 
-# Create the combined pipe for question encoding and answer retrieval
-combined_pipe = (
+# Initialize embedding pipeline (without Milvus search - we'll do that separately)
+embedding_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(**milvus_args))
-        .map('res', 'answer', lambda x: [id_answer[int(i[0])] for i in x])
-        .output('question', 'answer')
+    .map('question', 'vec', lambda x: x[:max_input_length])
+    .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')
 )
 
+def search_similar_questions(question: str) -> list:
+    """Search for similar questions using MilvusClient directly (Zilliz Serverless compatible)."""
+    # Get embedding for the question
+    result = embedding_pipe(question)
+    query_vector = result.get()[0].tolist()
+    
+    # Search using MilvusClient
+    search_results = search_client.search(
+        collection_name=COLLECTION_NAME,
+        data=[query_vector],
+        limit=1,
+        output_fields=["id"]
+    )
+    
+    # Extract answers from results
+    answers = []
+    for hits in search_results:
+        for hit in hits:
+            doc_id = hit['id']
+            if doc_id in id_answer:
+                answers.append(id_answer[doc_id])
+    
+    return answers
+
 # Step 3  - Custom LLM
 from openai import OpenAI
 def generate_stream(prompt, model="mixtral-8x7b"):
@@ -382,10 +399,8 @@ class CustomRetrieverLang(BaseRetriever):
         self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         # Perform the encoding and retrieval for a specific question
-        ans = combined_pipe(query)
-        ans = DataCollection(ans)
-        answer=ans[0]['answer']
-        answer_string = ' '.join(answer)
+        answers = search_similar_questions(query)
+        answer_string = ' '.join(answers) if answers else "No relevant information found."
         return [Document(page_content=answer_string)]   
 # Ensure correct VectorStoreRetriever usage
 retriever = CustomRetrieverLang()