crewai_db.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Install Dependencies**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install -qU  langchain-core\n",
    "# !pip install -qU langchain-community\n",
    "# !pip install -qU  \"crewai[tools]\"\n",
    "# !pip install -qU langchain-groq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Initial Setup**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import sqlite3\n",
    "from dataclasses import asdict, dataclass\n",
    "from datetime import datetime, timezone\n",
    "from pathlib import Path\n",
    "from textwrap import dedent\n",
    "from typing import Any, Dict, List, Tuple, Union\n",
    "\n",
    "import pandas as pd\n",
    "from crewai import Agent, Crew, Process, Task\n",
    "from crewai_tools import tool\n",
    "\n",
    "# from google.colab import userdata\n",
    "from langchain.schema import AgentFinish\n",
    "from langchain.schema.output import LLMResult\n",
    "from langchain_community.tools.sql_database.tool import (\n",
    "    InfoSQLDatabaseTool,\n",
    "    ListSQLDatabaseTool,\n",
    "    QuerySQLCheckerTool,\n",
    "    QuerySQLDataBaseTool,\n",
    ")\n",
    "from langchain_community.utilities.sql_database import SQLDatabase\n",
    "from langchain_core.callbacks.base import BaseCallbackHandler\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_groq import ChatGroq\n",
    "\n",
    "# from google.colab import userdata\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "os.environ[\"GROQ_API_KEY\"] = os.environ.get(\"GROQ_API_KEY\")\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.environ.get(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Import Local Data To Data Frame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(\"ds-salaries.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Convert The Data Frame To A Local SQLite DB File**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    os.remove(\"salaries.db\")\n",
    "except:\n",
    "    pass\n",
    "\n",
    "connection = sqlite3.connect(\"salaries.db\")\n",
    "df.to_sql(name=\"salaries\", con=connection)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set Up Callback Handler To Log Prompts And Responses From LLM**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Event:\n",
    "    event: str\n",
    "    timestamp: str\n",
    "    text: str\n",
    "\n",
    "\n",
    "def _current_time() -> str:\n",
    "    return datetime.now(timezone.utc).isoformat()\n",
    "\n",
    "\n",
    "class LLMCallbackHandler(BaseCallbackHandler):\n",
    "    def __init__(self, log_path: Path):\n",
    "        self.log_path = log_path\n",
    "\n",
    "    def on_llm_start(\n",
    "        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any\n",
    "    ) -> Any:\n",
    "        \"\"\"Run when LLM starts running.\"\"\"\n",
    "        assert len(prompts) == 1\n",
    "        event = Event(event=\"llm_start\", timestamp=_current_time(), text=prompts[0])\n",
    "        with self.log_path.open(\"a\", encoding=\"utf-8\") as file:\n",
    "            file.write(json.dumps(asdict(event)) + \"\\n\")\n",
    "\n",
    "    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:\n",
    "        \"\"\"Run when LLM ends running.\"\"\"\n",
    "        generation = response.generations[-1][-1].message.content\n",
    "        event = Event(event=\"llm_end\", timestamp=_current_time(), text=generation)\n",
    "        with self.log_path.open(\"a\", encoding=\"utf-8\") as file:\n",
    "            file.write(json.dumps(asdict(event)) + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Configure the Model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatGroq(\n",
    "    temperature=0,\n",
    "    model_name=\"llama3-70b-8192\",\n",
    "    # model_name=\"mixtral-8x7b-32768\",\n",
    "    callbacks=[LLMCallbackHandler(Path(\"prompts.jsonl\"))],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Open The DB Connection**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = SQLDatabase.from_uri(\"sqlite:///salaries.db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Create The Tools**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(\"list_tables\")\n",
    "def list_tables() -> str:\n",
    "    \"\"\"List the available tables in the database\"\"\"\n",
    "    return ListSQLDatabaseTool(db=db).invoke(\"\")\n",
    "\n",
    "\n",
    "list_tables.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(\"tables_schema\")\n",
    "def tables_schema(tables: str) -> str:\n",
    "    \"\"\"\n",
    "    Input is a comma-separated list of tables, output is the schema and sample rows\n",
    "    for those tables. Be sure that the tables actually exist by calling `list_tables` first!\n",
    "    Example Input: table1, table2, table3\n",
    "    \"\"\"\n",
    "    tool = InfoSQLDatabaseTool(db=db)\n",
    "    return tool.invoke(tables)\n",
    "\n",
    "\n",
    "print(tables_schema.run(\"salaries\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(\"execute_sql\")\n",
    "def execute_sql(sql_query: str) -> str:\n",
    "    \"\"\"Execute a SQL query against the database. Returns the result\"\"\"\n",
    "    return QuerySQLDataBaseTool(db=db).invoke(sql_query)\n",
    "\n",
    "\n",
    "execute_sql.run(\"SELECT * FROM salaries WHERE salary > 10000 LIMIT 5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(\"check_sql\")\n",
    "def check_sql(sql_query: str) -> str:\n",
    "    \"\"\"\n",
    "    Use this tool to double check if your query is correct before executing it. Always use this\n",
    "    tool before executing a query with `execute_sql`.\n",
    "    \"\"\"\n",
    "    return QuerySQLCheckerTool(db=db, llm=llm).invoke({\"query\": sql_query})\n",
    "\n",
    "\n",
    "check_sql.run(\"SELECT * WHERE salary > 10000 LIMIT 5 table = salaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Create The Agents**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_dev = Agent(\n",
    "    role=\"Senior Database Developer\",\n",
    "    goal=\"Construct and execute SQL queries based on a request\",\n",
    "    backstory=dedent(\n",
    "        \"\"\"\n",
    "        You are an experienced database engineer who is master at creating efficient and complex SQL queries.\n",
    "        You have a deep understanding of how different databases work and how to optimize queries.\n",
    "        Use the `list_tables` to find available tables.\n",
    "        Use the `tables_schema` to understand the metadata for the tables.\n",
    "        Use the `execute_sql` to check your queries for correctness.\n",
    "        Use the `check_sql` to execute queries against the database.\n",
    "    \"\"\"\n",
    "    ),\n",
    "    # llm=llm,\n",
    "    tools=[list_tables, tables_schema, execute_sql, check_sql],\n",
    "    allow_delegation=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_analyst = Agent(\n",
    "    role=\"Senior Data Analyst\",\n",
    "    goal=\"You receive data from the database developer and analyze it\",\n",
    "    backstory=dedent(\n",
    "        \"\"\"\n",
    "        You have deep experience with analyzing datasets using Python.\n",
    "        Your work is always based on the provided data and is clear,\n",
    "        easy-to-understand and to the point. You have attention\n",
    "        to detail and always produce very detailed work (as long as you need).\n",
    "    \"\"\"\n",
    "    ),\n",
    "    # llm=llm,\n",
    "    allow_delegation=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "report_writer = Agent(\n",
    "    role=\"Senior Report Editor\",\n",
    "    goal=\"Write an executive summary type of report based on the work of the analyst\",\n",
    "    backstory=dedent(\n",
    "        \"\"\"\n",
    "        Your writing still is well known for clear and effective communication.\n",
    "        You always summarize long texts into bullet points that contain the most\n",
    "        important details.\n",
    "        \"\"\"\n",
    "    ),\n",
    "    # llm=llm,\n",
    "    allow_delegation=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Create The Tasks**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extract_data = Task(\n",
    "    description=\"Extract data that is required for the query {query}.\",\n",
    "    expected_output=\"Database result for the query\",\n",
    "    agent=sql_dev,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analyze_data = Task(\n",
    "    description=\"Analyze the data from the database and write an analysis for {query}.\",\n",
    "    expected_output=\"Detailed analysis text\",\n",
    "    agent=data_analyst,\n",
    "    context=[extract_data],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_report = Task(\n",
    "    description=dedent(\n",
    "        \"\"\"\n",
    "        Write an executive summary of the report from the analysis. The report\n",
    "        must be less than 100 words.\n",
    "    \"\"\"\n",
    "    ),\n",
    "    expected_output=\"Markdown report\",\n",
    "    agent=report_writer,\n",
    "    context=[analyze_data],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Create The Crew**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "crew = Crew(\n",
    "    agents=[sql_dev, data_analyst, report_writer],\n",
    "    tasks=[extract_data, analyze_data, write_report],\n",
    "    process=Process.sequential,\n",
    "    verbose=True,\n",
    "    memory=False,\n",
    "    output_log_file=\"crew.log\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kickoff The Crew For Query 1:\n",
    "\n",
    "inputs = {\n",
    "    \"query\": \"Effects on salary (in USD) based on company location, size and employee experience\"\n",
    "}\n",
    "\n",
    "result = crew.kickoff(inputs=inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kickoff The Crew For Query 2:\n",
    "\n",
    "inputs = {\n",
    "    \"query\": \"How is the `Machine Learning Engineer` salary in USD is affected by remote positions\"\n",
    "}\n",
    "\n",
    "result = crew.kickoff(inputs=inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kickoff The Crew For Query 3:\n",
    "\n",
    "inputs = {\n",
    "    \"query\": \"How is the salary in USD based on employment type and experience level? \"\n",
    "}\n",
    "\n",
    "result = crew.kickoff(inputs=inputs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}