azure-ai-search-llamaindex-workflows.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Workflows in LlamaIndex\n",
    "A Workflow in LlamaIndex is an event-driven framework that allows you to chain together different computational steps to achieve complex tasks, such as building a retrieval-augmented generation (RAG) system. The workflow is composed of steps, where each step handles specific types of events and can emit new events.\n",
    "\n",
    "Workflows are designed to be asynchronous and event-driven, meaning that each step only runs when the appropriate event is triggered. This allows for the creation of complex, multi-step processes that can be easily managed and monitored."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Required libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install azure-identity\n",
    "!pip install azure-search-documents==11.4.0\n",
    "!pip install -U llama-index\n",
    "!pip install llama-index-embeddings-azure-openai\n",
    "!pip install llama-index-llms-azure-openai\n",
    "!pip install llama-index-vector-stores-azureaisearch\n",
    "!pip install nest-asyncio\n",
    "!pip install python-dotenv\n",
    "!pip install \"llama-index-core>=0.10.43\" \"openinference-instrumentation-llama-index>=2.2.2\" \"opentelemetry-proto>=1.12.0\" opentelemetry-exporter-otlp opentelemetry-sdk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from llama_index.core import (\n",
    "    StorageContext,\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    ")\n",
    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
    "from llama_index.llms.azure_openai import AzureOpenAI\n",
    "from llama_index.vector_stores.azureaisearch import (\n",
    "    AzureAISearchVectorStore,\n",
    "    IndexManagement,\n",
    ")\n",
    "from llama_index.core.response_synthesizers import CompactAndRefine\n",
    "from llama_index.core.postprocessor.llm_rerank import LLMRerank\n",
    "from llama_index.core.workflow.context import Context\n",
    "from llama_index.core.workflow.decorators import step\n",
    "from llama_index.core.workflow.events import Event, StartEvent, StopEvent\n",
    "from llama_index.core.workflow.workflow import Workflow\n",
    "from llama_index.core.schema import MetadataMode\n",
    "from llama_index.core.response.notebook_utils import display_response\n",
    "from llama_index.core.vector_stores.types import VectorStoreQueryMode\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Environment Variables\n",
    "AZURE_OPENAI_ENDPOINT = os.getenv(\"AZURE_OPENAI_ENDPOINT\")\n",
    "AZURE_OPENAI_API_KEY = os.getenv(\"AZURE_OPENAI_API_KEY\")\n",
    "AZURE_OPENAI_CHAT_COMPLETION_DEPLOYED_MODEL_NAME = os.getenv(\n",
    "    \"AZURE_OPENAI_CHAT_COMPLETION_DEPLOYED_MODEL_NAME\"\n",
    ")  # I'm using GPT-3.5-turbo\n",
    "AZURE_OPENAI_EMBEDDING_DEPLOYED_MODEL_NAME = os.getenv(\n",
    "    \"AZURE_OPENAI_EMBEDDING_DEPLOYED_MODEL_NAME\"\n",
    ")  # I'm using text-embedding-ada-002\n",
    "SEARCH_SERVICE_ENDPOINT = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")\n",
    "SEARCH_SERVICE_API_KEY = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "INDEX_NAME = \"llamindex-workflow-demo\"\n",
    "PHOENIX_API_KEY = os.getenv(\"PHOENIX_API_KEY\")\n",
    "\n",
    "# Initialize Azure OpenAI and embedding models\n",
    "llm = AzureOpenAI(\n",
    "    model=AZURE_OPENAI_CHAT_COMPLETION_DEPLOYED_MODEL_NAME,\n",
    "    deployment_name=AZURE_OPENAI_CHAT_COMPLETION_DEPLOYED_MODEL_NAME,\n",
    "    api_key=AZURE_OPENAI_API_KEY,\n",
    "    azure_endpoint=AZURE_OPENAI_ENDPOINT,\n",
    "    api_version=\"2024-02-01\",\n",
    ")\n",
    "\n",
    "embed_model = AzureOpenAIEmbedding(\n",
    "    model=AZURE_OPENAI_EMBEDDING_DEPLOYED_MODEL_NAME,\n",
    "    deployment_name=AZURE_OPENAI_EMBEDDING_DEPLOYED_MODEL_NAME,\n",
    "    api_key=AZURE_OPENAI_API_KEY,\n",
    "    azure_endpoint=AZURE_OPENAI_ENDPOINT,\n",
    "    api_version=\"2024-02-01\",\n",
    ")\n",
    "\n",
    "# Initialize search clients\n",
    "credential = AzureKeyCredential(SEARCH_SERVICE_API_KEY)\n",
    "index_client = SearchIndexClient(\n",
    "    endpoint=SEARCH_SERVICE_ENDPOINT, credential=credential\n",
    ")\n",
    "search_client = SearchClient(\n",
    "    endpoint=SEARCH_SERVICE_ENDPOINT, index_name=INDEX_NAME, credential=credential\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch Arize Phoenix for Observability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from opentelemetry.sdk import trace as trace_sdk\n",
    "from opentelemetry.sdk.trace.export import SimpleSpanProcessor\n",
    "from opentelemetry.exporter.otlp.proto.http.trace_exporter import (\n",
    "    OTLPSpanExporter as HTTPSpanExporter,\n",
    ")\n",
    "from openinference.instrumentation.llama_index import LlamaIndexInstrumentor\n",
    "\n",
    "\n",
    "# Add Phoenix API Key for tracing\n",
    "os.environ[\"OTEL_EXPORTER_OTLP_HEADERS\"] = f\"api_key={PHOENIX_API_KEY}\"\n",
    "\n",
    "# Add Phoenix\n",
    "span_phoenix_processor = SimpleSpanProcessor(\n",
    "    HTTPSpanExporter(endpoint=\"https://app.phoenix.arize.com/v1/traces\")\n",
    ")\n",
    "\n",
    "# Add them to the tracer\n",
    "tracer_provider = trace_sdk.TracerProvider()\n",
    "tracer_provider.add_span_processor(span_processor=span_phoenix_processor)\n",
    "\n",
    "# Instrument the application\n",
    "LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Key Components of a Workflow\n",
    "## Events\n",
    "Events are the fundamental objects that are passed between the steps of a workflow. Events can carry data and signal to other steps that certain actions need to be taken. There are special events like `StartEvent` and `StopEvent`, and you can also define custom events to carry specific types of data.\n",
    "\n",
    "In the example provided, two custom events are defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = AzureAISearchVectorStore(\n",
    "    search_or_index_client=index_client,\n",
    "    # filterable_metadata_field_keys=metadata_fields,\n",
    "    index_name=INDEX_NAME,\n",
    "    index_management=IndexManagement.CREATE_IF_NOT_EXISTS, # use VALIDATE to validate the index schema if using one that already exists\n",
    "    id_field_key=\"id\",\n",
    "    chunk_field_key=\"chunk\",\n",
    "    embedding_field_key=\"embedding\",\n",
    "    embedding_dimensionality=1536,  # Ensure this matches your embedding model\n",
    "    metadata_string_field_key=\"metadata\",\n",
    "    doc_id_field_key=\"doc_id\",\n",
    "    language_analyzer=\"en.lucene\",\n",
    "    vector_algorithm_type=\"exhaustiveKnn\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Existing Index\n",
    "If you have an existing index, use this code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    [],\n",
    "    storage_context=storage_context,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Key Components of a Workflow\n",
    "## Events\n",
    "Events are the fundamental objects that are passed between the steps of a workflow. Events can carry data and signal to other steps that certain actions need to be taken. There are special events like `StartEvent` and `StopEvent`, and you can also define custom events to carry specific types of data.\n",
    "\n",
    "In the example provided, two custom events are defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Event\n",
    "from llama_index.core.schema import NodeWithScore\n",
    "\n",
    "class RetrieverEvent(Event):\n",
    "    \"\"\"Result of running retrieval\"\"\"\n",
    "\n",
    "    nodes: list[NodeWithScore]\n",
    "\n",
    "class RerankEvent(Event):\n",
    "    \"\"\"Result of running reranking on retrieved nodes\"\"\"\n",
    "\n",
    "    nodes: list[NodeWithScore]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **RetrieverEvent**: This event carries the nodes retrieved from the vector store.\n",
    "- **RerankEvent**: This event carries the nodes after they have been reranked."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Steps are individual units of work within a workflow. Each step is defined as an asynchronous function and decorated with `@step()`. The decorator automatically handles the input and output types for validation and ensures that each step runs only when the appropriate event is ready.\n",
    "\n",
    "In the provided workflow, steps are defined to ingest documents, retrieve relevant information, rerank the retrieved nodes, and synthesize a final response:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RAGWorkflow(Workflow):\n",
    "    @step(pass_context=True)\n",
    "    async def ingest(self, ctx: Context, ev: StartEvent) -> StopEvent | None:\n",
    "        \"\"\"Entry point to ingest a document, triggered by a StartEvent with `dirname`.\"\"\"\n",
    "        dirname = ev.get(\"dirname\")\n",
    "        if not dirname:\n",
    "            return None\n",
    "\n",
    "        documents = SimpleDirectoryReader(dirname).load_data()\n",
    "        storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "        ctx.data[\"index\"] = VectorStoreIndex.from_documents(\n",
    "            documents=documents,\n",
    "            embed_model=embed_model,\n",
    "            storage_context=storage_context,\n",
    "        )\n",
    "        return StopEvent(result=f\"Indexed {len(documents)} documents.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Ingest Step** The ingest step handles the `StartEvent` to load and index documents from a specified directory. It uses the `SimpleDirectoryReader` to load the documents and then creates a `VectorStoreIndex`(in this case Azure AI Search) using these documents and the embedding model. This index is stored in the workflow's context for later retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@step(pass_context=True)\n",
    "async def retrieve(self, ctx: Context, ev: StartEvent) -> RetrieverEvent | None:\n",
    "    \"Entry point for RAG, triggered by a StartEvent with `query`.\"\n",
    "    query = ev.get(\"query\")\n",
    "    if not query:\n",
    "        return None\n",
    "\n",
    "    print(f\"Query the database with: {query}\")\n",
    "\n",
    "    # store the query in the global context\n",
    "    ctx.data[\"query\"] = query\n",
    "\n",
    "    # get the index from the global context\n",
    "    index = ctx.data.get(\"index\")\n",
    "    if index is None:\n",
    "        print(\"Index is empty, load some documents before querying!\")\n",
    "        return None\n",
    "\n",
    "    retriever = index.as_retriever(vector_store_query_mode=VectorStoreQueryMode.HYBRID, similarity_top_k=5)\n",
    "    nodes = retriever.retrieve(query)\n",
    "    print(f\"Retrieved {len(nodes)} nodes.\")\n",
    "    return RetrieverEvent(nodes=nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Retrieve Step**: The retrieve step is triggered by another `StartEvent` that contains a query. It retrieves relevant nodes from the indexed documents using a retriever configured with a similarity search `(similarity_top_k=2)`. The retrieved nodes are then passed along in a `RetrieverEvent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@step(pass_context=True)\n",
    "async def rerank(self, ctx: Context, ev: RetrieverEvent) -> RerankEvent:\n",
    "    # Rerank the nodes\n",
    "    ranker = LLMRerank(choice_batch_size=5, top_n=3, llm=llm)\n",
    "    print(ctx.data.get(\"query\"), flush=True)\n",
    "    new_nodes = ranker.postprocess_nodes(ev.nodes, query_str=ctx.data.get(\"query\"))\n",
    "    print(f\"Reranked nodes to {len(new_nodes)}\")\n",
    "    return RerankEvent(nodes=new_nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Rerank Step:** The rerank step takes the nodes from the `RetrieverEvent` and reranks them using a language model. The reranked nodes are then passed along in a `RerankEvent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@step(pass_context=True)\n",
    "async def synthesize(self, ctx: Context, ev: RerankEvent) -> StopEvent:\n",
    "    \"\"\"Return a streaming response using reranked nodes.\"\"\"\n",
    "    summarizer = CompactAndRefine(llm=llm, streaming=True, verbose=True)\n",
    "    query = ctx.data.get(\"query\")\n",
    "\n",
    "    response = await summarizer.asynthesize(query, nodes=ev.nodes)\n",
    "    return StopEvent(result=response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **Synthesize Step:** The synthesize step generates a final response by synthesizing the reranked nodes into a coherent answer using another language model. The result is then returned in a `StopEvent`, which stops the workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Context\n",
    "The **Context** object allows steps in the workflow to share data. In the example, the context is used to store the index created in the ingest step and the query provided in the retrieve step. This context ensures that the necessary data is available throughout the workflow's execution."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workflow Execution\n",
    "Once the workflow steps are defined, the workflow can be executed by creating an instance of the workflow and calling its run method. The run method is asynchronous and must be awaited. Each step in the workflow is executed in sequence, with data being passed between them using events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Workflow\n",
    "w = RAGWorkflow()\n",
    "\n",
    "# Ingest the documents (example with a directory 'data')\n",
    "await w.run(dirname=\"data/txt\")\n",
    "\n",
    "# Run a query\n",
    "result = await w.run(query=\"How was Llama2 trained?\")\n",
    "async for chunk in result.async_response_gen():\n",
    "    print(chunk, end=\"\", flush=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case:\n",
    "- **Ingestion:** The workflow is first run with a directory name to ingest documents into the vector index.\n",
    "- **Query Execution**: The workflow is then run again with a query to retrieve, rerank, and synthesize a response from the ingested documents.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "This workflow is a powerful abstraction that allows you to chain together complex tasks, such as retrieval-augmented generation, in a structured and manageable way. By leveraging custom events, context, and steps, you can build workflows that are both flexible and scalable, all while maintaining clear and concise code structure.\n",
    "\n",
    "Here is the full RAGWorklflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query the database with: does the president have a plan for covid-19\n",
      "Retrieved 5 nodes.\n",
      "does the president have a plan for covid-19\n",
      "Reranked nodes to 3\n",
      "=== GPT-4o-Generated Response ===\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "**`Final Response:`** Yes, the president has a comprehensive plan for COVID-19. The plan includes staying protected with vaccines and treatments, preparing for new variants, ending the shutdown of schools and businesses, and continuing to vaccinate the world. The plan emphasizes the importance of vaccines, treatments, and testing, as well as the need for readiness against new variants. It also focuses on reopening schools and businesses safely and providing global vaccine support."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "=== Details of Source Documents ===\n",
      "\n",
      "file_path: c:\\Dev\\azure-ai-search-python-playground\\data\\txt\\state_of_the_union.txt\n",
      "\n",
      "Because of the progress we’ve made, because of your resilience and the tools we have, tonight I can say  \n",
      "we are moving forward safely, back to more normal routines.  \n",
      "\n",
      "We’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July.  \n",
      "\n",
      "Just a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. \n",
      "\n",
      "Under these new guidelines, most Americans in most of the country can now be mask free.   \n",
      "\n",
      "And based on the projections, more of the country will reach that point across the next couple of weeks. \n",
      "\n",
      "Thanks to the progress we have made this past year, COVID-19 need no longer control our lives.  \n",
      "\n",
      "I know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. \n",
      "\n",
      "We will continue to combat the virus as we do other diseases. And because this is a virus that mutates and spreads, we will stay on guard. \n",
      "\n",
      "Here are four common sense steps as we move forward safely.  \n",
      "\n",
      "First, stay protected with vaccines and treatments. We know how incredibly effective vaccines are. If you’re vaccinated and boosted you have the highest degree of protection. \n",
      "\n",
      "We will never give up on vaccinating more Americans. Now, I know parents with kids under 5 are eager to see a vaccine authorized for their children. \n",
      "\n",
      "The scientists are working hard to get that done and we’ll be ready with plenty of vaccines when they do. \n",
      "\n",
      "We’re also ready with anti-viral treatments. If you get COVID-19, the Pfizer pill reduces your chances of ending up in the hospital by 90%.  \n",
      "\n",
      "We’ve ordered more of these pills than anyone in the world. And Pfizer is working overtime to get us 1 Million pills this month and more than double that next month.  \n",
      "\n",
      "And we’re launching the “Test to Treat” initiative so people can get tested at a pharmacy, and if they’re positive, receive antiviral pills on the spot at no cost.  \n",
      "\n",
      "If you’re immunocompromised or have some other vulnerability, we have treatments and free high-quality masks. \n",
      "\n",
      "We’re leaving no one behind or ignoring anyone’s needs as we move forward. \n",
      "\n",
      "And on testing, we have made hundreds of millions of tests available for you to order for free.   \n",
      "\n",
      "Even if you already ordered free tests tonight, I am announcing that you can order more from covidtests.gov starting next week. \n",
      "\n",
      "Second – we must prepare for new variants. Over the past year, we’ve gotten much better at detecting new variants. \n",
      "\n",
      "If necessary, we’ll be able to deploy new vaccines within 100 days instead of many more months or years.  \n",
      "\n",
      "And, if Congress provides the funds we need, we’ll have new stockpiles of tests, masks, and pills ready if needed. \n",
      "\n",
      "I cannot promise a new variant won’t come. But I can promise you we’ll do everything within our power to be ready if it does.  \n",
      "\n",
      "Third – we can end the shutdown of schools and businesses. We have the tools we need. \n",
      "\n",
      "It’s time for Americans to get back to work and fill our great downtowns again.  People working from home can feel safe to begin to return to the office.   \n",
      "\n",
      "We’re doing that here in the federal government. The vast majority of federal workers will once again work in person. \n",
      "\n",
      "Our schools are open. Let’s keep it that way. Our kids need to be in school. \n",
      "\n",
      "And with 75% of adult Americans fully vaccinated and hospitalizations down by 77%, most Americans can remove their masks, return to work, stay in the classroom, and move forward safely. \n",
      "\n",
      "We achieved this because we provided free vaccines, treatments, tests, and masks. \n",
      "\n",
      "Of course, continuing this costs money. \n",
      "\n",
      "I will soon send Congress a request. \n",
      "\n",
      "The vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly.   \n",
      "\n",
      "Fourth, we will continue vaccinating the world.     \n",
      "\n",
      "We’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. \n",
      "\n",
      "And we won’t stop. \n",
      "\n",
      "We have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. \n",
      "\n",
      "Let’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease.  \n",
      "\n",
      "Let’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans.  \n",
      "\n",
      "We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together.\n",
      "----------------------------------------\n",
      "\n",
      "file_path: c:\\Dev\\azure-ai-search-python-playground\\data\\txt\\state_of_the_union.txt\n",
      "\n",
      "Nobody.  \n",
      "\n",
      "The one thing all Americans agree on is that the tax system is not fair. We have to fix it.  \n",
      "\n",
      "I’m not looking to punish anyone. But let’s make sure corporations and the wealthiest Americans start paying their fair share. \n",
      "\n",
      "Just last year, 55 Fortune 500 corporations earned $40 billion in profits and paid zero dollars in federal income tax.  \n",
      "\n",
      "That’s simply not fair. That’s why I’ve proposed a 15% minimum tax rate for corporations. \n",
      "\n",
      "We got more than 130 countries to agree on a global minimum tax rate so companies can’t get out of paying their taxes at home by shipping jobs and factories overseas. \n",
      "\n",
      "That’s why I’ve proposed closing loopholes so the very wealthy don’t pay a lower tax rate than a teacher or a firefighter.  \n",
      "\n",
      "So that’s my plan. It will grow the economy and lower costs for families. \n",
      "\n",
      "So what are we waiting for? Let’s get this done. And while you’re at it, confirm my nominees to the Federal Reserve, which plays a critical role in fighting inflation.  \n",
      "\n",
      "My plan will not only lower costs to give families a fair shot, it will lower the deficit. \n",
      "\n",
      "The previous Administration not only ballooned the deficit with tax cuts for the very wealthy and corporations, it undermined the watchdogs whose job was to keep pandemic relief funds from being wasted. \n",
      "\n",
      "But in my administration, the watchdogs have been welcomed back. \n",
      "\n",
      "We’re going after the criminals who stole billions in relief money meant for small businesses and millions of Americans.  \n",
      "\n",
      "And tonight, I’m announcing that the Justice Department will name a chief prosecutor for pandemic fraud. \n",
      "\n",
      "By the end of this year, the deficit will be down to less than half what it was before I took office.  \n",
      "\n",
      "The only president ever to cut the deficit by more than one trillion dollars in a single year. \n",
      "\n",
      "Lowering your costs also means demanding more competition. \n",
      "\n",
      "I’m a capitalist, but capitalism without competition isn’t capitalism. \n",
      "\n",
      "It’s exploitation—and it drives up prices. \n",
      "\n",
      "When corporations don’t have to compete, their profits go up, your prices go up, and small businesses and family farmers and ranchers go under. \n",
      "\n",
      "We see it happening with ocean carriers moving goods in and out of America. \n",
      "\n",
      "During the pandemic, these foreign-owned companies raised prices by as much as 1,000% and made record profits. \n",
      "\n",
      "Tonight, I’m announcing a crackdown on these companies overcharging American businesses and consumers. \n",
      "\n",
      "And as Wall Street firms take over more nursing homes, quality in those homes has gone down and costs have gone up.  \n",
      "\n",
      "That ends on my watch. \n",
      "\n",
      "Medicare is going to set higher standards for nursing homes and make sure your loved ones get the care they deserve and expect. \n",
      "\n",
      "We’ll also cut costs and keep the economy going strong by giving workers a fair shot, provide more training and apprenticeships, hire them based on their skills not degrees. \n",
      "\n",
      "Let’s pass the Paycheck Fairness Act and paid leave.  \n",
      "\n",
      "Raise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. \n",
      "\n",
      "Let’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. \n",
      "\n",
      "And let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped.  \n",
      "\n",
      "When we invest in our workers, when we build the economy from the bottom up and the middle out together, we can do something we haven’t done in a long time: build a better America. \n",
      "\n",
      "For more than two years, COVID-19 has impacted every decision in our lives and the life of the nation. \n",
      "\n",
      "And I know you’re tired, frustrated, and exhausted. \n",
      "\n",
      "But I also know this. \n",
      "\n",
      "Because of the progress we’ve made, because of your resilience and the tools we have, tonight I can say  \n",
      "we are moving forward safely, back to more normal routines.  \n",
      "\n",
      "We’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July.  \n",
      "\n",
      "Just a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. \n",
      "\n",
      "Under these new guidelines, most Americans in most of the country can now be mask free.   \n",
      "\n",
      "And based on the projections, more of the country will reach that point across the next couple of weeks. \n",
      "\n",
      "Thanks to the progress we have made this past year, COVID-19 need no longer control our lives.  \n",
      "\n",
      "I know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. \n",
      "\n",
      "We will continue to combat the virus as we do other diseases.\n",
      "----------------------------------------\n",
      "\n",
      "file_path: c:\\Dev\\azure-ai-search-python-playground\\data\\txt\\state_of_the_union.txt\n",
      "\n",
      "I will soon send Congress a request. \n",
      "\n",
      "The vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly.   \n",
      "\n",
      "Fourth, we will continue vaccinating the world.     \n",
      "\n",
      "We’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. \n",
      "\n",
      "And we won’t stop. \n",
      "\n",
      "We have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. \n",
      "\n",
      "Let’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease.  \n",
      "\n",
      "Let’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans.  \n",
      "\n",
      "We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. \n",
      "\n",
      "I recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. \n",
      "\n",
      "They were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. \n",
      "\n",
      "Officer Mora was 27 years old. \n",
      "\n",
      "Officer Rivera was 22. \n",
      "\n",
      "Both Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. \n",
      "\n",
      "I spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. \n",
      "\n",
      "I’ve worked on these issues a long time. \n",
      "\n",
      "I know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. \n",
      "\n",
      "So let’s not abandon our streets. Or choose between safety and equal justice. \n",
      "\n",
      "Let’s come together to protect our communities, restore trust, and hold law enforcement accountable. \n",
      "\n",
      "That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. \n",
      "\n",
      "That’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope.  \n",
      "\n",
      "We should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. \n",
      "\n",
      "I ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe.  \n",
      "\n",
      "And I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. \n",
      "\n",
      "And I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? \n",
      "\n",
      "Ban assault weapons and high-capacity magazines. \n",
      "\n",
      "Repeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. \n",
      "\n",
      "These laws don’t infringe on the Second Amendment. They save lives. \n",
      "\n",
      "The most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. \n",
      "\n",
      "In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n",
      "\n",
      "We cannot let this happen. \n",
      "\n",
      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
      "\n",
      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
      "\n",
      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
      "\n",
      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. \n",
      "\n",
      "A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. \n",
      "\n",
      "And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. \n",
      "\n",
      "We can do both.\n",
      "----------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Define RAG Workflow with Reranking\n",
    "class RAGWorkflow(Workflow):\n",
    "    @step(pass_context=True)\n",
    "    async def ingest(self, ctx: Context, ev: StartEvent) -> StopEvent | None:\n",
    "        \"\"\"Entry point to ingest a document, triggered by a StartEvent with `dirname`.\"\"\"\n",
    "        dirname = ev.get(\"dirname\")\n",
    "        if not dirname:\n",
    "            return None\n",
    "\n",
    "        documents = SimpleDirectoryReader(dirname).load_data()\n",
    "        storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "        ctx.data[\"index\"] = VectorStoreIndex.from_documents(\n",
    "            documents=documents,\n",
    "            embed_model=embed_model,\n",
    "            storage_context=storage_context,\n",
    "        )\n",
    "        return StopEvent(result=f\"Indexed {len(documents)} documents.\")\n",
    "\n",
    "    @step(pass_context=True)\n",
    "    async def retrieve(self, ctx: Context, ev: StartEvent) -> RetrieverEvent | None:\n",
    "        \"Entry point for RAG, triggered by a StartEvent with `query`.\"\n",
    "        query = ev.get(\"query\")\n",
    "        if not query:\n",
    "            return None\n",
    "\n",
    "        print(f\"Query the database with: {query}\")\n",
    "\n",
    "        # store the query in the global context\n",
    "        ctx.data[\"query\"] = query\n",
    "\n",
    "        # get the index from the global context\n",
    "        index = ctx.data.get(\"index\")\n",
    "        if index is None:\n",
    "            print(\"Index is empty, load some documents before querying!\")\n",
    "            return None\n",
    "\n",
    "        retriever = index.as_retriever(vector_store_query_mode=VectorStoreQueryMode.HYBRID, similarity_top_k=5)\n",
    "        nodes = retriever.retrieve(query)\n",
    "        print(f\"Retrieved {len(nodes)} nodes.\")\n",
    "        return RetrieverEvent(nodes=nodes)\n",
    "\n",
    "    @step(pass_context=True)\n",
    "    async def rerank(self, ctx: Context, ev: RetrieverEvent) -> RerankEvent:\n",
    "        # Rerank the nodes\n",
    "        ranker = LLMRerank(choice_batch_size=5, top_n=3, llm=llm)\n",
    "        print(ctx.data.get(\"query\"), flush=True)\n",
    "        new_nodes = ranker.postprocess_nodes(ev.nodes, query_str=ctx.data.get(\"query\"))\n",
    "        print(f\"Reranked nodes to {len(new_nodes)}\")\n",
    "        return RerankEvent(nodes=new_nodes)\n",
    "\n",
    "    @step(pass_context=True)\n",
    "    async def synthesize(self, ctx: Context, ev: RerankEvent) -> StopEvent:\n",
    "        \"\"\"Return a streaming response using reranked nodes.\"\"\"\n",
    "        summarizer = CompactAndRefine(llm=llm, streaming=True, verbose=True)\n",
    "        query = ctx.data.get(\"query\")\n",
    "\n",
    "        response = await summarizer.asynthesize(query, nodes=ev.nodes)\n",
    "        return StopEvent(result=response)\n",
    "\n",
    "\n",
    "# Initialize Workflow\n",
    "w = RAGWorkflow()\n",
    "\n",
    "# Ingest the document (example with a specific file 'data/txt/state_of_the_union.txt')\n",
    "await w.run(dirname=\"data/txt\")\n",
    "\n",
    "# Run a query\n",
    "result = await w.run(query=\"does the president have a plan for covid-19\")\n",
    "\n",
    "\n",
    "# Function to display custom response\n",
    "def display_custom_response(response):\n",
    "    print(\"=== GPT-4o-Generated Response ===\")\n",
    "    display_response(response)\n",
    "    print(\"\\n=== Details of Source Documents ===\\n\")\n",
    "    for node in response.source_nodes:\n",
    "        print(node.get_content(metadata_mode=MetadataMode.LLM))\n",
    "        print(\"-\" * 40 + \"\\n\")\n",
    "\n",
    "\n",
    "# Await and collect the full response\n",
    "final_response = \"\"\n",
    "async for chunk in result.async_response_gen():\n",
    "    final_response += chunk\n",
    "\n",
    "\n",
    "# Create a mock response object for display\n",
    "class MockResponse:\n",
    "    def __init__(self, response, source_nodes):\n",
    "        self.response = response\n",
    "        self.source_nodes = source_nodes\n",
    "\n",
    "\n",
    "mock_response = MockResponse(final_response, result.source_nodes)\n",
    "\n",
    "# Display the response using the custom display function\n",
    "display_custom_response(mock_response)"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}