{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a6ab9a2-28a2-445d-8512-a0dc8d1b54e9",
   "metadata": {},
   "source": [
    "# Trading Code Generator\n",
    "\n",
    "A code generator that writes trading code to buy and sell equities in a simulated environment, based on a given API\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e610bf56-a46e-4aff-8de1-ab49d62b1ad3",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import io\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "from google import genai\n",
    "from google.genai import types\n",
    "import anthropic\n",
    "import ollama\n",
    "import gradio as gr\n",
    "import requests\n",
    "from typing import Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f672e1c-87e9-4865-b760-370fa605e614",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# environment\n",
    "\n",
    "load_dotenv(override=True)\n",
    "os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY', 'your-key-if-not-using-env')\n",
    "os.environ['GOOGLE_API_KEY'] = os.getenv('GOOGLE_API_KEY', 'your-key-if-not-using-env')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa149ed-9298-4d69-8fe2-8f5de0f667da",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "# initialize\n",
    "\n",
    "openai = OpenAI()\n",
    "claude = anthropic.Anthropic()\n",
    "client = genai.Client()\n",
    "\n",
    "\n",
    "OPENAI_MODEL = \"gpt-4o\"\n",
    "CLAUDE_MODEL = \"claude-sonnet-4-20250514\"\n",
    "GEMINI_MODEL = 'gemini-2.5-flash'\n",
    "LLAMA_MODEL = \"llama3.2\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b0a6f6",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "system_message = \"\"\"\n",
    "You are an effective programming assistant specialized to generate Python code based on the inputs.\n",
    "Respond only with Python code; use comments sparingly and do not provide any explanation other than occasional comments.\n",
    "Do not include Markdown formatting (```), language tags (python), or extra text.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7b3546-57aa-4c29-bc5d-f211970d04eb",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def get_user_prompt_for_trade_code_generation(currency, wallet_balance):\n",
    "    wallet_balance = str(wallet_balance)\n",
    "\n",
    "    user_prompt = \"\"\"\n",
    "        Create a simple Crypto trading engine Python code.\n",
    "        The engine will sell or buy the given crypto currency against USDT (Tether) based on the available wallet balance\n",
    "        This should be a simple Python code, not a function\n",
    "        The currency is: {}\n",
    "        The wallet balance is: {}\"\"\".format(currency, wallet_balance)\n",
    "    user_prompt += \"\"\"\n",
    "        Output will be a text containing the followings:\n",
    "            - advice to sell or buy\n",
    "            - amount in USDT\n",
    "        Rules you have to apply in the code:\n",
    "            - compose symbol: convert the input `crypto_currency` argument to upper case and concatenate it to string \"USDT\"\n",
    "            - compose url passing the previously composed symbol: `url = f\"https://data-api.binance.vision/api/v3/ticker/24hr?symbol={symbol}`\n",
    "            - call the api from with this url, expect to get the following json response, for example:\n",
    "                {'symbol': 'BTCUSDT',\n",
    "                'priceChange': '1119.99000000',\n",
    "                'priceChangePercent': '0.969',\n",
    "                'weightedAvgPrice': '116314.23644195',\n",
    "                'prevClosePrice': '115600.00000000',\n",
    "                'lastPrice': '116719.99000000',\n",
    "                'lastQty': '0.05368000',\n",
    "                'bidPrice': '116719.99000000',\n",
    "                'bidQty': '2.81169000',\n",
    "                'askPrice': '116720.00000000',\n",
    "                'askQty': '3.46980000',\n",
    "                'openPrice': '115600.00000000',\n",
    "                'highPrice': '117286.73000000',\n",
    "                'lowPrice': '114737.11000000',\n",
    "                'volume': '12500.51369000',\n",
    "                'quoteVolume': '1453987704.98443060',\n",
    "                'openTime': 1758015394001,\n",
    "                'closeTime': 1758101794001,\n",
    "                'firstId': 5236464586,\n",
    "                'lastId': 5238628513,\n",
    "                'count': 2163928}\n",
    "            - build a logic based on the retrieving values which can decide whether the engine should sell or buy he given crypto currency\n",
    "                - in the logic the code should also decide what is the confident level of selling or buying.\n",
    "                    - if the confident level is high the `amount` should be higher (closer to the `current_wallet_balance`)\n",
    "                    - if the confident level is lower then the amount should be lower as well\n",
    "            - error handling:\n",
    "                - if the api call returns with a json hving `code`, `msg` keys in it (eg. 'code': -1121, 'msg': 'Invalid symbol.') then handle this error message\n",
    "        Response rule: in your response do not include Markdown formatting (```), language tags (python), or extra text.\n",
    "    \"\"\"\n",
    "    return user_prompt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5030fdf5",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def get_user_prompt_for_docstring_generation(python_code):\n",
    "    return f\"\"\"\n",
    "        Consider the following Python code: \\n\\n\n",
    "        {python_code} \\n\\n\n",
    "\n",
    "        Generate a docstring comment around this code and it alongside with the Python code. \\n\n",
    "        Response rule: in your response do not include Markdown formatting (```), language tags (python), or extra text.\n",
    "\n",
    "        \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dc065c2",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def trade_gpt(currency, wallet_balance):\n",
    "    completion = openai.chat.completions.create(\n",
    "    model=OPENAI_MODEL,\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "        {\"role\": \"user\", \"content\": get_user_prompt_for_trade_code_generation(\n",
    "            currency,\n",
    "            wallet_balance\n",
    "            )\n",
    "            }\n",
    "    ]\n",
    "    )\n",
    "    return completion.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b402c67",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def code_comment_gpt(python_code):\n",
    "    completion = openai.chat.completions.create(\n",
    "    model=OPENAI_MODEL,\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": system_message},\n",
    "        {\"role\": \"user\", \"content\": get_user_prompt_for_docstring_generation(python_code)}\n",
    "        ]\n",
    "    )\n",
    "    return completion.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc80287",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def trade_cloude(currency, wallet_balance):\n",
    "    message = claude.messages.create(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        temperature=0.7,\n",
    "        system=system_message,\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\": get_user_prompt_for_trade_code_generation(\n",
    "                currency,\n",
    "                wallet_balance\n",
    "                )\n",
    "                },\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    return message.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90eb9547",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def code_comment_cloude(python_code):\n",
    "    message = claude.messages.create(\n",
    "        model=CLAUDE_MODEL,\n",
    "        max_tokens=2000,\n",
    "        temperature=0.7,\n",
    "        system=system_message,\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\": get_user_prompt_for_docstring_generation(python_code)\n",
    "                },\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    return message.content[0].text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b94fbd55",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "def trade_gemini(currency, wallet_balance):\n",
    "    response = client.models.generate_content(\n",
    "    model=GEMINI_MODEL,\n",
    "    config=types.GenerateContentConfig(\n",
    "        system_instruction=system_message),\n",
    "    contents=get_user_prompt_for_trade_code_generation(\n",
    "        currency,\n",
    "        wallet_balance\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return response.text\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f83ef7b8",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "def code_comment_gemini(python_code):\n",
    "    response = client.models.generate_content(\n",
    "    model=GEMINI_MODEL,\n",
    "    config=types.GenerateContentConfig(\n",
    "        system_instruction=system_message),\n",
    "    contents=get_user_prompt_for_docstring_generation(python_code)\n",
    "    )\n",
    "\n",
    "    return response.text\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6737962d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def trade_llama(currency, wallet_balance):\n",
    "    completion = ollama.chat(\n",
    "    model=LLAMA_MODEL,\n",
    "    messages=[\n",
    "                {\"role\": \"user\", \"content\": get_user_prompt_for_trade_code_generation(\n",
    "                    currency,\n",
    "                    wallet_balance\n",
    "                    )\n",
    "                    },\n",
    "            ],\n",
    "    )\n",
    "\n",
    "    return completion['message']['content']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b815aa07",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def code_comment_llama(python_code):\n",
    "    completion = ollama.chat(\n",
    "    model=LLAMA_MODEL,\n",
    "    messages=[\n",
    "                {\"role\": \"user\", \"content\": get_user_prompt_for_docstring_generation(python_code)},\n",
    "            ],\n",
    "    )\n",
    "\n",
    "    return completion['message']['content']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9e07437",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def generate_python_code(input_model, currency, wallet_balance):\n",
    "    model_mapping = {\"GPT\": trade_gpt(currency, wallet_balance),\n",
    "                     \"Claude\": trade_cloude(currency, wallet_balance),\n",
    "                     \"Gemini\": trade_gemini(currency, wallet_balance),\n",
    "                     \"Llama\": trade_llama(currency, wallet_balance)}\n",
    "    try:\n",
    "        return model_mapping[input_model]\n",
    "    except KeyError as e:\n",
    "        print(f\"{e}: {input_model} is not valid\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "016fed0e",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def comment_python_code(input_model, python_code):\n",
    "    model_mapping = {\"GPT\": code_comment_gpt(python_code),\n",
    "                     \"Claude\": code_comment_cloude(python_code),\n",
    "                     \"Gemini\": code_comment_gemini(python_code),\n",
    "                     \"Llama\": code_comment_llama(python_code)\n",
    "                     }\n",
    "    try:\n",
    "        return model_mapping[input_model]\n",
    "    except KeyError as e:\n",
    "        print(f\"{e}: {input_model} is not valid\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e224a715",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "def execute_python(code):\n",
    "    code = code.replace('```python\\n','').replace('```','')\n",
    "    try:\n",
    "        output = io.StringIO()\n",
    "        sys.stdout = output\n",
    "        exec(code)\n",
    "    finally:\n",
    "        sys.stdout = sys.__stdout__\n",
    "    return output.getvalue()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea96a88d",
   "metadata": {
    "vscode": {
     "languageId": "python"
    }
   },
   "outputs": [],
   "source": [
    "with gr.Blocks() as ui:\n",
    "    gr.Markdown(\"\"\"It generate Trading Python code, which will recommend you whether sell or buy a given crypto currency at its current price.\n",
    "                    Based on the confindency level of the prediction it will recommend what amount should be placed from your available wallet balance\"\"\")\n",
    "    with gr.Row():\n",
    "        crypto_currency = gr.Dropdown([\"BTC\", \"ETH\", \"SOL\"], label=\"The Crypto cyrrency\")\n",
    "        wallet_balance = gr.Number(label=\"Enter a number\")\n",
    "        model = gr.Dropdown([\"GPT\", \"Claude\", \"Gemini\", \"Llama\"], label=\"Select model\", value=\"GPT\")\n",
    "    with gr.Row():\n",
    "        generate_python_code_bt = gr.Button(\"Genarate Python code\")\n",
    "    with gr.Row():\n",
    "        with gr.Column():\n",
    "            python = gr.TextArea(label=\"Python Code\")\n",
    "            python_comment = gr.Button(\"Comment Python code\")\n",
    "            python_run = gr.Button(\"Run Python code\")\n",
    "    with gr.Row():\n",
    "        result_out = gr.TextArea(label=\"Trading advice\")\n",
    "\n",
    "    generate_python_code_bt.click(generate_python_code, inputs=[model, crypto_currency, wallet_balance], outputs=[python])\n",
    "    python_comment.click(comment_python_code, inputs=[model, python], outputs=python)\n",
    "    python_run.click(execute_python, inputs=[python], outputs=result_out)\n",
    "\n",
    "ui.launch(inbrowser=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}