Skip to content

voyageai_rerank_driver

VoyageAiRerankDriver

Bases: BaseRerankDriver

Voyage AI Rerank Driver.

Attributes:

Name Type Description
model str

Voyage AI rerank model name.
api_key str | None

API key to pass directly. Defaults to VOYAGE_API_KEY environment variable.
top_k int | None

Optional maximum number of results to return.
Source code in griptape/drivers/rerank/voyageai_rerank_driver.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@define(kw_only=True)
class VoyageAiRerankDriver(BaseRerankDriver):
    """Voyage AI Rerank Driver.

    Attributes:
        model: Voyage AI rerank model name.
        api_key: API key to pass directly. Defaults to `VOYAGE_API_KEY` environment variable.
        top_k: Optional maximum number of results to return.
    """

    model: str = field(default="rerank-2.5", metadata={"serializable": True})
    api_key: str | None = field(default=None, metadata={"serializable": False})
    top_k: int | None = field(default=None, metadata={"serializable": True})
    _client: Client | None = field(default=None, kw_only=True, alias="client", metadata={"serializable": False})

    @lazy_property()
    def client(self) -> Any:
        return import_optional_dependency("voyageai").Client(api_key=self.api_key)

    def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
        truthy_artifacts = [artifact for artifact in artifacts if artifact]

        if not truthy_artifacts:
            return []

        response = self.client.rerank(
            query=query,
            documents=[artifact.to_text() for artifact in truthy_artifacts],
            model=self.model,
            top_k=self.top_k,
        )

        return [truthy_artifacts[result.index] for result in response.results]

_client = field(default=None, kw_only=True, alias='client', metadata={'serializable': False}) class-attribute instance-attribute

api_key = field(default=None, metadata={'serializable': False}) class-attribute instance-attribute

model = field(default='rerank-2.5', metadata={'serializable': True}) class-attribute instance-attribute

top_k = field(default=None, metadata={'serializable': True}) class-attribute instance-attribute

client()

Source code in griptape/drivers/rerank/voyageai_rerank_driver.py 
32
33
34
@lazy_property()
def client(self) -> Any:
    return import_optional_dependency("voyageai").Client(api_key=self.api_key)

run(query, artifacts)

Source code in griptape/drivers/rerank/voyageai_rerank_driver.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def run(self, query: str, artifacts: list[TextArtifact]) -> list[TextArtifact]:
    truthy_artifacts = [artifact for artifact in artifacts if artifact]

    if not truthy_artifacts:
        return []

    response = self.client.rerank(
        query=query,
        documents=[artifact.to_text() for artifact in truthy_artifacts],
        model=self.model,
        top_k=self.top_k,
    )

    return [truthy_artifacts[result.index] for result in response.results]