CVE_assistant_api.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "880c83e5-599b-4763-a63b-7cd581b21ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "client = OpenAI()\n",
    "  \n",
    "assistant = client.beta.assistants.create(\n",
    "  name=\"CVE implementor\",\n",
    "  instructions=\"You are a cybersecurity analyist Write and run code to implement CVEs.\",\n",
    "  tools=[{\"type\": \"code_interpreter\"}],\n",
    "  model=\"gpt-4-turbo\",\n",
    ")\n",
    "\n",
    "thread = client.beta.threads.create()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "adcc8541-e525-4c6a-bd88-9ab8c28cc748",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant Received: {\"output\": \"Hello, World!\\n\", \"errors\": \"\", \"exit_status\": 0}\n"
     ]
    }
   ],
   "source": [
    "import subprocess\n",
    "import json\n",
    "\n",
    "def execute_bash_command(command):\n",
    "    \"\"\"\n",
    "    Executes a Bash command in a safe manner and returns the output, errors, and exit status.\n",
    "    \n",
    "    Args:\n",
    "        command (str): The Bash command to execute.\n",
    "    \n",
    "    Returns:\n",
    "        dict: A dictionary containing the output, errors, and exit status.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Execute the command, capture output and error\n",
    "        result = subprocess.run(command, shell=True, text=True, capture_output=True, check=False)\n",
    "        \n",
    "        return {\n",
    "            \"output\": result.stdout,\n",
    "            \"errors\": result.stderr,\n",
    "            \"exit_status\": result.returncode\n",
    "        }\n",
    "    except Exception as e:\n",
    "        return {\n",
    "            \"output\": \"\",\n",
    "            \"errors\": str(e),\n",
    "            \"exit_status\": -1\n",
    "        }\n",
    "\n",
    "def assistant_interact(command):\n",
    "    \"\"\"\n",
    "    This function simulates an interaction where the OpenAI Assistant sends a Bash command\n",
    "    and receives a response which it can then parse or use in conversation.\n",
    "    \n",
    "    Args:\n",
    "        command (str): A Bash command received from the OpenAI Assistant.\n",
    "    \n",
    "    Returns:\n",
    "        str: A JSON string containing the output, errors, and exit status.\n",
    "    \"\"\"\n",
    "    result = execute_bash_command(command)\n",
    "    return json.dumps(result)\n",
    "\n",
    "# Example interaction:\n",
    "if __name__ == \"__main__\":\n",
    "    # Simulating a command sent by the OpenAI Assistant\n",
    "    command_request = \"echo Hello, World!\"\n",
    "    response = assistant_interact(command_request)\n",
    "    print(\"Assistant Received:\", response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "287dd94e-68fe-4f31-a97f-993658f13613",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 1: Upgrade to Python SDK v1.2 with pip install --upgrade openai\n",
    "# Step 2: Install Tavily Python SDK with pip install tavily-python\n",
    "# Step 3: Build an OpenAI assistant with Python SDK documentation - https://platform.openai.com/docs/assistants/overview\n",
    "\n",
    "import os\n",
    "import json\n",
    "import time\n",
    "from openai import OpenAI\n",
    "from tavily import TavilyClient\n",
    "\n",
    "# Initialize clients with API keys\n",
    "client = OpenAI(api_key=os.environ[\"OPENAI_API_KEY\"])\n",
    "tavily_client = TavilyClient(api_key=os.environ[\"TAVILY_API_KEY\"])\n",
    "\n",
    "assistant_prompt_instruction = \"\"\"You are a finance expert. \n",
    "Your goal is to provide answers based on information from the internet. \n",
    "You must use the provided Tavily search API function to find relevant online information. \n",
    "You should never use your own knowledge to answer questions.\n",
    "Please include relevant url sources in the end of your answers.\n",
    "\"\"\"\n",
    "\n",
    "# Function to perform a Tavily search\n",
    "def tavily_search(query):\n",
    "    search_result = tavily_client.get_search_context(query, search_depth=\"advanced\", max_tokens=8000)\n",
    "    return search_result\n",
    "\n",
    "# Function to wait for a run to complete\n",
    "def wait_for_run_completion(thread_id, run_id):\n",
    "    while True:\n",
    "        time.sleep(1)\n",
    "        run = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run_id)\n",
    "        print(f\"Current run status: {run.status}\")\n",
    "        if run.status in ['completed', 'failed', 'requires_action']:\n",
    "            return run\n",
    "\n",
    "# Function to handle tool output submission\n",
    "def submit_tool_outputs(thread_id, run_id, tools_to_call):\n",
    "    tool_output_array = []\n",
    "    for tool in tools_to_call:\n",
    "        output = None\n",
    "        tool_call_id = tool.id\n",
    "        function_name = tool.function.name\n",
    "        function_args = tool.function.arguments\n",
    "\n",
    "        if function_name == \"tavily_search\":\n",
    "            output = tavily_search(query=json.loads(function_args)[\"query\"])\n",
    "\n",
    "        if output:\n",
    "            tool_output_array.append({\"tool_call_id\": tool_call_id, \"output\": output})\n",
    "\n",
    "    return client.beta.threads.runs.submit_tool_outputs(\n",
    "        thread_id=thread_id,\n",
    "        run_id=run_id,\n",
    "        tool_outputs=tool_output_array\n",
    "    )\n",
    "\n",
    "# Function to print messages from a thread\n",
    "def print_messages_from_thread(thread_id):\n",
    "    messages = client.beta.threads.messages.list(thread_id=thread_id)\n",
    "    for msg in messages:\n",
    "        print(f\"{msg.role}: {msg.content[0].text.value}\")\n",
    "\n",
    "# Create an assistant\n",
    "assistant = client.beta.assistants.create(\n",
    "    instructions=assistant_prompt_instruction,\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    tools=[{\"type\": \"code_interpreter\"},{\"type\": \"execute_bash_command\"} {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            \"name\": \"tavily_search\",\n",
    "            \"description\": \"Get information on recent events from the web.\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"query\": {\"type\": \"string\", \"description\": \"The search query to use. For example: 'Latest news on Nvidia stock performance'\"},\n",
    "                },\n",
    "                \"required\": [\"query\"]\n",
    "            }\n",
    "        }\n",
    "    }]\n",
    ")\n",
    "assistant_id = assistant.id\n",
    "print(f\"Assistant ID: {assistant_id}\")\n",
    "\n",
    "# Create a thread\n",
    "thread = client.beta.threads.create()\n",
    "print(f\"Thread: {thread}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f00410a-f859-4e10-ae89-d38f3b39ff39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ongoing conversation loop\n",
    "while True:\n",
    "    user_input = input(\"You: \")\n",
    "    if user_input.lower() == 'exit':\n",
    "        break\n",
    "\n",
    "    # Create a message\n",
    "    message = client.beta.threads.messages.create(\n",
    "        thread_id=thread.id,\n",
    "        role=\"user\",\n",
    "        content=user_input,\n",
    "    )\n",
    "\n",
    "    # Create a run\n",
    "    run = client.beta.threads.runs.create(\n",
    "        thread_id=thread.id,\n",
    "        assistant_id=assistant_id,\n",
    "    )\n",
    "    print(f\"Run ID: {run.id}\")\n",
    "\n",
    "    # Wait for run to complete\n",
    "    run = wait_for_run_completion(thread.id, run.id)\n",
    "\n",
    "    if run.status == 'failed':\n",
    "        print(run.error)\n",
    "        continue\n",
    "    elif run.status == 'requires_action':\n",
    "        run = submit_tool_outputs(thread.id, run.id, run.required_action.submit_tool_outputs.tool_calls)\n",
    "        run = wait_for_run_completion(thread.id, run.id)\n",
    "\n",
    "    # Print messages from the thread\n",
    "    print_messages_from_thread(thread.id)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

CVE_langchain_api.ipynb