lfx-nightly 0.1.13.dev0__py3-none-any.whl → 0.2.0.dev0__py3-none-any.whl

lfx/components/models/embedding_model.py

@@ -3,7 +3,9 @@ from typing import Any
 from langchain_openai import OpenAIEmbeddings
 
 from lfx.base.embeddings.model import LCEmbeddingsModel
+from lfx.base.models.ollama_constants import OLLAMA_EMBEDDING_MODELS
 from lfx.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
+from lfx.base.models.watsonx_constants import WATSONX_EMBEDDING_MODEL_NAMES
 from lfx.field_typing import Embeddings
 from lfx.io import (
     BoolInput,
@@ -29,11 +31,11 @@ class EmbeddingModelComponent(LCEmbeddingsModel):
         DropdownInput(
             name="provider",
             display_name="Model Provider",
-            options=["OpenAI"],
+            options=["OpenAI", "Ollama", "IBM watsonx.ai"],
             value="OpenAI",
             info="Select the embedding model provider",
             real_time_refresh=True,
-            options_metadata=[{"icon": "OpenAI"}],
+            options_metadata=[{"icon": "OpenAI"}, {"icon": "Ollama"}, {"icon": "WatsonxAI"}],
         ),
         DropdownInput(
             name="model",
@@ -56,6 +58,13 @@ class EmbeddingModelComponent(LCEmbeddingsModel):
             info="Base URL for the API. Leave empty for default.",
             advanced=True,
         ),
+        # Watson-specific inputs
+        MessageTextInput(
+            name="project_id",
+            display_name="Project ID",
+            info="IBM watsonx.ai Project ID (required for IBM watsonx.ai)",
+            show=False,
+        ),
         IntInput(
             name="dimensions",
             display_name="Dimensions",
@@ -102,13 +111,85 @@ class EmbeddingModelComponent(LCEmbeddingsModel):
                 show_progress_bar=show_progress_bar,
                 model_kwargs=model_kwargs,
             )
+
+        if provider == "Ollama":
+            try:
+                from langchain_ollama import OllamaEmbeddings
+            except ImportError:
+                try:
+                    from langchain_community.embeddings import OllamaEmbeddings
+                except ImportError:
+                    msg = "Please install langchain-ollama: pip install langchain-ollama"
+                    raise ImportError(msg) from None
+
+            return OllamaEmbeddings(
+                model=model,
+                base_url=api_base or "http://localhost:11434",
+                **model_kwargs,
+            )
+
+        if provider == "IBM watsonx.ai":
+            try:
+                from langchain_ibm import WatsonxEmbeddings
+            except ImportError:
+                msg = "Please install langchain-ibm: pip install langchain-ibm"
+                raise ImportError(msg) from None
+
+            if not api_key:
+                msg = "IBM watsonx.ai API key is required when using IBM watsonx.ai provider"
+                raise ValueError(msg)
+
+            project_id = self.project_id
+
+            if not project_id:
+                msg = "Project ID is required for IBM watsonx.ai provider"
+                raise ValueError(msg)
+
+            params = {
+                "model_id": model,
+                "url": api_base or "https://us-south.ml.cloud.ibm.com",
+                "apikey": api_key,
+            }
+
+            params["project_id"] = project_id
+
+            return WatsonxEmbeddings(**params)
+
         msg = f"Unknown provider: {provider}"
         raise ValueError(msg)
 
     def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None) -> dotdict:
-        if field_name == "provider" and field_value == "OpenAI":
-            build_config["model"]["options"] = OPENAI_EMBEDDING_MODEL_NAMES
-            build_config["model"]["value"] = OPENAI_EMBEDDING_MODEL_NAMES[0]
-            build_config["api_key"]["display_name"] = "OpenAI API Key"
-            build_config["api_base"]["display_name"] = "OpenAI API Base URL"
+        if field_name == "provider":
+            if field_value == "OpenAI":
+                build_config["model"]["options"] = OPENAI_EMBEDDING_MODEL_NAMES
+                build_config["model"]["value"] = OPENAI_EMBEDDING_MODEL_NAMES[0]
+                build_config["api_key"]["display_name"] = "OpenAI API Key"
+                build_config["api_key"]["required"] = True
+                build_config["api_key"]["show"] = True
+                build_config["api_base"]["display_name"] = "OpenAI API Base URL"
+                build_config["api_base"]["advanced"] = True
+                build_config["project_id"]["show"] = False
+
+            elif field_value == "Ollama":
+                build_config["model"]["options"] = OLLAMA_EMBEDDING_MODELS
+                build_config["model"]["value"] = OLLAMA_EMBEDDING_MODELS[0]
+                build_config["api_key"]["display_name"] = "API Key (Optional)"
+                build_config["api_key"]["required"] = False
+                build_config["api_key"]["show"] = False
+                build_config["api_base"]["display_name"] = "Ollama Base URL"
+                build_config["api_base"]["value"] = "http://localhost:11434"
+                build_config["api_base"]["advanced"] = True
+                build_config["project_id"]["show"] = False
+
+            elif field_value == "IBM watsonx.ai":
+                build_config["model"]["options"] = WATSONX_EMBEDDING_MODEL_NAMES
+                build_config["model"]["value"] = WATSONX_EMBEDDING_MODEL_NAMES[0]
+                build_config["api_key"]["display_name"] = "IBM watsonx.ai API Key"
+                build_config["api_key"]["required"] = True
+                build_config["api_key"]["show"] = True
+                build_config["api_base"]["display_name"] = "IBM watsonx.ai URL"
+                build_config["api_base"]["value"] = "https://us-south.ml.cloud.ibm.com"
+                build_config["api_base"]["advanced"] = False
+                build_config["project_id"]["show"] = True
+
         return build_config

lfx/components/ollama/ollama.py

@@ -1,4 +1,6 @@
 import asyncio
+import json
+from contextlib import suppress
 from typing import Any
 from urllib.parse import urljoin
 
@@ -8,11 +10,27 @@ from langchain_ollama import ChatOllama
 from lfx.base.models.model import LCModelComponent
 from lfx.field_typing import LanguageModel
 from lfx.field_typing.range_spec import RangeSpec
-from lfx.io import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, MessageTextInput, SliderInput
+from lfx.helpers.base_model import build_model_from_schema
+from lfx.io import (
+    BoolInput,
+    DictInput,
+    DropdownInput,
+    FloatInput,
+    IntInput,
+    MessageTextInput,
+    Output,
+    SecretStrInput,
+    SliderInput,
+    TableInput,
+)
 from lfx.log.logger import logger
+from lfx.schema.data import Data
+from lfx.schema.dataframe import DataFrame
+from lfx.schema.table import EditMode
 from lfx.utils.util import transform_localhost_url
 
 HTTP_STATUS_OK = 200
+TABLE_ROW_PLACEHOLDER = {"name": "field", "description": "description of field", "type": "str", "multiple": "False"}
 
 
 class ChatOllamaComponent(LCModelComponent):
@@ -28,11 +46,51 @@ class ChatOllamaComponent(LCModelComponent):
     DESIRED_CAPABILITY = "completion"
     TOOL_CALLING_CAPABILITY = "tools"
 
+    # Define the table schema for the format input
+    TABLE_SCHEMA = [
+        {
+            "name": "name",
+            "display_name": "Name",
+            "type": "str",
+            "description": "Specify the name of the output field.",
+            "default": "field",
+            "edit_mode": EditMode.INLINE,
+        },
+        {
+            "name": "description",
+            "display_name": "Description",
+            "type": "str",
+            "description": "Describe the purpose of the output field.",
+            "default": "description of field",
+            "edit_mode": EditMode.POPOVER,
+        },
+        {
+            "name": "type",
+            "display_name": "Type",
+            "type": "str",
+            "edit_mode": EditMode.INLINE,
+            "description": ("Indicate the data type of the output field (e.g., str, int, float, bool, dict)."),
+            "options": ["str", "int", "float", "bool", "dict"],
+            "default": "str",
+        },
+        {
+            "name": "multiple",
+            "display_name": "As List",
+            "type": "boolean",
+            "description": "Set to True if this output field should be a list of the specified type.",
+            "edit_mode": EditMode.INLINE,
+            "options": ["True", "False"],
+            "default": "False",
+        },
+    ]
+    default_table_row = {row["name"]: row.get("default", None) for row in TABLE_SCHEMA}
+    default_table_row_schema = build_model_from_schema([default_table_row]).model_json_schema()
+
     inputs = [
         MessageTextInput(
             name="base_url",
-            display_name="Base URL",
-            info="Endpoint of the Ollama API. Defaults to http://localhost:11434 .",
+            display_name="Ollama API URL",
+            info="Endpoint of the Ollama API. Defaults to http://localhost:11434.",
             value="http://localhost:11434",
             real_time_refresh=True,
         ),
@@ -44,6 +102,15 @@ class ChatOllamaComponent(LCModelComponent):
             refresh_button=True,
             real_time_refresh=True,
         ),
+        SecretStrInput(
+            name="api_key",
+            display_name="Ollama API Key",
+            info="Your Ollama API key.",
+            value=None,
+            required=False,
+            real_time_refresh=True,
+            advanced=True,
+        ),
         SliderInput(
             name="temperature",
             display_name="Temperature",
@@ -51,8 +118,13 @@ class ChatOllamaComponent(LCModelComponent):
             range_spec=RangeSpec(min=0, max=1, step=0.01),
             advanced=True,
         ),
-        MessageTextInput(
-            name="format", display_name="Format", info="Specify the format of the output (e.g., json).", advanced=True
+        TableInput(
+            name="format",
+            display_name="Format",
+            info="Specify the format of the output.",
+            advanced=False,
+            table_schema=TABLE_SCHEMA,
+            value=default_table_row,
         ),
         DictInput(name="metadata", display_name="Metadata", info="Metadata to add to the run trace.", advanced=True),
         DropdownInput(
@@ -112,7 +184,12 @@ class ChatOllamaComponent(LCModelComponent):
             name="top_k", display_name="Top K", info="Limits token selection to top K. (Default: 40)", advanced=True
         ),
         FloatInput(name="top_p", display_name="Top P", info="Works together with top-k. (Default: 0.9)", advanced=True),
-        BoolInput(name="verbose", display_name="Verbose", info="Whether to print out response text.", advanced=True),
+        BoolInput(
+            name="enable_verbose_output",
+            display_name="Ollama Verbose Output",
+            info="Whether to print out response text.",
+            advanced=True,
+        ),
         MessageTextInput(
             name="tags",
             display_name="Tags",
@@ -141,15 +218,22 @@ class ChatOllamaComponent(LCModelComponent):
         *LCModelComponent.get_base_inputs(),
     ]
 
+    outputs = [
+        Output(display_name="Text", name="text_output", method="text_response"),
+        Output(display_name="Language Model", name="model_output", method="build_model"),
+        Output(display_name="Data", name="data_output", method="build_data_output"),
+        Output(display_name="DataFrame", name="dataframe_output", method="build_dataframe_output"),
+    ]
+
     def build_model(self) -> LanguageModel:  # type: ignore[type-var]
         # Mapping mirostat settings to their corresponding values
         mirostat_options = {"Mirostat": 1, "Mirostat 2.0": 2}
 
-        # Default to 0 for 'Disabled'
-        mirostat_value = mirostat_options.get(self.mirostat, 0)
+        # Default to None for 'Disabled'
+        mirostat_value = mirostat_options.get(self.mirostat, None)
 
         # Set mirostat_eta and mirostat_tau to None if mirostat is disabled
-        if mirostat_value == 0:
+        if mirostat_value is None:
             mirostat_eta = None
             mirostat_tau = None
         else:
@@ -169,12 +253,18 @@ class ChatOllamaComponent(LCModelComponent):
                 "Learn more at https://docs.ollama.com/openai#openai-compatibility"
             )
 
+        try:
+            output_format = self._parse_format_field(self.format)
+        except Exception as e:
+            msg = f"Failed to parse the format field: {e}"
+            raise ValueError(msg) from e
+
         # Mapping system settings to their corresponding values
         llm_params = {
             "base_url": transformed_base_url,
             "model": self.model_name,
             "mirostat": mirostat_value,
-            "format": self.format,
+            "format": output_format,
             "metadata": self.metadata,
             "tags": self.tags.split(",") if self.tags else None,
             "mirostat_eta": mirostat_eta,
@@ -191,9 +281,12 @@ class ChatOllamaComponent(LCModelComponent):
             "timeout": self.timeout or None,
             "top_k": self.top_k or None,
             "top_p": self.top_p or None,
-            "verbose": self.verbose,
+            "verbose": self.enable_verbose_output or False,
             "template": self.template,
         }
+        headers = self.headers
+        if headers is not None:
+            llm_params["client_kwargs"] = {"headers": headers}
 
         # Remove parameters with None values
         llm_params = {k: v for k, v in llm_params.items() if v is not None}
@@ -219,7 +312,9 @@ class ChatOllamaComponent(LCModelComponent):
                 url = url.rstrip("/").removesuffix("/v1")
                 if not url.endswith("/"):
                     url = url + "/"
-                return (await client.get(urljoin(url, "api/tags"))).status_code == HTTP_STATUS_OK
+                return (
+                    await client.get(url=urljoin(url, "api/tags"), headers=self.headers)
+                ).status_code == HTTP_STATUS_OK
         except httpx.RequestError:
             return False
 
@@ -292,8 +387,9 @@ class ChatOllamaComponent(LCModelComponent):
             show_url = urljoin(base_url, "api/show")
 
             async with httpx.AsyncClient() as client:
+                headers = self.headers
                 # Fetch available models
-                tags_response = await client.get(tags_url)
+                tags_response = await client.get(url=tags_url, headers=headers)
                 tags_response.raise_for_status()
                 models = tags_response.json()
                 if asyncio.iscoroutine(models):
@@ -307,11 +403,12 @@ class ChatOllamaComponent(LCModelComponent):
                     await logger.adebug(f"Checking model: {model_name}")
 
                     payload = {"model": model_name}
-                    show_response = await client.post(show_url, json=payload)
+                    show_response = await client.post(url=show_url, json=payload, headers=headers)
                     show_response.raise_for_status()
                     json_data = show_response.json()
                     if asyncio.iscoroutine(json_data):
                         json_data = await json_data
+
                     capabilities = json_data.get(self.JSON_CAPABILITIES_KEY, [])
                     await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
 
@@ -325,3 +422,113 @@ class ChatOllamaComponent(LCModelComponent):
             raise ValueError(msg) from e
 
         return model_ids
+
+    def _parse_format_field(self, format_value: Any) -> Any:
+        """Parse the format field to handle both string and dict inputs.
+
+        The format field can be:
+        - A simple string like "json" (backward compatibility)
+        - A JSON string from NestedDictInput that needs parsing
+        - A dict/JSON schema (already parsed)
+        - None or empty
+
+        Args:
+            format_value: The raw format value from the input field
+
+        Returns:
+            Parsed format value as string, dict, or None
+        """
+        if not format_value:
+            return None
+
+        schema = format_value
+        if isinstance(format_value, list):
+            schema = build_model_from_schema(format_value).model_json_schema()
+            if schema == self.default_table_row_schema:
+                return None  # the rows are generic placeholder rows
+        elif isinstance(format_value, str):  # parse as json if string
+            with suppress(json.JSONDecodeError):  # e.g., literal "json" is valid for format field
+                schema = json.loads(format_value)
+
+        return schema or None
+
+    async def _parse_json_response(self) -> Any:
+        """Parse the JSON response from the model.
+
+        This method gets the text response and attempts to parse it as JSON.
+        Works with models that have format='json' or a JSON schema set.
+
+        Returns:
+            Parsed JSON (dict, list, or primitive type)
+
+        Raises:
+            ValueError: If the response is not valid JSON
+        """
+        message = await self.text_response()
+        text = message.text if hasattr(message, "text") else str(message)
+
+        if not text:
+            msg = "No response from model"
+            raise ValueError(msg)
+
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as e:
+            msg = f"Invalid JSON response. Ensure model supports JSON output. Error: {e}"
+            raise ValueError(msg) from e
+
+    async def build_data_output(self) -> Data:
+        """Build a Data output from the model's JSON response.
+
+        Returns:
+            Data: A Data object containing the parsed JSON response
+        """
+        parsed = await self._parse_json_response()
+
+        # If the response is already a dict, wrap it in Data
+        if isinstance(parsed, dict):
+            return Data(data=parsed)
+
+        # If it's a list, wrap in a results container
+        if isinstance(parsed, list):
+            if len(parsed) == 1:
+                return Data(data=parsed[0])
+            return Data(data={"results": parsed})
+
+        # For primitive types, wrap in a value container
+        return Data(data={"value": parsed})
+
+    async def build_dataframe_output(self) -> DataFrame:
+        """Build a DataFrame output from the model's JSON response.
+
+        Returns:
+            DataFrame: A DataFrame containing the parsed JSON response
+
+        Raises:
+            ValueError: If the response cannot be converted to a DataFrame
+        """
+        parsed = await self._parse_json_response()
+
+        # If it's a list of dicts, convert directly to DataFrame
+        if isinstance(parsed, list):
+            if not parsed:
+                return DataFrame()
+            # Ensure all items are dicts for proper DataFrame conversion
+            if all(isinstance(item, dict) for item in parsed):
+                return DataFrame(parsed)
+            msg = "List items must be dictionaries to convert to DataFrame"
+            raise ValueError(msg)
+
+        # If it's a single dict, wrap in a list to create a single-row DataFrame
+        if isinstance(parsed, dict):
+            return DataFrame([parsed])
+
+        # For primitive types, create a single-column DataFrame
+        return DataFrame([{"value": parsed}])
+
+    @property
+    def headers(self) -> dict[str, str] | None:
+        """Get the headers for the Ollama API."""
+        if self.api_key and self.api_key.strip():
+            return {"Authorization": f"Bearer {self.api_key}"}
+        return None