mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-23 19:41:54 +00:00
c2a3021bb0
Signed-off-by: ChengZi <chen.zhang@zilliz.com> Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com> Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com> Co-authored-by: Dan O'Donovan <dan.odonovan@gmail.com> Co-authored-by: Tom Daniel Grande <tomdgrande@gmail.com> Co-authored-by: Grande <Tom.Daniel.Grande@statsbygg.no> Co-authored-by: Bagatur <baskaryan@gmail.com> Co-authored-by: ccurme <chester.curme@gmail.com> Co-authored-by: Harrison Chase <hw.chase.17@gmail.com> Co-authored-by: Tomaz Bratanic <bratanic.tomaz@gmail.com> Co-authored-by: ZhangShenao <15201440436@163.com> Co-authored-by: Friso H. Kingma <fhkingma@gmail.com> Co-authored-by: ChengZi <chen.zhang@zilliz.com> Co-authored-by: Nuno Campos <nuno@langchain.dev> Co-authored-by: Morgante Pell <morgantep@google.com>
69 lines
2.7 KiB
Python
69 lines
2.7 KiB
Python
"""Test the voyageai reranker."""
|
||
|
||
import os
|
||
|
||
from langchain_core.documents import Document
|
||
|
||
from langchain_voyageai.rerank import VoyageAIRerank
|
||
|
||
|
||
def test_voyageai_reranker_init() -> None:
|
||
"""Test the voyageai reranker initializes correctly."""
|
||
VoyageAIRerank(voyage_api_key="foo", model="foo") # type: ignore[arg-type]
|
||
|
||
|
||
def test_sync() -> None:
|
||
rerank = VoyageAIRerank(
|
||
voyage_api_key=os.environ["VOYAGE_API_KEY"], # type: ignore[arg-type]
|
||
model="rerank-lite-1",
|
||
)
|
||
doc_list = [
|
||
"The Mediterranean diet emphasizes fish, olive oil, and vegetables"
|
||
", believed to reduce chronic diseases.",
|
||
"Photosynthesis in plants converts light energy into glucose and "
|
||
"produces essential oxygen.",
|
||
"20th-century innovations, from radios to smartphones, centered "
|
||
"on electronic advancements.",
|
||
"Rivers provide water, irrigation, and habitat for aquatic species, "
|
||
"vital for ecosystems.",
|
||
"Apple’s conference call to discuss fourth fiscal quarter results and "
|
||
"business updates is scheduled for Thursday, November 2, 2023 at 2:00 "
|
||
"p.m. PT / 5:00 p.m. ET.",
|
||
"Shakespeare's works, like 'Hamlet' and 'A Midsummer Night's Dream,' "
|
||
"endure in literature.",
|
||
]
|
||
documents = [Document(page_content=x) for x in doc_list]
|
||
|
||
result = rerank.compress_documents(
|
||
query="When is the Apple's conference call scheduled?", documents=documents
|
||
)
|
||
assert len(doc_list) == len(result)
|
||
|
||
|
||
async def test_async() -> None:
|
||
rerank = VoyageAIRerank(
|
||
voyage_api_key=os.environ["VOYAGE_API_KEY"], # type: ignore[arg-type]
|
||
model="rerank-lite-1",
|
||
)
|
||
doc_list = [
|
||
"The Mediterranean diet emphasizes fish, olive oil, and vegetables"
|
||
", believed to reduce chronic diseases.",
|
||
"Photosynthesis in plants converts light energy into glucose and "
|
||
"produces essential oxygen.",
|
||
"20th-century innovations, from radios to smartphones, centered "
|
||
"on electronic advancements.",
|
||
"Rivers provide water, irrigation, and habitat for aquatic species, "
|
||
"vital for ecosystems.",
|
||
"Apple’s conference call to discuss fourth fiscal quarter results and "
|
||
"business updates is scheduled for Thursday, November 2, 2023 at 2:00 "
|
||
"p.m. PT / 5:00 p.m. ET.",
|
||
"Shakespeare's works, like 'Hamlet' and 'A Midsummer Night's Dream,' "
|
||
"endure in literature.",
|
||
]
|
||
documents = [Document(page_content=x) for x in doc_list]
|
||
|
||
result = await rerank.acompress_documents(
|
||
query="When is the Apple's conference call scheduled?", documents=documents
|
||
)
|
||
assert len(doc_list) == len(result)
|