in databao/executors/lighthouse/graph.py [0:0]
def make_tools(self) -> list[BaseTool]:
@tool(parse_docstring=True)
def run_sql_query(sql: str, graph_state: Annotated[AgentState, InjectedState]) -> dict[str, Any]:
"""
Run a SELECT SQL query in the database. Returns the first 12 rows in csv format.
Args:
sql: SQL query
"""
try:
# TODO use ToolRuntime in LangChain v1.0
limit = graph_state["limit_max_rows"]
df = execute_duckdb_sql(sql, self._connection, limit=limit)
df_csv = df.head(self.MAX_TOOL_ROWS).to_csv(index=False)
df_markdown = dataframe_to_markdown(df.head(self.MAX_TOOL_ROWS), index=False)
if len(df) > self.MAX_TOOL_ROWS:
df_csv += f"\nResult is truncated from {len(df)} to {self.MAX_TOOL_ROWS} rows."
df_markdown += f"\nResult is truncated from {len(df)} to {self.MAX_TOOL_ROWS} rows."
return {"df": df, "sql": sql, "csv": df_csv, "markdown": df_markdown}
except Exception as e:
return {"error": exception_to_string(e)}
@tool(parse_docstring=True)
def submit_result(
query_id: str,
result_description: str,
visualization_prompt: str,
) -> str:
"""
Call this tool with the ID of the query you want to submit to the user.
This will return control to the user and must always be the last tool call.
The user will see the full query result, not just the first 12 rows. Returns a confirmation message.
Args:
query_id: The ID of the query to submit (query_ids are automatically generated when you run queries).
result_description: A comment to a final result. This will be included in the final result.
visualization_prompt: Optional visualization prompt. If not empty, a Vega-Lite visualization agent
will be asked to plot the submitted query data according to instructions in the prompt.
The instructions should be short and simple.
"""
return f"Query {query_id} submitted successfully. Your response is now visible to the user."
tools = [run_sql_query, submit_result]
return tools