diff --git a/examples/gmail/tools/chat.py b/examples/gmail/tools/chat.py index cff974c0..5687ed37 100644 --- a/examples/gmail/tools/chat.py +++ b/examples/gmail/tools/chat.py @@ -2,16 +2,17 @@ from toolserve.sdk import Param, tool, get_secret from toolserve.sdk.dataframe import get_df - +import pandas as pd import openai @tool -def summarize( - text: Param(str, "Text to summarize"), +async def summarize( + #text: Param(str, "Text to summarize"), + data_id: Param(int, "ID of the data 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. + """Summarize a piece of text using OpenAI Language models. Args: text (str): The text to summarize. @@ -20,7 +21,10 @@ def summarize( Returns: str: The summarized text. """ - api_key = get_secret("openai_api_key") + df = await get_df(data_id) + text = df.to_json(orient='records') + api_key = get_secret("openai_api_key", None) + model = get_secret("openai_model_summarize", "gpt-4-turbo") # Call the OpenAI model with the tools and messages messages = [ {"role": "system", "content": system_prompt}, @@ -28,8 +32,8 @@ def summarize( ] client = openai.Client(api_key=api_key) - completion = openai.chat.completions.create( - model=self.model, + completion = openai.chat.completions.create( + model=model, messages=messages, ) summary = completion.choices[0].message.content diff --git a/examples/sql-chat/agent.py b/examples/sql-chat/agent.py index 3615473e..13a9c6ca 100644 --- a/examples/sql-chat/agent.py +++ b/examples/sql-chat/agent.py @@ -32,7 +32,7 @@ class ToolClient: def __init__(self, base_url: str): self.base_url = base_url - self.client = httpx.Client() + self.client = httpx.Client(timeout=30) self.tools = self.__collect_tool_specs() def __collect_tool_specs(self) -> Dict[str, str]: @@ -115,6 +115,10 @@ class ToolRunner: def set_source(self, source: str): self._data_sources = self.__get_data_sources() + + if not source: + return + retries = 3 data_id = None while retries > 0: @@ -144,7 +148,9 @@ class ToolRunner: def __create_prompt(self, user_query: str, input_name: str, output_name: str) -> List[Dict[str, str]]: schema = self._data_schema - data_id = self._data_sources[input_name] + data_id = "No input" + if input_name: + data_id = self._data_sources[input_name] prompt = self.tool_prompt.format(schema=schema, data_id=data_id, output_name=output_name) messages = [ @@ -190,7 +196,7 @@ class ToolRunner: else: raise ValueError(f"Invalid params type: {type(params)}") - if "output_name" in args: + if "output_name" in args and output_name != "None": args["output_name"] = output_name if "data_id" in args: args["data_id"] = self._data_id @@ -206,6 +212,9 @@ class ToolRunner: :return: The result of the tool """ self.set_source(source) + print(f"Tool Name: {tool_name}") + print(f"Data ID: {self._data_id}") + print(f"Sourcing data from {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) @@ -248,9 +257,9 @@ class Edge(BaseModel): 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") + input_name: Optional[str] = Field(None, 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") + output_name: Optional[str] = Field(..., description="The name of the output data") class OutputType(Enum): DATA = "data" @@ -391,6 +400,8 @@ class ToolFlow: 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) + elif sink_output_type == OutputType.CHAT: + data = results[sink_node_id]["data"]["result"] else: data = results[sink_node_id] @@ -452,9 +463,22 @@ def summarize_flow_results(model_client, flow_results: Dict[str, Any], flow_sche +email_flow = FlowSchema( + nodes=[ + ToolNode(node_id=0, input_name=None, tool_name="ReadEmail", output_name="email_data_1"), + ToolNode(node_id=1, input_name="email_data_1", tool_name="Summarize", output_name=None), + ], + edges=[ + Edge(source=0, target=1) + ], + output_type=OutputType.CHAT +) +class Agent: + def __init__(self, flows: Dict[str, FlowSchema]): + self.flows = flows diff --git a/examples/sql-chat/main.py b/examples/sql-chat/main.py index bf3398b8..a4e41c0a 100644 --- a/examples/sql-chat/main.py +++ b/examples/sql-chat/main.py @@ -20,7 +20,7 @@ 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 ToolFlow, email_flow PROMPT = dedent("""Given a user query, construct a graph based representation of functions (nodes), and their data flow (edges) such that @@ -113,8 +113,9 @@ def submit(): with st.spinner(text="Wait for Agent..."): try: agent = get_agent() - flow = agent.infer_flow(submit_text) - json_flow = json.loads(flow) + #flow = agent.infer_flow(submit_text) + #json_flow = json.loads(flow) + json_flow = email_flow.dict() with st.expander("Show JSON Flow"): plot_flow(json_flow) res = agent.execute_flow(json_flow, submit_text) diff --git a/toolserve/pyproject.toml b/toolserve/pyproject.toml index a0afea25..c335a95b 100644 --- a/toolserve/pyproject.toml +++ b/toolserve/pyproject.toml @@ -3,7 +3,6 @@ name = "toolserve" version = "0.1.0" description = "" authors = ["Sam Partee "] -readme = "README.md" [tool.poetry.dependencies] python = "^3.10"