From 2e5e5ef7c864959d3eec85a505e74a84d3548ccb Mon Sep 17 00:00:00 2001 From: Sam Partee Date: Wed, 1 May 2024 23:53:53 -0700 Subject: [PATCH] more tools in sql example --- examples/gmail/pack.lock.toml | 2 + examples/gmail/pack.toml | 1 + examples/gmail/tools/chat.py | 36 ++ examples/gmail/tools/gmailer.py | 85 +++- examples/sql-chat/agent.py | 485 +++++++++++++++---- examples/sql-chat/main.py | 100 +++- toolserve/toolserve/apm/parse.py | 2 +- toolserve/toolserve/builtin/default/query.py | 4 +- toolserve/toolserve/utils/openai_tool.py | 29 +- 9 files changed, 601 insertions(+), 143 deletions(-) create mode 100644 examples/gmail/tools/chat.py diff --git a/examples/gmail/pack.lock.toml b/examples/gmail/pack.lock.toml index cfd84f29..ae60afcb 100644 --- a/examples/gmail/pack.lock.toml +++ b/examples/gmail/pack.lock.toml @@ -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" diff --git a/examples/gmail/pack.toml b/examples/gmail/pack.toml index f189689b..c5f9f76f 100644 --- a/examples/gmail/pack.toml +++ b/examples/gmail/pack.toml @@ -9,3 +9,4 @@ email = "sam@partee.io" [modules] gmailer = "0.1.0" +chat = "0.1.0" diff --git a/examples/gmail/tools/chat.py b/examples/gmail/tools/chat.py new file mode 100644 index 00000000..cff974c0 --- /dev/null +++ b/examples/gmail/tools/chat.py @@ -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 diff --git a/examples/gmail/tools/gmailer.py b/examples/gmail/tools/gmailer.py index 4acc25a3..5f264c76 100644 --- a/examples/gmail/tools/gmailer.py +++ b/examples/gmail/tools/gmailer.py @@ -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() + diff --git a/examples/sql-chat/agent.py b/examples/sql-chat/agent.py index 2ef8c53e..3615473e 100644 --- a/examples/sql-chat/agent.py +++ b/examples/sql-chat/agent.py @@ -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) - """ \ No newline at end of file + + + + + + + + + + + +#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) diff --git a/examples/sql-chat/main.py b/examples/sql-chat/main.py index 4e481e63..bf3398b8 100644 --- a/examples/sql-chat/main.py +++ b/examples/sql-chat/main.py @@ -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) diff --git a/toolserve/toolserve/apm/parse.py b/toolserve/toolserve/apm/parse.py index 7048facf..83e793a4 100644 --- a/toolserve/toolserve/apm/parse.py +++ b/toolserve/toolserve/apm/parse.py @@ -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) diff --git a/toolserve/toolserve/builtin/default/query.py b/toolserve/toolserve/builtin/default/query.py index 8368aee1..6872f38f 100644 --- a/toolserve/toolserve/builtin/default/query.py +++ b/toolserve/toolserve/builtin/default/query.py @@ -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) diff --git a/toolserve/toolserve/utils/openai_tool.py b/toolserve/toolserve/utils/openai_tool.py index 47f0405b..d4d90184 100644 --- a/toolserve/toolserve/utils/openai_tool.py +++ b/toolserve/toolserve/utils/openai_tool.py @@ -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