Tutorial - AI Tech Support Agent using PydanticAI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ee292622-e18f-4ca5-8160-a77960443826",
   "metadata": {},
   "source": [
    "# Tutorial - AI Tech Support Agent using PydanticAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ed71a46-b1ac-43c2-8db1-c75cfe102a35",
   "metadata": {},
   "source": [
    "Requirements for this tutorial:\n",
    "- Python 3.8 or higher\n",
    "- The pydantic-ai package\n",
    "- Access to OpenAI's API\n",
    "- The nest-asyncio package for handling async operations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "014d6490-27ab-4926-93e6-51a707736295",
   "metadata": {},
   "source": [
    "### Install Required Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "39f76ece-fe2c-4535-b4dc-e9fb2a3d6551",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pydantic-ai -q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c9ce263-86d5-4b7e-9248-9f375a9ae09b",
   "metadata": {},
   "source": [
    "### Load Environment Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "98f1d43f-c523-45af-ab5e-57235dd3bfae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv, find_dotenv\n",
    "load_dotenv(find_dotenv(), override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58d49925-aa6d-4f8d-a2f5-44cb62bd9b36",
   "metadata": {},
   "source": [
    "### AI Tech Support Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea551207-f032-4270-89ab-5d0533989ff7",
   "metadata": {},
   "source": [
    "#### Import Required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "860dffad-5f29-493c-a330-18c0d0d134ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "from typing import List, Optional, Dict\n",
    "from pydantic import BaseModel\n",
    "from pydantic_ai import Agent, RunContext\n",
    "import datetime\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e0fcb42-1b51-43c3-8a08-e50d94920a22",
   "metadata": {},
   "source": [
    "#### Define our data models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0d0121db-d097-4da2-9a12-39131b40d38c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SupportQuery(BaseModel):\n",
    "    issue: str\n",
    "    severity: str\n",
    "    product: str\n",
    "    user_id: str\n",
    "    timestamp: datetime.datetime = datetime.datetime.now()\n",
    "\n",
    "class SupportResponse(BaseModel):\n",
    "    solution: str\n",
    "    next_steps: List[str]\n",
    "    escalate: bool\n",
    "    priority_level: int\n",
    "    estimated_time: str\n",
    "\n",
    "class KnowledgeBase(BaseModel):\n",
    "    product: str\n",
    "    known_issues: Dict[str, str]\n",
    "    solutions: Dict[str, List[str]]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c9863d-160c-4d4b-aff6-dbd9acae69fe",
   "metadata": {},
   "source": [
    "### Initialize Knowledge base with sample data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8ab00acd-263e-4d9b-8766-3578f3adde2c",
   "metadata": {},
   "outputs": [],
   "source": [
    " kb = KnowledgeBase(\n",
    "        product=\"CloudDB\",\n",
    "        known_issues={\n",
    "            \"Can't connect to database\": \"Check connection string and firewall rules\",\n",
    "            \"Database crash\": \"Verify system resources and restart service\",\n",
    "            \"Slow queries\": \"Analyze query performance and optimize indexes\"\n",
    "        },\n",
    "        solutions={\n",
    "            \"connection\": [\"Check credentials\", \"Verify network access\", \"Test port availability\"],\n",
    "            \"performance\": [\"Run diagnostics\", \"Check resource usage\", \"Optimize queries\"],\n",
    "            \"crash\": [\"Collect logs\", \"Check error messages\", \"Restart service\"]\n",
    "        }\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc7d6bdf-a3cd-469a-8af8-028d5522205c",
   "metadata": {},
   "source": [
    "#### Initialize our Tech Support Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8c54cfad-1edd-4708-bc71-ae822c3b1b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "tech_support_agent = Agent(\n",
    "    'openai:gpt-4',\n",
    "    deps_type=KnowledgeBase,\n",
    "    result_type=SupportResponse,\n",
    "    system_prompt=\"\"\"\n",
    "    You are an expert technical support AI agent. Your role is to:\n",
    "    1. Analyze customer issues using the provided tools\n",
    "    2. Search the knowledge base for known solutions\n",
    "    3. Provide clear, actionable solutions\n",
    "    4. Determine if escalation is needed\n",
    "    5. Estimate resolution time\n",
    "    \n",
    "    Always use the available tools to verify information before responding.\n",
    "    \"\"\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3049f20f-0d88-4375-bf78-4b502034a7f8",
   "metadata": {},
   "source": [
    "#### Define tools for the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b167f47b-4e2d-4826-bd84-700e65c37142",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tech_support_agent.tool\n",
    "async def search_knowledge_base(ctx: RunContext[KnowledgeBase], issue: str, product: str) -> Dict[str, str]:\n",
    "    \"\"\"Search the knowledge base for known issues and solutions\"\"\"\n",
    "    kb = ctx.deps\n",
    "    if product != kb.product:\n",
    "        return {\"error\": \"Product not found in knowledge base\"}\n",
    "    \n",
    "    matches = {}\n",
    "    for known_issue, solution in kb.known_issues.items():\n",
    "        if issue.lower() in known_issue.lower():\n",
    "            matches[known_issue] = solution\n",
    "    return matches\n",
    "\n",
    "@tech_support_agent.tool\n",
    "async def check_severity(ctx: RunContext[KnowledgeBase], issue: str, severity: str) -> Dict[str, any]:\n",
    "    \"\"\"Analyze issue severity and recommend priority level\"\"\"\n",
    "    severity_levels = {\n",
    "        \"low\": 1,\n",
    "        \"medium\": 2,\n",
    "        \"high\": 3,\n",
    "        \"critical\": 4\n",
    "    }\n",
    "    \n",
    "    base_priority = severity_levels.get(severity.lower(), 1)\n",
    "    \n",
    "    # Check if issue contains critical keywords\n",
    "    critical_keywords = [\"crash\", \"data loss\", \"security\", \"breach\"]\n",
    "    if any(keyword in issue.lower() for keyword in critical_keywords):\n",
    "        base_priority = max(base_priority, 3)\n",
    "    \n",
    "    return {\n",
    "        \"priority_level\": base_priority,\n",
    "        \"needs_escalation\": base_priority >= 3,\n",
    "        \"estimated_time\": f\"{base_priority * 2}h\"\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dc06d77-ab66-47a9-ac1b-d6c8021c447a",
   "metadata": {},
   "source": [
    "#### Support Query Handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "3beed2e6-44b2-4c2b-b47f-ecf2d925c8a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def handle_support_query(query: SupportQuery) -> SupportResponse:\n",
    "    # Initialize knowledge base with sample data\n",
    "    kb = KnowledgeBase(\n",
    "        product=\"CloudDB\",\n",
    "        known_issues={\n",
    "            \"Can't connect to database\": \"Check connection string and firewall rules\",\n",
    "            \"Database crash\": \"Verify system resources and restart service\",\n",
    "            \"Slow queries\": \"Analyze query performance and optimize indexes\"\n",
    "        },\n",
    "        solutions={\n",
    "            \"connection\": [\"Check credentials\", \"Verify network access\", \"Test port availability\"],\n",
    "            \"performance\": [\"Run diagnostics\", \"Check resource usage\", \"Optimize queries\"],\n",
    "            \"crash\": [\"Collect logs\", \"Check error messages\", \"Restart service\"]\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    # Process the query using our agent\n",
    "    result = await tech_support_agent.run(\n",
    "        f\"\"\"\n",
    "        Customer Issue:\n",
    "        Product: {query.product}\n",
    "        Severity: {query.severity}\n",
    "        Issue: {query.issue}\n",
    "        \n",
    "        Please analyze this issue and provide a solution.\n",
    "        \"\"\",\n",
    "        deps=kb\n",
    "    )\n",
    "    \n",
    "    return result.data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5111ce1e-2245-42ad-8af4-aacd2d089d5f",
   "metadata": {},
   "source": [
    "#### Executing our Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "94e84ddd-9fd5-4f88-ab58-8912393560c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution: This seems to be a critical problem. The most probable cause could be an improper server restart or a networking issue. First, check the server logs for any initialisation errors. If the server logs report a healthy start, check the network connections and Firewall settings. Finally, verify the database connection string and the status of the database services.\n",
      "Next steps: ['Check server logs', 'Check network connections', 'Verify Firewall settings', 'Verify database connection settings', 'Confirm status of database services']\n",
      "Escalate: True\n",
      "Priority: 3\n",
      "Estimated time: 6h\n"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    \n",
    "    # Create a sample query\n",
    "    query = SupportQuery(\n",
    "        issue=\"Can't connect to database after server restart\",\n",
    "        severity=\"high\",\n",
    "        product=\"CloudDB\",\n",
    "        user_id=\"user123\"\n",
    "    )\n",
    "    \n",
    "    # Run the support agent\n",
    "    response = asyncio.run(handle_support_query(query))\n",
    "    print(f\"Solution: {response.solution}\")\n",
    "    print(f\"Next steps: {response.next_steps}\")\n",
    "    print(f\"Escalate: {response.escalate}\")\n",
    "    print(f\"Priority: {response.priority_level}\")\n",
    "    print(f\"Estimated time: {response.estimated_time}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6416afcd-4952-46e7-81e3-c7946ee1060f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5029534c-d14d-4400-89ce-249a74bb4a66",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b85581e5-d180-4670-8da2-e7d1b5ee48f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffaa4129-842d-46da-8585-e418ab0cc36d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f5169eb-cea2-4bb0-9900-b0cb00f10b43",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ae2964f-c611-4b87-b6d9-21197dabcb9b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac1488d-f5a5-4ebb-ab4f-1902b69e02c4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57f225fe-85fd-4c33-a4d7-79370c12f53e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c76cb72b-8c9c-4b01-8b61-2c77a53cc017",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad308675-41a6-48de-9dab-f12d2157ea92",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b85740e0-c570-4e85-8f52-8ee8fe60d950",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c2986e-37e0-4df1-9dcd-1ccd346f356a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c468f7-b9e5-4361-b0b1-113ef2002ff1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e26426-4a18-41f5-8e46-5b602e479a63",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c2612f-f7bc-478f-80df-57323fa70952",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e703541-93ea-487a-925f-574fd0f4e050",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b95399-cea6-4a80-9da0-327b735cde07",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6244f3b5-417f-4481-8cea-c4350a3b44e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c89061df-09a9-4a75-960e-f177c83f7e55",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df0bf3a7-1185-43b6-92dc-acc9859c67cb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51cb3575-54b3-4d94-a379-684f9105db06",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a422598-d977-427e-a2eb-ef5b9571773e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76fd3c7f-5627-4f66-8135-4b883264fed2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9f03cfd-b80f-4c3e-9066-630b6113368f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}