diff --git a/libs/arcade-cli/arcade_cli/display.py b/libs/arcade-cli/arcade_cli/display.py index abb3dc96..0d7e13bd 100644 --- a/libs/arcade-cli/arcade_cli/display.py +++ b/libs/arcade-cli/arcade_cli/display.py @@ -36,9 +36,13 @@ def display_tools_table(tools: list[ToolDefinition]) -> None: console.print(table) -def display_tool_details(tool: ToolDefinition) -> None: +def display_tool_details(tool: ToolDefinition, worker: bool = False) -> None: # noqa: C901 """ Display detailed information about a specific tool using multiple panels. + + Args: + tool: The tool definition to display + worker: If True, show full worker response structure. If False, show only value structure. """ # Description Panel description_panel = Panel( @@ -80,31 +84,147 @@ def display_tool_details(tool: ToolDefinition) -> None: border_style="green", ) - # Output Panel + # Output Panel - Show different levels based on worker flag output = tool.output - if output: - output_description = output.description or "No description available." - output_types = ", ".join(output.available_modes) - output_val_type = output.value_schema.val_type if output.value_schema else "N/A" - output_details = Text.assemble( - ("Description: ", "bold"), - (output_description, ""), - "\n", - ("Available Modes: ", "bold"), - (output_types, ""), - "\n", - ("Value Type: ", "bold"), - (output_val_type, ""), - ) + if output and output.value_schema: + output_table = Table(show_header=True, header_style="bold blue") + output_table.add_column("Field", style="cyan") + output_table.add_column("Type", style="magenta") + output_table.add_column("Description", style="white") + + if worker: + # Show full worker response structure + output_table.add_row( + "[bold]Response Structure[/bold]", + "", + "[dim]The tool response follows this structure:[/dim]", + ) + + # Available modes determine which fields can be present + modes = output.available_modes + + if "value" in modes: + # Show the value field with its schema + value_type: str = output.value_schema.val_type + display_type: str = value_type # Separate variable for display string + if value_type == "array" and output.value_schema.inner_val_type: + display_type = rf"array\[{output.value_schema.inner_val_type}]" + elif output.value_schema.enum: + display_type = f"{value_type} (enum: {', '.join(output.value_schema.enum)})" + + output_table.add_row( + " value", + display_type, + output.description or "The successful result from the tool", + ) + + # If the value is a json type with properties, show them + if ( + output.value_schema.val_type == "json" + and hasattr(output.value_schema, "properties") + and output.value_schema.properties + ): + _add_nested_properties(output_table, output.value_schema.properties, indent=2) + + if "error" in modes: + output_table.add_row( + " error", "object", "[dim]Error details if the tool fails[/dim]" + ) + output_table.add_row( + " message", "string", "[dim]User-facing error message[/dim]" + ) + output_table.add_row( + " developer_message", + "string?", + "[dim]Technical error details (optional)[/dim]", + ) + + if "null" in modes: + output_table.add_row(" value", "null", "[dim]Tool can return null/None[/dim]") + + # Additional fields that may be present + output_table.add_row("", "", "") + output_table.add_row( + "[bold]Additional Fields[/bold]", + "", + "[dim]May be present in any response:[/dim]", + ) + output_table.add_row( + " logs", "array?", "[dim]Optional warnings or info messages[/dim]" + ) + output_table.add_row( + " requires_authorization", + "object?", + "[dim]OAuth flow details if auth needed[/dim]", + ) + else: + # Show only the value structure (simplified view) + # Show the value type and description + display_type = _format_type_string(output.value_schema) + if output.value_schema.enum: + display_type = ( + f"{output.value_schema.val_type} (enum: {', '.join(output.value_schema.enum)})" + ) + + output_table.add_row( + "[bold]Value[/bold]", + display_type, + output.description or "The return value from the tool", + ) + + # If the value is a json type with properties, show them + if ( + output.value_schema.val_type == "json" + and hasattr(output.value_schema, "properties") + and output.value_schema.properties + ): + _add_nested_properties(output_table, output.value_schema.properties, indent=1) + + # Create subtitle with modes info + modes_text = Text() + modes_text.append("Response Modes: ", style="bold") + modes_text.append("One of { ", style="dim") + for i, mode in enumerate(output.available_modes): + if i > 0: + modes_text.append(", ", style="dim") + if mode == "value": + modes_text.append(mode, style="green") + elif mode == "error": + modes_text.append(mode, style="red") + elif mode == "null": + modes_text.append(mode, style="yellow") + else: + modes_text.append(mode, style="magenta") + modes_text.append(" }", style="dim") + output_panel = Panel( - output_details, - title="Expected Output", + output_table, + title="Output Schema", border_style="blue", + subtitle=modes_text, ) else: + # No schema defined + no_schema_table = Table(show_header=False) + no_schema_table.add_column() + + if worker: + no_schema_table.add_row( + "[dim]No output schema defined. The tool response will follow this structure:[/dim]" + ) + no_schema_table.add_row("") + no_schema_table.add_row("[cyan]Response Structure:[/cyan]") + no_schema_table.add_row(" • [bold]value[/bold]: null (when successful)") + no_schema_table.add_row(" • [bold]error[/bold]: object (when failed)") + no_schema_table.add_row(" • [bold]logs[/bold]: array? (optional warnings/info)") + else: + no_schema_table.add_row("[dim]No output schema defined.[/dim]") + no_schema_table.add_row("") + no_schema_table.add_row("The tool returns: [bold]null[/bold]") + output_panel = Panel( - "No output information available.", - title="Expected Output", + no_schema_table, + title="Output Schema", border_style="blue", ) @@ -114,6 +234,80 @@ def display_tool_details(tool: ToolDefinition) -> None: console.print(output_panel) +def _add_nested_properties( + table: Table, + properties: dict[str, Any], + indent: int = 0, + is_array_item: bool = False, +) -> None: + """ + Recursively add nested properties to the output table. + + Args: + table: The Rich table to add rows to + properties: Dictionary of property names to ValueSchema objects + indent: Current indentation level + is_array_item: Whether these properties are for array items + """ + indent_prefix = " " * indent + + # Show array item indicator if needed + if is_array_item and indent > 0: + table.add_row( + f"{indent_prefix[:-2]}[item]", + "", + "[dim]Each item in array:[/dim]", + ) + + for prop_name, prop_schema in properties.items(): + # Format the type string + type_str = _format_type_string(prop_schema) + + # Add the property row with better descriptions + description = "" + # For nested properties, we don't have descriptions yet, but we could add them + if hasattr(prop_schema, "description") and prop_schema.description: + description = prop_schema.description + + table.add_row( + f"{indent_prefix}{prop_name}", + type_str, + f"[dim]{description}[/dim]" if description else "", + ) + + # Recursively add nested properties if this is a json type with properties + if ( + prop_schema.val_type == "json" + and hasattr(prop_schema, "properties") + and prop_schema.properties + ): + _add_nested_properties(table, prop_schema.properties, indent + 1) + # Handle arrays with inner properties + elif ( + prop_schema.val_type == "array" + and hasattr(prop_schema, "inner_properties") + and prop_schema.inner_properties + ): + _add_nested_properties( + table, prop_schema.inner_properties, indent + 1, is_array_item=True + ) + + +def _format_type_string(schema: Any) -> str: + """Format type string for display.""" + type_str: str = schema.val_type + + if schema.val_type == "array": + if hasattr(schema, "inner_properties") and schema.inner_properties: + type_str = r"array\[object]" + elif schema.inner_val_type: + type_str = rf"array\[{schema.inner_val_type}]" + elif schema.enum: + type_str = f"{type_str} (enum)" + + return type_str + + def display_tool_messages(tool_messages: list[dict]) -> None: for message in tool_messages: if message["role"] == "assistant": @@ -124,7 +318,8 @@ def display_tool_messages(tool_messages: list[dict]) -> None: ) elif message["role"] == "tool": console.print( - f"[bold]'{message['name']}' tool returned:[/bold] {message['content']}", style="dim" + f"[bold]'{message['name']}' tool returned:[/bold] {message['content']}", + style="dim", ) diff --git a/libs/arcade-cli/arcade_cli/main.py b/libs/arcade-cli/arcade_cli/main.py index f71e6afe..54246509 100644 --- a/libs/arcade-cli/arcade_cli/main.py +++ b/libs/arcade-cli/arcade_cli/main.py @@ -55,7 +55,7 @@ cli = typer.Typer( cls=OrderCommands, add_completion=False, no_args_is_help=True, - pretty_exceptions_enable=False, + pretty_exceptions_enable=True, pretty_exceptions_show_locals=False, pretty_exceptions_short=True, rich_markup_mode="markdown", @@ -68,11 +68,16 @@ cli.add_typer( help="Manage deployments of tool servers (logs, list, etc)", rich_help_panel="Deployment", ) + + console = Console() def handle_cli_error( - message: str, error: Exception | None = None, debug: bool = True, should_exit: bool = True + message: str, + error: Optional[Exception] = None, + debug: bool = True, + should_exit: bool = True, ) -> None: """Handle CLI error reporting with optional debug traceback and exit.""" if error and debug: @@ -225,12 +230,34 @@ def show( "--no-tls", help="Whether to disable TLS for the connection to the Arcade Engine.", ), + worker: bool = typer.Option( + False, + "--worker", + "-w", + help="Show full worker response structure including error, logs, and authorization fields (only applies when used with -t/--tool).", + ), debug: bool = typer.Option(False, "--debug", "-d", help="Show debug information"), ) -> None: """ Show the available toolkits or detailed information about a specific tool. """ - show_logic(toolkit, tool, host, local, port, force_tls, force_no_tls, debug) + if worker and not tool: + console.print( + "⚠️ The -w/--worker flag only affects output when used with -t/--tool flag", + style="bold yellow", + ) + + show_logic( + toolkit=toolkit, + tool=tool, + host=host, + local=local, + port=port, + force_tls=force_tls, + force_no_tls=force_no_tls, + worker=worker, + debug=debug, + ) @cli.command( @@ -250,7 +277,7 @@ def chat( "--host", help="The Arcade Engine address to send chat requests to.", ), - port: int = typer.Option( + port: Optional[int] = typer.Option( None, "-p", "--port", @@ -388,7 +415,7 @@ def evals( "--cloud", help="Whether to run evaluations against the Arcade Cloud Engine. Overrides the 'host' option.", ), - port: int = typer.Option( + port: Optional[int] = typer.Option( None, "-p", "--port", @@ -509,10 +536,14 @@ def serve( show_default=True, ), port: int = typer.Option( - "8002", "-p", "--port", help="Port for the app, defaults to ", show_default=True + "8002", + "-p", + "--port", + help="Port for the app, defaults to ", + show_default=True, ), disable_auth: bool = typer.Option( - False, + True, "--no-auth", help="Disable authentication for the worker. Not recommended for production.", show_default=True, @@ -559,7 +590,9 @@ def serve( @cli.command( - help="Start a server with locally installed Arcade tools", rich_help_panel="Launch", hidden=True + help="Start a server with locally installed Arcade tools", + rich_help_panel="Launch", + hidden=True, ) def workerup( host: str = typer.Option( @@ -568,7 +601,11 @@ def workerup( show_default=True, ), port: int = typer.Option( - "8002", "-p", "--port", help="Port for the app, defaults to ", show_default=True + "8002", + "-p", + "--port", + help="Port for the app, defaults to ", + show_default=True, ), disable_auth: bool = typer.Option( False, @@ -610,7 +647,10 @@ def workerup( @cli.command(help="Deploy toolkits to Arcade Cloud", rich_help_panel="Deployment") def deploy( deployment_file: str = typer.Option( - "worker.toml", "--deployment-file", "-d", help="The deployment file to deploy." + "worker.toml", + "--deployment-file", + "-d", + help="The deployment file to deploy.", ), cloud_host: str = typer.Option( PROD_CLOUD_HOST, @@ -619,7 +659,7 @@ def deploy( help="The Arcade Cloud host to deploy to.", hidden=True, ), - cloud_port: int = typer.Option( + cloud_port: Optional[int] = typer.Option( None, "--cloud-port", "-cp", @@ -632,7 +672,7 @@ def deploy( "-h", help="The Arcade Engine host to register the worker to.", ), - port: int = typer.Option( + port: Optional[int] = typer.Option( None, "--port", "-p", @@ -674,7 +714,10 @@ def deploy( try: # Attempt to deploy worker worker.request().execute(cloud_client, engine_client) - console.log(f"✅ Worker '{worker.config.id}' deployed successfully.", style="dim") + console.log( + f"✅ Worker '{worker.config.id}' deployed successfully.", + style="dim", + ) except Exception as e: handle_cli_error(f"Failed to deploy worker '{worker.config.id}'", e, debug) diff --git a/libs/arcade-cli/arcade_cli/show.py b/libs/arcade-cli/arcade_cli/show.py index 310ef356..69d94455 100644 --- a/libs/arcade-cli/arcade_cli/show.py +++ b/libs/arcade-cli/arcade_cli/show.py @@ -1,3 +1,5 @@ +from typing import Optional + import typer from rich.console import Console from rich.markup import escape @@ -7,13 +9,14 @@ from arcade_cli.utils import create_cli_catalog, get_tools_from_engine def show_logic( - toolkit: str | None, - tool: str | None, + toolkit: Optional[str], + tool: Optional[str], host: str, local: bool, - port: int | None, + port: Optional[int], force_tls: bool, force_no_tls: bool, + worker: bool, debug: bool, ) -> None: """Wrapper function for the `arcade show` CLI command @@ -42,7 +45,7 @@ def show_logic( console.print(f"❌ Tool '{tool}' not found.", style="bold red") typer.Exit(code=1) else: - display_tool_details(tool_def) + display_tool_details(tool_def, worker=worker) else: # Display the list of tools as a table display_tools_table(tools) diff --git a/libs/arcade-cli/arcade_cli/utils.py b/libs/arcade-cli/arcade_cli/utils.py index 40b8e7ff..811c29f3 100644 --- a/libs/arcade-cli/arcade_cli/utils.py +++ b/libs/arcade-cli/arcade_cli/utils.py @@ -18,12 +18,20 @@ from arcade_core import ToolCatalog, Toolkit from arcade_core.config_model import Config from arcade_core.errors import ToolkitLoadError from arcade_core.schema import ToolDefinition -from arcadepy import NOT_GIVEN, APIConnectionError, APIStatusError, APITimeoutError, Arcade +from arcadepy import ( + NOT_GIVEN, + APIConnectionError, + APIStatusError, + APITimeoutError, + Arcade, +) from arcadepy.types import AuthorizationResponse from openai import OpenAI, Stream from openai.types.chat.chat_completion import Choice as ChatCompletionChoice from openai.types.chat.chat_completion_chunk import ChatCompletionChunk -from openai.types.chat.chat_completion_chunk import Choice as ChatCompletionChunkChoice +from openai.types.chat.chat_completion_chunk import ( + Choice as ChatCompletionChunkChoice, +) from pydantic import ValidationError from rich.console import Console from rich.live import Live @@ -231,7 +239,8 @@ def get_tools_from_engine( continue except APIConnectionError: console.print( - f"❌ Can't connect to Arcade Engine at {base_url}. (Is it running?)", style="bold red" + f"❌ Can't connect to Arcade Engine at {base_url}. (Is it running?)", + style="bold red", ) return tools @@ -326,7 +335,8 @@ def validate_and_get_config( if validate_user and (not config.user or not config.user.email): console.print( - "❌ User email not found in configuration. Please run `arcade login`.", style="bold red" + "❌ User email not found in configuration. Please run `arcade login`.", + style="bold red", ) raise typer.Exit(code=1) @@ -371,7 +381,11 @@ class ChatInteractionResult: def handle_chat_interaction( - client: OpenAI, model: str, history: list[dict], user_email: str | None, stream: bool = False + client: OpenAI, + model: str, + history: list[dict], + user_email: str | None, + stream: bool = False, ) -> ChatInteractionResult: """ Handle a single chat-request/chat-response interaction for both streamed and non-streamed responses. @@ -418,7 +432,8 @@ def handle_chat_interaction( elif role == "assistant": message_content = markdownify_urls(message_content) console.print( - f"\n[blue][bold]Assistant[/bold] ({model}):[/blue] ", Markdown(message_content) + f"\n[blue][bold]Assistant[/bold] ({model}):[/blue] ", + Markdown(message_content), ) else: console.print(f"\n[bold]{role}:[/bold] {message_content}") @@ -575,7 +590,10 @@ def load_eval_suites(eval_files: list[Path]) -> list[Callable]: ] if not eval_suite_funcs: - console.print(f"No @tool_eval functions found in {eval_file_path}", style="bold yellow") + console.print( + f"No @tool_eval functions found in {eval_file_path}", + style="bold yellow", + ) continue eval_suites.extend(eval_suite_funcs) @@ -628,7 +646,7 @@ def handle_user_command( user_input: str, history: list, host: str, - port: int, + port: int | None, force_tls: bool, force_no_tls: bool, show: Callable, @@ -658,6 +676,7 @@ def handle_user_command( force_tls=force_tls, force_no_tls=force_no_tls, debug=False, + worker=False, ) return True return False diff --git a/libs/arcade-core/arcade_core/catalog.py b/libs/arcade-core/arcade_core/catalog.py index 5fe38bb9..80c1e364 100644 --- a/libs/arcade-core/arcade_core/catalog.py +++ b/libs/arcade-core/arcade_core/catalog.py @@ -4,7 +4,7 @@ import logging import os import re import typing -from collections.abc import Iterator +from collections.abc import Callable, Iterator from dataclasses import dataclass from datetime import datetime from enum import Enum @@ -13,13 +13,12 @@ from types import ModuleType from typing import ( Annotated, Any, - Callable, Literal, - Optional, Union, cast, get_args, get_origin, + get_type_hints, ) from pydantic import BaseModel, Field, create_model @@ -62,6 +61,28 @@ InnerWireType = Literal["string", "integer", "number", "boolean", "json"] WireType = Union[InnerWireType, Literal["array"]] +def is_typeddict(tp: type) -> bool: + """ + Check if a type is a TypedDict. + Works with both typing.TypedDict and typing_extensions.TypedDict. + """ + try: + # TypedDict creates classes that inherit from dict + if not isinstance(tp, type) or not issubclass(tp, dict): + return False + + # Check for TypedDict-specific attributes + return ( + hasattr(tp, "__annotations__") + and hasattr(tp, "__total__") + and hasattr(tp, "__required_keys__") + and hasattr(tp, "__optional_keys__") + ) + except TypeError: + # Some special forms raise TypeError when checking issubclass + return False + + @dataclass class WireTypeInfo: """ @@ -71,6 +92,9 @@ class WireTypeInfo: wire_type: WireType inner_wire_type: InnerWireType | None = None enum_values: list[str] | None = None + properties: dict[str, "WireTypeInfo"] | None = None + inner_properties: dict[str, "WireTypeInfo"] | None = None + description: str | None = None class ToolMeta(BaseModel): @@ -79,9 +103,9 @@ class ToolMeta(BaseModel): """ module: str - toolkit: Optional[str] = None - package: Optional[str] = None - path: Optional[str] = None + toolkit: str | None = None + package: str | None = None + path: str | None = None date_added: datetime = Field(default_factory=datetime.now) date_updated: datetime = Field(default_factory=datetime.now) @@ -171,7 +195,7 @@ class ToolCatalog(BaseModel): def add_tool( self, tool_func: Callable, - toolkit_or_name: Union[str, Toolkit], + toolkit_or_name: str | Toolkit, module: ModuleType | None = None, ) -> None: """ @@ -289,7 +313,10 @@ class ToolCatalog(BaseModel): raise ValueError(f"Tool {func} not found in the catalog.") def get_tool_by_name( - self, name: str, version: Optional[str] = None, separator: str = TOOL_NAME_SEPARATOR + self, + name: str, + version: str | None = None, + separator: str = TOOL_NAME_SEPARATOR, ) -> MaterializedTool: """Get a tool from the catalog by name. @@ -353,8 +380,8 @@ class ToolCatalog(BaseModel): def create_tool_definition( tool: Callable, toolkit_name: str, - toolkit_version: Optional[str] = None, - toolkit_desc: Optional[str] = None, + toolkit_version: str | None = None, + toolkit_desc: str | None = None, ) -> ToolDefinition: """ Given a tool function, create a ToolDefinition @@ -431,16 +458,13 @@ def create_input_definition(func: Callable) -> ToolInput: description=tool_field_info.description, required=is_required, inferrable=tool_field_info.is_inferrable, - value_schema=ValueSchema( - val_type=tool_field_info.wire_type_info.wire_type, - inner_val_type=tool_field_info.wire_type_info.inner_wire_type, - enum=tool_field_info.wire_type_info.enum_values, - ), + value_schema=wire_type_info_to_value_schema(tool_field_info.wire_type_info), ) ) return ToolInput( - parameters=input_parameters, tool_context_parameter_name=tool_context_param_name + parameters=input_parameters, + tool_context_parameter_name=tool_context_param_name, ) @@ -478,11 +502,7 @@ def create_output_definition(func: Callable) -> ToolOutput: return ToolOutput( description=description, available_modes=available_modes, - value_schema=ValueSchema( - val_type=wire_type_info.wire_type, - inner_val_type=wire_type_info.inner_wire_type, - enum=wire_type_info.enum_values, - ), + value_schema=wire_type_info_to_value_schema(wire_type_info), ) @@ -669,12 +689,21 @@ def get_wire_type_info(_type: type) -> WireTypeInfo: # Is this a list type? # If so, get the inner (enclosed) type is_list = get_origin(_type) is list + inner_properties = None + if is_list: inner_type = get_args(_type)[0] - inner_wire_type = cast( - InnerWireType, - get_wire_type(str) if is_string_literal(inner_type) else get_wire_type(inner_type), - ) + + # Recursively get wire type info for inner type + inner_info = get_wire_type_info(inner_type) + inner_wire_type = cast(InnerWireType, inner_info.wire_type) + + # If inner type has properties (it's a complex object), propagate them + if inner_info.properties: + inner_properties = inner_info.properties + # If inner type is array (nested arrays), propagate inner_properties + elif inner_info.inner_properties: + inner_properties = inner_info.inner_properties else: inner_wire_type = None @@ -696,11 +725,133 @@ def get_wire_type_info(_type: type) -> WireTypeInfo: enum_values = [str(e) for e in get_args(type_to_check)] # Special case: Enum can be enumerated on the wire - elif issubclass(actual_type, Enum): + elif isinstance(actual_type, type) and issubclass(actual_type, Enum): is_enum = True - enum_values = [e.value for e in actual_type] # type: ignore[union-attr] + enum_values = [e.value for e in actual_type] - return WireTypeInfo(wire_type, inner_wire_type, enum_values if is_enum else None) + # Extract properties for complex types + properties = None + if wire_type == "json" and not is_list: + properties = extract_properties(type_to_check) + + return WireTypeInfo( + wire_type, + inner_wire_type, + enum_values if is_enum else None, + properties, + inner_properties, + ) + + +def _extract_typeddict_field_descriptions(typeddict_class: type) -> dict[str, str]: + """ + Extract field descriptions from TypedDict docstrings. + + TypedDict classes typically have field descriptions as docstrings after each field. + This function attempts to parse the source code to extract these descriptions. + """ + descriptions = {} + + try: + source = inspect.getsource(typeddict_class) + # Simple regex to match field: type pattern followed by a docstring + # This is a simplified approach - a full AST parser would be more robust + import re + + # Pattern to match field definition followed by docstring + pattern = r'(\w+):\s*[^"\n]+\n\s*"""([^"]+)"""' + matches = re.findall(pattern, source) + + for field_name, description in matches: + descriptions[field_name] = description.strip() + + except (OSError, TypeError): + # If we can't get the source, return empty descriptions + pass + + return descriptions + + +def extract_properties(type_to_check: type) -> dict[str, WireTypeInfo] | None: + """ + Extract properties from TypedDict, Pydantic models, or other structured types. + """ + properties = {} + + # Handle Pydantic BaseModel + if isinstance(type_to_check, type) and issubclass(type_to_check, BaseModel): + for field_name, field_info in type_to_check.model_fields.items(): + # Get the field type + field_type = field_info.annotation + if field_type is None: + continue + + # Handle Optional types (Union[T, None]) + if is_strict_optional(field_type): + # Extract the non-None type from Optional + field_type = next(arg for arg in get_args(field_type) if arg is not type(None)) + + # Get wire type info recursively + wire_info = get_wire_type_info(field_type) + properties[field_name] = wire_info + + # Handle TypedDict + elif is_typeddict(type_to_check): + # Get type hints for the TypedDict + type_hints = get_type_hints(type_to_check, include_extras=True) + + # Try to extract field descriptions from the class source + field_descriptions = _extract_typeddict_field_descriptions(type_to_check) + + for field_name, field_type in type_hints.items(): + # Handle Optional types (Union[T, None]) + if is_strict_optional(field_type): + # Extract the non-None type from Optional + field_type = next(arg for arg in get_args(field_type) if arg is not type(None)) + wire_info = get_wire_type_info(field_type) + + # Add description if available + if field_name in field_descriptions: + wire_info.description = field_descriptions[field_name] + + properties[field_name] = wire_info + + # Handle regular dict with type annotations (e.g., dict[str, Any]) + elif get_origin(type_to_check) is dict: + # For generic dicts, we can't extract specific properties + return None + + return properties if properties else None + + +def wire_type_info_to_value_schema(wire_info: WireTypeInfo) -> ValueSchema: + """ + Convert WireTypeInfo to ValueSchema, including nested properties. + """ + # Convert nested properties if they exist + properties = None + if wire_info.properties: + properties = { + name: wire_type_info_to_value_schema(nested_info) + for name, nested_info in wire_info.properties.items() + } + + # Convert inner properties for array items + inner_properties = None + if wire_info.inner_properties: + inner_properties = { + name: wire_type_info_to_value_schema(nested_info) + for name, nested_info in wire_info.inner_properties.items() + } + + return ValueSchema( + val_type=wire_info.wire_type, + inner_val_type=wire_info.inner_wire_type, + enum=wire_info.enum_values, + properties=properties, + inner_properties=inner_properties, + description=wire_info.description, + ) def extract_python_param_info(param: inspect.Parameter) -> ParamInfo: @@ -799,6 +950,9 @@ def get_wire_type( if isinstance(_type, type) and issubclass(_type, BaseModel): return "json" + if is_typeddict(_type): + return "json" + raise ToolDefinitionError(f"Unsupported parameter type: {_type}") @@ -831,7 +985,7 @@ def create_func_models(func: Callable) -> tuple[type[BaseModel], type[BaseModel] return input_model, output_model -def determine_output_model(func: Callable) -> type[BaseModel]: +def determine_output_model(func: Callable) -> type[BaseModel]: # noqa: C901 """ Determine the output model for a function based on its return annotation. """ @@ -845,6 +999,18 @@ def determine_output_model(func: Callable) -> type[BaseModel]: description = ( return_annotation.__metadata__[0] if return_annotation.__metadata__ else "" ) + + # Check if the field type is a TypedDict + if is_typeddict(field_type): + # Create a Pydantic model from TypedDict + typeddict_model = create_model_from_typeddict( + field_type, f"{output_model_name}TypedDict" + ) + return create_model( + output_model_name, + result=(typeddict_model, Field(description=str(description))), + ) + if description: return create_model( output_model_name, @@ -857,6 +1023,18 @@ def determine_output_model(func: Callable) -> type[BaseModel]: # TODO handle multiple non-None arguments. Raise error? for arg in get_args(return_annotation): if arg is not type(None): + # Check if the arg is a TypedDict + if is_typeddict(arg): + typeddict_model = create_model_from_typeddict( + arg, f"{output_model_name}TypedDict" + ) + return create_model( + output_model_name, + result=( + typeddict_model, + Field(description="No description provided."), + ), + ) return create_model( output_model_name, result=(arg, Field(description="No description provided.")), @@ -871,6 +1049,17 @@ def determine_output_model(func: Callable) -> type[BaseModel]: ), ) else: + # Check if return type is TypedDict + if is_typeddict(return_annotation): + typeddict_model = create_model_from_typeddict(return_annotation, output_model_name) + return create_model( + output_model_name, + result=( + typeddict_model, + Field(description="No description provided."), + ), + ) + # Handle simple return types (like str) return create_model( output_model_name, @@ -878,6 +1067,37 @@ def determine_output_model(func: Callable) -> type[BaseModel]: ) +def create_model_from_typeddict(typeddict_class: type, model_name: str) -> type[BaseModel]: + """ + Create a Pydantic model from a TypedDict class. + This enables runtime validation of TypedDict structures. + """ + # Get type hints for the TypedDict + type_hints = get_type_hints(typeddict_class, include_extras=True) + + # Build field definitions for the Pydantic model + field_definitions: dict[str, Any] = {} + for field_name, field_type in type_hints.items(): + # Check if field is required + is_required = field_name in getattr(typeddict_class, "__required_keys__", set()) + + # Handle nested TypedDict + if is_typeddict(field_type): + nested_model = create_model_from_typeddict(field_type, f"{model_name}_{field_name}") + if is_required: + field_definitions[field_name] = (nested_model, Field()) + else: + field_definitions[field_name] = (nested_model, Field(default=None)) + else: + if is_required: + field_definitions[field_name] = (field_type, Field()) + else: + field_definitions[field_name] = (field_type, Field(default=None)) + + # Create and return the Pydantic model + return create_model(model_name, **field_definitions) + + def to_tool_secret_requirements( secrets_requirement: list[str], ) -> list[ToolSecretRequirement]: diff --git a/libs/arcade-core/arcade_core/executor.py b/libs/arcade-core/arcade_core/executor.py index 58cf23d0..a9c67020 100644 --- a/libs/arcade-core/arcade_core/executor.py +++ b/libs/arcade-core/arcade_core/executor.py @@ -1,6 +1,7 @@ import asyncio import traceback -from typing import Any, Callable +from collections.abc import Callable +from typing import Any from pydantic import BaseModel, ValidationError @@ -12,7 +13,12 @@ from arcade_core.errors import ( ToolSerializationError, ) from arcade_core.output import output_factory -from arcade_core.schema import ToolCallLog, ToolCallOutput, ToolContext, ToolDefinition +from arcade_core.schema import ( + ToolCallLog, + ToolCallOutput, + ToolContext, + ToolDefinition, +) class ToolExecutor: @@ -34,7 +40,9 @@ class ToolExecutor: if definition.deprecation_message is not None: tool_call_logs.append( ToolCallLog( - message=definition.deprecation_message, level="warning", subtype="deprecation" + message=definition.deprecation_message, + level="warning", + subtype="deprecation", ) ) @@ -101,7 +109,8 @@ class ToolExecutor: except ValidationError as e: raise ToolInputError( - message="Error in tool input deserialization", developer_message=str(e) + message="Error in tool input deserialization", + developer_message=str(e), ) from e return inputs diff --git a/libs/arcade-core/arcade_core/output.py b/libs/arcade-core/arcade_core/output.py index bfd241b0..0b9b45b7 100644 --- a/libs/arcade-core/arcade_core/output.py +++ b/libs/arcade-core/arcade_core/output.py @@ -1,5 +1,7 @@ from typing import TypeVar +from pydantic import BaseModel + from arcade_core.schema import ToolCallError, ToolCallLog, ToolCallOutput from arcade_core.utils import coerce_empty_list_to_none @@ -17,9 +19,29 @@ class ToolOutputFactory: data: T | None = None, logs: list[ToolCallLog] | None = None, ) -> ToolCallOutput: - value = getattr(data, "result", "") if data else "" + # Extract the result value + """ + Extracts the result value for the tool output. + + The executor guarantees that `data` is either a string, a dict, or None. + """ + value: str | int | float | bool | dict | list[str] | None + if data is None: + value = "" + elif hasattr(data, "result"): + value = getattr(data, "result", "") + elif isinstance(data, BaseModel): + value = data.model_dump() + elif isinstance(data, (str, int, float, bool, list)): + value = data + else: + raise ValueError(f"Unsupported data output type: {type(data)}") + logs = coerce_empty_list_to_none(logs) - return ToolCallOutput(value=value, logs=logs) + return ToolCallOutput( + value=value, + logs=logs, + ) def fail( self, @@ -56,6 +78,7 @@ class ToolOutputFactory: can_retry=True, additional_prompt_content=additional_prompt_content, retry_after_ms=retry_after_ms, + traceback_info=traceback_info, ), logs=coerce_empty_list_to_none(logs), ) diff --git a/libs/arcade-core/arcade_core/parse.py b/libs/arcade-core/arcade_core/parse.py index f0d17122..148397fa 100644 --- a/libs/arcade-core/arcade_core/parse.py +++ b/libs/arcade-core/arcade_core/parse.py @@ -1,6 +1,5 @@ import ast from pathlib import Path -from typing import Optional, Union def load_ast_tree(filepath: str | Path) -> ast.AST: @@ -16,8 +15,8 @@ def load_ast_tree(filepath: str | Path) -> ast.AST: def get_function_name_if_decorated( - node: Union[ast.FunctionDef, ast.AsyncFunctionDef], -) -> Optional[str]: + node: ast.FunctionDef | ast.AsyncFunctionDef, +) -> str | None: """ Check if a function has a decorator. """ diff --git a/libs/arcade-core/arcade_core/schema.py b/libs/arcade-core/arcade_core/schema.py index 5a27ee0e..029dc198 100644 --- a/libs/arcade-core/arcade_core/schema.py +++ b/libs/arcade-core/arcade_core/schema.py @@ -1,7 +1,7 @@ import os from dataclasses import dataclass from enum import Enum -from typing import Any, Literal, Optional, Union +from typing import Any, Literal from pydantic import BaseModel, Field @@ -15,12 +15,21 @@ class ValueSchema(BaseModel): val_type: Literal["string", "integer", "number", "boolean", "json", "array"] """The type of the value.""" - inner_val_type: Optional[Literal["string", "integer", "number", "boolean", "json"]] = None + inner_val_type: Literal["string", "integer", "number", "boolean", "json"] | None = None """The type of the inner value, if the value is a list.""" - enum: Optional[list[str]] = None + enum: list[str] | None = None """The list of possible values for the value, if it is a closed list.""" + properties: dict[str, "ValueSchema"] | None = None + """For object types (json), the schema of nested properties.""" + + inner_properties: dict[str, "ValueSchema"] | None = None + """For array types with json items, the schema of properties for each array item.""" + + description: str | None = None + """Optional description of the value.""" + class InputParameter(BaseModel): """A parameter that can be passed to a tool.""" @@ -30,8 +39,9 @@ class InputParameter(BaseModel): ..., description="Whether this parameter is required (true) or optional (false).", ) - description: Optional[str] = Field( - None, description="A descriptive, human-readable explanation of the parameter." + description: str | None = Field( + None, + description="A descriptive, human-readable explanation of the parameter.", ) value_schema: ValueSchema = Field( ..., @@ -59,14 +69,14 @@ class ToolInput(BaseModel): class ToolOutput(BaseModel): """The output of a tool.""" - description: Optional[str] = Field( + description: str | None = Field( None, description="A descriptive, human-readable explanation of the output." ) available_modes: list[str] = Field( default_factory=lambda: ["value", "error", "null"], description="The available modes for the output.", ) - value_schema: Optional[ValueSchema] = Field( + value_schema: ValueSchema | None = Field( None, description="The schema of the value of the output." ) @@ -74,7 +84,7 @@ class ToolOutput(BaseModel): class OAuth2Requirement(BaseModel): """Indicates that the tool requires OAuth 2.0 authorization.""" - scopes: Optional[list[str]] = None + scopes: list[str] | None = None """The scope(s) needed for the authorized action.""" @@ -90,16 +100,16 @@ class ToolAuthRequirement(BaseModel): # # The Arcade SDK translates these into the appropriate provider ID (Google) and type (OAuth2). # The only time the developer will set these is if they are using a custom auth provider. - provider_id: Optional[str] = None + provider_id: str | None = None """The provider ID configured in Arcade that acts as an alias to well-known configuration.""" provider_type: str """The type of the authorization provider.""" - id: Optional[str] = None + id: str | None = None """A provider's unique identifier, allowing the tool to specify a specific authorization provider. Recommended for private tools only.""" - oauth2: Optional[OAuth2Requirement] = None + oauth2: OAuth2Requirement | None = None """The OAuth 2.0 requirement, if any.""" @@ -133,13 +143,13 @@ class ToolMetadataRequirement(BaseModel): class ToolRequirements(BaseModel): """The requirements for a tool to run.""" - authorization: Union[ToolAuthRequirement, None] = None + authorization: ToolAuthRequirement | None = None """The authorization requirements for the tool, if any.""" - secrets: Union[list[ToolSecretRequirement], None] = None + secrets: list[ToolSecretRequirement] | None = None """The secret requirements for the tool, if any.""" - metadata: Union[list[ToolMetadataRequirement], None] = None + metadata: list[ToolMetadataRequirement] | None = None """The metadata requirements for the tool, if any.""" @@ -149,10 +159,10 @@ class ToolkitDefinition(BaseModel): name: str """The name of the toolkit.""" - description: Optional[str] = None + description: str | None = None """The description of the toolkit.""" - version: Optional[str] = None + version: str | None = None """The version identifier of the toolkit.""" @@ -166,7 +176,7 @@ class FullyQualifiedName: toolkit_name: str """The name of the toolkit containing the tool.""" - toolkit_version: Optional[str] = None + toolkit_version: str | None = None """The version of the toolkit containing the tool.""" def __str__(self) -> str: @@ -225,7 +235,7 @@ class ToolDefinition(BaseModel): requirements: ToolRequirements """The requirements (e.g. authorization) for the tool to run.""" - deprecation_message: Optional[str] = None + deprecation_message: str | None = None """The message to display when the tool is deprecated.""" def get_fully_qualified_name(self) -> FullyQualifiedName: @@ -241,7 +251,7 @@ class ToolReference(BaseModel): toolkit: str """The name of the toolkit containing the tool.""" - version: Optional[str] = None + version: str | None = None """The version of the toolkit containing the tool.""" def get_fully_qualified_name(self) -> FullyQualifiedName: @@ -313,7 +323,10 @@ class ToolContext(BaseModel): return self._get_item(key, self.metadata, "metadata") def _get_item( - self, key: str, items: list[ToolMetadataItem] | list[ToolSecretItem] | None, item_name: str + self, + key: str, + items: list[ToolMetadataItem] | list[ToolSecretItem] | None, + item_name: str, ) -> str: if not key or not key.strip(): raise ValueError( @@ -368,7 +381,7 @@ class ToolCallLog(BaseModel): ] """The level of severity for the log.""" - subtype: Optional[Literal["deprecation"]] = None + subtype: Literal["deprecation"] | None = None """Optional field for further categorization of the log.""" @@ -405,7 +418,7 @@ class ToolCallRequiresAuthorization(BaseModel): class ToolCallOutput(BaseModel): """The output of a tool invocation.""" - value: Union[str, int, float, bool, dict, list[str]] | None = None + value: str | int | float | bool | dict | list[str] | None = None """The value returned by the tool.""" logs: list[ToolCallLog] | None = None """The logs that occurred during the tool invocation.""" diff --git a/libs/arcade-core/arcade_core/utils.py b/libs/arcade-core/arcade_core/utils.py index d9c58696..12bbd3df 100644 --- a/libs/arcade-core/arcade_core/utils.py +++ b/libs/arcade-core/arcade_core/utils.py @@ -3,9 +3,9 @@ from __future__ import annotations import ast import inspect import re -from collections.abc import Iterable +from collections.abc import Callable, Iterable from types import UnionType -from typing import Any, Callable, Literal, TypeVar, Union, get_args, get_origin +from typing import Any, Literal, TypeVar, Union, get_args, get_origin T = TypeVar("T") diff --git a/libs/arcade-serve/arcade_serve/core/components.py b/libs/arcade-serve/arcade_serve/core/components.py index ae9274e1..a85d8db2 100644 --- a/libs/arcade-serve/arcade_serve/core/components.py +++ b/libs/arcade-serve/arcade_serve/core/components.py @@ -82,17 +82,17 @@ class HealthCheckComponent(WorkerComponent): "health", self, method="GET", - require_auth=False, response_type=HealthCheckResponse, operation_id="health_check", - description="Check the health of the worker", - summary="Check the health of the worker", + description="Health check", + summary="Health check", tags=["Arcade"], + require_auth=False, ) async def __call__(self, request: RequestData) -> HealthCheckResponse: """ - Handle the request for a health check. + Handle the request to check the health of the worker. """ tracer = trace.get_tracer(__name__) with tracer.start_as_current_span("HealthCheck"): diff --git a/libs/tests/cli/test_show.py b/libs/tests/cli/test_show.py index 15a7cbdf..ad8c33eb 100644 --- a/libs/tests/cli/test_show.py +++ b/libs/tests/cli/test_show.py @@ -14,6 +14,7 @@ def test_show_logic_local_false(): port=None, force_tls=False, force_no_tls=False, + worker=False, debug=False, ) @@ -33,6 +34,7 @@ def test_show_logic_local_true(): port=None, force_tls=False, force_no_tls=False, + worker=False, debug=False, ) diff --git a/libs/tests/core/test_executor.py b/libs/tests/core/test_executor.py index 488d0a99..c5136652 100644 --- a/libs/tests/core/test_executor.py +++ b/libs/tests/core/test_executor.py @@ -179,7 +179,7 @@ def check_output(output: ToolCallOutput, expected_output: ToolCallOutput): output_logs = output.logs or [] expected_logs = expected_output.logs or [] assert len(output_logs) == len(expected_logs) - for output_log, expected_log in zip(output_logs, expected_logs): + for output_log, expected_log in zip(output_logs, expected_logs, strict=False): assert output_log.message == expected_log.message assert output_log.level == expected_log.level assert output_log.subtype == expected_log.subtype diff --git a/libs/tests/tool/test_create_tool_definition.py b/libs/tests/tool/test_create_tool_definition.py index 34e4cc45..d46809d0 100644 --- a/libs/tests/tool/test_create_tool_definition.py +++ b/libs/tests/tool/test_create_tool_definition.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Annotated, Literal, Optional, Union +from typing import Annotated, Literal import pytest from arcade_core.catalog import ToolCatalog @@ -186,20 +186,20 @@ def func_with_param_with_default(param1: Annotated[str, "First param"] = "defaul @tool(desc="A function with an optional input parameter") -def func_with_optional_param(param1: Annotated[Optional[str], "First param"]): +def func_with_optional_param(param1: Annotated[str | None, "First param"]): pass @tool(desc="A function with an optional input parameter (default: None)") def func_with_optional_param_with_default_None( - param1: Annotated[Optional[str], "First param"] = None, + param1: Annotated[str | None, "First param"] = None, ): pass @tool(desc="A function with an optional input parameter with default value") def func_with_optional_param_with_default_value( - param1: Annotated[Optional[str], "First param"] = "default", + param1: Annotated[str | None, "First param"] = "default", ): pass @@ -220,14 +220,14 @@ def func_with_optional_param_with_bar_syntax_2( @tool(desc="A function with an optional input parameter with union syntax") def func_with_optional_param_with_union_syntax_1( - param1: Annotated[Union[str, None], "First param"] = None, + param1: Annotated[str | None, "First param"] = None, ): pass @tool(desc="A function with an optional input parameter with union syntax") def func_with_optional_param_with_union_syntax_2( - param1: Annotated[Union[None, str], "First param"] = None, + param1: Annotated[None | str, "First param"] = None, ): pass @@ -290,7 +290,7 @@ def func_with_annotated_return() -> Annotated[str, "Annotated return description @tool(desc="A function with an optional return type") -def func_with_optional_return() -> Optional[str]: +def func_with_optional_return() -> str | None: return "maybe output" @@ -305,12 +305,12 @@ def func_with_optional_return_with_bar_syntax_2() -> None | str: @tool(desc="A function with an optional return type that uses union syntax") -def func_with_optional_return_with_union_syntax_1() -> Union[str, None]: +def func_with_optional_return_with_union_syntax_1() -> str | None: return "maybe output" @tool(desc="A function with an optional return type that uses union syntax") -def func_with_optional_return_with_union_syntax_2() -> Union[None, str]: +def func_with_optional_return_with_union_syntax_2() -> None | str: return "maybe output" diff --git a/libs/tests/tool/test_create_tool_definition_errors.py b/libs/tests/tool/test_create_tool_definition_errors.py index 23458ea3..da90337d 100644 --- a/libs/tests/tool/test_create_tool_definition_errors.py +++ b/libs/tests/tool/test_create_tool_definition_errors.py @@ -1,4 +1,4 @@ -from typing import Annotated, Union +from typing import Annotated import pytest from arcade_core.catalog import ToolCatalog @@ -18,7 +18,7 @@ def func_with_missing_return_type(): @tool(desc="A function with a union return type (illegal)") -def func_with_union_return_type_1() -> Union[str, int]: +def func_with_union_return_type_1() -> str | int: return "hello world" @@ -48,7 +48,7 @@ def func_with_union_param_1(param1: str | int): @tool(desc="A function with a union parameter (illegal)") -def func_with_union_param_2(param1: Union[str, int]): +def func_with_union_param_2(param1: str | int): pass diff --git a/libs/tests/tool/test_create_tool_definition_pydantic.py b/libs/tests/tool/test_create_tool_definition_pydantic.py index 38b3c816..c50a5cdc 100644 --- a/libs/tests/tool/test_create_tool_definition_pydantic.py +++ b/libs/tests/tool/test_create_tool_definition_pydantic.py @@ -1,4 +1,4 @@ -from typing import Annotated, Optional, Union +from typing import Annotated import pytest from arcade_core.catalog import ToolCatalog @@ -12,15 +12,33 @@ from arcade_tdk import tool from pydantic import BaseModel, Field -class ProductOutput(BaseModel): - product_name: str = Field(..., description="The name of the product") - price: int = Field(..., description="The price of the product") - stock_quantity: int = Field(..., description="The stock quantity of the product") +class ProductOutputModel(BaseModel): + product_name: str + """The name of the product""" + price: int + """The price of the product""" + stock_quantity: int + """The stock quantity of the product""" + + class Config: + extra = "forbid" @tool(desc="A function that returns a Pydantic model") -def func_returns_pydantic_model() -> Annotated[ProductOutput, "The product, price, and quantity"]: - return ProductOutput( +def func_returns_pydantic_model() -> Annotated[ + ProductOutputModel, "The product, price, and quantity" +]: + """ + Returns a ProductOutput Pydantic model with sample data. + + Returns: + ProductOutput: The product, price, and quantity. + + Example: + >>> func_returns_pydantic_model() + ProductOutput(product_name='Product 1', price=100, stock_quantity=1000) + """ + return ProductOutputModel( product_name="Product 1", price=100, stock_quantity=1000, @@ -36,23 +54,23 @@ def func_takes_pydantic_field_with_description( @tool(desc="A function that accepts an optional Pydantic Field") def func_takes_pydantic_field_optional( - product_name: Optional[str] = Field(None, description="The name of the product"), + product_name: str | None = Field(None, description="The name of the product"), ) -> str: - return product_name + return product_name if product_name is not None else "Product 1" @tool(desc="A function that accepts an optional Pydantic Field with bar syntax") def func_takes_pydantic_field_optional_bar_syntax( product_name: str | None = Field(None, description="The name of the product"), -) -> str: - return product_name +) -> str | None: + return product_name if product_name is not None else None @tool(desc="A function that accepts an optional Pydantic Field with union syntax") def func_takes_pydantic_field_optional_union_syntax( - product_name: Union[str, None] = Field(None, description="The name of the product"), + product_name: str | None = Field(None, description="The name of the product"), ) -> str: - return product_name + return product_name if product_name is not None else "Product 1" # Annotated[] takes precedence over Field() properties @@ -85,9 +103,22 @@ def func_takes_pydantic_field_default( @tool(desc="A function that accepts a Pydantic Field with a default value factory") def func_takes_pydantic_field_default_factory( product_name: str = Field( - ..., description="The name of the product", default_factory=lambda: "Product 1" + default_factory=lambda: "Product 1", description="The name of the product" ), ) -> str: + """ + Accepts a product name with a default value provided by a factory. + + Parameters: + product_name: The name of the product. Defaults to "Product 1" if not provided. + + Returns: + str: The product name. + + Example: + >>> func_takes_pydantic_field_default_factory() + 'Product 1' + """ return product_name @@ -114,9 +145,18 @@ class FilterPriceLessThan(ProductFilter): class ProductSearch(BaseModel): - column: str = Field("Product Name", description="The column to search in") + column: str = Field(..., description="The column to search in") query: str = Field(..., description="The query to search for") - filter_operation: Union[FilterRating, FilterPriceGreaterThan, FilterPriceLessThan] = None + filter_operation: FilterRating | None = Field( + default=None, + description="The filter operation to apply (rating or price filter).", + ) + highest_price: FilterPriceGreaterThan | None = Field( + default=None, description="The highest price to filter by" + ) + lowest_price: FilterPriceLessThan | None = Field( + default=None, description="The lowest price to filter by" + ) class ProductOutput(BaseModel): @@ -129,14 +169,31 @@ class ProductOutput(BaseModel): def read_products( action: Annotated[ProductSearch, "The search query to perform"], cols: list[str] = Field( - ..., - description="The columns to return", default_factory=lambda: ["Product Name", "Price", "Stock Quantity"], + description="The columns to return", ), ) -> Annotated[list[ProductOutput], "Data with the selected columns"]: - """Used to search through products by name and filter by rating or price.""" + """ + Used to search through products by name and filter by rating or price. - pass + Parameters: + action: The search query to perform, as a ProductSearch model. + cols: The columns to return. Defaults to ["Product Name", "Price", "Stock Quantity"]. + + Returns: + list[ProductOutput]: Data with the selected columns. + + Raises: + None + + Example: + >>> await read_products(ProductSearch(query="Widget"), ["Product Name", "Price"]) + """ + # This is a stub implementation for testing; in real code, this would query a database or service. + return [ + ProductOutput(product_name="Widget", price=100, stock_quantity=50), + ProductOutput(product_name="Gadget", price=150, stock_quantity=20), + ] @pytest.mark.parametrize( @@ -146,7 +203,15 @@ def read_products( func_returns_pydantic_model, { "output": ToolOutput( - value_schema=ValueSchema(val_type="json", enum=None), + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + }, + ), available_modes=["value", "error"], description="The product, price, and quantity", ) @@ -299,12 +364,46 @@ def read_products( description="The search query to perform", required=True, inferrable=True, - value_schema=ValueSchema(val_type="json", enum=None), + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "column": ValueSchema(val_type="string", enum=None), + "query": ValueSchema(val_type="string", enum=None), + "filter_operation": ValueSchema( + val_type="json", + enum=None, + properties={ + "column": ValueSchema(val_type="string", enum=None), + "greater_than": ValueSchema( + val_type="integer", enum=None + ), + }, + ), + "highest_price": ValueSchema( + val_type="json", + enum=None, + properties={ + "column": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + }, + ), + "lowest_price": ValueSchema( + val_type="json", + enum=None, + properties={ + "column": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + }, + ), + }, + ), ), InputParameter( name="cols", description="The columns to return", required=False, + inferrable=True, value_schema=ValueSchema( val_type="array", inner_val_type="string", enum=None ), @@ -312,7 +411,16 @@ def read_products( ] ), "output": ToolOutput( - value_schema=ValueSchema(val_type="array", inner_val_type="json", enum=None), + value_schema=ValueSchema( + val_type="array", + inner_val_type="json", + enum=None, + inner_properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + }, + ), available_modes=["value", "error"], description="Data with the selected columns", ), diff --git a/libs/tests/tool/test_create_tool_definition_pydantic_errors.py b/libs/tests/tool/test_create_tool_definition_pydantic_errors.py index 478b041e..554ce71f 100644 --- a/libs/tests/tool/test_create_tool_definition_pydantic_errors.py +++ b/libs/tests/tool/test_create_tool_definition_pydantic_errors.py @@ -1,4 +1,4 @@ -from typing import Annotated, Union +from typing import Annotated import pytest from arcade_core.catalog import ToolCatalog @@ -29,7 +29,7 @@ def func_takes_pydantic_field_non_strict_optional_bar_syntax( @tool(desc="A function that accepts an optional Pydantic Field with non-strict optional syntax") def func_takes_pydantic_field_non_strict_optional_union_syntax( - product_name: Union[str, int, None] = Field(None, description="The name of the product"), + product_name: str | int | None = Field(None, description="The name of the product"), ) -> str: return product_name diff --git a/libs/tests/tool/test_create_tool_definition_typeddict.py b/libs/tests/tool/test_create_tool_definition_typeddict.py new file mode 100644 index 00000000..1ac8ffc7 --- /dev/null +++ b/libs/tests/tool/test_create_tool_definition_typeddict.py @@ -0,0 +1,323 @@ +from typing import Annotated + +import pytest +from arcade_core.catalog import ToolCatalog +from arcade_core.schema import ( + InputParameter, + ToolInput, + ToolOutput, + ValueSchema, +) +from arcade_tdk import tool +from typing_extensions import TypedDict + + +class ProductOutputDict(TypedDict): + """A product with its details.""" + + product_name: str + price: int + stock_quantity: int + + +@tool(desc="A function that returns a TypedDict") +def func_returns_typeddict() -> Annotated[ProductOutputDict, "The product, price, and quantity"]: + """Returns a ProductOutput TypedDict with sample data.""" + return ProductOutputDict( + product_name="Product 1", + price=100, + stock_quantity=1000, + ) + + +@tool(desc="A function that returns a list of TypedDict") +def func_returns_list_of_typeddict() -> Annotated[ + list[ProductOutputDict], "The product, price, and quantity" +]: + """Returns a list of ProductOutput TypedDict with sample data.""" + return [ + ProductOutputDict( + product_name="Product 1", + price=100, + stock_quantity=1000, + ), + ProductOutputDict( + product_name="Product 2", + price=200, + stock_quantity=2000, + ), + ] + + +@tool(desc="A function that accepts an optional TypedDict parameter") +def func_takes_optional_typeddict_param( + product: Annotated[ProductOutputDict | None, "The product information"] = None, +) -> str: + if product is None: + return "No product provided" + return f"{product['product_name']} for price {product['price']}" + + +class ProductOutputDictWithOptional(TypedDict): + """A product with its details.""" + + product_name: str + price: int + stock_quantity: int + description: str | None + + +@tool(desc="A function that returns a TypedDict with an optional field") +def func_returns_typeddict_with_optional_field() -> Annotated[ + ProductOutputDictWithOptional, "The product, price, and quantity" +]: + """ + Returns a ProductOutput TypedDict with sample data. + """ + return ProductOutputDictWithOptional( + product_name="Product 1", + price=100, + stock_quantity=1000, + ) + + +class ProductListDict(TypedDict): + """A collection of products.""" + + category: str + products: list[str] + + +@tool(desc="A function that accepts a TypedDict with list fields") +def func_takes_typeddict_with_list_field( + product_list: Annotated[ProductListDict | None, "The product collection"] = None, +) -> Annotated[list[str], "The product names."]: + """Accepts a product list with category information.""" + if product_list is None: + return ["Laptop", "Phone"] + return product_list["products"] + + +### TypedDict with total=False for optional fields +class OptionalFieldsDict(TypedDict, total=False): + """A TypedDict with all optional fields.""" + + name: str + description: str + price: int + + +@tool(desc="A function that returns a TypedDict with optional fields") +def func_returns_typeddict_optional_fields() -> Annotated[ + OptionalFieldsDict, "Product info with optional fields" +]: + """Returns a TypedDict where some fields may be missing.""" + return OptionalFieldsDict(name="Product 1") + + +### Nested TypedDict example +class AddressDict(TypedDict): + """Address information.""" + + street: str + city: str + zip_code: str + + +class CustomerDict(TypedDict): + """Customer information with nested address.""" + + name: str + email: str + address: AddressDict + + +@tool(desc="A function that returns nested Typedicts") +def func_returns_nested_typedicts() -> Annotated[CustomerDict, "Customer information with address"]: + """Returns a nested TypedDict structure.""" + return CustomerDict( + name="John Doe", + email="john@example.com", + address=AddressDict( + street="123 Main St", + city="Anytown", + zip_code="12345", + ), + ) + + +@pytest.mark.parametrize( + "func_under_test, expected_tool_def_fields", + [ + pytest.param( + func_returns_typeddict, + { + "output": ToolOutput( + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + }, + ), + available_modes=["value", "error"], + description="The product, price, and quantity", + ) + }, + id="func_returns_typeddict", + ), + pytest.param( + func_returns_list_of_typeddict, + { + "output": ToolOutput( + value_schema=ValueSchema( + val_type="array", + inner_val_type="json", + enum=None, + inner_properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + }, + ), + available_modes=["value", "error"], + description="The product, price, and quantity", + ) + }, + id="func_returns_list_of_typeddict", + ), + pytest.param( + func_takes_optional_typeddict_param, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="product", + description="The product information", + required=False, + inferrable=True, + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + }, + ), + ) + ] + ) + }, + id="func_takes_optional_typeddict_param", + ), + pytest.param( + func_returns_typeddict_with_optional_field, + { + "output": ToolOutput( + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "product_name": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + "stock_quantity": ValueSchema(val_type="integer", enum=None), + "description": ValueSchema(val_type="string", enum=None, nullable=True), + }, + ), + available_modes=["value", "error"], + description="The product, price, and quantity", + ) + }, + id="func_returns_typeddict_with_optional_field", + ), + pytest.param( + func_takes_typeddict_with_list_field, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="product_list", + description="The product collection", + required=False, + inferrable=True, + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "category": ValueSchema(val_type="string", enum=None), + "products": ValueSchema( + val_type="array", inner_val_type="string", enum=None + ), + }, + ), + ) + ] + ), + "output": ToolOutput( + value_schema=ValueSchema( + val_type="array", + inner_val_type="string", + enum=None, + ), + available_modes=["value", "error"], + description="The product names.", + ), + }, + id="func_takes_typeddict_with_list_field", + ), + pytest.param( + func_returns_typeddict_optional_fields, + { + "output": ToolOutput( + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "name": ValueSchema(val_type="string", enum=None), + "description": ValueSchema(val_type="string", enum=None), + "price": ValueSchema(val_type="integer", enum=None), + }, + ), + available_modes=["value", "error"], + description="Product info with optional fields", + ) + }, + id="func_returns_typeddict_optional_fields", + ), + pytest.param( + func_returns_nested_typedicts, + { + "output": ToolOutput( + value_schema=ValueSchema( + val_type="json", + enum=None, + properties={ + "name": ValueSchema(val_type="string", enum=None), + "email": ValueSchema(val_type="string", enum=None), + "address": ValueSchema( + val_type="json", + enum=None, + properties={ + "street": ValueSchema(val_type="string", enum=None), + "city": ValueSchema(val_type="string", enum=None), + "zip_code": ValueSchema(val_type="string", enum=None), + }, + ), + }, + ), + available_modes=["value", "error"], + description="Customer information with address", + ) + }, + id="func_returns_nested_typedicts", + ), + ], +) +def test_create_tool_def_from_typeddict(func_under_test, expected_tool_def_fields): + tool_def = ToolCatalog.create_tool_definition(func_under_test, "1.0") + + for field, expected_value in expected_tool_def_fields.items(): + assert getattr(tool_def, field) == expected_value diff --git a/libs/tests/tool/test_create_tool_definition_typeddict_errors.py b/libs/tests/tool/test_create_tool_definition_typeddict_errors.py new file mode 100644 index 00000000..2d2c711d --- /dev/null +++ b/libs/tests/tool/test_create_tool_definition_typeddict_errors.py @@ -0,0 +1,73 @@ +from typing import Annotated + +import pytest +from arcade_core.catalog import ToolCatalog +from arcade_core.errors import ToolDefinitionError +from arcade_tdk import tool +from typing_extensions import NotRequired, TypedDict + + +class ProductWithNotRequired(TypedDict): + """Product with optional field using NotRequired.""" + + name: str + price: float + description: NotRequired[str] # NotRequired in TypedDict field is not supported + + +@tool +def func_takes_typeddict_with_notrequired( + product: Annotated[ProductWithNotRequired, "Product information"], +) -> Annotated[str, "Product summary"]: + """Process a product with NotRequired field.""" + return f"Product: {product['name']}" + + +class ProductWithUnionField(TypedDict): + """Product with union type field.""" + + name: str + price: float | int # Union type in TypedDict field is not supported + stock: int + + +@tool +def func_takes_typeddict_with_union_field( + product: Annotated[ProductWithUnionField, "Product with union price field"], +) -> Annotated[str, "Product info"]: + """Process a product with union type field.""" + return f"Product: {product['name']}, Price: {product['price']}" + + +@tool +def func_takes_optional_typeddict_non_strict( + config: ProductWithNotRequired | None = None, +) -> Annotated[str, "Configuration status"]: + """Process optional TypedDict with non-strict syntax.""" + return "processed" if config else "no config" + + +@pytest.mark.parametrize( + "func_under_test, exception_type", + [ + pytest.param( + func_takes_typeddict_with_notrequired, + ToolDefinitionError, + id="typeddict_with_notrequired", + ), + pytest.param( + func_takes_typeddict_with_union_field, + ToolDefinitionError, + id="typeddict_with_union_field", + ), + pytest.param( + func_takes_optional_typeddict_non_strict, + ToolDefinitionError, + id="optional_typeddict_non_strict", + ), + ], +) +def test_typeddict_errors_raise_tool_definition_error(func_under_test, exception_type): + """Test that various TypedDict error scenarios raise ToolDefinitionError.""" + with pytest.raises(exception_type): + ToolCatalog.create_tool_definition(func_under_test, "1.0") diff --git a/toolkits/google_news/arcade_google_news/__init__.py b/toolkits/google_news/arcade_google_news/__init__.py index e371d62b..42ae5b0c 100644 --- a/toolkits/google_news/arcade_google_news/__init__.py +++ b/toolkits/google_news/arcade_google_news/__init__.py @@ -1,3 +1,19 @@ from arcade_google_news.tools import search_news_stories +from arcade_google_news.types import ( + CountryCode, + GoogleNewsResponse, + LanguageCode, + SearchNewsOutput, + SearchNewsParams, + SimplifiedNewsResult, +) -__all__ = ["search_news_stories"] +__all__ = [ + "search_news_stories", + "CountryCode", + "GoogleNewsResponse", + "LanguageCode", + "SearchNewsOutput", + "SearchNewsParams", + "SimplifiedNewsResult", +] diff --git a/toolkits/google_news/arcade_google_news/google_data.py b/toolkits/google_news/arcade_google_news/google_data.py index 789e3183..22cdec15 100644 --- a/toolkits/google_news/arcade_google_news/google_data.py +++ b/toolkits/google_news/arcade_google_news/google_data.py @@ -1,4 +1,4 @@ -COUNTRY_CODES = { +COUNTRY_CODES: dict[str, str] = { "af": "Afghanistan", "al": "Albania", "dz": "Algeria", @@ -246,7 +246,7 @@ COUNTRY_CODES = { } -LANGUAGE_CODES = { +LANGUAGE_CODES: dict[str, str] = { "ar": "Arabic", "bn": "Bengali", "da": "Danish", diff --git a/toolkits/google_news/arcade_google_news/tools/google_news.py b/toolkits/google_news/arcade_google_news/tools/google_news.py index 73c81a40..71f76a70 100644 --- a/toolkits/google_news/arcade_google_news/tools/google_news.py +++ b/toolkits/google_news/arcade_google_news/tools/google_news.py @@ -1,12 +1,20 @@ -from typing import Annotated, Any +from typing import Annotated from arcade_tdk import ToolContext, tool from arcade_tdk.errors import ToolExecutionError -from arcade_google_news.constants import DEFAULT_GOOGLE_NEWS_COUNTRY, DEFAULT_GOOGLE_NEWS_LANGUAGE +from arcade_google_news.constants import ( + DEFAULT_GOOGLE_NEWS_COUNTRY, + DEFAULT_GOOGLE_NEWS_LANGUAGE, +) from arcade_google_news.exceptions import CountryNotFoundError, LanguageNotFoundError from arcade_google_news.google_data import COUNTRY_CODES, LANGUAGE_CODES -from arcade_google_news.utils import call_serpapi, extract_news_results, prepare_params +from arcade_google_news.types import CountryCode, LanguageCode, SearchNewsOutput +from arcade_google_news.utils import ( + call_serpapi, + extract_news_results, + prepare_params, +) @tool(requires_secrets=["SERP_API_KEY"]) @@ -17,12 +25,13 @@ async def search_news_stories( "Keywords to search for news articles. E.g. 'Apple launches new iPhone'.", ], country_code: Annotated[ - str | None, - "2-character country code to search for news articles. E.g. 'us' (United States). " + CountryCode | None, + "2-character country code to search for news articles. " + "E.g. 'us' (United States). " f"Defaults to '{DEFAULT_GOOGLE_NEWS_COUNTRY}'.", ] = None, language_code: Annotated[ - str, + LanguageCode, "2-character language code to search for news articles. E.g. 'en' (English). " f"Defaults to '{DEFAULT_GOOGLE_NEWS_LANGUAGE}'.", ] = DEFAULT_GOOGLE_NEWS_LANGUAGE, @@ -31,7 +40,7 @@ async def search_news_stories( "Maximum number of news articles to return. Defaults to None " "(returns all results found by the API).", ] = None, -) -> Annotated[dict[str, list[dict[str, Any]]], "News results."]: +) -> Annotated[SearchNewsOutput, "News search results with article details."]: """Search for news articles related to a given query.""" if not keywords: raise ToolExecutionError("Keywords are required to search for news articles.") @@ -44,4 +53,4 @@ async def search_news_stories( params = prepare_params("google_news", q=keywords, gl=country_code, hl=language_code) results = call_serpapi(context, params) - return {"news_results": extract_news_results(results, limit=limit)} + return SearchNewsOutput(news_results=extract_news_results(results, limit=limit)) diff --git a/toolkits/google_news/arcade_google_news/types.py b/toolkits/google_news/arcade_google_news/types.py new file mode 100644 index 00000000..5a6c4c22 --- /dev/null +++ b/toolkits/google_news/arcade_google_news/types.py @@ -0,0 +1,196 @@ +"""Type definitions for Google News API responses and parameters.""" + +from typing_extensions import TypedDict + +# For now, we'll use str type alias to maintain compatibility +# In the future, these could be converted to proper Literal types +CountryCode = str +LanguageCode = str + + +class SearchNewsParams(TypedDict): + """Input parameters for searching news articles.""" + + keywords: str + """Search query terms to find relevant news articles \ + (e.g., 'Apple launches new iPhone').""" + + country_code: CountryCode | None + """Optional 2-letter country code to filter news by region \ + (e.g., 'us' for United States, 'uk' for United Kingdom).""" + + language_code: LanguageCode | None + """Optional 2-letter language code to filter news by language \ + (e.g., 'en' for English, 'es' for Spanish).""" + + limit: int | None + """Optional maximum number of news articles to return. \ + If not specified, returns all results from the API.""" + + +class SourceInfo(TypedDict, total=False): + """Information about the news source/publication.""" + + name: str + """Name of the publication (e.g., 'CNN', 'BBC News', 'The New York Times').""" + + icon: str + """URL to the source's favicon or logo image.""" + + authors: list[str] + """List of author names for the article, if available.""" + + +class NewsResult(TypedDict, total=False): + """Individual news article from the Google News API response.""" + + position: int + """Ranking position of this result in the search results.""" + + title: str + """Headline or title of the news article.""" + + link: str + """Full URL to the original news article.""" + + source: SourceInfo + """Information about the publication source.""" + + date: str + """Publication date and time (e.g., '2 hours ago', 'Dec 15, 2023').""" + + snippet: str + """Brief excerpt or summary from the article content.""" + + thumbnail: str + """URL to a high-resolution thumbnail image for the article.""" + + thumbnail_small: str + """URL to a low-resolution thumbnail image for the article.""" + + story_token: str + """Token for accessing full coverage of this news story across multiple sources.""" + + stories: list["NewsResult"] + """Related news stories from other sources covering the same topic.""" + + highlight: dict + """Additional highlighted information about the story.""" + + +class SearchMetadata(TypedDict, total=False): + """Metadata about the search request and processing.""" + + id: str + """Unique identifier for this search request within SerpApi.""" + + status: str + """Current processing status ('Processing', 'Success', or 'Error').""" + + json_endpoint: str + """URL to retrieve the JSON results for this search.""" + + created_at: str + """Timestamp when the search request was created.""" + + processed_at: str + """Timestamp when the search request was processed.""" + + google_news_url: str + """Original Google News URL that would return these results.""" + + total_time_taken: float + """Total time in seconds taken to process this search.""" + + +class SearchParameters(TypedDict, total=False): + """Parameters used for the search request.""" + + engine: str + """Search engine used (always 'google_news' for this API).""" + + q: str + """Search query string.""" + + gl: str + """Country code used for geographic filtering.""" + + hl: str + """Language code used for language filtering.""" + + topic_token: str + """Token for accessing specific news topics (e.g., 'World', 'Business', 'Technology').""" + + publication_token: str + """Token for accessing news from specific publishers.""" + + +class MenuLink(TypedDict): + """Navigation link for news categories or topics.""" + + title: str + """Display text for the menu item (e.g., 'Technology', 'Sports', 'Business').""" + + topic_token: str + """Token to access this specific topic or category.""" + + serpapi_link: str + """SerpApi URL to search within this topic.""" + + +class TopStoriesLink(TypedDict): + """Link to top stories section.""" + + topic_token: str + """Token to access top stories.""" + + serpapi_link: str + """SerpApi URL to retrieve top stories.""" + + +class GoogleNewsResponse(TypedDict, total=False): + """Complete response from the Google News API.""" + + search_metadata: SearchMetadata + """Metadata about the search request and processing.""" + + search_parameters: SearchParameters + """Parameters that were used for this search.""" + + news_results: list[NewsResult] + """List of news articles matching the search criteria.""" + + menu_links: list[MenuLink] + """Navigation links to different news categories and topics.""" + + top_stories_link: TopStoriesLink + """Link to access top stories.""" + + title: str + """Title of the page or topic being displayed.""" + + +class SimplifiedNewsResult(TypedDict): + """Simplified news article format for tool output.""" + + title: str + """Headline of the news article.""" + + link: str + """URL to the full article.""" + + source: str | None + """Name of the publication source.""" + + date: str | None + """When the article was published.""" + + snippet: str | None + """Brief excerpt from the article.""" + + +class SearchNewsOutput(TypedDict): + """Output format for the search_news_stories tool.""" + + news_results: list[SimplifiedNewsResult] + """List of news articles in simplified format.""" diff --git a/toolkits/google_news/arcade_google_news/utils.py b/toolkits/google_news/arcade_google_news/utils.py index a401b7eb..458faae4 100644 --- a/toolkits/google_news/arcade_google_news/utils.py +++ b/toolkits/google_news/arcade_google_news/utils.py @@ -5,6 +5,8 @@ from arcade_tdk import ToolContext from arcade_tdk.errors import ToolExecutionError from serpapi import Client as SerpClient +from arcade_google_news.types import GoogleNewsResponse, SimplifiedNewsResult + def prepare_params(engine: str, **kwargs: Any) -> dict[str, Any]: """ @@ -23,7 +25,7 @@ def prepare_params(engine: str, **kwargs: Any) -> dict[str, Any]: return params -def call_serpapi(context: ToolContext, params: dict) -> dict: +def call_serpapi(context: ToolContext, params: dict[str, Any]) -> GoogleNewsResponse: """ Execute a search query using the SerpAPI client and return the results as a dictionary. @@ -38,7 +40,7 @@ def call_serpapi(context: ToolContext, params: dict) -> dict: client = SerpClient(api_key=api_key) try: search = client.search(params) - return cast(dict[str, Any], search.as_dict()) + return cast(GoogleNewsResponse, search.as_dict()) except Exception as e: # SerpAPI error messages sometimes contain the API key, so we need to sanitize it sanitized_e = re.sub(r"(api_key=)[^ &]+", r"\1***", str(e)) @@ -48,16 +50,20 @@ def call_serpapi(context: ToolContext, params: dict) -> dict: ) -def extract_news_results(results: dict[str, Any], limit: int | None = None) -> list[dict[str, Any]]: - news_results = [] +def extract_news_results( + results: GoogleNewsResponse, limit: int | None = None +) -> list[SimplifiedNewsResult]: + news_results: list[SimplifiedNewsResult] = [] for result in results.get("news_results", []): - news_results.append({ - "title": result.get("title"), - "snippet": result.get("snippet"), - "link": result.get("link"), - "date": result.get("date"), - "source": result.get("source", {}).get("name"), - }) + news_results.append( + SimplifiedNewsResult( + title=result.get("title", ""), + link=result.get("link", ""), + source=result.get("source", {}).get("name"), + date=result.get("date"), + snippet=result.get("snippet"), + ) + ) if limit: return news_results[:limit]