mirror of
https://github.com/run-llama/workflows-py.git
synced 2026-07-18 16:14:58 -04:00
Adrian Lyjak
829 lines
40 KiB
Plaintext
829 lines
40 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Iterative Extraction with Workflows + LlamaCloud\n",
|
|
"\n",
|
|
"In this example, we'll build a workflow that can\n",
|
|
"1. Parse a document (using `LlamaParse`)\n",
|
|
"2. Use an LLM to generate a JSON schema for the data we want to extract (using `OpenAI`). This includes error handling and retries for when the JSON schema is invalid!\n",
|
|
"3. Use human-in-the-loop to either accept or provide feedback on the proposed schema.\n",
|
|
"4. Send the finalized schema and parsed content to `LlamaExtract` to extract the data\n",
|
|
"\n",
|
|
"To use `LlamaCloud` and `OpenAI`, you'll need a few API keys:\n",
|
|
"1. [LlamaCloud](https://cloud.llamaindex.ai)\n",
|
|
"2. [OpenAI](https://platform.openai.com/account/api-keys)\n",
|
|
"\n",
|
|
"You'll also need to install the some required libraries:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-workflows llama-cloud-services jsonschema, openai"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup\n",
|
|
"\n",
|
|
"### Workflow Resources\n",
|
|
"\n",
|
|
"First, we'll define some resource functions for our workflow.\n",
|
|
"\n",
|
|
"This will be called once per workflow run, and will be used to return clients for various services."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_cloud_services.extract import (\n",
|
|
" ExtractConfig,\n",
|
|
" ExtractMode,\n",
|
|
" LlamaExtract,\n",
|
|
" SourceText,\n",
|
|
")\n",
|
|
"from llama_cloud_services.parse import LlamaParse\n",
|
|
"from openai import AsyncOpenAI\n",
|
|
"\n",
|
|
"\n",
|
|
"async def get_parse_client(**kwargs):\n",
|
|
" return LlamaParse(api_key=\"llx-...\", parse_mode=\"parse_page_with_agent\")\n",
|
|
"\n",
|
|
"\n",
|
|
"async def get_extract_client(**kwargs):\n",
|
|
" return LlamaExtract(api_key=\"llx-...\")\n",
|
|
"\n",
|
|
"\n",
|
|
"async def get_openai_client(**kwargs):\n",
|
|
" return AsyncOpenAI(api_key=\"sk-...\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Workflow State\n",
|
|
"\n",
|
|
"As our workflow runs, we can store global state between steps. While we could skip this step, providing a state class allows us to have more type safety, and also control over how the state is serialized and deserialized (using optional pydantic serializers and validators)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from pydantic import BaseModel, Field\n",
|
|
"\n",
|
|
"\n",
|
|
"class WorkflowState(BaseModel):\n",
|
|
" file_path: str | None = Field(default=None)\n",
|
|
" file_content: str | None = Field(default=None)\n",
|
|
" current_schema: dict | None = Field(default=None)\n",
|
|
" current_feedback: str | None = Field(default=None)\n",
|
|
" original_prompt: str | None = Field(default=None)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Schema Extraction Prompt\n",
|
|
"\n",
|
|
"Our prompt for generating a JSON schema will be doing double duty:\n",
|
|
"1. Handles the initial generation of the schema\n",
|
|
"2. Placeholders for handling the feedback from past generations/human feedback.\n",
|
|
"\n",
|
|
"To make the output easy to parse, we'll instruct the LLM to use the `<schema>` and `</schema>` tags to wrap the schema.\n",
|
|
"\n",
|
|
"`LlamaExtract` expects a JSON schema that has a root node with \"type\": \"object\" and fields inside \"properties\", so we'll instruct the LLM to output a schema that matches this format."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"extract_prompt = \"\"\"\\\n",
|
|
"<file>\n",
|
|
"{file_content}\n",
|
|
"</file>\n",
|
|
"\n",
|
|
"<user_prompt>\n",
|
|
"{prompt}\n",
|
|
"</user_prompt>\n",
|
|
"{past_attempt}\n",
|
|
"\n",
|
|
"Given the file content above, and the user prompt, output a JSON schema that will be used to extract the data from the file.\n",
|
|
"\n",
|
|
"Your JSON schema should have a root node with \"type\": \"object\" and fields inside \"properties\".\n",
|
|
"\n",
|
|
"Wrap your schema in <schema>...</schema> tags.\n",
|
|
"\"\"\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Workflow Events\n",
|
|
"\n",
|
|
"We'll define a few events for our workflow.\n",
|
|
"\n",
|
|
"- `InputEvent`: Start the workflow by providing a file path and a prompt. Subclass of `StartEvent`, so that the workflow knows to start with this event.\n",
|
|
"- `ParsedContent`: Parses the content and carry forward the content and prompt. Will also be used to re-trigger the schema generation if the user provides feedback on the current proposed schema.\n",
|
|
"- `RunExtraction`: Runs the extraction on the provided file content and schema.\n",
|
|
"- `ProposedSchema`: Proposes a schema for the extraction. Needs human approval. Subclass of `InputRequiredEvent`, so that the workflow knows to wait for the human approval before continuing.\n",
|
|
"- `ApprovedSchema`: Handles the human approval of the proposed schema. Subclass of `HumanResponseEvent`, so that the workflow knows to wait for the human approval before continuing.\n",
|
|
"- `ExtractedData`: Outputs the extracted data and the agent ID from the workflow run. Subclass of `StopEvent`, so that the workflow knows to stop when this event is received.\n",
|
|
"\n",
|
|
"### Workflow Design\n",
|
|
"\n",
|
|
"Using all the pieces above, we can now define our workflow.\n",
|
|
"\n",
|
|
"This workflow has a few key features that make it unique:\n",
|
|
"1. It uses a human-in-the-loop to iteratively improve the schema for the extraction.\n",
|
|
"2. The generated schema is validated using `jsonschema` before being sent to `LlamaExtract`.\n",
|
|
"3. It loops between the schema generation and human approval until the schema is approved."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from workflows.events import (\n",
|
|
" Event,\n",
|
|
" HumanResponseEvent,\n",
|
|
" InputRequiredEvent,\n",
|
|
" StartEvent,\n",
|
|
" StopEvent,\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"class InputEvent(StartEvent):\n",
|
|
" \"\"\"Start the workflow by providing a file path and a prompt.\"\"\"\n",
|
|
"\n",
|
|
" file_path: str\n",
|
|
" prompt: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class ParsedContent(Event):\n",
|
|
" \"\"\"Parses the content and carry forward the content and prompt.\"\"\"\n",
|
|
"\n",
|
|
" file_content: str\n",
|
|
" prompt: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class RunExtraction(Event):\n",
|
|
" \"\"\"Runs the extraction on the provided file content and schema.\"\"\"\n",
|
|
"\n",
|
|
" generated_schema: dict\n",
|
|
" file_content: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class ProposedSchema(InputRequiredEvent):\n",
|
|
" \"\"\"Proposes a schema for the extraction. Needs human approval.\"\"\"\n",
|
|
"\n",
|
|
" generated_schema: dict\n",
|
|
"\n",
|
|
"\n",
|
|
"class ApprovedSchema(HumanResponseEvent):\n",
|
|
" \"\"\"Handles the human approval of the proposed schema.\"\"\"\n",
|
|
"\n",
|
|
" approved: bool\n",
|
|
" feedback: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class ExtractedData(StopEvent):\n",
|
|
" \"\"\"Outputs the extracted data and the agent ID from the workflow run.\"\"\"\n",
|
|
"\n",
|
|
" data: dict\n",
|
|
" agent_id: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class ProgressEvent(Event):\n",
|
|
" \"\"\"Propagates a progress message to the user as the workflow runs.\"\"\"\n",
|
|
"\n",
|
|
" msg: str"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Workflow Design\n",
|
|
"\n",
|
|
"Using all the pieces above, we can now define our workflow.\n",
|
|
"\n",
|
|
"This workflow has a few key features that make it unique:\n",
|
|
"1. It uses a human-in-the-loop to iteratively improve the schema for the extraction.\n",
|
|
"2. The generated schema is validated using `jsonschema` before being sent to `LlamaExtract`.\n",
|
|
"3. It loops between the schema generation and human approval until the schema is approved."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import json\n",
|
|
"import re\n",
|
|
"import uuid\n",
|
|
"from typing import Annotated\n",
|
|
"\n",
|
|
"from jsonschema import Draft202012Validator\n",
|
|
"from workflows import Context, Workflow, step\n",
|
|
"from workflows.resource import Resource\n",
|
|
"\n",
|
|
"\n",
|
|
"class IterativeExtractionWorkflow(Workflow):\n",
|
|
" @step\n",
|
|
" async def parse_file(\n",
|
|
" self,\n",
|
|
" ev: InputEvent,\n",
|
|
" ctx: Context[WorkflowState],\n",
|
|
" parser: Annotated[LlamaParse, Resource(get_parse_client)],\n",
|
|
" ) -> ParsedContent:\n",
|
|
" ctx.write_event_to_stream(ProgressEvent(msg=f\"Parsing file: {ev.file_path}\"))\n",
|
|
" result = await parser.aparse(ev.file_path)\n",
|
|
" ctx.write_event_to_stream(ProgressEvent(msg=\"File parsed successfully\"))\n",
|
|
"\n",
|
|
" # Update the state with the file content and path\n",
|
|
" async with ctx.store.edit_state() as state:\n",
|
|
" state.original_prompt = ev.prompt\n",
|
|
" state.file_path = ev.file_path\n",
|
|
" state.file_content = \"\\n\\n\".join([page.md for page in result.pages])\n",
|
|
"\n",
|
|
" return ParsedContent(file_content=state.file_content, prompt=ev.prompt)\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def propose_schema(\n",
|
|
" self,\n",
|
|
" ev: ParsedContent,\n",
|
|
" ctx: Context[WorkflowState],\n",
|
|
" client: Annotated[AsyncOpenAI, Resource(get_openai_client)],\n",
|
|
" ) -> ProposedSchema:\n",
|
|
" ctx.write_event_to_stream(ProgressEvent(msg=\"Proposing schema\"))\n",
|
|
"\n",
|
|
" # Inject feedback from previous attempts if available\n",
|
|
" state = await ctx.store.get_state()\n",
|
|
" if state.current_feedback and state.current_schema:\n",
|
|
" past_attempt_str = f\"\\n<past_attempt>\\n<feedback>{state.current_feedback}</feedback>\\n<schema>{str(state.current_schema)}</schema>\\n</past_attempt>\\n\"\n",
|
|
" else:\n",
|
|
" past_attempt_str = \"\"\n",
|
|
"\n",
|
|
" # Start the extraction process with a fresh chat history\n",
|
|
" prompt = extract_prompt.format(\n",
|
|
" file_content=ev.file_content,\n",
|
|
" prompt=ev.prompt,\n",
|
|
" past_attempt=past_attempt_str,\n",
|
|
" )\n",
|
|
"\n",
|
|
" history = [{\"role\": \"user\", \"content\": prompt}]\n",
|
|
"\n",
|
|
" # Generate a new schema using OpenAI\n",
|
|
" response = await client.responses.create(\n",
|
|
" input=history,\n",
|
|
" model=\"gpt-4.1\",\n",
|
|
" temperature=1.0,\n",
|
|
" store=False,\n",
|
|
" )\n",
|
|
" history.append({\"role\": \"assistant\", \"content\": response.output_text})\n",
|
|
"\n",
|
|
" # Try to parse the schema from the response. If it fails, try again, using chat history\n",
|
|
" # to keep track of failed attempts.\n",
|
|
" attempts = 1\n",
|
|
" schema = {}\n",
|
|
" while attempts <= 3 and not schema:\n",
|
|
" try:\n",
|
|
" ctx.write_event_to_stream(\n",
|
|
" ProgressEvent(\n",
|
|
" msg=f\"Attempting to parse schema string from:\\n{response.output_text}\"\n",
|
|
" )\n",
|
|
" )\n",
|
|
" json_str = re.sub(\n",
|
|
" r\"<schema>([\\s\\S]*)<\\/schema>\", r\"\\1\", response.output_text\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Validate the schema\n",
|
|
" schema = json.loads(json_str)\n",
|
|
" Draft202012Validator.check_schema(schema)\n",
|
|
"\n",
|
|
" async with ctx.store.edit_state() as state:\n",
|
|
" state.current_schema = schema\n",
|
|
"\n",
|
|
" break\n",
|
|
" except Exception as e:\n",
|
|
" ctx.write_event_to_stream(\n",
|
|
" ProgressEvent(msg=f\"Schema parsing failed:\\n{e}\\n\\nTrying again...\")\n",
|
|
" )\n",
|
|
" history.append(\n",
|
|
" {\"role\": \"user\", \"content\": f\"Error: {e}\\n\\nPlease try again.\"}\n",
|
|
" )\n",
|
|
" response = await client.responses.create(\n",
|
|
" input=history,\n",
|
|
" model=\"gpt-4.1\",\n",
|
|
" temperature=1.0,\n",
|
|
" store=False,\n",
|
|
" )\n",
|
|
" history.append({\"role\": \"assistant\", \"content\": response.output_text})\n",
|
|
" attempts += 1\n",
|
|
"\n",
|
|
" if attempts > 3:\n",
|
|
" raise Exception(\"Failed to propose a valid schema after 3 attempts!\")\n",
|
|
"\n",
|
|
" ctx.write_event_to_stream(ProgressEvent(msg=\"Schema proposed successfully\"))\n",
|
|
" return ProposedSchema(generated_schema=schema)\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def handle_schema_approval(\n",
|
|
" self,\n",
|
|
" ev: ApprovedSchema,\n",
|
|
" ctx: Context[WorkflowState],\n",
|
|
" ) -> ParsedContent | RunExtraction:\n",
|
|
" async with ctx.store.edit_state() as state:\n",
|
|
" state.current_feedback = ev.feedback\n",
|
|
"\n",
|
|
" # If the schema is approved, run the extraction. Otherwise, go back to the start and try again.\n",
|
|
" if ev.approved:\n",
|
|
" return RunExtraction(\n",
|
|
" generated_schema=state.current_schema, file_content=state.file_content\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" return ParsedContent(\n",
|
|
" file_content=state.file_content, prompt=state.original_prompt\n",
|
|
" )\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def run_extraction(\n",
|
|
" self,\n",
|
|
" ev: RunExtraction,\n",
|
|
" ctx: Context[WorkflowState],\n",
|
|
" extract: Annotated[LlamaExtract, Resource(get_extract_client)],\n",
|
|
" ) -> ExtractedData:\n",
|
|
" ctx.write_event_to_stream(ProgressEvent(msg=\"Running extraction\"))\n",
|
|
"\n",
|
|
" # Persist an extraction agent + schema to llama-cloud\n",
|
|
" agent = extract.create_agent(\n",
|
|
" name=f\"extraction_workflow_{uuid.uuid4()}\",\n",
|
|
" data_schema=ev.generated_schema,\n",
|
|
" config=ExtractConfig(\n",
|
|
" extraction_mode=ExtractMode.BALANCED,\n",
|
|
" ),\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Run the extraction\n",
|
|
" file_path = await ctx.store.get(\"file_path\")\n",
|
|
" result = await agent.aextract(\n",
|
|
" files=SourceText(text_content=ev.file_content, filename=file_path),\n",
|
|
" )\n",
|
|
"\n",
|
|
" return ExtractedData(data=result.data, agent_id=agent.id)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Running the Workflow\n",
|
|
"\n",
|
|
"First, let's download some sample data to work with."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!wget https://arxiv.org/pdf/2506.05176 -O qwen3_embed_paper.pdf"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"When running the workflow, since we have human-in-the-loop, we have two options for how to run the workflow:\n",
|
|
"1. Assume the workflow will run to completion while waiting for human input. This setup is useful for \"online\" applications like websockets, running in a CLI, or similar environments where you expect a response from the user quickly.\n",
|
|
"2. Assuming the workflow will need to pause and restart once human input is recieved. In this case, we need to serialize the workflow context and restart the workflow from the point of the pause when the human input is recieved. This setup is useful for more asynchronous applications, like a REST API, where you expect the user to take some time to respond."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Option 1: Running the Workflow to Completion\n",
|
|
"\n",
|
|
"In this option, we'll run the workflow to completion while waiting for human input. This setup is useful for \"online\" applications like websockets, running in a CLI, or similar environments where you expect a response from the user quickly."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Parsing file: ./qwen3_embed_paper.pdf\n",
|
|
"Started parsing the file under job_id 9315e44f-6f5e-439f-a088-0e8aeacc1d56\n",
|
|
"File parsed successfully\n",
|
|
"Proposing schema\n",
|
|
"Attempting to parse schema string from:\n",
|
|
"<schema>\n",
|
|
"{\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"title\": {\n",
|
|
" \"type\": \"string\",\n",
|
|
" \"description\": \"The title of the paper.\"\n",
|
|
" },\n",
|
|
" \"authors\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"description\": \"A list of the authors of the paper.\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"key_takeaways\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"description\": \"A list of the main findings or key insights from the paper.\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" }\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"required\": [\"title\", \"authors\", \"key_takeaways\"]\n",
|
|
"}\n",
|
|
"</schema>\n",
|
|
"Schema proposed successfully\n",
|
|
"Proposed schema: {'type': 'object', 'properties': {'title': {'type': 'string', 'description': 'The title of the paper.'}, 'authors': {'type': 'array', 'description': 'A list of the authors of the paper.', 'items': {'type': 'string'}}, 'key_takeaways': {'type': 'array', 'description': 'A list of the main findings or key insights from the paper.', 'items': {'type': 'string'}}}, 'required': ['title', 'authors', 'key_takeaways']}\n",
|
|
"Approved? can you add a section about the datasets used in the paper?\n",
|
|
"Proposing schema\n",
|
|
"Attempting to parse schema string from:\n",
|
|
"<schema>\n",
|
|
"{\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"title\": {\n",
|
|
" \"type\": \"string\",\n",
|
|
" \"description\": \"The title of the paper.\"\n",
|
|
" },\n",
|
|
" \"authors\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"description\": \"A list of the authors of the paper.\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"key_takeaways\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"description\": \"A list of the main findings or key insights from the paper.\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"datasets\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"description\": \"A list of datasets used in the paper, including synthetic and publicly available datasets.\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"name\": { \"type\": \"string\", \"description\": \"The name of the dataset.\" },\n",
|
|
" \"type\": { \"type\": \"string\", \"description\": \"The type/category of the dataset, e.g. 'synthetic', 'benchmark', 'retrieval', etc.\" },\n",
|
|
" \"description\": { \"type\": \"string\", \"description\": \"A brief description of the dataset and its role in the study.\" },\n",
|
|
" \"size\": { \"type\": \"string\", \"description\": \"Approximate size/count, if available.\" }\n",
|
|
" },\n",
|
|
" \"required\": [\"name\"]\n",
|
|
" }\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"required\": [\"title\", \"authors\", \"key_takeaways\", \"datasets\"]\n",
|
|
"}\n",
|
|
"</schema>\n",
|
|
"\n",
|
|
"Schema proposed successfully\n",
|
|
"Proposed schema: {'type': 'object', 'properties': {'title': {'type': 'string', 'description': 'The title of the paper.'}, 'authors': {'type': 'array', 'description': 'A list of the authors of the paper.', 'items': {'type': 'string'}}, 'key_takeaways': {'type': 'array', 'description': 'A list of the main findings or key insights from the paper.', 'items': {'type': 'string'}}, 'datasets': {'type': 'array', 'description': 'A list of datasets used in the paper, including synthetic and publicly available datasets.', 'items': {'type': 'object', 'properties': {'name': {'type': 'string', 'description': 'The name of the dataset.'}, 'type': {'type': 'string', 'description': \"The type/category of the dataset, e.g. 'synthetic', 'benchmark', 'retrieval', etc.\"}, 'description': {'type': 'string', 'description': 'A brief description of the dataset and its role in the study.'}, 'size': {'type': 'string', 'description': 'Approximate size/count, if available.'}}, 'required': ['name']}}}, 'required': ['title', 'authors', 'key_takeaways', 'datasets']}\n",
|
|
"Approved? y\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading files: 0%| | 0/1 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Running extraction\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading files: 100%|██████████| 1/1 [00:01<00:00, 1.04s/it]\n",
|
|
"Creating extraction jobs: 100%|██████████| 1/1 [00:01<00:00, 1.12s/it]\n",
|
|
"Extracting files: 100%|██████████| 1/1 [00:15<00:00, 15.05s/it]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Agent ID: 374730cc-974e-42f6-a64b-b5f163ad8818\n",
|
|
"Extracted data: {'title': 'Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models', 'authors': ['Yanzhao Zhang', 'Mingxin Li', 'Dingkun Long', 'Xin Zhang', 'Huan Lin', 'Baosong Yang', 'Pengjun Xie', 'An Yang', 'Dayiheng Liu', 'Junyang Lin', 'Fei Huang', 'Jingren Zhou'], 'key_takeaways': ['Qwen3 Embedding series introduces advanced text embedding and reranking models based on Qwen3 foundation models, supporting multiple model sizes (0.6B, 4B, 8B) for diverse deployment needs.', 'The models leverage a multi-stage training pipeline: large-scale weakly supervised pre-training on synthetic data, supervised fine-tuning on high-quality datasets, and model merging for robustness and generalization.', 'Qwen3 LLMs are used not only as backbones but also to synthesize high-quality, diverse, multilingual training data, enhancing the training process.', 'The Qwen3 Embedding models achieve state-of-the-art results on multiple benchmarks, including MTEB (Multilingual, English, Code), CMTEB, MMTEB, and various retrieval tasks, outperforming or matching leading proprietary and open-source models.', 'Ablation studies show that both large-scale synthetic data pre-training and model merging are critical for achieving superior performance.', 'The models and code are open-sourced under Apache 2.0 to promote reproducibility and community research.'], 'datasets': [{'name': 'Synthetic Data (for weakly supervised pre-training)', 'type': 'synthetic', 'description': 'Large-scale synthetic text pair data generated using Qwen3-32B for tasks such as retrieval, bitext mining, classification, and semantic textual similarity. Data is generated with diverse prompts to ensure variety in task, language, length, and difficulty.', 'size': '~150M pairs'}, {'name': 'High-quality Synthetic Data (for supervised fine-tuning)', 'type': 'synthetic', 'description': 'Filtered subset of synthetic data with cosine similarity > 0.7, used for supervised fine-tuning to enhance model performance and generalization.', 'size': '~12M pairs'}, {'name': 'Labeled Data (for supervised fine-tuning)', 'type': 'benchmark', 'description': 'A collection of high-quality labeled datasets used in supervised fine-tuning, including MS MARCO, NQ, HotpotQA, NLI, Dureader, T²-Ranking, SimCLUE, MIRACL, MLDR, Mr.TyDi, Multi-CPR, CodeSearchNet, etc.', 'size': '~7M pairs'}, {'name': 'MTEB (Massive Text Embedding Benchmark)', 'type': 'benchmark', 'description': 'A large-scale benchmark for evaluating text embedding models, including multilingual, English, and code tasks. Used for comprehensive evaluation of model performance.', 'size': 'MTEB Multilingual: 131 tasks; MTEB English v2: 41 tasks; MTEB Code: 12 tasks'}, {'name': 'MMTEB (Massive Multilingual Text Embedding Benchmark)', 'type': 'benchmark', 'description': 'An expansion of MTEB, covering over 500 quality-controlled evaluation tasks across multiple languages and domains.', 'size': '216 tasks used in evaluation (from MTEB Multilingual, English, CMTEB, and Code)'}, {'name': 'CMTEB (Chinese Massive Text Embedding Benchmark)', 'type': 'benchmark', 'description': 'A benchmark for evaluating Chinese text embedding models, used in the evaluation of Qwen3 models.', 'size': '32 tasks'}, {'name': 'MLDR (Multilingual Long Document Retrieval)', 'type': 'benchmark', 'description': 'A benchmark for multilingual long document retrieval tasks, used in reranking evaluation.', 'size': None}, {'name': 'FollowIR', 'type': 'benchmark', 'description': 'A benchmark for complex instruction retrieval tasks, used in reranking evaluation.', 'size': None}]}\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"wf = IterativeExtractionWorkflow(timeout=None)\n",
|
|
"\n",
|
|
"handler = wf.run(\n",
|
|
" file_path=\"./qwen3_embed_paper.pdf\",\n",
|
|
" prompt=\"Extract the title, authors, and key takeaways from the paper.\",\n",
|
|
")\n",
|
|
"async for ev in handler.stream_events():\n",
|
|
" if isinstance(ev, ProgressEvent):\n",
|
|
" print(ev.msg, flush=True)\n",
|
|
" elif isinstance(ev, ProposedSchema):\n",
|
|
" print(f\"Proposed schema: {ev.generated_schema}\", flush=True)\n",
|
|
" approved = input(\"Approve? (y/<reason>): \").strip().lower()\n",
|
|
" print(f\"Approved? {approved}\", flush=True)\n",
|
|
" if approved == \"y\":\n",
|
|
" handler.ctx.send_event(ApprovedSchema(approved=True, feedback=\"Approved\"))\n",
|
|
" else:\n",
|
|
" handler.ctx.send_event(ApprovedSchema(approved=False, feedback=approved))\n",
|
|
"\n",
|
|
"result = await handler\n",
|
|
"print(f\"Agent ID: {result.agent_id}\", flush=True)\n",
|
|
"print(f\"Extracted data: {result.data}\", flush=True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Option 2: Serializing the Workflow Context\n",
|
|
"\n",
|
|
"In this option, we'll serialize the workflow context and restart the workflow from the point of the pause when the human input is recieved. This setup is useful for more asynchronous applications, like a REST API, where you expect the user to take some time to respond.\n",
|
|
"\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Faking some DB for storing the workflow context between runs\n",
|
|
"in_memory_store = {}\n",
|
|
"\n",
|
|
"# Define the workflow instance\n",
|
|
"wf = IterativeExtractionWorkflow(timeout=None)\n",
|
|
"\n",
|
|
"\n",
|
|
"# Define a function to run the workflow\n",
|
|
"async def run_workflow(\n",
|
|
" file_path: str | None = None,\n",
|
|
" prompt: str | None = None,\n",
|
|
" response_ev: ApprovedSchema | None = None,\n",
|
|
") -> ProposedSchema | ExtractedData:\n",
|
|
" # Check if there's existing context from a previous run\n",
|
|
" existing_context = in_memory_store.get(\"workflow_context\")\n",
|
|
" if existing_context:\n",
|
|
" ctx = Context.from_dict(wf, existing_context)\n",
|
|
" handler = wf.run(ctx=ctx)\n",
|
|
" handler.ctx.send_event(response_ev)\n",
|
|
" else:\n",
|
|
" handler = wf.run(file_path=file_path, prompt=prompt)\n",
|
|
"\n",
|
|
" # Stream events until we get to the end of the workflow or hit a input required event\n",
|
|
" async for ev in handler.stream_events():\n",
|
|
" if isinstance(ev, ProgressEvent):\n",
|
|
" print(ev.msg, flush=True)\n",
|
|
" elif isinstance(ev, ProposedSchema):\n",
|
|
" # Here, you would serialize the workflow context and save it to some DB or storage\n",
|
|
" # And resume the workflow from the point of the pause when the human input is recieved.\n",
|
|
" in_memory_store[\"workflow_context\"] = handler.ctx.to_dict()\n",
|
|
" await handler.cancel_run()\n",
|
|
" return ev\n",
|
|
"\n",
|
|
" return await handler"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Parsing file: ./qwen3_embed_paper.pdf\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Started parsing the file under job_id cc7807e8-bc53-43ec-a467-7b5df172ae83\n",
|
|
"File parsed successfully\n",
|
|
"Proposing schema\n",
|
|
"Attempting to parse schema string from:\n",
|
|
"<schema>\n",
|
|
"{\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"title\": {\n",
|
|
" \"type\": \"string\",\n",
|
|
" \"description\": \"The full title of the paper\"\n",
|
|
" },\n",
|
|
" \"authors\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" },\n",
|
|
" \"description\": \"A list of authors as full names\"\n",
|
|
" },\n",
|
|
" \"key_takeaways\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" },\n",
|
|
" \"description\": \"A list of the main contributions, findings, or conclusions drawn from the paper\"\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"required\": [\"title\", \"authors\", \"key_takeaways\"]\n",
|
|
"}\n",
|
|
"</schema>\n",
|
|
"\n",
|
|
"Schema proposed successfully\n",
|
|
"Proposed schema: {'type': 'object', 'properties': {'title': {'type': 'string', 'description': 'The full title of the paper'}, 'authors': {'type': 'array', 'items': {'type': 'string'}, 'description': 'A list of authors as full names'}, 'key_takeaways': {'type': 'array', 'items': {'type': 'string'}, 'description': 'A list of the main contributions, findings, or conclusions drawn from the paper'}}, 'required': ['title', 'authors', 'key_takeaways']}\n",
|
|
"Approved? can you add a section about the datasets used in the paper?\n",
|
|
"Proposing schema\n",
|
|
"Attempting to parse schema string from:\n",
|
|
"<schema>\n",
|
|
"{\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"title\": {\n",
|
|
" \"type\": \"string\",\n",
|
|
" \"description\": \"The full title of the paper\"\n",
|
|
" },\n",
|
|
" \"authors\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" },\n",
|
|
" \"description\": \"A list of authors as full names\"\n",
|
|
" },\n",
|
|
" \"key_takeaways\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"string\"\n",
|
|
" },\n",
|
|
" \"description\": \"A list of the main contributions, findings, or conclusions drawn from the paper\"\n",
|
|
" },\n",
|
|
" \"datasets\": {\n",
|
|
" \"type\": \"array\",\n",
|
|
" \"items\": {\n",
|
|
" \"type\": \"object\",\n",
|
|
" \"properties\": {\n",
|
|
" \"name\": { \"type\": \"string\" },\n",
|
|
" \"type\": { \"type\": \"string\", \"description\": \"Type or purpose of the dataset (e.g., Weakly Supervised, Synthetic, Supervised Fine Tuning)\" },\n",
|
|
" \"size\": { \"type\": \"string\", \"description\": \"Size or number of examples\" },\n",
|
|
" \"notes\": { \"type\": \"string\", \"description\": \"Additional details about the dataset, such as source, languages, or tasks covered\" }\n",
|
|
" },\n",
|
|
" \"required\": [\"name\"]\n",
|
|
" },\n",
|
|
" \"description\": \"A list of datasets used or introduced in the paper, including key attributes for each dataset.\"\n",
|
|
" }\n",
|
|
" },\n",
|
|
" \"required\": [\"title\", \"authors\", \"key_takeaways\", \"datasets\"]\n",
|
|
"}\n",
|
|
"</schema>\n",
|
|
"\n",
|
|
"Schema proposed successfully\n",
|
|
"Proposed schema: {'type': 'object', 'properties': {'title': {'type': 'string', 'description': 'The full title of the paper'}, 'authors': {'type': 'array', 'items': {'type': 'string'}, 'description': 'A list of authors as full names'}, 'key_takeaways': {'type': 'array', 'items': {'type': 'string'}, 'description': 'A list of the main contributions, findings, or conclusions drawn from the paper'}, 'datasets': {'type': 'array', 'items': {'type': 'object', 'properties': {'name': {'type': 'string'}, 'type': {'type': 'string', 'description': 'Type or purpose of the dataset (e.g., Weakly Supervised, Synthetic, Supervised Fine Tuning)'}, 'size': {'type': 'string', 'description': 'Size or number of examples'}, 'notes': {'type': 'string', 'description': 'Additional details about the dataset, such as source, languages, or tasks covered'}}, 'required': ['name']}, 'description': 'A list of datasets used or introduced in the paper, including key attributes for each dataset.'}}, 'required': ['title', 'authors', 'key_takeaways', 'datasets']}\n",
|
|
"Approved? y\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading files: 0%| | 0/1 [00:00<?, ?it/s]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Running extraction\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Uploading files: 100%|██████████| 1/1 [00:01<00:00, 1.21s/it]\n",
|
|
"Creating extraction jobs: 100%|██████████| 1/1 [00:00<00:00, 1.01it/s]\n",
|
|
"Extracting files: 100%|██████████| 1/1 [00:14<00:00, 14.94s/it]"
|
|
]
|
|
},
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Extracted data: {'title': 'Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models', 'authors': ['Yanzhao Zhang', 'Mingxin Li', 'Dingkun Long', 'Xin Zhang', 'Huan Lin', 'Baosong Yang', 'Pengjun Xie', 'An Yang', 'Dayiheng Liu', 'Junyang Lin', 'Fei Huang', 'Jingren Zhou'], 'key_takeaways': ['Qwen3 Embedding series introduces advanced text embedding and reranking models based on Qwen3 foundation models, significantly improving over the previous GTE-Qwen series.', 'The models leverage a multi-stage training pipeline: large-scale weakly supervised pre-training on synthetic data, supervised fine-tuning on high-quality datasets, and model merging for robustness and generalization.', 'Qwen3 LLMs are used both as backbone models and as generators of high-quality, diverse, multilingual synthetic training data, enhancing the training process.', 'The series includes models of various sizes (0.6B, 4B, 8B parameters) for both embedding and reranking, supporting flexible deployment for efficiency or effectiveness.', 'Empirical results show state-of-the-art performance on multiple benchmarks, including MTEB (multilingual, English, Chinese, code), MMTEB, and various retrieval tasks, outperforming or matching leading proprietary and open-source models.', 'The synthetic dataset generation is highly controllable, allowing for task, language, length, and difficulty customization, resulting in 150M weakly supervised pairs and 12M high-quality pairs for fine-tuning.', 'Model merging via spherical linear interpolation (slerp) of checkpoints further boosts robustness and generalization.', 'The models and code are open-sourced under Apache 2.0 to promote reproducibility and community research.'], 'datasets': [{'name': 'Synthetic Data (for Weakly Supervised Pre-Training)', 'type': 'Weakly Supervised / Synthetic', 'size': '~150M pairs', 'notes': 'Generated using Qwen3-32B; covers retrieval, bitext mining, classification, and semantic textual similarity tasks; multilingual and cross-lingual; data synthesis is highly controllable (task, language, length, difficulty).'}, {'name': 'High-quality Synthetic Data (for Supervised Fine Tuning)', 'type': 'Synthetic (filtered)', 'size': '~12M pairs', 'notes': 'Filtered from the 150M synthetic pairs using cosine similarity > 0.7; used in the second stage of supervised training.'}, {'name': 'Labeled Data (for Supervised Fine Tuning)', 'type': 'Supervised Fine Tuning', 'size': '~7M pairs', 'notes': 'Includes MS MARCO, NQ, HotpotQA, NLI, Dureader, T²-Ranking, SimCLUE, MIRACL, MLDR, Mr.TyDi, Multi-CPR, CodeSearchNet, etc.; covers multiple languages and tasks.'}, {'name': 'MMTEB (Massive Multilingual Text Embedding Benchmark)', 'type': 'Evaluation Benchmark', 'size': '216 tasks (131 MTEB Multilingual, 41 MTEB English v2, 32 CMTEB, 12 MTEB Code)', 'notes': 'Covers over 500 evaluation tasks in total; used for comprehensive evaluation of embedding models.'}, {'name': 'MTEB (Massive Text Embedding Benchmark)', 'type': 'Evaluation Benchmark', 'size': 'Not specified (part of MMTEB)', 'notes': 'Includes MTEB Multilingual, MTEB English v2, MTEB Code; used for evaluation.'}, {'name': 'CMTEB (Chinese Massive Text Embedding Benchmark)', 'type': 'Evaluation Benchmark', 'size': '32 tasks (as part of MMTEB)', 'notes': 'Chinese language evaluation.'}, {'name': 'MTEB-Code', 'type': 'Evaluation Benchmark', 'size': '12 code retrieval tasks', 'notes': 'Code-related retrieval evaluation.'}, {'name': 'MLDR, MIRACL, Mr.TyDi, Multi-CPR, CodeSearchNet, T²-Ranking, SimCLUE, Dureader, NQ, HotpotQA, NLI, MS MARCO', 'type': 'Supervised Fine Tuning / Evaluation', 'size': 'Included in ~7M labeled data', 'notes': 'Used for supervised fine-tuning and/or evaluation; covers various languages and tasks.'}, {'name': 'FollowIR', 'type': 'Evaluation Benchmark', 'size': 'Not specified', 'notes': 'Used for complex instruction retrieval evaluation.'}]}\n",
|
|
"Agent ID: a4c1a25c-c00e-4a0d-b89d-1ac7099355db\n"
|
|
]
|
|
},
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Run the workflow in a loop, waiting for human input between schema generations.\n",
|
|
"# Loop is over once we get to the ExtractedData event.\n",
|
|
"ev = await run_workflow(\n",
|
|
" file_path=\"./qwen3_embed_paper.pdf\",\n",
|
|
" prompt=\"Extract the title, authors, and key takeaways from the paper.\",\n",
|
|
")\n",
|
|
"while not isinstance(ev, ExtractedData):\n",
|
|
" if isinstance(ev, ProposedSchema):\n",
|
|
" print(f\"Proposed schema: {ev.generated_schema}\", flush=True)\n",
|
|
" approved = input(\"Approve? (y/<reason>): \").strip()\n",
|
|
" print(f\"Approved? {approved}\", flush=True)\n",
|
|
" if approved.lower() == \"y\":\n",
|
|
" ev = await run_workflow(\n",
|
|
" response_ev=ApprovedSchema(approved=True, feedback=\"Approved\")\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" ev = await run_workflow(\n",
|
|
" response_ev=ApprovedSchema(approved=False, feedback=approved)\n",
|
|
" )\n",
|
|
" else:\n",
|
|
" break\n",
|
|
"\n",
|
|
"print(f\"Extracted data: {ev.data}\", flush=True)\n",
|
|
"print(f\"Agent ID: {ev.agent_id}\", flush=True)"
|
|
]
|
|
}
|
|
],
|
|
"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.3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|