{ "cells": [ { "cell_type": "markdown", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "source": [ "# Dataiku: voilà an agent application!\n", "\n", "## Defining the tools" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "from typing import Type\n", "\n", "from dataiku.langchain.dku_llm import DKUChatModel\n", "from dataiku.sql import Constant, toSQL, Dialects\n", "from ddgs import DDGS\n", "from langchain.agents import create_agent\n", "from langchain.tools import BaseTool\n", "from pydantic import BaseModel, Field\n", "\n", "import dataiku\n", "from dataiku import SQLExecutor2\n", "\n", "LLM_ID = \"\" # Fill with a VALID LLM id\n", "DATASET_NAME = \"pro_customers_sql\"\n", "llm = DKUChatModel(llm_id=LLM_ID)\n", "\n", "\n", "class CustomerInfo(BaseModel):\n", " \"\"\"Parameter for GetCustomerInfo\"\"\"\n", " id: str = Field(description=\"customer ID\")\n", "\n", "\n", "class GetCustomerInfo(BaseTool):\n", " \"\"\"Gathering customer information\"\"\"\n", "\n", " name: str = \"GetCustomerInfo\"\n", " description: str = \"Provide a name, job title and company of a customer, given the customer's ID\"\n", " args_schema: Type[BaseModel] = CustomerInfo\n", "\n", " def _run(self, id: str):\n", " dataset = dataiku.Dataset(DATASET_NAME)\n", " table_name = dataset.get_location_info().get('info', {}).get('quotedResolvedTableName')\n", " executor = SQLExecutor2(dataset=dataset)\n", " cid = Constant(str(id))\n", " escaped_cid = toSQL(cid, dialect=Dialects.POSTGRES) # Replace by your DB\n", " query_reader = executor.query_to_iter(\n", " f\"\"\"SELECT \"name\", \"job\", \"company\" FROM {table_name} WHERE \"id\" = {escaped_cid}\"\"\")\n", " for (name, job, company) in query_reader.iter_tuples():\n", " return f\"The customer's name is \\\"{name}\\\", holding the position \\\"{job}\\\" at the company named {company}\"\n", " return f\"No information can be found about the customer {id}\"\n", "\n", " def _arun(self, id: str):\n", " raise NotImplementedError(\"This tool does not support async\")\n", "\n", "\n", "class CompanyInfo(BaseModel):\n", " \"\"\"Parameter for the GetCompanyInfo\"\"\"\n", " name: str = Field(description=\"Company's name\")\n", "\n", "\n", "class GetCompanyInfo(BaseTool):\n", " \"\"\"Class for gathering in the company information\"\"\"\n", "\n", " name: str = \"GetCompanyInfo\"\n", " description: str = \"Provide general information about a company, given the company's name.\"\n", " args_schema: Type[BaseModel] = CompanyInfo\n", "\n", " def _run(self, name: str):\n", " results = DDGS().text(name + \" (company)\", max_results=1)\n", " result = \"Information found about \" + name + \": \" + results[0][\"body\"] + \"\\n\" \\\n", " if len(results) > 0 and \"body\" in results[0] \\\n", " else None\n", " if not result:\n", " results = DDGS().text(name, max_results=1)\n", " result = \"Information found about \" + name + \": \" + results[0][\"body\"] + \"\\n\" \\\n", " if len(results) > 0 and \"body\" in results[0] \\\n", " else \"No information can be found about the company \" + name\n", " return result\n", "\n", " def _arun(self, name: str):\n", " raise NotImplementedError(\"This tool does not support async\")\n", "\n", "\n", "tools = [GetCustomerInfo(), GetCompanyInfo()]\n", "tool_names = [tool.name for tool in tools]\n", "\n", "# Initializes the agent\n", "prompt = \"\"\"Answer the following questions as best you can. You have only access to the following tools:\n", "{tools}\n", "Use the following format:\n", "Question: the input question you must answer\n", "Thought: you should always think about what to do\n", "Action: the action to take should be one of [{tool_names}]\n", "Action Input: the input to the action\n", "Observation: the result of the action\n", "... (this Thought/Action/Action Input/Observation can repeat N times)\n", "Thought: I now know the final answer\n", "Final Answer: the final answer to the original input question\n", "Begin!\n", "Question: {input}\n", "Thought:{agent_scratchpad}\"\"\"\n", "\n", "agent = create_agent(model=llm, tools=tools, system_prompt=prompt)" ] }, { "cell_type": "markdown", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "source": [ "## Voilà application" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "import ipywidgets as widgets\n", "\n", "label = widgets.Label(value=\"Enter the customer ID\")\n", "status = widgets.HTML(value=\"\")\n", "text = widgets.Text(placeholder=\"fdouetteau\", continuous_update=False)\n", "\n", "result = widgets.Output()\n", "\n", "\n", "def search(customer_id):\n", " \"\"\"\n", " Search information about a customer\n", " Args:\n", " customer_id: customer ID\n", " Returns:\n", " the agent result\n", " \"\"\"\n", " return agent.invoke({\n", " \"messages\": [{\n", " \"role\": \"user\",\n", " \"content\": f\"\"\"Give all the professional information you can about the customer with ID: {customer_id}.\n", " Also include information about the company if you can.\"\"\"\n", " }]\n", " })[\"messages\"][-1].content\n", "\n", "\n", "def callback(customerId):\n", " \"\"\"\n", " Callback function for calling the agent\n", " Args:\n", " customerId: customer ID\n", " \"\"\"\n", " result.clear_output()\n", " status.value = \"Processing request\"\n", " with result:\n", " print(search(customerId.get('new', '')))\n", " status.value = \"\"\n", "\n", "\n", "text.observe(callback, 'value')\n", "display(widgets.VBox([widgets.HBox([label, text]), status, result]))\n", "\n" ] } ], "metadata": { "kernelspec": { "display_name": "DSS Codeenv - pyenv-voila", "language": "python", "name": "py-dku-venv-pyenv-voila" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.14" } }, "nbformat": 4, "nbformat_minor": 4 }