diff --git a/src/backend/langflow/components/embeddings/OpenAIEmbeddings.py b/src/backend/langflow/components/embeddings/OpenAIEmbeddings.py index 7faf4fd07..6303389c8 100644 --- a/src/backend/langflow/components/embeddings/OpenAIEmbeddings.py +++ b/src/backend/langflow/components/embeddings/OpenAIEmbeddings.py @@ -1,7 +1,8 @@ - from langflow import CustomComponent -from typing import Optional, Set, Dict, Any, Union, Callable -from langchain.embeddings import OpenAIEmbeddings +from langflow.field_typing import NestedDict +from typing import List, Optional, Dict, Any, Union, Callable +from langchain_openai.embeddings.base import OpenAIEmbeddings + class OpenAIEmbeddingsComponent(CustomComponent): display_name = "OpenAIEmbeddings" @@ -9,40 +10,75 @@ class OpenAIEmbeddingsComponent(CustomComponent): def build_config(self): return { - "allowed_special": {"display_name": "Allowed Special", "advanced": True}, - "disallowed_special": {"display_name": "Disallowed Special", "advanced": True}, + "allowed_special": { + "display_name": "Allowed Special", + "advanced": True, + "field_type": "str", + "is_list": True, + }, + "default_headers": { + "display_name": "Default Headers", + "advanced": True, + "field_type": "dict", + }, + "default_query": { + "display_name": "Default Query", + "advanced": True, + "field_type": "NestedDict", + }, + "disallowed_special": { + "display_name": "Disallowed Special", + "advanced": True, + "field_type": "str", + "is_list": True, + }, "chunk_size": {"display_name": "Chunk Size", "advanced": True}, "client": {"display_name": "Client", "advanced": True}, "deployment": {"display_name": "Deployment", "advanced": True}, - "embedding_ctx_length": {"display_name": "Embedding Context Length", "advanced": True}, + "embedding_ctx_length": { + "display_name": "Embedding Context Length", + "advanced": True, + }, "max_retries": {"display_name": "Max Retries", "advanced": True}, "model": {"display_name": "Model", "advanced": True}, "model_kwargs": {"display_name": "Model Kwargs", "advanced": True}, - "openai_api_base": {"display_name": "OpenAI API Base", "advanced": True}, - "openai_api_key": {"display_name": "OpenAI API Key"}, - "openai_api_type": {"display_name": "OpenAI API Type", "advanced": True}, - "openai_api_version": {"display_name": "OpenAI API Version", "advanced": True}, - "openai_organization": {"display_name": "OpenAI Organization", "advanced": True}, + "openai_api_base": {"display_name": "OpenAI API Base","password":True, "advanced": True}, + "openai_api_key": {"display_name": "OpenAI API Key","password":True}, + "openai_api_type": {"display_name": "OpenAI API Type", "advanced": True,"password":True}, + "openai_api_version": { + "display_name": "OpenAI API Version", + "advanced": True, + }, + "openai_organization": { + "display_name": "OpenAI Organization", + "advanced": True, + }, "openai_proxy": {"display_name": "OpenAI Proxy", "advanced": True}, "request_timeout": {"display_name": "Request Timeout", "advanced": True}, - "show_progress_bar": {"display_name": "Show Progress Bar", "advanced": True}, + "show_progress_bar": { + "display_name": "Show Progress Bar", + "advanced": True, + }, "skip_empty": {"display_name": "Skip Empty", "advanced": True}, "tiktoken_model_name": {"display_name": "TikToken Model Name"}, + "tikToken_enable": {"display_name": "TikToken Enable"}, } def build( self, - allowed_special: Optional[Set[str]] = set(), - disallowed_special: str = "all", + default_headers: Optional[Dict[str, str]] = None, + default_query: Optional[NestedDict] = {}, + allowed_special: Optional[List[str]] = [], + disallowed_special: List[str] = ["all"], chunk_size: Optional[int] = 1000, client: Optional[Any] = None, deployment: str = "text-embedding-ada-002", embedding_ctx_length: Optional[int] = 8191, max_retries: Optional[int] = 6, model: str = "text-embedding-ada-002", - model_kwargs: Optional[Dict[str, Any]] = None, + model_kwargs: NestedDict = {}, openai_api_base: Optional[str] = None, - openai_api_key: Optional[str] = '', + openai_api_key: Optional[str] = "", openai_api_type: Optional[str] = None, openai_api_version: Optional[str] = None, openai_organization: Optional[str] = None, @@ -50,11 +86,15 @@ class OpenAIEmbeddingsComponent(CustomComponent): request_timeout: Optional[float] = None, show_progress_bar: Optional[bool] = False, skip_empty: Optional[bool] = False, + tikToken_enable: Optional[bool] = True, tiktoken_model_name: Optional[str] = None, ) -> Union[OpenAIEmbeddings, Callable]: return OpenAIEmbeddings( - allowed_special=allowed_special, - disallowed_special=disallowed_special, + tiktoken_enabled=tikToken_enable, + default_headers=default_headers, + default_query=default_query, + allowed_special=set(allowed_special), + disallowed_special=set(disallowed_special), chunk_size=chunk_size, client=client, deployment=deployment,