137_PDFToStrategy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2f75653",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#000;\"><img src=\"pqn.png\"></img></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdd371da",
   "metadata": {},
   "source": [
    "These libraries enable us to work with advanced financial modeling and asynchronous programming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09065da2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import asyncio\n",
    "from dotenv import load_dotenv\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.workflow import Context\n",
    "from llama_index.core.agent.workflow import (\n",
    "    FunctionAgent,\n",
    "    AgentWorkflow,\n",
    "    AgentOutput,\n",
    "    ToolCall,\n",
    "    ToolCallResult,\n",
    ")\n",
    "from llama_index.core import SimpleDirectoryReader, GPTVectorStoreIndex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fa74bfc",
   "metadata": {},
   "source": [
    "We start by loading environment variables and initializing our language model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0da3c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf26db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3da4196",
   "metadata": {},
   "source": [
    "This code sets up our environment and initializes the language model. We load environment variables, which likely include API keys. Then we create an instance of the OpenAI language model, specifically using the GPT-4 model. This prepares us for natural language processing tasks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a65ad8",
   "metadata": {},
   "source": [
    "## Define our tool functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1591f6ac",
   "metadata": {},
   "source": [
    "We define asynchronous functions that will serve as tools for our agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee9ba65",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def read_pdf_tool(ctx: Context) -> str:\n",
    "    documents = SimpleDirectoryReader(input_files=[\"pairs.pdf\"]).load_data()\n",
    "    index = GPTVectorStoreIndex(documents)\n",
    "    query_engine = index.as_query_engine()\n",
    "    query = (\n",
    "        \"Extract a detailed description of the pairs trading strategy implementation from this PDF. \"\n",
    "        \"Ensure the description is detailed enough to reproduce the strategy in code.\"\n",
    "    )\n",
    "    response = query_engine.query(query)\n",
    "    return str(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c87c969",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def build_plan_tool(ctx: Context, plan: str) -> str:\n",
    "    current_state = await ctx.get(\"state\")\n",
    "    current_state[\"implementation_plan\"] = plan\n",
    "    await ctx.set(\"state\", current_state)\n",
    "    return \"Implementation plan recorded.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eeb94f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def write_code_tool(ctx: Context, plan: str) -> str:\n",
    "    current_state = await ctx.get(\"state\")\n",
    "    current_state[\"python_code\"] = plan\n",
    "    await ctx.set(\"state\", current_state)\n",
    "    return \"Python code recorded.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5f8826d",
   "metadata": {},
   "source": [
    "These functions are tools that our agents will use. The read_pdf_tool extracts information from a PDF file about a pairs trading strategy. The build_plan_tool and write_code_tool update the workflow state with an implementation plan and Python code, respectively. These tools allow our agents to process information and generate outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2129958a",
   "metadata": {},
   "source": [
    "## Create our function agents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "620ca627",
   "metadata": {},
   "source": [
    "We define three function agents to handle different tasks in our workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d1caf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_reader_agent = FunctionAgent(\n",
    "    name=\"PDFReaderAgent\",\n",
    "    description=(\n",
    "        \"Reads a PDF file containing a pairs trading strategy and extracts a detailed description \"\n",
    "        \"of the strategy implementation.\"\n",
    "    ),\n",
    "    system_prompt=(\n",
    "        \"You are the PDFReaderAgent that can read PDFs containing implementation details of pairs trading strategies \"\n",
    "        \"and describe the strategy in detail. Once you read the PDF and describe the implementation details and are \"\n",
    "        \"satisfied, you should hand off control to the PlanBuilderAgent to develop an implementation plan. \"\n",
    "        \"You should have a detailed description of the strategy before handing off control to the PlanBuilderAgent.\"\n",
    "    ),\n",
    "    llm=llm,\n",
    "    tools=[read_pdf_tool],\n",
    "    can_handoff_to=[\"PlanBuilderAgent\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70b36149",
   "metadata": {},
   "outputs": [],
   "source": [
    "plan_builder_agent = FunctionAgent(\n",
    "    name=\"PlanBuilderAgent\",\n",
    "    description=(\n",
    "        \"Takes the detailed strategy description and builds a detailed plan to implement the strategy in Python.\"\n",
    "    ),\n",
    "    system_prompt=(\n",
    "        \"You are the PlanBuilderAgent. Your task is to analyze the strategy description from the state and generate \"\n",
    "        \"a detailed plan outlining the steps, functions, and code structure required to implement the pairs trading strategy. \"\n",
    "        \"Include suggested Python libraries for the implementation. Your plan should be in markdown format. Once the plan \"\n",
    "        \"is written, you should hand off control to the CodeWriterAgent. Output your plan with no preamble. Just output the plan.\"\n",
    "    ),\n",
    "    llm=llm,\n",
    "    tools=[build_plan_tool],\n",
    "    can_handoff_to=[\"CodeWriterAgent\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a502afb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "code_writer_agent = FunctionAgent(\n",
    "    name=\"CodeWriterAgent\",\n",
    "    description=(\n",
    "        \"Takes the implementation plan and writes complete Python code for the pairs trading strategy.\"\n",
    "    ),\n",
    "    system_prompt=(\n",
    "        \"You are the CodeWriterAgent. Your task is to use the implementation plan from the state to write complete \"\n",
    "        \"and executable Python code for the pairs trading strategy. Output your Python code with no preamble. Just \"\n",
    "        \"output the Python code.\"\n",
    "    ),\n",
    "    llm=llm,\n",
    "    tools=[write_code_tool],\n",
    "    can_handoff_to=[],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c8846bc",
   "metadata": {},
   "source": [
    "We create three function agents: PDFReaderAgent, PlanBuilderAgent, and CodeWriterAgent. Each agent has a specific role in our workflow. They use the language model and their assigned tools to perform tasks like reading PDFs, building implementation plans, and writing code. These agents work together to process information and generate outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e11a7006",
   "metadata": {},
   "source": [
    "## Set up our agent workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2965b25",
   "metadata": {},
   "source": [
    "We create an agent workflow to orchestrate the interaction between our agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c65f0e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_workflow = AgentWorkflow(\n",
    "    agents=[pdf_reader_agent, plan_builder_agent, code_writer_agent],\n",
    "    root_agent=pdf_reader_agent.name,\n",
    "    initial_state={\n",
    "        \"strategy_description\": \"Not generated yet.\",\n",
    "        \"implementation_plan\": \"Not generated yet.\",\n",
    "        \"python_code\": \"Not generated yet.\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb05ad8",
   "metadata": {},
   "source": [
    "This code sets up our agent workflow. It defines the sequence of agents that will work on our task, starting with the PDFReaderAgent. We also set an initial state for our workflow, which will be updated as the agents perform their tasks. This workflow structure allows our agents to collaborate effectively."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ac88ad0",
   "metadata": {},
   "source": [
    "## Execute our workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1614aa83",
   "metadata": {},
   "source": [
    "We run our agent workflow and process its output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b8b5a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_msg = (\n",
    "    \"Please process the PDF file 'pairs.pdf'. \"\n",
    "    \"Extract a detailed description of the pairs trading strategy, build an implementation plan in Python, \"\n",
    "    \"and finally generate the complete Python code for the strategy.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f875cdde",
   "metadata": {},
   "outputs": [],
   "source": [
    "handler = agent_workflow.run(user_msg=user_msg)\n",
    "current_agent = None\n",
    "async for event in handler.stream_events():\n",
    "    if hasattr(event, \"current_agent_name\") and event.current_agent_name != current_agent:\n",
    "        current_agent = event.current_agent_name\n",
    "        print(f\"\\n{'='*50}\")\n",
    "        print(f\"🤖 Agent: {current_agent}\")\n",
    "        print(f\"{'='*50}\\n\")\n",
    "    elif isinstance(event, AgentOutput):\n",
    "        if event.response.content:\n",
    "            print(\"📤 Output:\", event.response.content)\n",
    "        if event.tool_calls:\n",
    "            print(\"🛠️  Planning to use tools:\", [call.tool_name for call in event.tool_calls])\n",
    "    elif isinstance(event, ToolCallResult):\n",
    "        print(f\"🔧 Tool Result ({event.tool_name}):\")\n",
    "        print(\"  Arguments:\", event.tool_kwargs)\n",
    "        print(\"  Output:\", event.tool_output)\n",
    "    elif isinstance(event, ToolCall):\n",
    "        print(f\"🔨 Calling Tool: {event.tool_name}\")\n",
    "        print(\"  With arguments:\", event.tool_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6950123d",
   "metadata": {},
   "source": [
    "We execute our agent workflow by providing a user message that outlines the task. The code then processes the events streamed by the workflow. It prints information about which agent is currently active, their outputs, tool usage, and results. This allows us to see the step-by-step progress of our workflow as it processes the PDF, builds a plan, and generates code for the pairs trading strategy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0841bcfb",
   "metadata": {},
   "source": [
    "<a href=\"https://pyquantnews.com/\">PyQuant News</a> is where finance practitioners level up with Python for quant finance, algorithmic trading, and market data analysis. Looking to get started? Check out the fastest growing, top-selling course to <a href=\"https://gettingstartedwithpythonforquantfinance.com/\">get started with Python for quant finance</a>. For educational purposes. Not investment advise. Use at your own risk."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}