mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-24 20:09:01 +00:00
fireworks[patch]: Add Fireworks partner packages (#17694)
--------- Co-authored-by: Erick Friis <erick@langchain.dev>
This commit is contained in:
committed by
GitHub
parent
11cf95e810
commit
ee6a773456
52
libs/partners/fireworks/langchain_fireworks/embeddings.py
Normal file
52
libs/partners/fireworks/langchain_fireworks/embeddings.py
Normal file
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator
|
||||
from langchain_core.utils import convert_to_secret_str
|
||||
from openai import OpenAI # type: ignore
|
||||
|
||||
|
||||
class FireworksEmbeddings(BaseModel, Embeddings):
|
||||
"""FireworksEmbeddings embedding model.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from langchain_fireworks import FireworksEmbeddings
|
||||
|
||||
model = FireworksEmbeddings(
|
||||
model='nomic-ai/nomic-embed-text-v1.5'
|
||||
)
|
||||
"""
|
||||
|
||||
_client: OpenAI = Field(default=None)
|
||||
fireworks_api_key: SecretStr = convert_to_secret_str("")
|
||||
model: str = "nomic-ai/nomic-embed-text-v1.5"
|
||||
|
||||
@root_validator()
|
||||
def validate_environment(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Validate environment variables."""
|
||||
fireworks_api_key = convert_to_secret_str(
|
||||
values.get("fireworks_api_key") or os.getenv("FIREWORKS_API_KEY") or ""
|
||||
)
|
||||
values["fireworks_api_key"] = fireworks_api_key
|
||||
|
||||
# note this sets it globally for module
|
||||
# there isn't currently a way to pass it into client
|
||||
api_key = fireworks_api_key.get_secret_value()
|
||||
values["_client"] = OpenAI(
|
||||
api_key=api_key, base_url="https://api.fireworks.ai/inference/v1"
|
||||
)
|
||||
return values
|
||||
|
||||
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]
|
Reference in New Issue
Block a user