admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

more tools in sql example

Browse source
This commit is contained in:
Sam Partee 2024-05-01 23:53:53 -07:00
parent 41b783ef2e
commit 2e5e5ef7c8
9 changed files with 601 additions and 143 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
examples/gmail/pack.lock.toml
View file

@ -10,3 +10,5 @@ email = "sam@partee.io"
[tools]
SendEmail = "gmailer.send_email@0.1.0"
ReadEmail = "gmailer.read_email@0.1.0"
PlotDataframe = "gmailer.plot_dataframe@0.1.0"
Summarize = "chat.summarize@0.1.0"

1
examples/gmail/pack.toml
View file

@ -9,3 +9,4 @@ email = "sam@partee.io"
[modules]
gmailer = "0.1.0"
chat = "0.1.0"

36
examples/gmail/tools/chat.py Normal file
View file

@ -0,0 +1,36 @@
from toolserve.sdk import Param, tool, get_secret
from toolserve.sdk.dataframe import get_df
import openai
@tool
def summarize(
text: Param(str, "Text to summarize"),
system_prompt: Param(str, "System prompt to use") = "Summarize the following text",
max_tokens: Param(int, "Maximum number of tokens to generate") = 1000,
) -> Param(str, "Summarized text"):
"""Summarize a piece of text using OpenAI's GPT-3 model.
Args:
text (str): The text to summarize.
max_tokens (int): The maximum number of tokens to generate.
Returns:
str: The summarized text.
"""
api_key = get_secret("openai_api_key")
# Call the OpenAI model with the tools and messages
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": text},
]
client = openai.Client(api_key=api_key)
completion = openai.chat.completions.create(
model=self.model,
messages=messages,
)
summary = completion.choices[0].message.content
return summary

85
examples/gmail/tools/gmailer.py
View file

@ -7,13 +7,16 @@ import email
from email.header import decode_header
from pydantic import BaseModel
import pandas as pd
import plotly.express as px
from bs4 import BeautifulSoup
import re
from toolserve.sdk import Param, tool, get_secret
from toolserve.sdk.dataframe import get_df, save_df
@tool
def send_email(
async def send_email(
sender_email: Param(str, "Email address of the sender"),
recipient_email: Param(str, "Email address of the recipient"),
subject: Param(str, "Subject of the email"),
@ -44,12 +47,13 @@ def send_email(
@tool
def read_email(
email_address: Param(str, "Email address of the recipient"),
async def read_email(
output_name: Param(str, "Name of the output data"),
n_emails: Param(int, "Number of emails to read") = 5,
) -> Param(str, "JSON dataframe of List of emails"):
"""Read emails from a Gmail account"""
):
"""Read emails from a Gmail account and extract plain text content, removing any HTML."""
email_address = get_secret("gmail_email")
password = get_secret("gmail_password")
server = get_secret("gmail_stmp_server", "smtp.gmail.com")
port = get_secret("gmail_smtp_port", 587)
@ -73,23 +77,82 @@ def read_email(
email_details = {
"from": msg["From"],
"to": msg["To"],
#"subject": decode_header(msg["Subject"])[0][0],
"date": msg["Date"]
}
if msg.is_multipart():
for part in msg.walk():
if part.get_content_type() == "text/plain":
email_details["body"] = part.get_payload(decode=True)
body = part.get_payload(decode=True).decode('utf-8')
email_details["body"] = clean_email_body(body)
else:
email_details["body"] = msg.get_payload(decode=True)
body = msg.get_payload(decode=True).decode('utf-8')
email_details["body"] = clean_email_body(body)
emails.append(email_details)
mail.close()
mail.logout()
return pd.DataFrame(emails).to_json()
df = pd.DataFrame(emails)
await save_df(df, output_name)
def clean_email_body(body: str) -> str:
"""Remove HTML tags and non-sentence elements from email body text."""
# Remove HTML tags using BeautifulSoup
soup = BeautifulSoup(body, "html.parser")
text = soup.get_text(separator=' ')
# Remove any non-sentence elements (e.g., URLs, email addresses, etc.)
text = re.sub(r'\S*@\S*\s?', '', text) # Remove emails
text = re.sub(r'http\S+', '', text) # Remove URLs
text = re.sub(r'[^.!?a-zA-Z0-9\s]', '', text) # Remove non-sentence characters
text = ' '.join(text.split()) # Remove extra whitespace
return text
@tool
async def plot_dataframe(
data_id: Param(int, "Data ID of the dataframe"),
x: Param(str, "Column to use as x-axis"),
y: Param(str, "Column to use as y-axis"),
kind: Param(str, "Type of plot") = "line",
title: Param(str, "Title of the plot") = "Plot",
xlabel: Param(str, "Label for x-axis") = "X",
ylabel: Param(str, "Label for y-axis") = "Y",
) -> Param(str, "JSON representation of the plot"):
"""
Asynchronously generates a plot from a dataframe using Plotly and returns the plot as a JSON string.
Args:
data_id (int): The ID of the dataframe to plot.
x (str): The column name to use as the x-axis.
y (str): The column name to use as the y-axis.
kind (str): The type of plot to generate (e.g., 'line', 'scatter', 'bar').
title (str): The title of the plot.
xlabel (str): The label for the x-axis.
ylabel (str): The label for the y-axis.
Returns:
str: The JSON representation of the plot.
"""
import plotly.express as px
df = await get_df(data_id)
if kind == 'line':
fig = px.line(df, x=x, y=y, title=title)
elif kind == 'scatter':
fig = px.scatter(df, x=x, y=y, title=title)
elif kind == 'bar':
fig = px.bar(df, x=x, y=y, title=title)
else:
raise ValueError(f"Unsupported plot type: {kind}")
fig.update_layout(xaxis_title=xlabel, yaxis_title=ylabel)
return fig.to_json()

485
examples/sql-chat/agent.py
View file

@ -1,21 +1,38 @@
import httpx
import json
import time
import openai
from typing import Any, Dict, List, Optional
class Toolchain:
from pydantic import BaseModel
from typing import List, Dict
from textwrap import dedent
from pydantic import BaseModel, Field
from enum import Enum
from typing import Type
from toolserve.utils.openai_tool import model_to_json_schema
from typing import Dict, Any, Optional
import json
from collections import deque
class ToolClient:
available_tools = {
"query_sql": "/tool/query/query_sql",
"list_data_sources": "/tool/query/list_data_sources",
"get_data_schema": "/tool/query/get_data_schema"
"get_data_schema": "/tool/query/get_data_schema",
"PlotDataframe": "/tool/gmailer/PlotDataframe",
"ReadEmail": "/tool/gmailer/ReadEmail",
"Summarize": "/tool/chat/Summarize",
}
def __init__(self, base_url: str, openai_api_key: str, model: str = "gpt-4-turbo"):
def __init__(self, base_url: str):
self.base_url = base_url
self.client = httpx.Client()
self.openai_client = openai.Client(api_key=openai_api_key)
self.model = model
self.tools = self.__collect_tool_specs()
def __collect_tool_specs(self) -> Dict[str, str]:
@ -26,6 +43,18 @@ class Toolchain:
return tools
def call_api(self, method: str, endpoint: str, params: dict = {}, data: dict = {}, json_data: dict = {}) -> Dict[str, Any]:
"""Call the Darkstar Toolserver API with the given parameters.
Args:
method (str): The HTTP method to use for the request.
endpoint (str): The endpoint to call.
params (dict): The query parameters for the request.
data (dict): The data to send in the request body.
json_data (dict): The JSON data to send in the request body.
Returns:
Dict[str, Any]: The response from the API.
"""
url = f"{self.base_url}{endpoint}"
response = self.client.request(method, url, params=params, json=json_data, data=data)
try:
@ -35,32 +64,6 @@ class Toolchain:
result = response.json()
return result
def get_tool_args(self, tool_name: str, messages: List[Dict[str, str]]) -> Dict[str, Any]:
"""
Retrieves the required arguments for an tool from the Darkstar Toolserver API and
uses them to call an OpenAI model with predefined tools and messages.
:param tool_name: The name of the tool to execute.
:param messages: A list of messages to provide to the model.
:return: The result of the OpenAI model call.
"""
func_spec = self.tools.get(tool_name, {})
if not func_spec:
raise ValueError(f"Tool '{tool_name}' not found in available tools.")
tool = json.loads(func_spec)
# Call the OpenAI model with the tools and messages
completion = self.openai_client.chat.completions.create(
model="gpt-4-turbo",
messages=messages,
tools=[tool],
tool_choice="auto"
)
predicted_args = completion.choices[0].message.tool_calls[0].function.arguments
print(predicted_args)
print("-----")
return predicted_args
def execute_tool(self, tool_name: str, tool_args: Optional[Dict[str, Any]]) -> Dict[str, Any]:
"""
Executes an tool using the Darkstar Toolserver API and an OpenAI model.
@ -75,114 +78,394 @@ class Toolchain:
return result
from pydantic import BaseModel
from typing import List, Dict
from textwrap import dedent
class Agent:
prompt = dedent("""Given a user query and a schema of a table, generate the SQL query to answer the user query.
The generated SQL query should only refer to columns in the table schema list below. The table schema is as follows:
class ToolRunner:
tool_prompt = dedent("""
Given a user query and the schema of the fields in a dataframe, generate the arguments for a tool to execute.
YOU MUST CALL THE TOOL.
The schema of the fields in the dataframe is as follows:
{schema}
The data_id of this source is: {data_id}
If needed, the data_id for the source is: {data_id}
If needed, the output_name should be: {output_name}
""")
def __init__(self, base_url: str, model: str, api_key: str):
"""
Initialize the ToolRunner with necessary configurations.
def __init__(self, toolchain: Toolchain):
self.toolchain = toolchain
self.data_sources = self.__get_data_sources()
Args:
base_url (str): The base URL for the API calls.
model (str): The model identifier to be used for queries.
api_key (str): The API key for authentication.
"""
self._client = ToolClient(base_url)
self._model = model
self._openai_client = openai.Client(api_key=api_key)
self._data_sources = self.__get_data_sources()
self._source = None
self._data_schema = None
self._data_id = None
def set_source(self, source: str):
if source not in self.data_sources.keys():
raise ValueError(f"Data source '{source}' not found.")
else:
data_id = self.data_sources[source]
# get the schema
schema = self.toolchain.call_api("POST", "/tool/query/get_data_schema", json_data={"data_id": data_id})
self._source = source
self._data_schema = schema
self._data_sources = self.__get_data_sources()
retries = 3
data_id = None
while retries > 0:
try:
data_id = self._data_sources[source]
break
except KeyError:
retries -= 1
time.sleep(1)
self._data_sources = self.__get_data_sources()
def get_source(self) -> str:
return self._source
if data_id is None:
raise ValueError(f"Data source '{source}' not found.")
# get the schema
schema = self._client.call_api("POST", "/tool/query/get_data_schema", json_data={"data_id": data_id})
self._source = source
self._data_schema = schema
self._data_id = data_id
def __get_data_sources(self) -> Dict[str, Dict[str, str]]:
response = self.toolchain.call_api("POST", "/tool/query/list_data_sources")
response = self._client.call_api("POST", "/tool/query/list_data_sources")
sources = {}
for _id, source_data in response["data"]["result"].items():
sources[source_data["file_name"]] = _id
return sources
def query(self, user_query: str) -> str:
if not self._source:
raise ValueError("Data source not set. Please set a data source before querying.")
def __create_prompt(self, user_query: str, input_name: str, output_name: str) -> List[Dict[str, str]]:
schema = self._data_schema
prompt = self.prompt.format(schema=schema, data_id=self.data_sources[self._source])
data_id = self._data_sources[input_name]
prompt = self.tool_prompt.format(schema=schema, data_id=data_id, output_name=output_name)
# Prepare the input message for the OpenAI model
messages = [
{"role": "system", "content": prompt},
{"role": "user", "content": user_query}]
{"role": "user", "content": user_query}
]
return messages
tool_args = self.toolchain.get_tool_args("query_sql", messages)
args = json.loads(tool_args)
params = args.get("params", [])
if params:
if isinstance(params, dict):
args["params"] = list(params.values())
elif isinstance(params, str):
args["params"] = [params]
elif isinstance(params, list):
args["params"] = params
else:
raise ValueError(f"Invalid params type: {type(params)}")
def get_tool_args(self, tool_name: str, messages: List[Dict[str, str]], output_name: str) -> Dict[str, Any]:
"""
Retrieves the required arguments for an tool from the Darkstar Toolserver API and
uses them to call an OpenAI model with predefined tools and messages.
:param tool_name: The name of the tool to execute.
:param messages: A list of messages to provide to the model.
:return: The result of the OpenAI model call.
"""
func_spec = self._client.tools.get(tool_name, {})
if not func_spec:
raise ValueError(f"Tool '{tool_name}' not found in available tools.")
tool = json.loads(func_spec)
print(tool)
# Call the OpenAI model with the tools and messages
completion = self._openai_client.chat.completions.create(
model="gpt-4-turbo",
messages=messages,
tools=[tool],
tool_choice="required"
)
predicted_args = completion.choices[0].message.tool_calls[0].function.arguments
args = json.loads(predicted_args)
if "params" in args:
params = args.get("params", [])
if params:
if isinstance(params, dict):
args["params"] = list(params.values())
elif isinstance(params, str):
args["params"] = [params]
elif isinstance(params, list):
args["params"] = params
else:
raise ValueError(f"Invalid params type: {type(params)}")
if "output_name" in args:
args["output_name"] = output_name
if "data_id" in args:
args["data_id"] = self._data_id
return args
def run_tool(self, tool_name: str, user_query: str, source: str, output_name: str) -> Any:
"""
Executes an tool using the Darkstar Toolserver API and an OpenAI model.
:param tool_name: The name of the tool to execute.
:param user_query: The user query to provide to the model.
:return: The result of the tool
"""
self.set_source(source)
messages = self.__create_prompt(user_query, source, output_name)
tool_args = self.get_tool_args(tool_name, messages, output_name)
result = self._client.execute_tool(tool_name, tool_args)
return result
def get_data_object(self, data_id: int) -> Dict[str, Any]:
"""
Retrieves a data object from the Darkstar Toolserver API.
:param data_id: The ID of the data object to retrieve.
:return: The data object.
"""
return self._client.call_api("GET", f"/api/v1/data/object/{data_id}")["data"]["json_blob"]
response = self.toolchain.execute_tool("query_sql", args)
if response["code"] != 200:
raise ValueError(f"Error executing tool: {response['message']}")
data_id = response["data"]["result"]["data_id"]
def pydantic_to_openai_tool(model: Type[BaseModel]) -> str:
"""
Convert a Pydantic model to an OpenAI tool schema.
# get the data
data_response = self.toolchain.call_api("GET", f"/api/v1/data/object/{data_id}")
if data_response["code"] != 200:
raise ValueError(f"Error retrieving data: {data_response['message']}")
data = data_response["data"]["json_blob"]
return data
Args:
model (Type[BaseModel]): The Pydantic model to convert.
Returns:
str: The OpenAI tool schema.
"""
schema = model_to_json_schema(model)
tool_schema = {
"type": "function",
"function": {
"name": model.__name__,
"description": model.__doc__ or "",
"parameters": schema
}
}
return json.dumps(tool_schema)
from pydantic import BaseModel, Field
from enum import Enum
class Edge(BaseModel):
source: int = Field(..., description="The ID of the source node")
target: int = Field(..., description="The ID of the target node")
class ToolNode:
pass
class ToolNode(BaseModel):
node_id: int = Field(..., description="The ID of the node", ge=0)
input_name: str = Field(..., description="The name of the input data")
tool_name: str = Field(..., description="The name of the tool to execute")
output_name: str = Field(..., description="The name of the output data")
class OutputType(Enum):
DATA = "data"
CHAT = "chat"
ARTIFACT = "artifact"
class FlowSchema(BaseModel):
"""A graph based representation of functions (nodes), and their data flow (edges)"""
nodes: List[ToolNode] = Field(..., description="The nodes in the flow")
edges: List[Edge] = Field([], description="The IDs of the adjacent nodes")
output_type: OutputType = Field(OutputType.CHAT, description="The type of the output")
class Config:
arbitrary_types_allowed = True
use_enum_values = True
class ToolFlow:
tools = {
"query_sql": (OutputType.DATA, True, False),
"PlotDataframe": (OutputType.ARTIFACT, False, True),
"ReadEmail": (OutputType.CHAT, True, False),
"Summarize": (OutputType.CHAT, False, True),
}
def __init__(
self,
name,
description,
sources,
name: str,
description: str,
prompt: str,
base_url: str = "http://localhost:8000",
model: str = "gpt-4-turbo",
model_api_key: Optional[str] = None
):
pass
self.name = name
self.description = description
self.prompt = prompt
self.runner = ToolRunner(base_url, model, model_api_key)
self.model = model
self.openai_client = openai.Client(api_key=model_api_key)
def __create_prompt(self, user_query: str) -> List[Dict[str, str]]:
tool_list = ""
for tool, spec in self.tools.items():
tool_list += f"- Name: {tool}\n"
tool_list += f" - Output Type: {spec[0].value}\n"
tool_list += f" - Can be source node: {spec[1]}\n"
tool_list += f" - Can be sink node: {spec[2]}\n"
source_list = "\n".join(self.runner._data_sources.keys())
prompt = self.prompt.format(nodes=tool_list, sources=source_list)
messages = [
{"role": "system", "content": prompt},
{"role": "user", "content": user_query}
]
return messages
def infer_flow(self, user_query: str) -> FlowSchema:
"""
Infer the tool flow based on the user query.
Args:
user_query (str): The user's query string.
Returns:
FlowSchema: The inferred tool flow schema.
"""
messages = self.__create_prompt(user_query)
func_spec = pydantic_to_openai_tool(FlowSchema)
tool = json.loads(func_spec)
# Call the OpenAI model with the tools and messages
completion = self.openai_client.chat.completions.create(
model=self.model,
messages=messages,
tools=[tool],
tool_choice="required"
)
predicted_args = completion.choices[0].message.tool_calls[0].function.arguments
print(predicted_args)
return predicted_args
def execute_flow(self, flow_schema: Dict[str, Any], user_query: str) -> Any:
"""
Executes the tool flow based on the provided schema. This method performs a breadth-first search (BFS)
on the graph defined by the flow schema and executes each node according to the order determined by the BFS.
Args:
flow_schema (Dict[str, Any]): The schema representing the tool flow to be executed.
user_query (str): The user's query string that may influence tool execution.
Returns:
Any: The result of executing the tool flow.
"""
# Initialize a queue for BFS
execution_queue = deque([flow_schema['nodes'][0]]) # Start BFS from the source node
visited = set()
results = {}
while execution_queue:
current_node = execution_queue.popleft()
node_id = current_node['node_id']
if node_id in visited:
continue
visited.add(node_id)
# Execute the current node's operation using runner.run_tool
operation_result = self.runner.run_tool(
current_node['tool_name'],
user_query,
current_node['input_name'],
current_node['output_name']
)
results[node_id] = operation_result
# Enqueue all adjacent nodes
for edge in flow_schema.get('edges', []):
if edge['source'] == node_id:
target_node_id = edge['target']
target_node = next(node for node in flow_schema['nodes'] if node['node_id'] == target_node_id)
if target_node_id not in visited:
execution_queue.append(target_node)
# Assuming the last node processed is the sink node
sink_node = flow_schema['nodes'][-1]
sink_tool_name = sink_node['tool_name']
sink_node_id = sink_node['node_id']
sink_output_type = self.tools[sink_tool_name][0]
if sink_output_type == OutputType.DATA:
data = self.runner.get_data_object(self.runner._data_id)
else:
data = results[sink_node_id]
return (data, results, sink_output_type)
def summarize_flow_results(model_client, flow_results: Dict[str, Any], flow_schema) -> str:
"""
Summarizes the results of a tool flow execution using an OpenAI model to generate a chat response.
Args:
model_client (openai.Client): The OpenAI client to use for generating chat responses.
flow_results (Dict[str, Any]): The results of the tool flow execution.
flow_schema (Dict[str, Any]): The schema representing the tool flow.
Returns:
Dict[str, str]: A dictionary containing the chat response under the key "data".
"""
try:
# Check if flow_results is already a JSON string, otherwise convert it
if isinstance(flow_results, str):
flow_summary = flow_results
else:
flow_summary = json.dumps(flow_results, indent=2)
# Construct a concise and informative prompt for the chat model
prompt_content = dedent(f"""
Please review the tool execution results and the flow schema provided below.
Use the results of the final tool to describe the outcomes. Be concise and only use the provided information.
If the results seem incorrect or incomplete, kindly ask the user to reformulate their query for better accuracy.
The execution path, expressed a a JSON object where nodes represent tools and edges represent data flow:
{flow_schema}
The results of the execution, expressed as a JSON object:
{flow_summary}
""")
messages = [
{"role": "system", "content": prompt_content}
]
# Call the OpenAI chat model
response = model_client.chat.completions.create(
model="gpt-4-turbo",
messages=messages
)
# Extract the chat response
chat_response = response.choices[0].message.content
return chat_response
except Exception as e:
print(f"Error in summarizing flow results: {e}")
return "Error: Failed to generate summary due to an internal error."
""" # Example usage:
oai_key = "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"
toolchain = Toolchain(base_url="http://localhost:8000", model="gpt-4-turbo", openai_api_key=oai_key)
agent = Agent(toolchain)
agent.set_source("users_db")
while True:
user_query = input("Enter a query: ")
result = agent.query(user_query)
print(result)
"""
#flow_schema = tf.infer_flow("Plot the users' age distribution")
#from pprint import pprint
#flow = json.loads(flow_schema)
#pprint(flow)
#result = tf.execute_flow(flow, "Plot the users' age distribution")
#print(result)

100
examples/sql-chat/main.py
View file

@ -11,20 +11,73 @@ import time
import traceback
import os
from typing import Dict, Any
import networkx as nx
import matplotlib.pyplot as plt
import pandas as pd
import streamlit as st
from pydantic import BaseModel
from streamlit_chat import message
from textwrap import dedent
import plotly.express as px
from agent import ToolFlow
from agent import Agent, Toolchain
PROMPT = dedent("""Given a user query, construct a graph based representation of functions (nodes), and their data flow (edges) such that
the graph can be executed to supply the user query enough information to answer their query.
You must construct the graph with the following constraints:
- There can only be 1 source node and 1 sink node.
- There should be no leaf nodes besides the sink node.
- The source and sink can be the same node.
Only use the available nodes and their output types as edges. Create unique ids for each node starting from 0.
The available nodes are:
{nodes}
The available input names for the source are:
{sources}
""")
oai_key = "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"
def plot_flow(data: Dict[str, Any]):
# Create a directed graph
G = nx.DiGraph()
# Add nodes
for node in data['nodes']:
G.add_node(node['node_id'], label=node['tool_name'])
# Add edges
if 'edges' in data:
for edge in data['edges']:
G.add_edge(edge['source'], edge['target'])
# Node labels with specific formatting
labels = {node['node_id']: f"{node['tool_name']}\n({node['input_name']} -> {node['output_name']})" for node in data['nodes']}
# Position nodes using the spring layout
pos = nx.spring_layout(G)
plt.figure(figsize=(4, 3))
nx.draw(G, pos, with_labels=False, node_size=3000, node_color='skyblue', font_size=9, font_weight='bold')
nx.draw_networkx_labels(G, pos, labels, font_size=8)
st.write("Graph of the data flow:")
# Use Streamlit's function to display the plot
st.pyplot(plt, use_container_width=False)
@st.cache_resource()
def get_agent():
toolchain = Toolchain(base_url="http://localhost:8000", model="gpt-4-turbo", openai_api_key=oai_key)
agent = Agent(toolchain)
agent.set_source("users_db")
return agent
AnalysisTool = ToolFlow(
name="data_analysis",
description="A tool flow for data analysis",
prompt=PROMPT,
model_api_key=oai_key
)
return AnalysisTool
# From here down is all the StreamLit UI.
@ -60,9 +113,15 @@ def submit():
with st.spinner(text="Wait for Agent..."):
try:
agent = get_agent()
res = agent.query(submit_text)
flow = agent.infer_flow(submit_text)
json_flow = json.loads(flow)
with st.expander("Show JSON Flow"):
plot_flow(json_flow)
res = agent.execute_flow(json_flow, submit_text)
except Exception:
res = traceback.format_exc()
st.error("Error executing the flow:")
st.error(traceback.format_exc())
return
st.session_state.past.append(submit_text)
st.session_state.generated.append(res)
@ -80,24 +139,21 @@ if st.session_state["generated"]:
): # range(len(st.session_state["generated"]) - 1, -1, -1):
message(st.session_state["past"][i], is_user=True, key=str(i) + "_user")
res = st.session_state["generated"][i]
result = st.session_state["generated"][i]
res, all_results, output_type = result
try:
json_res = json.loads(res)["data"]
print(json_res)
except Exception:
json_res = None
if json_res:
try:
res = pd.DataFrame(json_res)
except Exception:
res = json_res
if isinstance(res, str):
output_type = output_type.value
if output_type == "artifact":
# plot the json returned in res
fig_json = res["data"]["result"]
# plot the json with ploylu atream lit
st.plotly_chart(json.loads(fig_json))
elif output_type == "chat":
st.write(res)
elif isinstance(res, pd.DataFrame):
st.dataframe(res)
elif output_type == "data":
json_res = json.loads(res)["data"]
st.dataframe(json_res)
else:
st.error("Returned result:")
st.error(res)

2
toolserve/toolserve/apm/parse.py
View file

@ -76,7 +76,7 @@ def get_tools_from_file(filepath: str) -> List[str]:
tree = load_ast_tree(filepath)
tools = []
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
tool_name = get_function_name_if_decorated(node)
if tool_name:
tools.append(tool_name)

4
toolserve/toolserve/builtin/default/query.py
View file

@ -47,6 +47,7 @@ async def get_data_schema(
async def query_sql(
data_id: Param(int, "id of the data source"),
sql: Param(str, "parameterized SQL query to execute"),
output_name: Param(str, "name of the output data to save"),
params: Param(Optional[List[Union[str, int]]], "parameters to pass to the SQL query") = None,
) -> Dict[str, Union[int, str]]:
"""Query a data source using SQL
@ -62,6 +63,7 @@ async def query_sql(
Args:
data_id (int): The id of the data source to query.
sql (str): The parameterized SQL query to execute.
output_name (str): The name of the output data to save.
params (Optional[Dict[str, Any]]): Parameters to pass to the SQL query.
Returns:
@ -79,7 +81,7 @@ async def query_sql(
result_df = con.execute(sql).fetchdf()
# Save the resulting DataFrame and create a new data source
result = await save_df(result_df, f"query_result_{data_id}")
result = await save_df(result_df, output_name)
result_id = result["id"]
# Retrieve and return the schema of the new data source
return get_df_info(result_df, data_id=result_id)

29
toolserve/toolserve/utils/openai_tool.py
View file

@ -1,6 +1,9 @@
import json
from typing import Any, Dict, Type
from pydantic import BaseModel
from pydantic_core import PydanticUndefined
from enum import Enum
from toolserve.server.core.catalog import ToolSchema
@ -25,13 +28,18 @@ def python_type_to_json_type(python_type: Type) -> Dict[str, Any]:
"""
if hasattr(python_type, '__origin__'):
origin = python_type.__origin__
if origin is list:
item_type = python_type_to_json_type(python_type.__args__[0])
return {'type': 'array', 'items': item_type}
elif origin is dict:
key_type = python_type_to_json_type(python_type.__args__[0])
value_type = python_type_to_json_type(python_type.__args__[1])
return {'type': 'object', 'additionalProperties': value_type}
elif issubclass(python_type, BaseModel):
return model_to_json_schema(python_type)
return PYTHON_TO_JSON_TYPES.get(python_type, "string")
def model_to_json_schema(model: Type[BaseModel]) -> Dict[str, Any]:
@ -47,12 +55,19 @@ def model_to_json_schema(model: Type[BaseModel]) -> Dict[str, Any]:
properties = {}
required = []
for field_name, model_field in model.model_fields.items():
field_schema = {
"type": python_type_to_json_type(model_field.annotation),
"description": model_field.description or "",
}
if model_field.default is not None:
field_schema["default"] = model_field.default
type_json = python_type_to_json_type(model_field.annotation)
if isinstance(type_json, dict):
field_schema = type_json
else:
field_schema = {
"type": type_json,
"description": model_field.description or "",
}
if model_field.default not in [None, PydanticUndefined]:
if isinstance(model_field.default, Enum):
field_schema["default"] = model_field.default.value
else:
field_schema["default"] = model_field.default
if model_field.is_required():
required.append(field_name)
properties[field_name] = field_schema