{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customer Information Assistant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d8170be280d44b0d87b8af209b52d05f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(Label(value='Enter your query about a customer'),)), HBox(children=(Text(value='…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import dataiku\n",
    "import ipywidgets as widgets\n",
    "import json\n",
    "import os\n",
    "from utils import get_customer_details, search_company_info, process_tool_calls, create_chat_session\n",
    "\n",
    "# LLM setup\n",
    "LLM_ID = \"openai:yv-openai-only35:gpt-4o-mini\"  # LLM ID for the LLM Mesh connection + model goes here\n",
    "client = dataiku.api_client()\n",
    "project = client.get_default_project()\n",
    "llm = project.get_llm(LLM_ID)\n",
    "\n",
    "def process_agent_response(chat, query):\n",
    "    \"\"\"Process the agent's response and handle any tool calls\"\"\"\n",
    "    chat.with_message(query, role=\"user\")\n",
    "    \n",
    "    while True:\n",
    "        response = chat.execute()\n",
    "        \n",
    "        if not response.tool_calls:\n",
    "            # Final answer received\n",
    "            chat.with_message(response.text, role=\"assistant\")\n",
    "            chat = create_chat_session(llm, project) # refresh chat\n",
    "            return response.text\n",
    "            \n",
    "        # Handle tool calls\n",
    "        chat.with_tool_calls(response.tool_calls, role=\"assistant\")\n",
    "        tool_name = response.tool_calls[0][\"function\"][\"name\"]\n",
    "        tool_args = response.tool_calls[0][\"function\"][\"arguments\"]\n",
    "        \n",
    "        # Process tool call and get result\n",
    "        tool_call_result = process_tool_calls(response.tool_calls)\n",
    "        chat.with_tool_output(tool_call_result, tool_call_id=response.tool_calls[0][\"id\"])\n",
    "\n",
    "\n",
    "# Create widgets\n",
    "label = widgets.Label(value=\"Enter your query about a customer\")\n",
    "query_input = widgets.Text(\n",
    "    placeholder=\"Tell me about customer fdouetteau\",\n",
    "    continuous_update=False,\n",
    "    layout=widgets.Layout(width='50%')\n",
    ")\n",
    "result = widgets.HTML(value=\"\")\n",
    "button = widgets.Button(description=\"Ask\")\n",
    "\n",
    "# Create the chat session\n",
    "chat = create_chat_session(llm, project)\n",
    "\n",
    "def on_button_click(b):\n",
    "    \"\"\"Handle button click event\"\"\"\n",
    "    query = query_input.value\n",
    "    if query:\n",
    "        try:\n",
    "            response = process_agent_response(chat, query)\n",
    "            result.value = f\"<div style='white-space: pre-wrap;'>{response}</div>\"\n",
    "        except Exception as e:\n",
    "            result.value = f\"<div style='color: red'>Error: {str(e)}</div>\"\n",
    "            \n",
    "button.on_click(on_button_click)\n",
    "\n",
    "# Layout\n",
    "display(widgets.VBox([\n",
    "    widgets.HBox([label]),\n",
    "    widgets.HBox([query_input, button]),\n",
    "    widgets.HBox([result])\n",
    "], layout=widgets.Layout(padding='20px')))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}