mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-07 05:30:39 +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>
109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
from typing import List
|
|
|
|
from langchain_core.embeddings import Embeddings
|
|
from langchain_core.utils import secret_from_env
|
|
from openai import OpenAI
|
|
from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator
|
|
from typing_extensions import Self
|
|
|
|
# type: ignore
|
|
|
|
|
|
class FireworksEmbeddings(BaseModel, Embeddings):
|
|
"""Fireworks embedding model integration.
|
|
|
|
Setup:
|
|
Install ``langchain_fireworks`` and set environment variable
|
|
``FIREWORKS_API_KEY``.
|
|
|
|
.. code-block:: bash
|
|
|
|
pip install -U langchain_fireworks
|
|
export FIREWORKS_API_KEY="your-api-key"
|
|
|
|
Key init args — completion params:
|
|
model: str
|
|
Name of Fireworks model to use.
|
|
|
|
Key init args — client params:
|
|
fireworks_api_key: SecretStr
|
|
Fireworks API key.
|
|
|
|
See full list of supported init args and their descriptions in the params section.
|
|
|
|
Instantiate:
|
|
.. code-block:: python
|
|
|
|
from langchain_fireworks import FireworksEmbeddings
|
|
|
|
model = FireworksEmbeddings(
|
|
model='nomic-ai/nomic-embed-text-v1.5'
|
|
# Use FIREWORKS_API_KEY env var or pass it in directly
|
|
# fireworks_api_key="..."
|
|
)
|
|
|
|
Embed multiple texts:
|
|
.. code-block:: python
|
|
|
|
vectors = embeddings.embed_documents(['hello', 'goodbye'])
|
|
# Showing only the first 3 coordinates
|
|
print(len(vectors))
|
|
print(vectors[0][:3])
|
|
|
|
.. code-block:: python
|
|
|
|
2
|
|
[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
|
|
|
|
|
|
Embed single text:
|
|
.. code-block:: python
|
|
|
|
input_text = "The meaning of life is 42"
|
|
vector = embeddings.embed_query('hello')
|
|
print(vector[:3])
|
|
|
|
.. code-block:: python
|
|
|
|
[-0.024603435769677162, -0.007543657906353474, 0.0039630369283258915]
|
|
"""
|
|
|
|
client: OpenAI = Field(default=None, exclude=True) #: :meta private:
|
|
fireworks_api_key: SecretStr = Field(
|
|
alias="api_key",
|
|
default_factory=secret_from_env(
|
|
"FIREWORKS_API_KEY",
|
|
default="",
|
|
),
|
|
)
|
|
"""Fireworks API key.
|
|
|
|
Automatically read from env variable `FIREWORKS_API_KEY` if not provided.
|
|
"""
|
|
model: str = "nomic-ai/nomic-embed-text-v1.5"
|
|
|
|
model_config = ConfigDict(
|
|
populate_by_name=True,
|
|
arbitrary_types_allowed=True,
|
|
)
|
|
|
|
@model_validator(mode="after")
|
|
def validate_environment(self) -> Self:
|
|
"""Validate environment variables."""
|
|
self.client = OpenAI(
|
|
api_key=self.fireworks_api_key.get_secret_value(),
|
|
base_url="https://api.fireworks.ai/inference/v1",
|
|
)
|
|
return self
|
|
|
|
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
|
"""Embed search docs."""
|
|
return [
|
|
i.embedding
|
|
for i in self.client.embeddings.create(input=texts, model=self.model).data
|
|
]
|
|
|
|
def embed_query(self, text: str) -> List[float]:
|
|
"""Embed query text."""
|
|
return self.embed_documents([text])[0]
|