Skip to content

Commit 4eaec79

Browse files
eyurtsevolgamurraft
authored andcommitted
ai21[patch]: Update API reference documentation (langchain-ai#25302)
Issue: langchain-ai#24856
1 parent 880c7c4 commit 4eaec79

File tree

2 files changed

+60
-6
lines changed

2 files changed

+60
-6
lines changed

libs/partners/ai21/langchain_ai21/ai21_base.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,17 @@ class Config:
1616

1717
client: Any = Field(default=None, exclude=True) #: :meta private:
1818
api_key: Optional[SecretStr] = None
19+
"""API key for AI21 API."""
1920
api_host: Optional[str] = None
21+
"""Host URL"""
2022
timeout_sec: Optional[float] = None
23+
"""Timeout in seconds.
24+
25+
If not set, it will default to the value of the environment
26+
variable `AI21_TIMEOUT_SEC` or 300 seconds.
27+
"""
2128
num_retries: Optional[int] = None
29+
"""Maximum number of retries for API requests before giving up."""
2230

2331
@root_validator()
2432
def validate_environment(cls, values: Dict) -> Dict:

libs/partners/ai21/langchain_ai21/embeddings.py

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,64 @@ def _split_texts_into_batches(texts: List[str], batch_size: int) -> Iterator[Lis
1515

1616

1717
class AI21Embeddings(Embeddings, AI21Base):
18-
"""AI21 embedding model.
18+
"""AI21 embedding model integration.
1919
20-
To use, you should have the 'AI21_API_KEY' environment variable set
21-
or pass as a named parameter to the constructor.
20+
Install ``langchain_ai21`` and set environment variable ``AI21_API_KEY``.
2221
23-
Example:
22+
.. code-block:: bash
23+
24+
pip install -U langchain_ai21
25+
export AI21_API_KEY="your-api-key"
26+
27+
Key init args — client params:
28+
api_key: Optional[SecretStr]
29+
batch_size: int
30+
The number of texts that will be sent to the API in each batch.
31+
Use larger batch sizes if working with many short texts. This will reduce
32+
the number of API calls made, and can improve the time it takes to embed
33+
a large number of texts.
34+
num_retries: Optional[int]
35+
Maximum number of retries for API requests before giving up.
36+
timeout_sec: Optional[float]
37+
Timeout in seconds for API requests. If not set, it will default to the
38+
value of the environment variable `AI21_TIMEOUT_SEC` or 300 seconds.
39+
40+
See full list of supported init args and their descriptions in the params section.
41+
42+
Instantiate:
2443
.. code-block:: python
2544
2645
from langchain_ai21 import AI21Embeddings
2746
28-
embeddings = AI21Embeddings()
29-
query_result = embeddings.embed_query("Hello embeddings world!")
47+
embed = AI21Embeddings(
48+
# api_key="...",
49+
# batch_size=128,
50+
)
51+
52+
Embed single text:
53+
.. code-block:: python
54+
55+
input_text = "The meaning of life is 42"
56+
vector = embed.embed_query(input_text)
57+
print(vector[:3])
58+
59+
.. code-block:: python
60+
61+
[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
62+
63+
Embed multiple texts:
64+
.. code-block:: python
65+
66+
input_texts = ["Document 1...", "Document 2..."]
67+
vectors = embed.embed_documents(input_texts)
68+
print(len(vectors))
69+
# The first 3 coordinates for the first vector
70+
print(vectors[0][:3])
71+
72+
.. code-block:: python
73+
74+
2
75+
[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
3076
"""
3177

3278
batch_size: int = _DEFAULT_BATCH_SIZE

0 commit comments

Comments
 (0)