llama-index

1
"""Voyage embeddings file."""
2

3
from typing import Any, List, Optional
4

5
from llama_index.legacy.bridge.pydantic import PrivateAttr
6
from llama_index.legacy.callbacks.base import CallbackManager
7
from llama_index.legacy.embeddings.base import BaseEmbedding
8

9
DEFAULT_VOYAGE_BATCH_SIZE = 8
10

11

12
class VoyageEmbedding(BaseEmbedding):
13
    """Class for Voyage embeddings.
14

15
    Args:
16
        model_name (str): Model for embedding.
17
            Defaults to "voyage-01".
18

19
        voyage_api_key (Optional[str]): Voyage API key. Defaults to None.
20
            You can either specify the key here or store it as an environment variable.
21
    """
22

23
    _model: Any = PrivateAttr()
24

25
    def __init__(
26
        self,
27
        model_name: str = "voyage-01",
28
        voyage_api_key: Optional[str] = None,
29
        embed_batch_size: int = DEFAULT_VOYAGE_BATCH_SIZE,
30
        callback_manager: Optional[CallbackManager] = None,
31
        **kwargs: Any,
32
    ):
33
        try:
34
            import voyageai
35
        except ImportError:
36
            raise ImportError(
37
                "voyageai package not found, install with" "'pip install voyageai'"
38
            )
39
        if voyage_api_key:
40
            voyageai.api_key = voyage_api_key
41
        self._model = voyageai
42

43
        super().__init__(
44
            model_name=model_name,
45
            embed_batch_size=embed_batch_size,
46
            callback_manager=callback_manager,
47
            **kwargs,
48
        )
49

50
    @classmethod
51
    def class_name(cls) -> str:
52
        return "VoyageEmbedding"
53

54
    def _get_query_embedding(self, query: str) -> List[float]:
55
        """Get query embedding."""
56
        return self._model.get_embedding(
57
            query, model=self.model_name, input_type="query"
58
        )
59

60
    async def _aget_query_embedding(self, query: str) -> List[float]:
61
        """The asynchronous version of _get_query_embedding."""
62
        return await self._model.aget_embedding(
63
            query, model=self.model_name, input_type="query"
64
        )
65

66
    def _get_text_embedding(self, text: str) -> List[float]:
67
        """Get text embedding."""
68
        return self._model.get_embedding(
69
            text, model=self.model_name, input_type="document"
70
        )
71

72
    async def _aget_text_embedding(self, text: str) -> List[float]:
73
        """Asynchronously get text embedding."""
74
        return await self._model.aget_embedding(
75
            text, model=self.model_name, input_type="document"
76
        )
77

78
    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
79
        """Get text embeddings."""
80
        return self._model.get_embeddings(
81
            texts, model=self.model_name, input_type="document"
82
        )
83

84
    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
85
        """Asynchronously get text embeddings."""
86
        return await self._model.aget_embeddings(
87
            texts, model=self.model_name, input_type="document"
88
        )
89

90
    def get_general_text_embedding(
91
        self, text: str, input_type: Optional[str] = None
92
    ) -> List[float]:
93
        """Get general text embedding with input_type."""
94
        return self._model.get_embedding(
95
            text, model=self.model_name, input_type=input_type
96
        )
97

98
    async def aget_general_text_embedding(
99
        self, text: str, input_type: Optional[str] = None
100
    ) -> List[float]:
101
        """Asynchronously get general text embedding with input_type."""
102
        return await self._model.aget_embedding(
103
            text, model=self.model_name, input_type=input_type
104
        )
105