src/backend/base/langflow/components/embeddings/text_embedder.py

import logging
from typing import TYPE_CHECKING

from langflow.custom import Component
from langflow.io import HandleInput, MessageInput, Output
from langflow.schema import Data

if TYPE_CHECKING:
    from langflow.field_typing import Embeddings
    from langflow.schema.message import Message


class TextEmbedderComponent(Component):
    display_name: str = "Text Embedder"
    description: str = "Generate embeddings for a given message using the specified embedding model."
    icon = "binary"
    inputs = [
        HandleInput(
            name="embedding_model",
            display_name="Embedding Model",
            info="The embedding model to use for generating embeddings.",
            input_types=["Embeddings"],
        ),
        MessageInput(
            name="message",
            display_name="Message",
            info="The message to generate embeddings for.",
        ),
    ]
    outputs = [
        Output(display_name="Embedding Data", name="embeddings", method="generate_embeddings"),
    ]

    def generate_embeddings(self) -> Data:
        try:
            embedding_model: Embeddings = self.embedding_model
            message: Message = self.message

            # Validate embedding model
            if not embedding_model:
                msg = "Embedding model not provided"
                raise ValueError(msg)

            # Extract the text content from the message
            text_content = message.text if message and message.text else ""
            if not text_content:
                msg = "No text content found in message"
                raise ValueError(msg)

            # Check if the embedding model has the required attributes
            if not hasattr(embedding_model, "client") or not embedding_model.client:
                msg = "Embedding model client not properly initialized"
                raise ValueError(msg)

            # Ensure the base URL has proper protocol
            if hasattr(embedding_model.client, "base_url"):
                base_url = embedding_model.client.base_url
                if not base_url.startswith(("http://", "https://")):
                    embedding_model.client.base_url = f"https://{base_url}"

            # Generate embeddings using the provided embedding model
            embeddings = embedding_model.embed_documents([text_content])

            # Validate embeddings output
            if not embeddings or not isinstance(embeddings, list):
                msg = "Invalid embeddings generated"
                raise ValueError(msg)

            embedding_vector = embeddings[0]

        except Exception as e:
            logging.exception("Error generating embeddings")
            # Return empty data with error status
            error_data = Data(data={"text": "", "embeddings": [], "error": str(e)})
            self.status = {"error": str(e)}
            return error_data

        # Create a Data object to encapsulate the results
        result_data = Data(data={"text": text_content, "embeddings": embedding_vector})
        self.status = {"text": text_content, "embeddings": embedding_vector}
        return result_data