[Bug fix] Bring test_vector_database.py up to date.

kyolee415 · kyolee415 · commit 51a8f38b165c · 2024-02-02T06:28:14.000Z
diff --git a/demos/palm/python/docs-agent/scripts/test_vector_database.py b/demos/palm/python/docs-agent/scripts/test_vector_database.py
@@ -108,8 +108,12 @@ def embed_palm(texts: Documents) -> Embeddings:
         model_name=LOCAL_LLM
     )
 
+embedding_function_gemini_retrieval = embedding_functions.GoogleGenerativeAiEmbeddingFunction(
+    api_key=API_KEY, model_name=EMBEDDING_MODEL,
+    task_type="RETRIEVAL_QUERY")
+
 collection = chroma_client.get_collection(
-    name=COLLECTION_NAME, embedding_function=emb_fn
+    name=COLLECTION_NAME, embedding_function=embedding_function_gemini_retrieval
 )
 
 results = collection.query(query_texts=[QUESTION], n_results=NUM_RETURNS)

Original file line number	Diff line number	Diff line change
`@@ -108,8 +108,12 @@ def embed_palm(texts: Documents) -> Embeddings:`
`108`	`108`	`model_name=LOCAL_LLM`
`109`	`109`	`)`
`110`	`110`
	`111`	`+embedding_function_gemini_retrieval = embedding_functions.GoogleGenerativeAiEmbeddingFunction(`
	`112`	`+ api_key=API_KEY, model_name=EMBEDDING_MODEL,`
	`113`	`+ task_type="RETRIEVAL_QUERY")`
	`114`	`+`
`111`	`115`	`collection = chroma_client.get_collection(`
`112`		`- name=COLLECTION_NAME, embedding_function=emb_fn`
	`116`	`+ name=COLLECTION_NAME, embedding_function=embedding_function_gemini_retrieval`
`113`	`117`	`)`
`114`	`118`
`115`	`119`	`results = collection.query(query_texts=[QUESTION], n_results=NUM_RETURNS)`