more tools in sql example

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"

examples/gmail/pack.toml

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

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

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(
+async def read_email(
-    email_address: Param(str, "Email address of the recipient"),
+    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()
-
+    df = pd.DataFrame(emails)
-    return pd.DataFrame(emails).to_json()
+    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()
+

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):
+        Args:
-        self.toolchain = toolchain
+            base_url (str): The base URL for the API calls.
-        self.data_sources = self.__get_data_sources()
+            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():
+        self._data_sources = self.__get_data_sources()
-            raise ValueError(f"Data source '{source}' not found.")
+        retries = 3
-        else:
+        data_id = None
-            data_id = self.data_sources[source]
+        while retries > 0:
-            # get the schema
+            try:
-            schema = self.toolchain.call_api("POST", "/tool/query/get_data_schema", json_data={"data_id": data_id})
+                data_id = self._data_sources[source]
-            self._source = source
+                break
-            self._data_schema = schema
+            except KeyError:
+                retries -= 1
+                time.sleep(1)
+                self._data_sources = self.__get_data_sources()
 
-    def get_source(self) -> str:
+        if data_id is None:
-        return self._source
+            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 __create_prompt(self, user_query: str, input_name: str, output_name: str) -> List[Dict[str, str]]:
-    def query(self, user_query: str) -> str:
-        if not self._source:
-            raise ValueError("Data source not set. Please set a data source before querying.")
         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)
+    def get_tool_args(self, tool_name: str, messages: List[Dict[str, str]], output_name: str) -> Dict[str, Any]:
-        args = json.loads(tool_args)
+        """
-        params = args.get("params", [])
+        Retrieves the required arguments for an tool from the Darkstar Toolserver API and
-        if params:
+        uses them to call an OpenAI model with predefined tools and messages.
-            if isinstance(params, dict):
+
-                args["params"] = list(params.values())
+        :param tool_name: The name of the tool to execute.
-            elif isinstance(params, str):
+        :param messages: A list of messages to provide to the model.
-                args["params"] = [params]
+        :return: The result of the OpenAI model call.
-            elif isinstance(params, list):
+        """
-                args["params"] = params
+        func_spec = self._client.tools.get(tool_name, {})
-            else:
+        if not func_spec:
-                raise ValueError(f"Invalid params type: {type(params)}")
+            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)
+def pydantic_to_openai_tool(model: Type[BaseModel]) -> str:
-        if response["code"] != 200:
+    """
-            raise ValueError(f"Error executing tool: {response['message']}")
+    Convert a Pydantic model to an OpenAI tool schema.
-        data_id = response["data"]["result"]["data_id"]
 
-        # get the data
+    Args:
-        data_response = self.toolchain.call_api("GET", f"/api/v1/data/object/{data_id}")
+        model (Type[BaseModel]): The Pydantic model to convert.
-        if data_response["code"] != 200:
-            raise ValueError(f"Error retrieving data: {data_response['message']}")
-        data = data_response["data"]["json_blob"]
-        return data
 
+    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
+class Edge(BaseModel):
-from enum import Enum
+    source: int = Field(..., description="The ID of the source node")
+    target: int = Field(..., description="The ID of the target node")
 
-class ToolNode:
+class ToolNode(BaseModel):
-    pass
+    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,
+        name: str,
-        description,
+        description: str,
-        sources,
+        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)

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)
+    AnalysisTool = ToolFlow(
-    agent = Agent(toolchain)
+        name="data_analysis",
-    agent.set_source("users_db")
+        description="A tool flow for data analysis",
-    return agent
+        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:
+            output_type = output_type.value
-                try:
+            if output_type == "artifact":
-                    res = pd.DataFrame(json_res)
+                # plot the json returned in res
-                except Exception:
+                fig_json = res["data"]["result"]
-                    res = json_res
+                # plot the json with ploylu atream lit
-
+                st.plotly_chart(json.loads(fig_json))
-            if isinstance(res, str):
+            elif output_type == "chat":
                 st.write(res)
-            elif isinstance(res, pd.DataFrame):
+            elif output_type == "data":
-                st.dataframe(res)
+                json_res = json.loads(res)["data"]
+                st.dataframe(json_res)
             else:
                 st.error("Returned result:")
                 st.error(res)

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)

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)

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_json = python_type_to_json_type(model_field.annotation)
-            "type": python_type_to_json_type(model_field.annotation),
+        if isinstance(type_json, dict):
-            "description": model_field.description or "",
+            field_schema = type_json
-        }
+        else:
-        if model_field.default is not None:
+            field_schema = {
-            field_schema["default"] = model_field.default
+                "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