chore: Add LCEmbeddingsModel class for handling embeddings in langflow

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-19 19:01:17 -03:00
commit 2f04ac3e44
3 changed files with 26 additions and 7 deletions

View file

@ -0,0 +1,23 @@
from langflow.custom import Component
from langflow.field_typing import Embeddings
from langflow.io import Output
class LCEmbeddingsModel(Component):
trace_type = "embedding"
outputs = [
Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
]
def _validate_outputs(self):
required_output_methods = ["build_embeddings"]
output_names = [output.name for output in self.outputs]
for method_name in required_output_methods:
if method_name not in output_names:
raise ValueError(f"Output with name '{method_name}' must be defined.")
elif not hasattr(self, method_name):
raise ValueError(f"Method '{method_name}' must be defined.")
def build_embeddings(self) -> Embeddings:
raise NotImplementedError("You must implement the build_embeddings method in your class.")

View file

@ -1,11 +1,11 @@
from langchain_openai.embeddings.base import OpenAIEmbeddings
from langflow.base.models.model import LCModelComponent
from langflow.base.embeddings.model import LCEmbeddingsModel
from langflow.field_typing import Embeddings
from langflow.io import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, Output, SecretStrInput, TextInput
from langflow.io import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, SecretStrInput, TextInput
class OpenAIEmbeddingsComponent(LCModelComponent):
class OpenAIEmbeddingsComponent(LCEmbeddingsModel):
display_name = "OpenAI Embeddings"
description = "Generate embeddings using OpenAI models."
icon = "OpenAI"
@ -66,10 +66,6 @@ class OpenAIEmbeddingsComponent(LCModelComponent):
),
]
outputs = [
Output(display_name="Embeddings", name="embeddings", method="build_embeddings"),
]
def build_embeddings(self) -> Embeddings:
return OpenAIEmbeddings(
tiktoken_enabled=self.tiktoken_enable,