Skip to content

Commit 89128b7

Browse files
authored
community[patch]: add detailed paragraph and example for BaichuanTextEmbeddings (#22031)
- **Description:** add detailed paragraph and example for BaichuanTextEmbeddings - **Issue:** the issue #21983
1 parent 4e676a6 commit 89128b7

File tree

1 file changed

+33
-23
lines changed
  • libs/community/langchain_community/embeddings

1 file changed

+33
-23
lines changed

libs/community/langchain_community/embeddings/baichuan.py

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from langchain_core.embeddings import Embeddings
55
from langchain_core.pydantic_v1 import BaseModel, SecretStr, root_validator
66
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env
7+
from requests import RequestException
78

89
BAICHUAN_API_URL: str = "http://api.baichuan-ai.com/v1/embeddings"
910

@@ -22,11 +23,23 @@
2223
# NOTE!! BaichuanTextEmbeddings only supports Chinese text embedding.
2324
# Multi-language support is coming soon.
2425
class BaichuanTextEmbeddings(BaseModel, Embeddings):
25-
"""Baichuan Text Embedding models."""
26+
"""Baichuan Text Embedding models.
27+
28+
To use, you should set the environment variable ``BAICHUAN_API_KEY`` to
29+
your API key or pass it as a named parameter to the constructor.
30+
31+
Example:
32+
.. code-block:: python
33+
34+
from langchain_community.embeddings import BaichuanTextEmbeddings
35+
36+
baichuan = BaichuanTextEmbeddings(baichuan_api_key="my-api-key")
37+
"""
2638

2739
session: Any #: :meta private:
2840
model_name: str = "Baichuan-Text-Embedding"
2941
baichuan_api_key: Optional[SecretStr] = None
42+
"""Automatically inferred from env var `BAICHUAN_API_KEY` if not provided."""
3043

3144
@root_validator(allow_reuse=True)
3245
def validate_environment(cls, values: Dict) -> Dict:
@@ -65,29 +78,26 @@ def _embed(self, texts: List[str]) -> Optional[List[List[float]]]:
6578
A list of list of floats representing the embeddings, or None if an
6679
error occurs.
6780
"""
68-
try:
69-
response = self.session.post(
70-
BAICHUAN_API_URL, json={"input": texts, "model": self.model_name}
81+
response = self.session.post(
82+
BAICHUAN_API_URL, json={"input": texts, "model": self.model_name}
83+
)
84+
# Raise exception if response status code from 400 to 600
85+
response.raise_for_status()
86+
# Check if the response status code indicates success
87+
if response.status_code == 200:
88+
resp = response.json()
89+
embeddings = resp.get("data", [])
90+
# Sort resulting embeddings by index
91+
sorted_embeddings = sorted(embeddings, key=lambda e: e.get("index", 0))
92+
# Return just the embeddings
93+
return [result.get("embedding", []) for result in sorted_embeddings]
94+
else:
95+
# Log error or handle unsuccessful response appropriately
96+
# Handle 100 <= status_code < 400, not include 200
97+
raise RequestException(
98+
f"Error: Received status code {response.status_code} from "
99+
"`BaichuanEmbedding` API"
71100
)
72-
# Check if the response status code indicates success
73-
if response.status_code == 200:
74-
resp = response.json()
75-
embeddings = resp.get("data", [])
76-
# Sort resulting embeddings by index
77-
sorted_embeddings = sorted(embeddings, key=lambda e: e.get("index", 0))
78-
# Return just the embeddings
79-
return [result.get("embedding", []) for result in sorted_embeddings]
80-
else:
81-
# Log error or handle unsuccessful response appropriately
82-
print( # noqa: T201
83-
f"Error: Received status code {response.status_code} from "
84-
"embedding API"
85-
)
86-
return None
87-
except Exception as e:
88-
# Log the exception or handle it as needed
89-
print(f"Exception occurred while trying to get embeddings: {str(e)}") # noqa: T201
90-
return None
91101

92102
def embed_documents(self, texts: List[str]) -> Optional[List[List[float]]]: # type: ignore[override]
93103
"""Public method to get embeddings for a list of documents.

0 commit comments

Comments
 (0)