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

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,
         ),
@@ -43,6 +101,16 @@ class ChatOllamaComponent(LCModelComponent):
             info="Refer to https://ollama.com/library for more models.",
             refresh_button=True,
             real_time_refresh=True,
+            required=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",
@@ -51,8 +119,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.",
+            table_schema=TABLE_SCHEMA,
+            value=default_table_row,
+            show=False,
         ),
         DictInput(name="metadata", display_name="Metadata", info="Metadata to add to the run trace.", advanced=True),
         DropdownInput(
@@ -112,7 +185,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",
@@ -138,18 +216,33 @@ class ChatOllamaComponent(LCModelComponent):
         MessageTextInput(
             name="template", display_name="Template", info="Template to use for generating text.", advanced=True
         ),
+        BoolInput(
+            name="enable_structured_output",
+            display_name="Enable Structured Output",
+            info="Whether to enable structured output in the model.",
+            value=False,
+            advanced=False,
+            real_time_refresh=True,
+        ),
         *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 +262,18 @@ class ChatOllamaComponent(LCModelComponent):
                 "Learn more at https://docs.ollama.com/openai#openai-compatibility"
             )
 
+        try:
+            output_format = self._parse_format_field(self.format) if self.enable_structured_output else None
+        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 or None,
             "metadata": self.metadata,
             "tags": self.tags.split(",") if self.tags else None,
             "mirostat_eta": mirostat_eta,
@@ -191,9 +290,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,11 +321,16 @@ 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
 
     async def update_build_config(self, build_config: dict, field_value: Any, field_name: str | None = None):
+        if field_name == "enable_structured_output":  # bind enable_structured_output boolean to format show value
+            build_config["format"]["show"] = field_value
+
         if field_name == "mirostat":
             if field_value == "Disabled":
                 build_config["mirostat_eta"]["advanced"] = True
@@ -243,6 +350,8 @@ class ChatOllamaComponent(LCModelComponent):
                     build_config["mirostat_tau"]["value"] = 5
 
         if field_name in {"model_name", "base_url", "tool_model_enabled"}:
+            logger.warning(f"Fetching Ollama models from updated URL: {build_config['base_url']}")
+
             if await self.is_valid_ollama_url(self.base_url):
                 tool_model_enabled = build_config["tool_model_enabled"].get("value", False) or self.tool_model_enabled
                 build_config["model_name"]["options"] = await self.get_models(
@@ -292,8 +401,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 +417,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 +436,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

lfx/components/openrouter/openrouter.py

@@ -1,6 +1,3 @@
-from collections import defaultdict
-from typing import Any
-
 import httpx
 from langchain_openai import ChatOpenAI
 from pydantic.v1 import SecretStr
@@ -8,13 +5,7 @@ from pydantic.v1 import SecretStr
 from lfx.base.models.model import LCModelComponent
 from lfx.field_typing import LanguageModel
 from lfx.field_typing.range_spec import RangeSpec
-from lfx.inputs.inputs import (
-    DropdownInput,
-    IntInput,
-    SecretStrInput,
-    SliderInput,
-    StrInput,
-)
+from lfx.inputs.inputs import DropdownInput, IntInput, SecretStrInput, SliderInput, StrInput
 
 
 class OpenRouterComponent(LCModelComponent):
@@ -28,36 +19,13 @@ class OpenRouterComponent(LCModelComponent):
 
     inputs = [
         *LCModelComponent.get_base_inputs(),
-        SecretStrInput(
-            name="api_key", display_name="OpenRouter API Key", required=True, info="Your OpenRouter API key"
-        ),
-        StrInput(
-            name="site_url",
-            display_name="Site URL",
-            info="Your site URL for OpenRouter rankings",
-            advanced=True,
-        ),
-        StrInput(
-            name="app_name",
-            display_name="App Name",
-            info="Your app name for OpenRouter rankings",
-            advanced=True,
-        ),
-        DropdownInput(
-            name="provider",
-            display_name="Provider",
-            info="The AI model provider",
-            options=["Loading providers..."],
-            value="Loading providers...",
-            real_time_refresh=True,
-            required=True,
-        ),
+        SecretStrInput(name="api_key", display_name="API Key", required=True),
         DropdownInput(
             name="model_name",
             display_name="Model",
-            info="The model to use for chat completion",
-            options=["Select a provider first"],
-            value="Select a provider first",
+            options=[],
+            value="",
+            refresh_button=True,
             real_time_refresh=True,
             required=True,
         ),
@@ -66,137 +34,71 @@ class OpenRouterComponent(LCModelComponent):
             display_name="Temperature",
             value=0.7,
             range_spec=RangeSpec(min=0, max=2, step=0.01),
-            info="Controls randomness. Lower values are more deterministic, higher values are more creative.",
-            advanced=True,
-        ),
-        IntInput(
-            name="max_tokens",
-            display_name="Max Tokens",
-            info="Maximum number of tokens to generate",
             advanced=True,
         ),
+        IntInput(name="max_tokens", display_name="Max Tokens", advanced=True),
+        StrInput(name="site_url", display_name="Site URL", advanced=True),
+        StrInput(name="app_name", display_name="App Name", advanced=True),
     ]
 
-    def fetch_models(self) -> dict[str, list]:
-        """Fetch available models from OpenRouter API and organize them by provider."""
-        url = "https://openrouter.ai/api/v1/models"
-
+    def fetch_models(self) -> list[dict]:
+        """Fetch available models from OpenRouter."""
         try:
-            with httpx.Client() as client:
-                response = client.get(url)
-                response.raise_for_status()
-
-                models_data = response.json().get("data", [])
-                provider_models = defaultdict(list)
-
-                for model in models_data:
-                    model_id = model.get("id", "")
-                    if "/" in model_id:
-                        provider = model_id.split("/")[0].title()
-                        provider_models[provider].append(
-                            {
-                                "id": model_id,
-                                "name": model.get("name", ""),
-                                "description": model.get("description", ""),
-                                "context_length": model.get("context_length", 0),
-                            }
-                        )
-
-                return dict(provider_models)
-
-        except httpx.HTTPError as e:
-            self.log(f"Error fetching models: {e!s}")
-            return {"Error": [{"id": "error", "name": f"Error fetching models: {e!s}"}]}
+            response = httpx.get("https://openrouter.ai/api/v1/models", timeout=10.0)
+            response.raise_for_status()
+            models = response.json().get("data", [])
+            return sorted(
+                [
+                    {
+                        "id": m["id"],
+                        "name": m.get("name", m["id"]),
+                        "context": m.get("context_length", 0),
+                    }
+                    for m in models
+                    if m.get("id")
+                ],
+                key=lambda x: x["name"],
+            )
+        except (httpx.RequestError, httpx.HTTPStatusError) as e:
+            self.log(f"Error fetching models: {e}")
+            return []
+
+    def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None) -> dict:  # noqa: ARG002
+        """Update model options."""
+        models = self.fetch_models()
+        if models:
+            build_config["model_name"]["options"] = [m["id"] for m in models]
+            build_config["model_name"]["tooltips"] = {m["id"]: f"{m['name']} ({m['context']:,} tokens)" for m in models}
+        else:
+            build_config["model_name"]["options"] = ["Failed to load models"]
+            build_config["model_name"]["value"] = "Failed to load models"
+        return build_config
 
     def build_model(self) -> LanguageModel:
-        """Build and return the OpenRouter language model."""
-        model_not_selected = "Please select a model"
-        api_key_required = "API key is required"
-
-        if not self.model_name or self.model_name == "Select a provider first":
-            raise ValueError(model_not_selected)
-
+        """Build the OpenRouter model."""
         if not self.api_key:
-            raise ValueError(api_key_required)
-
-        api_key = SecretStr(self.api_key).get_secret_value()
+            msg = "API key is required"
+            raise ValueError(msg)
+        if not self.model_name or self.model_name == "Loading...":
+            msg = "Please select a model"
+            raise ValueError(msg)
 
-        # Build base configuration
-        kwargs: dict[str, Any] = {
+        kwargs = {
             "model": self.model_name,
-            "openai_api_key": api_key,
+            "openai_api_key": SecretStr(self.api_key).get_secret_value(),
             "openai_api_base": "https://openrouter.ai/api/v1",
             "temperature": self.temperature if self.temperature is not None else 0.7,
         }
 
-        # Add optional parameters
         if self.max_tokens:
-            kwargs["max_tokens"] = self.max_tokens
+            kwargs["max_tokens"] = int(self.max_tokens)
 
         headers = {}
         if self.site_url:
             headers["HTTP-Referer"] = self.site_url
         if self.app_name:
             headers["X-Title"] = self.app_name
-
         if headers:
             kwargs["default_headers"] = headers
 
-        try:
-            return ChatOpenAI(**kwargs)
-        except (ValueError, httpx.HTTPError) as err:
-            error_msg = f"Failed to build model: {err!s}"
-            self.log(error_msg)
-            raise ValueError(error_msg) from err
-
-    def _get_exception_message(self, e: Exception) -> str | None:
-        """Get a message from an OpenRouter exception.
-
-        Args:
-            e (Exception): The exception to get the message from.
-
-        Returns:
-            str | None: The message from the exception, or None if no specific message can be extracted.
-        """
-        try:
-            from openai import BadRequestError
-
-            if isinstance(e, BadRequestError):
-                message = e.body.get("message")
-                if message:
-                    return message
-        except ImportError:
-            pass
-        return None
-
-    def update_build_config(self, build_config: dict, field_value: str, field_name: str | None = None) -> dict:
-        """Update build configuration based on field updates."""
-        try:
-            if field_name is None or field_name == "provider":
-                provider_models = self.fetch_models()
-                build_config["provider"]["options"] = sorted(provider_models.keys())
-                if build_config["provider"]["value"] not in provider_models:
-                    build_config["provider"]["value"] = build_config["provider"]["options"][0]
-
-            if field_name == "provider" and field_value in self.fetch_models():
-                provider_models = self.fetch_models()
-                models = provider_models[field_value]
-
-                build_config["model_name"]["options"] = [model["id"] for model in models]
-                if models:
-                    build_config["model_name"]["value"] = models[0]["id"]
-
-                tooltips = {
-                    model["id"]: (f"{model['name']}\nContext Length: {model['context_length']}\n{model['description']}")
-                    for model in models
-                }
-                build_config["model_name"]["tooltips"] = tooltips
-
-        except httpx.HTTPError as e:
-            self.log(f"Error updating build config: {e!s}")
-            build_config["provider"]["options"] = ["Error loading providers"]
-            build_config["provider"]["value"] = "Error loading providers"
-            build_config["model_name"]["options"] = ["Error loading models"]
-            build_config["model_name"]["value"] = "Error loading models"
-
-        return build_config
+        return ChatOpenAI(**kwargs)