AI Agents¶
lumen.ai.agents
¶
__all__ = ['Agent', 'AnalysisAgent', 'AnalystAgent', 'ChatAgent', 'DbtslAgent', 'DocumentListAgent', 'hvPlotAgent', 'SourceAgent', 'SQLAgent', 'TableListAgent', 'ValidationAgent', 'VegaLiteAgent']
module-attribute
¶
Agent
¶
Bases:
Agents are actors responsible for taking a user query and performing a particular task, either by adding context or generating outputs.
Agents have access to an LLM and are given context and can solve tasks by executing a series of prompts or by rendering contents such as forms or widgets to gather user input.
agents = param.List(doc='\n List of agents this agent can invoke.')
class-attribute
instance-attribute
¶
debug = param.Boolean(default=False, doc='\n Whether to enable verbose error reporting.')
class-attribute
instance-attribute
¶
llm = param.ClassSelector(class_=Llm, doc='\n The LLM implementation to query.')
class-attribute
instance-attribute
¶
user = param.String(default='Agent', doc='\n The name of the user that will be respond to the user query.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
Additional checks to determine if the agent should be used.
respond(messages, context, step_title=None)
async
¶
Provides a response to the user query.
The type of the response may be a simple string or an object.
Arguments
messages: list[Message] The list of messages corresponding to the user query and any other system messages to be included. context: TContext A mapping containing context for the agent to perform its task. step_title: str | None If the Agent response is part of a longer query this describes the step currently being processed.
AnalysisAgent
¶
Bases:
analyses = param.List([])
class-attribute
instance-attribute
¶
conditions = param.List(default=['Use for custom analysis, advanced analytics, or domain-specific methods', "Use when the user query matches one of the available analyses' name or description below.", "Include the selected analysis' required cols in the instructions", 'NOT for simple queries or basic visualizations'])
class-attribute
instance-attribute
¶
input_schema = AnalysisInputs
class-attribute
instance-attribute
¶
output_schema = AnalysisOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'AnalysisAgent' / 'main.jinja2', 'response_model': make_analysis_model}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Perform custom analyses that are reliable and repeatable.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
AnalystAgent
¶
Bases:
conditions = param.List(default=['Use for interpreting and analyzing results from executed queries', 'NOT for initial data queries, data exploration, or technical programming questions'])
class-attribute
instance-attribute
¶
input_schema = AnalystInputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'AnalystAgent' / 'main.jinja2', 'tools': []}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Responsible for analyzing results and providing clear, concise, and actionable insights.\n Focuses on breaking down complex data findings into understandable points for\n high-level decision-making. Emphasizes detailed interpretation, trends, and\n relationships within the data, while avoiding general overviews or\n superficial descriptions.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
ChatAgent
¶
Bases:
ChatAgent provides general information about available data and other topics to the user.
conditions = param.List(default=['Use for high-level data information or general conversation', "Use for technical questions about programming, functions, methods, libraries, APIs, software tools, or 'how to' code usage", 'NOT for data-specific questions that require querying data'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'ChatAgent' / 'main.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Engages in conversations about high-level data topics, programming questions,\n technical documentation, and general conversation. Handles questions about\n specific functions, methods, libraries, and provides coding guidance and\n technical explanations.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
DbtslAgent
¶
Bases:
Responsible for creating and executing queries against a dbt Semantic Layer to answer user questions about business metrics.
conditions = param.List(default=['Always use this when dbtsl_metaset is available'])
class-attribute
instance-attribute
¶
output_schema = DbtslOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': DbtslQueryParams, 'template': PROMPTS_DIR / 'DbtslAgent' / 'main.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'BaseLumenAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Responsible for displaying data to answer user queries about\n business metrics using dbt Semantic Layers. This agent can compile\n and execute metric queries against a dbt Semantic Layer.')
class-attribute
instance-attribute
¶
requires = param.List(default=['source', 'dbtsl_metaset'], readonly=True)
class-attribute
instance-attribute
¶
source = param.ClassSelector(class_=BaseSQLSource, doc='\n The source associated with the dbt Semantic Layer.')
class-attribute
instance-attribute
¶
user = param.String(default='DBT')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
Responds to user messages by generating and executing a dbt Semantic Layer query.
DocumentListAgent
¶
Bases:
The DocumentListAgent lists all available documents provided by the user.
conditions = param.List(default=['Use when user asks to list or see all available documents', 'NOT when user asks about specific document content'])
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Displays a list of all available documents.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
SQLAgent
¶
Bases:
conditions = param.List(default=['Use for displaying, examining, or querying data resulting in a data pipeline', "Use for calculations that require data (e.g., 'calculate average', 'sum by category')", 'Do not sample or show limited rows of the data, unless explicitly requested', 'NOT for non-data questions or technical programming help', 'NOT useful if the user is using the same data for plotting'])
class-attribute
instance-attribute
¶
exclusions = param.List(default=['dbtsl_metaset'])
class-attribute
instance-attribute
¶
exploration_enabled = param.Boolean(default=True, doc='\n Whether to enable SQL exploration mode. When False, only attempts oneshot SQL generation.')
class-attribute
instance-attribute
¶
input_schema = SQLInputs
class-attribute
instance-attribute
¶
not_with = param.List(default=['DbtslAgent', 'TableLookup', 'TableListAgent'])
class-attribute
instance-attribute
¶
output_schema = SQLOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': make_sql_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'main.jinja2'}, 'select_discoveries': {'response_model': make_discovery_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'select_discoveries.jinja2'}, 'check_sufficiency': {'response_model': make_discovery_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'check_sufficiency.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'SQLAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Handles the display of data and the creation, modification, and execution\n of SQL queries to address user queries about the data. Executes queries in\n a single step, encompassing tasks such as table joins, filtering, aggregations,\n and calculations. If additional columns are required, SQLAgent can join the\n current table with other tables to fulfill the query requirements.')
class-attribute
instance-attribute
¶
user = param.String(default='SQL')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
Execute SQL generation with one-shot attempt first, then exploration if needed.
revise(feedback, messages, context, view=None, spec=None, language=None, errors=None, **kwargs)
async
¶
SourceAgent
¶
Bases:
SourceAgent renders a form that allows a user to upload or provide a URI to one or more datasets.
conditions = param.List(default=['Use ONLY when user explicitly asks to upload or connect to NEW data sources (NOT if data sources already exist)'])
class-attribute
instance-attribute
¶
output_schema = SourceOutputs
class-attribute
instance-attribute
¶
purpose = param.String(default='Allows a user to upload new datasets, data, or documents.')
class-attribute
instance-attribute
¶
source_controls = SourceControls
class-attribute
¶
respond(messages, context, step_title=None)
async
¶
TableListAgent
¶
Bases:
The TableListAgent lists all available data and lets the user pick one.
conditions = param.List(default=["Use when user explicitly asks to 'list data', 'show available data', or 'what data do you have'", 'NOT for showing actual data contents, querying, or analyzing data'])
class-attribute
instance-attribute
¶
input_schema = TableListInputs
class-attribute
instance-attribute
¶
not_with = param.List(default=['DbtslAgent', 'SQLAgent'])
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Displays a list of all available data & datasets. Not useful for identifying which dataset to use for analysis.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
ValidationAgent
¶
Bases:
ValidationAgent focuses solely on validating whether the executed plan fully answered the user's original query. It identifies missing elements and suggests next steps when validation fails.
conditions = param.List(default=['Use to validate whether executed plans fully answered user queries', 'Use to identify missing elements from the original user request', 'NOT for data analysis, pattern identification, or technical programming questions'])
class-attribute
instance-attribute
¶
input_schema = ValidationInputs
class-attribute
instance-attribute
¶
output_schema = ValidationOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'ValidationAgent' / 'main.jinja2', 'response_model': QueryCompletionValidation, 'tools': []}})
class-attribute
instance-attribute
¶
purpose = param.String(default="\n Validates whether executed plans fully answered the user's original query.\n Identifies missing elements, assesses completeness, and suggests next steps\n when validation fails. Acts as a quality gate for plan execution.")
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
VegaLiteAgent
¶
Bases:
conditions = param.List(default=['Use for publication-ready visualizations or when user specifically requests Vega-Lite charts', 'Use for polished charts intended for presentation or sharing'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': VegaLiteSpec, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'main.jinja2'}, 'interaction_polish': {'response_model': VegaLiteSpecUpdate, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'interaction_polish.jinja2'}, 'annotate_plot': {'response_model': VegaLiteSpecUpdate, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'annotate_plot.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Generates a vega-lite plot specification from the input data pipeline.')
class-attribute
instance-attribute
¶
user = param.String(default='Vega')
class-attribute
instance-attribute
¶
vector_store_path = param.Path(default=None, check_exists=False, doc='\n Path to a custom vector store for storing and retrieving Vega-Lite examples;\n if not provided a default store will be used depending on the LLM--\n OpenAIEmbeddings for OpenAI LLM or NumpyEmbeddings for all others.')
class-attribute
instance-attribute
¶
view_type = VegaLiteView
class-attribute
instance-attribute
¶
annotate(instruction, messages, context, spec)
async
¶
Apply annotations based on user request.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instruction
|
|
User's description of what to annotate |
required |
messages
|
|
Chat history for context |
required |
context
|
|
Session context |
required |
spec
|
|
The current VegaLite specification (full dict with 'spec' key) |
required |
Returns:
| Type | Description |
|---|---|
|
Updated specification with annotations |
respond(messages, context, step_title=None)
async
¶
Generates a VegaLite visualization using progressive building approach with real-time updates.
revise(feedback, messages, context, view=None, spec=None, language=None, errors=None, **kwargs)
async
¶
hvPlotAgent
¶
Bases:
conditions = param.List(default=['Use for exploratory data analysis, interactive plots, and dynamic filtering', 'Use for quick, iterative data visualization during analysis'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'hvPlotAgent' / 'main.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Generates a plot of the data given a user prompt.')
class-attribute
instance-attribute
¶
view_type = hvPlotUIView
class-attribute
instance-attribute
¶
analysis
¶
AnalysisAgent
¶
Bases:
analyses = param.List([])
class-attribute
instance-attribute
¶
conditions = param.List(default=['Use for custom analysis, advanced analytics, or domain-specific methods', "Use when the user query matches one of the available analyses' name or description below.", "Include the selected analysis' required cols in the instructions", 'NOT for simple queries or basic visualizations'])
class-attribute
instance-attribute
¶
input_schema = AnalysisInputs
class-attribute
instance-attribute
¶
output_schema = AnalysisOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'AnalysisAgent' / 'main.jinja2', 'response_model': make_analysis_model}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Perform custom analyses that are reliable and repeatable.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
AnalysisOutputs
¶
make_analysis_model(analyses)
¶
analyst
¶
AnalystAgent
¶
Bases:
conditions = param.List(default=['Use for interpreting and analyzing results from executed queries', 'NOT for initial data queries, data exploration, or technical programming questions'])
class-attribute
instance-attribute
¶
input_schema = AnalystInputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'AnalystAgent' / 'main.jinja2', 'tools': []}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Responsible for analyzing results and providing clear, concise, and actionable insights.\n Focuses on breaking down complex data findings into understandable points for\n high-level decision-making. Emphasizes detailed interpretation, trends, and\n relationships within the data, while avoiding general overviews or\n superficial descriptions.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
base
¶
Agent
¶
Bases:
Agents are actors responsible for taking a user query and performing a particular task, either by adding context or generating outputs.
Agents have access to an LLM and are given context and can solve tasks by executing a series of prompts or by rendering contents such as forms or widgets to gather user input.
agents = param.List(doc='\n List of agents this agent can invoke.')
class-attribute
instance-attribute
¶
debug = param.Boolean(default=False, doc='\n Whether to enable verbose error reporting.')
class-attribute
instance-attribute
¶
llm = param.ClassSelector(class_=Llm, doc='\n The LLM implementation to query.')
class-attribute
instance-attribute
¶
user = param.String(default='Agent', doc='\n The name of the user that will be respond to the user query.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
Additional checks to determine if the agent should be used.
respond(messages, context, step_title=None)
async
¶
Provides a response to the user query.
The type of the response may be a simple string or an object.
Arguments
messages: list[Message] The list of messages corresponding to the user query and any other system messages to be included. context: TContext A mapping containing context for the agent to perform its task. step_title: str | None If the Agent response is part of a longer query this describes the step currently being processed.
base_list
¶
base_lumen
¶
BaseLumenAgent
¶
Bases:
prompts = param.Dict(default={'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'BaseLumenAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
user = param.String(default='Lumen')
class-attribute
instance-attribute
¶
revise(instruction, messages, context, view=None, spec=None, language=None, errors=None, **kwargs)
async
¶
Retry the output by line, allowing the user to provide instruction on why the output was not satisfactory, or an error.
base_view
¶
BaseViewAgent
¶
Bases:
input_schema = ViewInputs
class-attribute
instance-attribute
¶
output_schema = ViewOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'BaseViewAgent' / 'main.jinja2'}})
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
Generates a visualization based on user messages and the current data pipeline.
ViewInputs
¶
chat
¶
ChatAgent
¶
Bases:
ChatAgent provides general information about available data and other topics to the user.
conditions = param.List(default=['Use for high-level data information or general conversation', "Use for technical questions about programming, functions, methods, libraries, APIs, software tools, or 'how to' code usage", 'NOT for data-specific questions that require querying data'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'ChatAgent' / 'main.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Engages in conversations about high-level data topics, programming questions,\n technical documentation, and general conversation. Handles questions about\n specific functions, methods, libraries, and provides coding guidance and\n technical explanations.')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
dbtsl
¶
DbtslAgent
¶
Bases:
Responsible for creating and executing queries against a dbt Semantic Layer to answer user questions about business metrics.
conditions = param.List(default=['Always use this when dbtsl_metaset is available'])
class-attribute
instance-attribute
¶
output_schema = DbtslOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': DbtslQueryParams, 'template': PROMPTS_DIR / 'DbtslAgent' / 'main.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'BaseLumenAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Responsible for displaying data to answer user queries about\n business metrics using dbt Semantic Layers. This agent can compile\n and execute metric queries against a dbt Semantic Layer.')
class-attribute
instance-attribute
¶
requires = param.List(default=['source', 'dbtsl_metaset'], readonly=True)
class-attribute
instance-attribute
¶
source = param.ClassSelector(class_=BaseSQLSource, doc='\n The source associated with the dbt Semantic Layer.')
class-attribute
instance-attribute
¶
user = param.String(default='DBT')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
Responds to user messages by generating and executing a dbt Semantic Layer query.
DbtslQueryParams
¶
Bases:
Model for dbtsl.client.query() parameters.
chain_of_thought = Field(description="You are a world-class dbt Semantic Layer expert. Think step by step about\n what metrics are needed, what dimensions to group by, what time granularity\n to use, and any filters that should be applied; if filters are applied, include those\n filtered dimensions in group_by. If there are errors, mention how you'll address the errors.\n ")
class-attribute
instance-attribute
¶
expr_slug = Field(description='Give the query a concise, but descriptive, slug that includes the metrics\n and dimensions used, e.g. monthly_revenue_by_region. The slug must be unique.')
class-attribute
instance-attribute
¶
group_by = Field(default_factory=list, description="A list of dimensions to group by, e.g. ['metric_time__month'], must include dimensions from where.")
class-attribute
instance-attribute
¶
limit = Field(default=None, description='The maximum number of rows to return.')
class-attribute
instance-attribute
¶
metrics = Field(default_factory=list, description="A list of metrics to include in the query, e.g. ['revenue']")
class-attribute
instance-attribute
¶
order_by = Field(default_factory=list, description="A list of columns or expressions to order the results by, e.g. ['metric_time__month']")
class-attribute
instance-attribute
¶
where = Field(default_factory=list, description="A list of conditions to filter the results; dimensions referenced here must also be in group_by, e.g. ['metric_time__month >= date_trunc('month', '2024-09-30'::date)']")
class-attribute
instance-attribute
¶
document_list
¶
DocumentListAgent
¶
Bases:
The DocumentListAgent lists all available documents provided by the user.
conditions = param.List(default=['Use when user asks to list or see all available documents', 'NOT when user asks about specific document content'])
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Displays a list of all available documents.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
hvplot
¶
hvPlotAgent
¶
Bases:
conditions = param.List(default=['Use for exploratory data analysis, interactive plots, and dynamic filtering', 'Use for quick, iterative data visualization during analysis'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'hvPlotAgent' / 'main.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Generates a plot of the data given a user prompt.')
class-attribute
instance-attribute
¶
view_type = hvPlotUIView
class-attribute
instance-attribute
¶
source
¶
SourceAgent
¶
Bases:
SourceAgent renders a form that allows a user to upload or provide a URI to one or more datasets.
conditions = param.List(default=['Use ONLY when user explicitly asks to upload or connect to NEW data sources (NOT if data sources already exist)'])
class-attribute
instance-attribute
¶
output_schema = SourceOutputs
class-attribute
instance-attribute
¶
purpose = param.String(default='Allows a user to upload new datasets, data, or documents.')
class-attribute
instance-attribute
¶
source_controls = SourceControls
class-attribute
¶
respond(messages, context, step_title=None)
async
¶
sql
¶
DistinctQuery
¶
Bases:
Universal column analysis with optional pattern matching - handles join keys, categories, date ranges.
column = Field(description='Column to analyze unique values')
class-attribute
instance-attribute
¶
offset = Field(default=0, description='Number of distinct values to skip (0 for initial, 10 for follow-up)')
class-attribute
instance-attribute
¶
pattern = Field(default='', description="Optional pattern to search for (e.g., 'chin' for China/Chinese variations). Leave empty for all distinct values.")
class-attribute
instance-attribute
¶
query = Field(default="SELECT DISTINCT {column} FROM {slug[table]} WHERE {column} ILIKE '%{pattern}%' LIMIT 10 OFFSET {offset}")
class-attribute
instance-attribute
¶
model_post_init(__context)
¶
Adjust query template based on whether pattern is provided.
SQLAgent
¶
Bases:
conditions = param.List(default=['Use for displaying, examining, or querying data resulting in a data pipeline', "Use for calculations that require data (e.g., 'calculate average', 'sum by category')", 'Do not sample or show limited rows of the data, unless explicitly requested', 'NOT for non-data questions or technical programming help', 'NOT useful if the user is using the same data for plotting'])
class-attribute
instance-attribute
¶
exclusions = param.List(default=['dbtsl_metaset'])
class-attribute
instance-attribute
¶
exploration_enabled = param.Boolean(default=True, doc='\n Whether to enable SQL exploration mode. When False, only attempts oneshot SQL generation.')
class-attribute
instance-attribute
¶
input_schema = SQLInputs
class-attribute
instance-attribute
¶
not_with = param.List(default=['DbtslAgent', 'TableLookup', 'TableListAgent'])
class-attribute
instance-attribute
¶
output_schema = SQLOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': make_sql_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'main.jinja2'}, 'select_discoveries': {'response_model': make_discovery_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'select_discoveries.jinja2'}, 'check_sufficiency': {'response_model': make_discovery_model, 'template': PROMPTS_DIR / 'SQLAgent' / 'check_sufficiency.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'SQLAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Handles the display of data and the creation, modification, and execution\n of SQL queries to address user queries about the data. Executes queries in\n a single step, encompassing tasks such as table joins, filtering, aggregations,\n and calculations. If additional columns are required, SQLAgent can join the\n current table with other tables to fulfill the query requirements.')
class-attribute
instance-attribute
¶
user = param.String(default='SQL')
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
Execute SQL generation with one-shot attempt first, then exploration if needed.
revise(feedback, messages, context, view=None, spec=None, language=None, errors=None, **kwargs)
async
¶
SQLInputs
¶
SQLOutputs
¶
SQLQuery
¶
Bases:
A single SQL query with its associated metadata.
query = Field(description="\n One, correct, valid SQL query that answers the user's question;\n should only be one query and do NOT add extraneous comments; no multiple semicolons")
class-attribute
instance-attribute
¶
table_slug = Field(description='\n Provide a unique, descriptive table slug for the SQL expression that clearly indicates the key transformations and source tables involved.\n Include 1 or 2 elements of data lineage in the slug, such as the main transformation and original table names,\n e.g. top_5_athletes_in_2020 or distinct_years_from_wx_table.\n Ensure the slug does not duplicate any existing table names or slugs.\n ')
class-attribute
instance-attribute
¶
SampleQuery
¶
Bases:
See actual data content - reveals format issues and patterns.
query = Field(default='SELECT * FROM {slug[table]} LIMIT 5')
class-attribute
instance-attribute
¶
make_discovery_model(sources)
¶
make_source_table(sources)
¶
make_sql_model(sources)
¶
table_list
¶
TableListAgent
¶
Bases:
The TableListAgent lists all available data and lets the user pick one.
conditions = param.List(default=["Use when user explicitly asks to 'list data', 'show available data', or 'what data do you have'", 'NOT for showing actual data contents, querying, or analyzing data'])
class-attribute
instance-attribute
¶
input_schema = TableListInputs
class-attribute
instance-attribute
¶
not_with = param.List(default=['DbtslAgent', 'SQLAgent'])
class-attribute
instance-attribute
¶
purpose = param.String(default='\n Displays a list of all available data & datasets. Not useful for identifying which dataset to use for analysis.')
class-attribute
instance-attribute
¶
applies(context)
async
classmethod
¶
validation
¶
QueryCompletionValidation
¶
Bases:
Validation of whether the executed plan answered the user's query
chain_of_thought = Field(description='Restate intent and results succinctly; then explain your reasoning as to why you will be answering yes or no.')
class-attribute
instance-attribute
¶
correct = Field(description='True if query correctly solves user request, otherwise False.')
class-attribute
instance-attribute
¶
missing_elements = Field(default_factory=list, description="List of specific elements from the user's query that weren't addressed")
class-attribute
instance-attribute
¶
suggestions = Field(default_factory=list, description='Suggestions for additional steps that could complete the query if not fully answered')
class-attribute
instance-attribute
¶
ValidationAgent
¶
Bases:
ValidationAgent focuses solely on validating whether the executed plan fully answered the user's original query. It identifies missing elements and suggests next steps when validation fails.
conditions = param.List(default=['Use to validate whether executed plans fully answered user queries', 'Use to identify missing elements from the original user request', 'NOT for data analysis, pattern identification, or technical programming questions'])
class-attribute
instance-attribute
¶
input_schema = ValidationInputs
class-attribute
instance-attribute
¶
output_schema = ValidationOutputs
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'template': PROMPTS_DIR / 'ValidationAgent' / 'main.jinja2', 'response_model': QueryCompletionValidation, 'tools': []}})
class-attribute
instance-attribute
¶
purpose = param.String(default="\n Validates whether executed plans fully answered the user's original query.\n Identifies missing elements, assesses completeness, and suggests next steps\n when validation fails. Acts as a quality gate for plan execution.")
class-attribute
instance-attribute
¶
respond(messages, context, step_title=None)
async
¶
ValidationInputs
¶
vega_lite
¶
VegaLiteAgent
¶
Bases:
conditions = param.List(default=['Use for publication-ready visualizations or when user specifically requests Vega-Lite charts', 'Use for polished charts intended for presentation or sharing'])
class-attribute
instance-attribute
¶
prompts = param.Dict(default={'main': {'response_model': VegaLiteSpec, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'main.jinja2'}, 'interaction_polish': {'response_model': VegaLiteSpecUpdate, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'interaction_polish.jinja2'}, 'annotate_plot': {'response_model': VegaLiteSpecUpdate, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'annotate_plot.jinja2'}, 'revise_output': {'response_model': RetrySpec, 'template': PROMPTS_DIR / 'VegaLiteAgent' / 'revise_output.jinja2'}})
class-attribute
instance-attribute
¶
purpose = param.String(default='Generates a vega-lite plot specification from the input data pipeline.')
class-attribute
instance-attribute
¶
user = param.String(default='Vega')
class-attribute
instance-attribute
¶
vector_store_path = param.Path(default=None, check_exists=False, doc='\n Path to a custom vector store for storing and retrieving Vega-Lite examples;\n if not provided a default store will be used depending on the LLM--\n OpenAIEmbeddings for OpenAI LLM or NumpyEmbeddings for all others.')
class-attribute
instance-attribute
¶
view_type = VegaLiteView
class-attribute
instance-attribute
¶
annotate(instruction, messages, context, spec)
async
¶
Apply annotations based on user request.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
instruction
|
|
User's description of what to annotate |
required |
messages
|
|
Chat history for context |
required |
context
|
|
Session context |
required |
spec
|
|
The current VegaLite specification (full dict with 'spec' key) |
required |
Returns:
| Type | Description |
|---|---|
|
Updated specification with annotations |
respond(messages, context, step_title=None)
async
¶
Generates a VegaLite visualization using progressive building approach with real-time updates.
revise(feedback, messages, context, view=None, spec=None, language=None, errors=None, **kwargs)
async
¶
VegaLiteSpec
¶
Bases:
chain_of_thought = Field(description='Explain your design choices based on visualization theory:\n - What story does this data tell?\n - Which visual encodings (position, color, size) best reveal patterns?\n - Should color highlight specific insights or remain neutral?\n - What makes this plot engaging and useful for the user?\n Then describe the basic plot structure.')
class-attribute
instance-attribute
¶
yaml_spec = Field(description='A basic vega-lite YAML specification with core plot elements only (data, mark, basic x/y encoding).')
class-attribute
instance-attribute
¶
VegaLiteSpecUpdate
¶
Bases: