Skip to content

Agents and tools

Registry and contracts

A registered strategy declares name, task_type, output_type, and config_cls. It provides from_config_async(...) and the predict_async(...) method appropriate for its task family.

The pipeline creates one agent per task. Pass the class directly to predict_async(...), or register it with agent_registry.register(YourAgent) for name-based lookup in the current Python process.

Return pred_query=None for an intentional abstention in a simple task. Let unexpected prediction exceptions propagate so the pipeline logs them and records empty outputs. Use extra_pred_info to retain predictions from before postprocessing. Reference queries remain in task outputs for evaluation; include only question context and schema in model prompts.

See Build a custom agent for an implementation.

agent_registry module-attribute 

agent_registry = ClassRegistry[Any]('agent')

SimpleAgentProtocol

Bases: Protocol

Strategy that predicts one query for a simple task.

name class-attribute 

name: str

task_type class-attribute 

task_type: str

output_type class-attribute 

output_type: str

config_cls class-attribute 

config_cls: Any

predict_async async 

predict_async(
    task: SimpleNL2QTask, db_connector: DataConnector
) -> SimpleNL2QTaskOutput

AmbigSQLAgentProtocol

Bases: Protocol

Strategy that resolves and predicts queries for an ambiguous SQL task.

name class-attribute 

name: str

task_type class-attribute 

task_type: str

output_type class-attribute 

output_type: str

config_cls class-attribute 

config_cls: Any

predict_async async 

predict_async(
    task: AmbigNL2QTask,
    db_connector: SQLConnector,
    user_simulator: UserSimulatorProtocol,
) -> (
    SimpleAmbigNL2QTaskOutput
    | FlatAmbigNL2QTaskOutput
    | StructuredAmbigNL2QTaskOutput
)

DbtAgentProtocol

Bases: Protocol

Strategy that produces a transformed dbt project.

name class-attribute 

name: str

task_type class-attribute 

task_type: str

output_type class-attribute 

output_type: str

config_cls class-attribute 

config_cls: Any

predict_async async 

predict_async(
    task: DbtTask, db_connector: SQLConnector
) -> DbtTaskOutput

BasicAgentConfig

Bases: BaseModel

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

Simple strategies

Compare built-in methods in Agents.

DirectPromptAgent 

DirectPromptAgent(config: BasicAgentConfig)

name class-attribute 

name: str = 'direct_prompting'

task_type class-attribute 

task_type: str = 'simple'

output_type class-attribute 

output_type: str = 'simple'

config_cls class-attribute 

config_cls: type[BasicAgentConfig] = BasicAgentConfig

config instance-attribute 

config = config

from_config_async async classmethod 

from_config_async(
    config: BasicAgentConfig,
) -> DirectPromptAgent

predict_async async 

predict_async(
    task: SimpleNL2QTask, db_connector: DataConnector
) -> SimpleNL2QTaskOutput

FullSchemaAgent 

FullSchemaAgent(config: BasicAgentConfig)

name class-attribute 

name: str = 'full_schema'

task_type class-attribute 

task_type: str = 'simple'

output_type class-attribute 

output_type: str = 'simple'

config_cls class-attribute 

config_cls: type[BasicAgentConfig] = BasicAgentConfig

config instance-attribute 

config = config

from_config_async async classmethod 

from_config_async(
    config: BasicAgentConfig,
) -> FullSchemaAgent

predict_async async 

predict_async(
    task: SimpleNL2QTask, db_connector: DataConnector
) -> SimpleNL2QTaskOutput

SchemaLinkingAgent 

SchemaLinkingAgent(
    config: SchemaLinkingAgentConfig,
    few_shot_dataset: NL2QDataset | None = None,
    few_shot_embeddings: NDArray[Any] | None = None,
)

name class-attribute 

name: str = 'schema_linking'

task_type class-attribute 

task_type: str = 'simple'

output_type class-attribute 

output_type: str = 'simple'

config_cls class-attribute 

config_cls: type[SchemaLinkingAgentConfig] = (
    SchemaLinkingAgentConfig
)

config instance-attribute 

config = config

few_shot_dataset instance-attribute 

few_shot_dataset = few_shot_dataset

few_shot_embeddings instance-attribute 

few_shot_embeddings = few_shot_embeddings

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

schema_linker instance-attribute 

schema_linker = (
    SchemaLinker(config)
    if config.do_schema_linking
    else None
)

postprocessor instance-attribute 

postprocessor = (
    Postprocessor(config)
    if config.do_postprocessing
    else None
)

from_config_async async classmethod 

from_config_async(
    config: SchemaLinkingAgentConfig,
    few_shot_dataset: NL2QDataset | None = None,
) -> SchemaLinkingAgent

predict_async async 

predict_async(
    task: SimpleNL2QTask, db_connector: DataConnector
) -> SimpleNL2QTaskOutput

SchemaLinkingAgentConfig

Bases: BasicAgentConfig

min_columns_for_schema_linking class-attribute instance-attribute 

min_columns_for_schema_linking: int = 20

num_few_shot_examples class-attribute instance-attribute 

num_few_shot_examples: int = 0

do_schema_linking class-attribute instance-attribute 

do_schema_linking: bool = True

do_postprocessing class-attribute instance-attribute 

do_postprocessing: bool = True

question_embedder_embedding_llm class-attribute instance-attribute 

question_embedder_embedding_llm: str = (
    "openai:text-embedding-3-small"
)

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

SchemaDiscoveryAgent 

SchemaDiscoveryAgent(config: SchemaDiscoveryAgentConfig)

name class-attribute 

name: str = 'schema_discovery'

task_type class-attribute 

task_type: str = 'simple'

output_type class-attribute 

output_type: str = 'simple'

config_cls class-attribute 

config_cls: type[SchemaDiscoveryAgentConfig] = (
    SchemaDiscoveryAgentConfig
)

config instance-attribute 

config = config

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

from_config_async async classmethod 

from_config_async(
    config: SchemaDiscoveryAgentConfig,
) -> SchemaDiscoveryAgent

predict_async async 

predict_async(
    task: SimpleNL2QTask, db_connector: DataConnector
) -> SimpleNL2QTaskOutput

SchemaDiscoveryAgentConfig

Bases: BasicAgentConfig

db_summarizer_llm class-attribute instance-attribute 

db_summarizer_llm: str = 'openai:gpt-5.6-sol'

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

Schema-linking components

SchemaLinker 

SchemaLinker(config: SchemaLinkingAgentConfig)

config instance-attribute 

config = config

expand_schema_async async 

expand_schema_async(
    ctx: SchemaLinkingContext,
    schema_to_expand: SQLSchema,
    task: SimpleNL2QTask,
    batch_size: int = 5,
) -> SQLSchema

link_schema_async(
    ctx: SchemaLinkingContext, task: SimpleNL2QTask
) -> SQLSchema

Postprocessor 

Postprocessor(config: SchemaLinkingAgentConfig)

config instance-attribute 

config = config

postprocess_async async 

postprocess_async(
    ctx: SchemaLinkingContext,
    task: SimpleNL2QTask,
    pred_query: PredQuery,
) -> PredQuery

SchemaLinkingContext dataclass 

SchemaLinkingContext(
    task: NL2QTask,
    db_connector: SQLConnector,
    preprocessed_schema: SQLSchema,
    schema_formatter: SQLSchemaFormatter,
    usage: Usage,
    tools: dict[str, AgentTool],
    trajectories: list[Trajectory],
    er_diagram: ERDiagram | None = None,
    er_diagram_formatter: MermaidERDiagramFormatter
    | None = None,
    few_shot_examples: list[SimpleNL2QTask] = list(),
)

Bases: TaskRunContext

db_connector instance-attribute 

db_connector: SQLConnector

er_diagram class-attribute instance-attribute 

er_diagram: ERDiagram | None = None

er_diagram_formatter class-attribute instance-attribute 

er_diagram_formatter: MermaidERDiagramFormatter | None = (
    None
)

few_shot_examples class-attribute instance-attribute 

few_shot_examples: list[SimpleNL2QTask] = field(
    default_factory=list
)

task instance-attribute 

task: NL2QTask

preprocessed_schema instance-attribute 

preprocessed_schema: SQLSchema

schema_formatter instance-attribute 

schema_formatter: SQLSchemaFormatter

usage instance-attribute 

usage: Usage

tools instance-attribute 

tools: dict[str, AgentTool]

trajectories instance-attribute 

trajectories: list[Trajectory]

TaskRunContext dataclass 

TaskRunContext(
    task: NL2QTask,
    db_connector: DataConnector,
    preprocessed_schema: SQLSchema,
    schema_formatter: SQLSchemaFormatter,
    usage: Usage,
    tools: dict[str, AgentTool],
    trajectories: list[Trajectory],
)

task instance-attribute 

task: NL2QTask

db_connector instance-attribute 

db_connector: DataConnector

preprocessed_schema instance-attribute 

preprocessed_schema: SQLSchema

schema_formatter instance-attribute 

schema_formatter: SQLSchemaFormatter

usage instance-attribute 

usage: Usage

tools instance-attribute 

tools: dict[str, AgentTool]

trajectories instance-attribute 

trajectories: list[Trajectory]

Ambiguity-aware strategies

AmbigSimpleSQLAgent 

AmbigSimpleSQLAgent(config: AmbigSimpleSQLAgentConfig)

name class-attribute 

name: str = 'ambig_simple_sql_agent'

task_type class-attribute 

task_type: str = 'ambig'

output_type class-attribute 

output_type: str = 'ambig-simple'

config_cls class-attribute 

config_cls: type[AmbigSimpleSQLAgentConfig] = (
    AmbigSimpleSQLAgentConfig
)

config instance-attribute 

config = config

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

from_config_async async classmethod 

from_config_async(
    config: AmbigSimpleSQLAgentConfig,
) -> AmbigSimpleSQLAgent

predict_async async 

predict_async(
    task: AmbigNL2QTask,
    db_connector: SQLConnector,
    user_simulator: UserSimulatorProtocol,
) -> SimpleAmbigNL2QTaskOutput

AmbigSimpleSQLAgentConfig

Bases: BasicAgentConfig

user_patience class-attribute instance-attribute 

user_patience: int | Literal["NUM_AMBIG_POINTS"] | None = (
    None
)

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

AmbigFlatSQLAgent 

AmbigFlatSQLAgent(config: AmbigFlatSQLAgentConfig)

name class-attribute 

name: str = 'ambig_flat_sql_agent'

task_type class-attribute 

task_type: str = 'ambig'

output_type class-attribute 

output_type: str = 'ambig-flat'

config_cls class-attribute 

config_cls: type[AmbigFlatSQLAgentConfig] = (
    AmbigFlatSQLAgentConfig
)

config instance-attribute 

config = config

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

from_config_async async classmethod 

from_config_async(
    config: AmbigFlatSQLAgentConfig,
) -> AmbigFlatSQLAgent

predict_async async 

predict_async(
    task: AmbigNL2QTask,
    db_connector: SQLConnector,
    user_simulator: UserSimulatorProtocol,
) -> FlatAmbigNL2QTaskOutput

AmbigFlatSQLAgentConfig

Bases: BasicAgentConfig

query_for_intended_only class-attribute instance-attribute 

query_for_intended_only: bool = True

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

AmbigStructuredSQLAgent 

AmbigStructuredSQLAgent(
    config: AmbigStructuredSQLAgentConfig,
)

name class-attribute 

name: str = 'ambig_structured_sql_agent'

task_type class-attribute 

task_type: str = 'ambig'

output_type class-attribute 

output_type: str = 'ambig-structured'

config_cls class-attribute 

config_cls: type[AmbigStructuredSQLAgentConfig] = (
    AmbigStructuredSQLAgentConfig
)

config instance-attribute 

config = config

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

from_config_async async classmethod 

from_config_async(
    config: AmbigStructuredSQLAgentConfig,
) -> AmbigStructuredSQLAgent

predict_async async 

predict_async(
    task: AmbigNL2QTask,
    db_connector: SQLConnector,
    user_simulator: UserSimulatorProtocol,
) -> StructuredAmbigNL2QTaskOutput

AmbigStructuredSQLAgentConfig

Bases: BasicAgentConfig

query_for_intended_only class-attribute instance-attribute 

query_for_intended_only: bool = True

use_gold_phrases class-attribute instance-attribute 

use_gold_phrases: bool = False

use_gold_ambiguity_points class-attribute instance-attribute 

use_gold_ambiguity_points: bool = False

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

dbt strategy

DbtAgent works on Spider 2.0 dbt projects and produces transformed tables. Give Spider2DbtDatasetLoader a distinct workspace_dir for each run so it can create isolated project copies and connectors before prediction. Evaluate with Spider2DuckdbMatch; the query execution stage does not execute dbt projects.

DbtAgent 

DbtAgent(config: DbtAgentConfig)

name class-attribute 

name: str = 'dbt_agent'

task_type class-attribute 

task_type: str = 'dbt'

output_type class-attribute 

output_type: str = 'dbt'

config_cls class-attribute 

config_cls: type[DbtAgentConfig] = DbtAgentConfig

config instance-attribute 

config = config

formatter instance-attribute 

formatter = config.create_schema_formatter('sql')

from_config_async async classmethod 

from_config_async(config: DbtAgentConfig) -> DbtAgent

predict_async async 

predict_async(
    task: DbtTask, db_connector: SQLConnector
) -> DbtTaskOutput

DbtAgentConfig

Bases: BasicAgentConfig

db_summarizer_llm class-attribute instance-attribute 

db_summarizer_llm: str = 'openai:gpt-5.6-sol'

use_bash_tool class-attribute instance-attribute 

use_bash_tool: bool = False

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-5.6-sol'

schema_formatter class-attribute instance-attribute 

schema_formatter: str | None = None

compact_table_families class-attribute instance-attribute 

compact_table_families: bool = True

temperature class-attribute instance-attribute 

temperature: float | None = None

max_steps class-attribute instance-attribute 

max_steps: int = Field(default=50, ge=1)

formatter_max_total_columns class-attribute instance-attribute 

formatter_max_total_columns: int | None = 5000

use_column_descriptions class-attribute instance-attribute 

use_column_descriptions: bool = True

reasoning class-attribute instance-attribute 

reasoning: ReasoningLevel | None = None

service_tier class-attribute instance-attribute 

service_tier: ServiceTier | None = None

create_schema_formatter 

create_schema_formatter(
    kind: Literal["sql"],
) -> SQLSchemaFormatter

create_schema_formatter(
    kind: Literal["property_graph"],
) -> PropertyGraphSchemaFormatter

create_schema_formatter(
    kind: Literal["sql", "property_graph"],
) -> SQLSchemaFormatter | PropertyGraphSchemaFormatter

Create a compatible formatter with this agent's schema options.

to_model_settings 

to_model_settings() -> dict[str, Any]

User simulation

UserSimulator 

UserSimulator(config: UserSimulatorConfig)

config instance-attribute 

config = config

control_agent instance-attribute 

control_agent = make_agent(
    self.config.llm,
    tools=[],
    instructions=control_agent_system_prompt,
    model_settings={"temperature": self.config.temperature},
)

usage 

usage() -> Usage

trajectory 

trajectory() -> Trajectory

user_effort 

user_effort() -> float

from_ambig_nl2q_task classmethod 

from_ambig_nl2q_task(
    task: AmbigNL2QTask,
    llm: str = "openai:gpt-4.1-2025-04-14",
    temperature: float = 0.0,
    include_history: bool = True,
    answer_with_multiple_ambig_points: bool = False,
) -> UserSimulator

ask_free_text_async async 

ask_free_text_async(
    question: UserFreeTextQuestion,
) -> UserFreeTextAnswer | None

ask_multiple_choice_async async 

ask_multiple_choice_async(
    question: UserMultipleChoiceQuestion,
) -> UserMultipleChoiceAnswer | None

ask_value_async async 

ask_value_async(
    question: UserValueQuestion,
) -> UserValueAnswer | None

ask_async async 

ask_async(question: UserQuestion) -> UserAnswer | None

UserSimulatorConfig

Bases: BaseModel

task instance-attribute 

task: str

ambig_points instance-attribute 

ambig_points: list[NLAmbigPoint]

llm class-attribute instance-attribute 

llm: str = 'openai:gpt-4.1-2025-04-14'

temperature class-attribute instance-attribute 

temperature: float = 0.0

include_history class-attribute instance-attribute 

include_history: bool = True

answer_with_multiple_ambig_points class-attribute instance-attribute 

answer_with_multiple_ambig_points: bool = False

NLAmbigPoint

Bases: BaseModel

id instance-attribute 

id: str

phrase instance-attribute 

phrase: str

intended_interpretation instance-attribute 

intended_interpretation: str

all_interpretations instance-attribute 

all_interpretations: list[str] | None

Research tools

For general data, browser, and filesystem tools, see the library reference.

AskUserTool 

AskUserTool(
    user_simulator: UserSimulatorProtocol,
    patience: int | None = None,
)

name class-attribute instance-attribute 

name: ClassVar = 'ask_user'

user_simulator instance-attribute 

user_simulator = user_simulator

patience instance-attribute 

patience = patience

__call__ async 

__call__(question: str) -> str

Ask the user a question and get a response.

Parameters:

Name Type Description Default
question str

The question to ask the user.

 required

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> Tool

metrics 

metrics() -> AskUserToolMetrics

FinishTool 

FinishTool()

name class-attribute instance-attribute 

name: ClassVar = 'finish'

__call__ 

__call__(trajectory: Trajectory) -> None

Finish the task. The last executed query will be considered as the final answer. No parameters needed.

Example: 

finish()

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> ToolOutput[None]

metrics 

metrics() -> FinishToolMetrics

GetSchemaTool 

GetSchemaTool(
    schema: SQLSchema, formatter: SQLSchemaFormatter
)

Tool that retrieves the full database schema.

Formats and returns the complete schema using the configured formatter.

Attributes:

Name Type Description
schema

The physical SQL schema containing all available tables.
formatter

The formatter used to render the schema as text.

name class-attribute instance-attribute 

name: ClassVar = 'get_schema'

schema instance-attribute 

schema = schema

formatter instance-attribute 

formatter = formatter

__call__ async 

__call__() -> str

Get the schema of the database.

Example: 

get_schema()

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> Tool

metrics 

metrics() -> GetSchemaToolMetrics

GetColumnDescriptionTool 

GetColumnDescriptionTool(schema: SQLSchema)

Tool that retrieves the description of a specific column in a table.

Looks up a column by schema name, table name, and column name, then returns its description if available.

Attributes:

Name Type Description
schema

The physical SQL schema containing all available tables.

name class-attribute instance-attribute 

name: ClassVar = 'get_column_description'

schema instance-attribute 

schema = schema

__call__ async 

__call__(
    schema_name: str | None,
    table_name: str,
    column_name: str,
) -> str

Get the description of a column of a table.

Parameters:

Name Type Description Default
schema_name str | None

The name of the schema, or None if schema is not applicable.

 required
table_name str

The name of the table.

 required
column_name str

The name of the column.

 required

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> Tool

metrics 

metrics() -> GetColumnDescriptionToolMetrics

SearchKeywordsTool 

SearchKeywordsTool(
    db_connector: SQLConnector,
    max_visible_results: int = 40,
)

name class-attribute instance-attribute 

name: ClassVar = 'search_keywords'

db_connector instance-attribute 

db_connector = db_connector

max_visible_results instance-attribute 

max_visible_results = max_visible_results

__call__ async 

__call__(
    schema_name: str | None,
    table_name: str,
    column_name: str,
    keywords: list[str],
) -> str

Search for values in a column of a table that match any of the keywords.

Parameters:

Name Type Description Default
schema_name str | None

The name of the schema to which the table belongs, or None if schema is not applicable.

 required
table_name str

The name of the table to which the column belongs.

 required
column_name str

The name of the column to search in. The datatype of the column must be text-like.

 required
keywords list[str]

A list of keywords to search for. A value is considered a match if it contains any of the keywords.

 required

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> Tool

metrics 

metrics() -> SearchKeywordsToolMetrics

RunDbtTool 

RunDbtTool(
    working_dir: str,
    pre_run_hook: Callable[[], Awaitable[None]]
    | None = None,
)

Execute dbt CLI commands in a project working directory.

The tool automatically sets --project-dir and --profiles-dir to the working directory, ensuring dbt always operates on the correct project. Only a fixed set of subcommands is allowed.

An optional pre_run_hook can be supplied (e.g. to restore a pristine database before each build). The hook is invoked only before run and build commands, which re-materialise all models; read-only commands such as test, ls, and compile skip the hook so they operate on the database state left by the most recent build.

Attributes:

Name Type Description
working_dir

Path to the dbt project directory.
pre_run_hook

Optional async callback invoked before run and build commands (e.g. to restore a pristine database).

name class-attribute instance-attribute 

name: ClassVar = 'run_dbt'

__call__ async 

__call__(
    command: DbtCommand,
    select: str | None = None,
    exclude: str | None = None,
) -> str

Run a dbt CLI command in the project directory.

Parameters:

Name Type Description Default
command DbtCommand

The dbt subcommand to run. One of "run", "build", "test", "compile", "debug", "ls", "deps".

 required
select str | None

Optional --select node selector (e.g. "my_model" or "tag:daily"). Applies to run, build, test, compile, and ls.

 None
exclude str | None

Optional --exclude node selector. Same commands as select.

 None

as_pydantic_ai_tool 

as_pydantic_ai_tool() -> Tool

metrics 

metrics() -> RunDbtToolMetrics

DbtCommand module-attribute 

DbtCommand = Literal[
    "run", "build", "test", "compile", "debug", "ls", "deps"
]

Tool metrics

AskUserToolMetrics

Bases: BaseModel

num_calls class-attribute instance-attribute 

num_calls: int = 0

user_refused_to_answer class-attribute instance-attribute 

user_refused_to_answer: int = 0

FinishToolMetrics

Bases: BaseModel

num_calls class-attribute instance-attribute 

num_calls: int = 0

error_no_query_executed class-attribute instance-attribute 

error_no_query_executed: int = 0

GetSchemaToolMetrics

Bases: BaseModel

num_calls class-attribute instance-attribute 

num_calls: int = 0

GetColumnDescriptionToolMetrics

Bases: BaseModel

num_calls class-attribute instance-attribute 

num_calls: int = 0

error_table_not_found class-attribute instance-attribute 

error_table_not_found: int = 0

error_column_not_found class-attribute instance-attribute 

error_column_not_found: int = 0

SearchKeywordsToolMetrics

Bases: BaseModel

num_calls class-attribute instance-attribute 

num_calls: int = 0

error_table_not_found class-attribute instance-attribute 

error_table_not_found: int = 0

error_column_not_found class-attribute instance-attribute 

error_column_not_found: int = 0

error_column_not_string class-attribute instance-attribute 

error_column_not_string: int = 0

RunDbtToolMetrics

Bases: BaseModel

num_run class-attribute instance-attribute 

num_run: int = 0

num_build class-attribute instance-attribute 

num_build: int = 0

num_test class-attribute instance-attribute 

num_test: int = 0

num_compile class-attribute instance-attribute 

num_compile: int = 0

num_debug class-attribute instance-attribute 

num_debug: int = 0

num_ls class-attribute instance-attribute 

num_ls: int = 0

num_deps class-attribute instance-attribute 

num_deps: int = 0

error_count class-attribute instance-attribute 

error_count: int = 0

num_run_success class-attribute instance-attribute 

num_run_success: int = 0

num_run_failure class-attribute instance-attribute 

num_run_failure: int = 0

last_run_success class-attribute instance-attribute 

last_run_success: bool | None = None

Tracing

Group model activity by task QID. See tracing setup for provider settings.

configure_research_observability 

configure_research_observability() -> None

Configure research tracing and enable agent instrumentation once.

trace_prediction 

trace_prediction(
    predict_async_fn: Callable[..., Any],
) -> Callable[..., Any]

Trace a top-level research prediction without nesting duplicate spans.