\",\n",
+ " \"tasks\": [\n",
+ " {{\n",
+ " \"type\": \"formal\",\n",
+ " \"description\": \"Write a precise, technical version that emphasizes specifications\"\n",
+ " }},\n",
+ " {{\n",
+ " \"type\": \"conversational\",\n",
+ " \"description\": \"Write an engaging, friendly version that connects with readers\"\n",
+ " }}\n",
+ " ]\n",
+ " }}\n",
+ " \"\"\",\n",
+ " \"response_format\": {\n",
+ " \"type\": \"json_schema\",\n",
+ " \"json_schema\": OrchestratorOutputSchema.model_json_schema()\n",
+ " }\n",
+ "}\n",
+ "\n",
+ "worker_agent_config = {\n",
+ " **base_agent_config,\n",
+ " \"instructions\": \"\"\"You will be given a task guideline. Generate content based on the provided\n",
+ " task, following the style and guideline descriptions. \n",
+ "\n",
+ " Return your response in this format:\n",
+ "\n",
+ " Response: Your content here, maintaining the specified style and fully addressing requirements.\n",
+ " \"\"\",\n",
+ "}\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 104,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def orchestrator_worker_workflow(task, context):\n",
+ " # single orchestrator agent\n",
+ " orchestrator_agent = Agent(client, **orchestrator_agent_config)\n",
+ " orchestrator_session_id = orchestrator_agent.create_session(session_name=f\"orchestrator_agent_{uuid.uuid4()}\")\n",
+ "\n",
+ " orchestrator_response = orchestrator_agent.create_turn(\n",
+ " messages=[{\"role\": \"user\", \"content\": f\"Your task is to {task}. Here is some context: {context}\"}],\n",
+ " stream=False,\n",
+ " session_id=orchestrator_session_id,\n",
+ " )\n",
+ "\n",
+ " orchestrator_result = json.loads(orchestrator_response.output_message.content)\n",
+ " rich.print(f\"[bold cyan] Orchestrator Analysis: [/bold cyan]\")\n",
+ " pprint(orchestrator_result)\n",
+ "\n",
+ " workers = {}\n",
+ " # spawn multiple worker agents\n",
+ " for task in orchestrator_result[\"tasks\"]:\n",
+ " worker_agent = Agent(client, **worker_agent_config)\n",
+ " worker_session_id = worker_agent.create_session(session_name=f\"worker_agent_{uuid.uuid4()}\")\n",
+ " workers[task[\"type\"]] = worker_agent\n",
+ " \n",
+ " worker_response = worker_agent.create_turn(\n",
+ " messages=[{\"role\": \"user\", \"content\": f\"Your task is to {task['description']}.\"}],\n",
+ " stream=False,\n",
+ " session_id=worker_session_id,\n",
+ " )\n",
+ " rich.print(f\"[bold yellow] >>> Worker {task['type']} <<< [/bold yellow]\")\n",
+ " rich.print(worker_response.output_message.content)\n",
+ " \n",
+ " return orchestrator_agent, workers"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 105,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ " Orchestrator Analysis: \n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36m Orchestrator Analysis: \u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ "β 'analysis': \"To create an effective product description for the new eco-friendly water bottle, it's essential to consider the target audience of environmentally conscious millennials. This demographic values sustainability and is likely to be drawn to products that not only reduce waste but also offer long-term durability. The key features of the water bottle, including its plastic-free construction, insulated design, and lifetime warranty, should be highlighted in a way that resonates with this audience. Different approaches can serve various aspects of the task, such as emphasizing the technical specifications for a formal tone or focusing on the environmental benefits and user experience for a more conversational tone.\",\n",
+ "β 'tasks': [\n",
+ "β β {\n",
+ "β β β 'type': 'formal',\n",
+ "β β β 'description': 'Write a detailed, technical product description that outlines the specifications and features of the eco-friendly water bottle, including its plastic-free materials, insulation properties, and lifetime warranty.'\n",
+ "β β },\n",
+ "β β {\n",
+ "β β β 'type': 'conversational',\n",
+ "β β β 'description': \"Craft an engaging product description that speaks directly to environmentally conscious millennials, highlighting how the water bottle's eco-friendly design, insulated performance, and lifetime warranty align with their values and lifestyle.\"\n",
+ "β β },\n",
+ "β β {\n",
+ "β β β 'type': 'creative',\n",
+ "β β β 'description': 'Develop a compelling narrative around the eco-friendly water bottle, incorporating storytelling elements that illustrate the positive impact of choosing a plastic-free, insulated, and durable product on both personal health and the environment.'\n",
+ "β β }\n",
+ "β ]\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'analysis'\u001b[0m: \u001b[32m\"To create an effective product description for the new eco-friendly water bottle, it's essential to consider the target audience of environmentally conscious millennials. This demographic values sustainability and is likely to be drawn to products that not only reduce waste but also offer long-term durability. The key features of the water bottle, including its plastic-free construction, insulated design, and lifetime warranty, should be highlighted in a way that resonates with this audience. Different approaches can serve various aspects of the task, such as emphasizing the technical specifications for a formal tone or focusing on the environmental benefits and user experience for a more conversational tone.\"\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'tasks'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'type'\u001b[0m: \u001b[32m'formal'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'description'\u001b[0m: \u001b[32m'Write a detailed, technical product description that outlines the specifications and features of the eco-friendly water bottle, including its plastic-free materials, insulation properties, and lifetime warranty.'\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'type'\u001b[0m: \u001b[32m'conversational'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'description'\u001b[0m: \u001b[32m\"Craft an engaging product description that speaks directly to environmentally conscious millennials, highlighting how the water bottle's eco-friendly design, insulated performance, and lifetime warranty align with their values and lifestyle.\"\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'type'\u001b[0m: \u001b[32m'creative'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'description'\u001b[0m: \u001b[32m'Develop a compelling narrative around the eco-friendly water bottle, incorporating storytelling elements that illustrate the positive impact of choosing a plastic-free, insulated, and durable product on both personal health and the environment.'\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ " >>> Worker formal <<< \n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33m >>> Worker formal <<< \u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Response: \n",
+ "\n",
+ "**Introduction to the EcoPro Water Bottle**\n",
+ "\n",
+ "The EcoPro Water Bottle is a revolutionary, eco-friendly hydration solution designed for the environmentally \n",
+ "conscious consumer. This premium water bottle is crafted from high-quality, plastic-free materials that not only \n",
+ "reduce waste but also provide superior insulation and durability. With its innovative design and commitment to \n",
+ "sustainability, the EcoPro Water Bottle is the perfect accessory for outdoor enthusiasts, commuters, and anyone \n",
+ "seeking a reliable and guilt-free drinking experience.\n",
+ "\n",
+ "**Plastic-Free Materials**\n",
+ "\n",
+ "The EcoPro Water Bottle is made from a unique blend of 18/8 stainless steel and natural, non-toxic materials. The \n",
+ "bottle's body is constructed from a single piece of stainless steel, ensuring a seamless and leak-proof design. The\n",
+ "lid and cap are crafted from a plant-based, bioplastic material derived from renewable resources such as corn \n",
+ "starch and sugarcane. This eco-friendly material is not only compostable but also resistant to extreme temperatures\n",
+ "and UV light.\n",
+ "\n",
+ "**Insulation Properties**\n",
+ "\n",
+ "The EcoPro Water Bottle features advanced insulation technology that keeps drinks hot or cold for hours. The \n",
+ "bottle's double-walled design, combined with a proprietary insulation material, provides exceptional thermal \n",
+ "performance. This means that your beverage will remain at the optimal temperature, whether you're sipping hot \n",
+ "coffee on a chilly morning or enjoying a refreshing cold drink on a sweltering summer day. The insulation \n",
+ "properties of the EcoPro Water Bottle are as follows:\n",
+ "\n",
+ "* Keeps drinks hot for up to 12 hours\n",
+ "* Keeps drinks cold for up to 24 hours\n",
+ "* Resistant to condensation and sweating\n",
+ "\n",
+ "**Lifetime Warranty**\n",
+ "\n",
+ "At EcoPro, we stand behind the quality and durability of our water bottles. That's why we offer a lifetime warranty\n",
+ "on all our products. If your EcoPro Water Bottle ever leaks, cracks, or fails to perform as expected, we will \n",
+ "replace it free of charge. This warranty is a testament to our commitment to producing high-quality, sustainable \n",
+ "products that will last a lifetime.\n",
+ "\n",
+ "**Additional Features**\n",
+ "\n",
+ "The EcoPro Water Bottle boasts a range of innovative features that make it a joy to use. These include:\n",
+ "\n",
+ "* **Wide Mouth**: The bottle's wide mouth makes it easy to clean and fill with ice or your favorite beverage.\n",
+ "* **Spout Lid**: The spout lid allows for easy sipping and is designed to prevent spills and leaks.\n",
+ "* **Carry Loop**: The carry loop provides a secure and comfortable way to transport your bottle on-the-go.\n",
+ "* **Measurement Markings**: The bottle features measurement markings, making it easy to track\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "Response: \n",
+ "\n",
+ "**Introduction to the EcoPro Water Bottle**\n",
+ "\n",
+ "The EcoPro Water Bottle is a revolutionary, eco-friendly hydration solution designed for the environmentally \n",
+ "conscious consumer. This premium water bottle is crafted from high-quality, plastic-free materials that not only \n",
+ "reduce waste but also provide superior insulation and durability. With its innovative design and commitment to \n",
+ "sustainability, the EcoPro Water Bottle is the perfect accessory for outdoor enthusiasts, commuters, and anyone \n",
+ "seeking a reliable and guilt-free drinking experience.\n",
+ "\n",
+ "**Plastic-Free Materials**\n",
+ "\n",
+ "The EcoPro Water Bottle is made from a unique blend of \u001b[1;36m18\u001b[0m/\u001b[1;36m8\u001b[0m stainless steel and natural, non-toxic materials. The \n",
+ "bottle's body is constructed from a single piece of stainless steel, ensuring a seamless and leak-proof design. The\n",
+ "lid and cap are crafted from a plant-based, bioplastic material derived from renewable resources such as corn \n",
+ "starch and sugarcane. This eco-friendly material is not only compostable but also resistant to extreme temperatures\n",
+ "and UV light.\n",
+ "\n",
+ "**Insulation Properties**\n",
+ "\n",
+ "The EcoPro Water Bottle features advanced insulation technology that keeps drinks hot or cold for hours. The \n",
+ "bottle's double-walled design, combined with a proprietary insulation material, provides exceptional thermal \n",
+ "performance. This means that your beverage will remain at the optimal temperature, whether you're sipping hot \n",
+ "coffee on a chilly morning or enjoying a refreshing cold drink on a sweltering summer day. The insulation \n",
+ "properties of the EcoPro Water Bottle are as follows:\n",
+ "\n",
+ "* Keeps drinks hot for up to \u001b[1;36m12\u001b[0m hours\n",
+ "* Keeps drinks cold for up to \u001b[1;36m24\u001b[0m hours\n",
+ "* Resistant to condensation and sweating\n",
+ "\n",
+ "**Lifetime Warranty**\n",
+ "\n",
+ "At EcoPro, we stand behind the quality and durability of our water bottles. That's why we offer a lifetime warranty\n",
+ "on all our products. If your EcoPro Water Bottle ever leaks, cracks, or fails to perform as expected, we will \n",
+ "replace it free of charge. This warranty is a testament to our commitment to producing high-quality, sustainable \n",
+ "products that will last a lifetime.\n",
+ "\n",
+ "**Additional Features**\n",
+ "\n",
+ "The EcoPro Water Bottle boasts a range of innovative features that make it a joy to use. These include:\n",
+ "\n",
+ "* **Wide Mouth**: The bottle's wide mouth makes it easy to clean and fill with ice or your favorite beverage.\n",
+ "* **Spout Lid**: The spout lid allows for easy sipping and is designed to prevent spills and leaks.\n",
+ "* **Carry Loop**: The carry loop provides a secure and comfortable way to transport your bottle on-the-go.\n",
+ "* **Measurement Markings**: The bottle features measurement markings, making it easy to track\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ " >>> Worker conversational <<< \n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33m >>> Worker conversational <<< \u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Response:\n",
+ "\n",
+ "**Introducing the Ultimate Eco-Friendly Companion for the Conscious Adventurer**\n",
+ "\n",
+ "Are you tired of contributing to the staggering 8 million tons of plastic waste that enter our oceans every year? \n",
+ "Do you believe that staying hydrated on-the-go shouldn't come at the cost of the planet? Look no further! Our \n",
+ "eco-friendly water bottle is designed specifically with you, the environmentally conscious millennial, in mind.\n",
+ "\n",
+ "**Designed with the Planet in Mind**\n",
+ "\n",
+ "Our water bottle is crafted from high-quality, BPA-free materials that are not only durable but also fully \n",
+ "recyclable. The sleek and modern design is inspired by nature, with a minimalist aesthetic that reflects your \n",
+ "values of simplicity and sustainability. By choosing our water bottle, you're reducing your reliance on single-use \n",
+ "plastics and helping to minimize the staggering amount of waste that ends up in our landfills and oceans.\n",
+ "\n",
+ "**Performance that Keeps Up with Your Active Lifestyle**\n",
+ "\n",
+ "But our water bottle is more than just a pretty face. Its insulated design keeps your drinks hot or cold for hours,\n",
+ "whether you're hiking through the mountains, exploring the city, or simply need a refreshing pick-me-up at your \n",
+ "desk. The double-walled insulation ensures that your hands stay cool and dry, even when filled with scorching hot \n",
+ "coffee or icy cold water.\n",
+ "\n",
+ "**A Lifetime of Hydration, Guaranteed**\n",
+ "\n",
+ "We're so confident in the quality and durability of our water bottle that we're backing it with a lifetime \n",
+ "warranty. That's right - if your bottle ever breaks or malfunctions, we'll replace it free of charge. This means \n",
+ "you can enjoy years of hassle-free hydration, without worrying about the environmental or financial costs of \n",
+ "constantly replacing disposable water bottles.\n",
+ "\n",
+ "**Join a Community of Like-Minded Individuals**\n",
+ "\n",
+ "By choosing our eco-friendly water bottle, you're not just making a statement - you're joining a movement. You're \n",
+ "part of a community that values sustainability, simplicity, and the great outdoors. You're a conscious consumer who\n",
+ "demands more from the products you use and the companies you support. And we're proud to be a part of that journey \n",
+ "with you.\n",
+ "\n",
+ "**Upgrade to a Better Way of Hydrating**\n",
+ "\n",
+ "So why wait? Ditch the disposable water bottles and upgrade to a hydration companion that aligns with your values \n",
+ "and lifestyle. Our eco-friendly water bottle is the perfect accessory for any conscious adventurer, whether you're \n",
+ "a busy professional, an outdoor enthusiast, or simply someone who cares about the planet. Join the movement and \n",
+ "experience the freedom of hydration that's as sustainable as it is stylish.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "Response:\n",
+ "\n",
+ "**Introducing the Ultimate Eco-Friendly Companion for the Conscious Adventurer**\n",
+ "\n",
+ "Are you tired of contributing to the staggering \u001b[1;36m8\u001b[0m million tons of plastic waste that enter our oceans every year? \n",
+ "Do you believe that staying hydrated on-the-go shouldn't come at the cost of the planet? Look no further! Our \n",
+ "eco-friendly water bottle is designed specifically with you, the environmentally conscious millennial, in mind.\n",
+ "\n",
+ "**Designed with the Planet in Mind**\n",
+ "\n",
+ "Our water bottle is crafted from high-quality, BPA-free materials that are not only durable but also fully \n",
+ "recyclable. The sleek and modern design is inspired by nature, with a minimalist aesthetic that reflects your \n",
+ "values of simplicity and sustainability. By choosing our water bottle, you're reducing your reliance on single-use \n",
+ "plastics and helping to minimize the staggering amount of waste that ends up in our landfills and oceans.\n",
+ "\n",
+ "**Performance that Keeps Up with Your Active Lifestyle**\n",
+ "\n",
+ "But our water bottle is more than just a pretty face. Its insulated design keeps your drinks hot or cold for hours,\n",
+ "whether you're hiking through the mountains, exploring the city, or simply need a refreshing pick-me-up at your \n",
+ "desk. The double-walled insulation ensures that your hands stay cool and dry, even when filled with scorching hot \n",
+ "coffee or icy cold water.\n",
+ "\n",
+ "**A Lifetime of Hydration, Guaranteed**\n",
+ "\n",
+ "We're so confident in the quality and durability of our water bottle that we're backing it with a lifetime \n",
+ "warranty. That's right - if your bottle ever breaks or malfunctions, we'll replace it free of charge. This means \n",
+ "you can enjoy years of hassle-free hydration, without worrying about the environmental or financial costs of \n",
+ "constantly replacing disposable water bottles.\n",
+ "\n",
+ "**Join a Community of Like-Minded Individuals**\n",
+ "\n",
+ "By choosing our eco-friendly water bottle, you're not just making a statement - you're joining a movement. You're \n",
+ "part of a community that values sustainability, simplicity, and the great outdoors. You're a conscious consumer who\n",
+ "demands more from the products you use and the companies you support. And we're proud to be a part of that journey \n",
+ "with you.\n",
+ "\n",
+ "**Upgrade to a Better Way of Hydrating**\n",
+ "\n",
+ "So why wait? Ditch the disposable water bottles and upgrade to a hydration companion that aligns with your values \n",
+ "and lifestyle. Our eco-friendly water bottle is the perfect accessory for any conscious adventurer, whether you're \n",
+ "a busy professional, an outdoor enthusiast, or simply someone who cares about the planet. Join the movement and \n",
+ "experience the freedom of hydration that's as sustainable as it is stylish.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ " >>> Worker creative <<< \n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33m >>> Worker creative <<< \u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Response:\n",
+ "\n",
+ "In a world where single-use plastics have become an epidemic, threatening the very foundations of our ecosystems, a\n",
+ "hero emerges in the form of an eco-friendly water bottle. This isn't just any water bottle; it's a symbol of a \n",
+ "movement, a beacon of hope for a healthier planet and a healthier you. Let's dive into the story of how this \n",
+ "simple, yet powerful, product can change your life and the lives of those around you.\n",
+ "\n",
+ "Meet Emma, a young professional who, like many of us, was accustomed to grabbing a plastic water bottle on the go. \n",
+ "Every day, she'd use one, sometimes two, without giving it a second thought. But Emma began to notice the toll this\n",
+ "habit was taking. Her body wasn't retaining heat well, and she found herself constantly buying new bottles, \n",
+ "contributing to the plastic waste that was polluting her beloved local park and, ultimately, the oceans. The guilt \n",
+ "was creeping in, but the convenience was hard to give up.\n",
+ "\n",
+ "That was until Emma discovered the eco-friendly water bottle. Made from durable, BPA-free materials and designed \n",
+ "with insulation that keeps drinks hot or cold for hours, this bottle quickly became her constant companion. Not \n",
+ "only did it reduce her reliance on single-use plastics, but it also improved her hydration habits. The insulation \n",
+ "meant her drinks stayed at the perfect temperature, encouraging her to drink more throughout the day. Her energy \n",
+ "levels soared, and she noticed an improvement in her overall health.\n",
+ "\n",
+ "But the impact didn't stop there. Emma soon realized that her choice was part of a larger movement. By opting for a\n",
+ "plastic-free, insulated, and durable water bottle, she was contributing to a reduction in plastic waste. It's \n",
+ "estimated that if we don't change our ways, there will be more plastic than fish in the ocean by 2050. Emma's small\n",
+ "action, multiplied by millions of others making the same choice, could significantly alter this grim forecast.\n",
+ "\n",
+ "As word of her eco-friendly water bottle spread, Emma found herself at the forefront of a local initiative to \n",
+ "reduce plastic use in her community. Together with friends, family, and like-minded individuals, they organized \n",
+ "clean-up events, spread awareness about the dangers of single-use plastics, and encouraged others to make the \n",
+ "switch to reusable products. The community began to flourish, not just environmentally, but socially as well. \n",
+ "People from all walks of life came together, united by a common goal: to protect their home, the Earth.\n",
+ "\n",
+ "The story of Emma and her eco-friendly water bottle serves as a powerful reminder that our daily choices have the\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "Response:\n",
+ "\n",
+ "In a world where single-use plastics have become an epidemic, threatening the very foundations of our ecosystems, a\n",
+ "hero emerges in the form of an eco-friendly water bottle. This isn't just any water bottle; it's a symbol of a \n",
+ "movement, a beacon of hope for a healthier planet and a healthier you. Let's dive into the story of how this \n",
+ "simple, yet powerful, product can change your life and the lives of those around you.\n",
+ "\n",
+ "Meet Emma, a young professional who, like many of us, was accustomed to grabbing a plastic water bottle on the go. \n",
+ "Every day, she'd use one, sometimes two, without giving it a second thought. But Emma began to notice the toll this\n",
+ "habit was taking. Her body wasn't retaining heat well, and she found herself constantly buying new bottles, \n",
+ "contributing to the plastic waste that was polluting her beloved local park and, ultimately, the oceans. The guilt \n",
+ "was creeping in, but the convenience was hard to give up.\n",
+ "\n",
+ "That was until Emma discovered the eco-friendly water bottle. Made from durable, BPA-free materials and designed \n",
+ "with insulation that keeps drinks hot or cold for hours, this bottle quickly became her constant companion. Not \n",
+ "only did it reduce her reliance on single-use plastics, but it also improved her hydration habits. The insulation \n",
+ "meant her drinks stayed at the perfect temperature, encouraging her to drink more throughout the day. Her energy \n",
+ "levels soared, and she noticed an improvement in her overall health.\n",
+ "\n",
+ "But the impact didn't stop there. Emma soon realized that her choice was part of a larger movement. By opting for a\n",
+ "plastic-free, insulated, and durable water bottle, she was contributing to a reduction in plastic waste. It's \n",
+ "estimated that if we don't change our ways, there will be more plastic than fish in the ocean by \u001b[1;36m2050\u001b[0m. Emma's small\n",
+ "action, multiplied by millions of others making the same choice, could significantly alter this grim forecast.\n",
+ "\n",
+ "As word of her eco-friendly water bottle spread, Emma found herself at the forefront of a local initiative to \n",
+ "reduce plastic use in her community. Together with friends, family, and like-minded individuals, they organized \n",
+ "clean-up events, spread awareness about the dangers of single-use plastics, and encouraged others to make the \n",
+ "switch to reusable products. The community began to flourish, not just environmentally, but socially as well. \n",
+ "People from all walks of life came together, united by a common goal: to protect their home, the Earth.\n",
+ "\n",
+ "The story of Emma and her eco-friendly water bottle serves as a powerful reminder that our daily choices have the\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "orchestrator_agent, workers = orchestrator_worker_workflow(\n",
+ " task=\"Write a product description for a new eco-friendly water bottle\",\n",
+ " context={\n",
+ " \"target_audience\": \"environmentally conscious millennials\",\n",
+ " \"key_features\": [\"plastic-free\", \"insulated\", \"lifetime warranty\"]\n",
+ " }\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### 3.2. Monitor Orchestrator-Workers Workflow's Internals\n",
+ "\n",
+ "Let's see what happened with the orchestrator agent and worker agents it spawn up. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 91,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ "β 'session_id': '8e765c0f-e71d-4c0c-9986-ee729d73966e',\n",
+ "β 'session_name': 'orchestrator_agent_976ef2f2-911c-47ac-9860-1c38d9038a91',\n",
+ "β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 28, 669769),\n",
+ "β 'turns': [\n",
+ "β β {\n",
+ "β β β 'input_messages': [\n",
+ "β β β β {\n",
+ "β β β β β 'content': \"Your task is to Write a product description for a new eco-friendly water bottle. Here is some context: {'target_audience': 'environmentally conscious millennials', 'key_features': ['plastic-free', 'insulated', 'lifetime warranty']}\",\n",
+ "β β β β β 'role': 'user',\n",
+ "β β β β β 'context': None\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'output_message': {\n",
+ "β β β β 'content': '{\\n\"analysis\": \"The task of writing a product description for a new eco-friendly water bottle requires a deep understanding of the target audience, which is environmentally conscious millennials. To effectively connect with this audience, the description should highlight the key features of the product, such as being plastic-free, insulated, and having a lifetime warranty. A valuable approach would be to emphasize the eco-friendly aspects of the product, as this aligns with the values and concerns of the target audience. Additionally, emphasizing the practical benefits of the product, such as its insulation and durability, would also be effective. Lastly, using a tone that is both informative and engaging would help to capture the reader\\'s attention and convey the product\\'s value.\",\\n\"tasks\": [\\n{\\n\"type\": \"formal\",\\n\"description\": \"Write a precise, technical description that highlights the product\\'s key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product\\'s specifications.\"\\n},\\n{\\n\"type\": \"conversational\",\\n\"description\": \"Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values.\"\\n},\\n{\\n\"type\": \"creative\",\\n\"description\": \"Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader\\'s attention and leaving a lasting impression.\"\\n}\\n]\\n}',\n",
+ "β β β β 'role': 'assistant',\n",
+ "β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β 'tool_calls': []\n",
+ "β β β },\n",
+ "β β β 'session_id': '8e765c0f-e71d-4c0c-9986-ee729d73966e',\n",
+ "β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 28, 687648, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=57600))),\n",
+ "β β β 'steps': [\n",
+ "β β β β {\n",
+ "β β β β β 'model_response': {\n",
+ "β β β β β β 'content': '{\\n\"analysis\": \"The task of writing a product description for a new eco-friendly water bottle requires a deep understanding of the target audience, which is environmentally conscious millennials. To effectively connect with this audience, the description should highlight the key features of the product, such as being plastic-free, insulated, and having a lifetime warranty. A valuable approach would be to emphasize the eco-friendly aspects of the product, as this aligns with the values and concerns of the target audience. Additionally, emphasizing the practical benefits of the product, such as its insulation and durability, would also be effective. Lastly, using a tone that is both informative and engaging would help to capture the reader\\'s attention and convey the product\\'s value.\",\\n\"tasks\": [\\n{\\n\"type\": \"formal\",\\n\"description\": \"Write a precise, technical description that highlights the product\\'s key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product\\'s specifications.\"\\n},\\n{\\n\"type\": \"conversational\",\\n\"description\": \"Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values.\"\\n},\\n{\\n\"type\": \"creative\",\\n\"description\": \"Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader\\'s attention and leaving a lasting impression.\"\\n}\\n]\\n}',\n",
+ "β β β β β β 'role': 'assistant',\n",
+ "β β β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β β β 'tool_calls': []\n",
+ "β β β β β },\n",
+ "β β β β β 'step_id': 'd340d9ae-3aed-4042-aefd-9d9ce9448bee',\n",
+ "β β β β β 'step_type': 'inference',\n",
+ "β β β β β 'turn_id': '0ceb314a-82e0-4728-9b08-0dbb89ee6f25',\n",
+ "β β β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 32, 72702, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 28, 698909, tzinfo=TzInfo(-08:00))\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'turn_id': '0ceb314a-82e0-4728-9b08-0dbb89ee6f25',\n",
+ "β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 32, 86428, tzinfo=TzInfo(-08:00)),\n",
+ "β β β 'output_attachments': []\n",
+ "β β }\n",
+ "β ]\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'8e765c0f-e71d-4c0c-9986-ee729d73966e'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_name'\u001b[0m: \u001b[32m'orchestrator_agent_976ef2f2-911c-47ac-9860-1c38d9038a91'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m28\u001b[0m, \u001b[1;36m669769\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'turns'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'input_messages'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Your task is to Write a product description for a new eco-friendly water bottle. Here is some context: \u001b[0m\u001b[32m{\u001b[0m\u001b[32m'target_audience': 'environmentally conscious millennials', 'key_features': \u001b[0m\u001b[32m[\u001b[0m\u001b[32m'plastic-free', 'insulated', 'lifetime warranty'\u001b[0m\u001b[32m]\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'user'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'context'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_message'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m'\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"analysis\": \"The task of writing a product description for a new eco-friendly water bottle requires a deep understanding of the target audience, which is environmentally conscious millennials. To effectively connect with this audience, the description should highlight the key features of the product, such as being plastic-free, insulated, and having a lifetime warranty. A valuable approach would be to emphasize the eco-friendly aspects of the product, as this aligns with the values and concerns of the target audience. Additionally, emphasizing the practical benefits of the product, such as its insulation and durability, would also be effective. Lastly, using a tone that is both informative and engaging would help to capture the reader\\'s attention and convey the product\\'s value.\",\\n\"tasks\": \u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"formal\",\\n\"description\": \"Write a precise, technical description that highlights the product\\'s key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product\\'s specifications.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m,\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"conversational\",\\n\"description\": \"Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m,\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"creative\",\\n\"description\": \"Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader\\'s attention and leaving a lasting impression.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'8e765c0f-e71d-4c0c-9986-ee729d73966e'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m28\u001b[0m, \u001b[1;36m687648\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.timezone\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mdatetime.timedelta\u001b[0m\u001b[1m(\u001b[0m\u001b[33mdays\u001b[0m=\u001b[1;36m-1\u001b[0m, \u001b[33mseconds\u001b[0m=\u001b[1;36m57600\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'model_response'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m'\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"analysis\": \"The task of writing a product description for a new eco-friendly water bottle requires a deep understanding of the target audience, which is environmentally conscious millennials. To effectively connect with this audience, the description should highlight the key features of the product, such as being plastic-free, insulated, and having a lifetime warranty. A valuable approach would be to emphasize the eco-friendly aspects of the product, as this aligns with the values and concerns of the target audience. Additionally, emphasizing the practical benefits of the product, such as its insulation and durability, would also be effective. Lastly, using a tone that is both informative and engaging would help to capture the reader\\'s attention and convey the product\\'s value.\",\\n\"tasks\": \u001b[0m\u001b[32m[\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"formal\",\\n\"description\": \"Write a precise, technical description that highlights the product\\'s key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product\\'s specifications.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m,\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"conversational\",\\n\"description\": \"Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m,\\n\u001b[0m\u001b[32m{\u001b[0m\u001b[32m\\n\"type\": \"creative\",\\n\"description\": \"Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader\\'s attention and leaving a lasting impression.\"\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\\n\u001b[0m\u001b[32m}\u001b[0m\u001b[32m'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_id'\u001b[0m: \u001b[32m'd340d9ae-3aed-4042-aefd-9d9ce9448bee'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_type'\u001b[0m: \u001b[32m'inference'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'0ceb314a-82e0-4728-9b08-0dbb89ee6f25'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m32\u001b[0m, \u001b[1;36m72702\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m28\u001b[0m, \u001b[1;36m698909\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'0ceb314a-82e0-4728-9b08-0dbb89ee6f25'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m32\u001b[0m, \u001b[1;36m86428\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_attachments'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Worker formal Session:\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ "β 'session_id': '30a5e169-2aeb-4e20-99b9-f060349b6b55',\n",
+ "β 'session_name': 'worker_agent_2824b8d3-3059-4862-966d-12ce895d6c0b',\n",
+ "β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 32, 154138),\n",
+ "β 'turns': [\n",
+ "β β {\n",
+ "β β β 'input_messages': [\n",
+ "β β β β {\n",
+ "β β β β β 'content': \"Your task is to Write a precise, technical description that highlights the product's key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product's specifications..\",\n",
+ "β β β β β 'role': 'user',\n",
+ "β β β β β 'context': None\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'output_message': {\n",
+ "β β β β 'content': \"Response: \\n\\nThe product in question is a cutting-edge, eco-friendly solution designed to provide superior performance while minimizing environmental impact. Its key features include a plastic-free construction, leveraging high-quality, sustainable materials that not only reduce waste but also ensure durability and longevity. \\n\\nOne of the standout aspects of this product is its exceptional insulation capabilities. Engineered with advanced technology, it effectively retains heat in colder conditions and keeps warmth at bay in hotter environments, thereby optimizing energy efficiency and comfort. This feature is particularly beneficial for applications where temperature control is crucial, making it an ideal choice for a wide range of uses.\\n\\nFurthermore, the product comes with a comprehensive lifetime warranty, reflecting the manufacturer's confidence in its quality and performance. This warranty provides users with peace of mind, knowing that they are protected against defects and functional failures for the entire lifespan of the product. It underscores the commitment to customer satisfaction and the dedication to delivering products that meet the highest standards of excellence.\\n\\nIn terms of specifications, the product boasts a robust design that is both lightweight and easy to use, making it versatile and adaptable to various settings. Its plastic-free construction not only supports eco-friendly initiatives but also contributes to a healthier indoor air quality by eliminating the potential for plastic off-gassing.\\n\\nThe insulation properties are further enhanced by a unique design that minimizes thermal bridging, ensuring consistent and reliable performance. Whether used in residential, commercial, or industrial applications, this product is designed to deliver consistent results, combining sustainability with functional superiority.\\n\\nOverall, the product represents a significant advancement in eco-friendly technology, combining a plastic-free construction, superior insulation capabilities, and a lifetime warranty to offer a solution that is as environmentally responsible as it is effective. It is an exemplary model of innovative design and manufacturing excellence, catering to the evolving needs of consumers who prioritize both performance and sustainability.\",\n",
+ "β β β β 'role': 'assistant',\n",
+ "β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β 'tool_calls': []\n",
+ "β β β },\n",
+ "β β β 'session_id': '30a5e169-2aeb-4e20-99b9-f060349b6b55',\n",
+ "β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 32, 161464, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=57600))),\n",
+ "β β β 'steps': [\n",
+ "β β β β {\n",
+ "β β β β β 'model_response': {\n",
+ "β β β β β β 'content': \"Response: \\n\\nThe product in question is a cutting-edge, eco-friendly solution designed to provide superior performance while minimizing environmental impact. Its key features include a plastic-free construction, leveraging high-quality, sustainable materials that not only reduce waste but also ensure durability and longevity. \\n\\nOne of the standout aspects of this product is its exceptional insulation capabilities. Engineered with advanced technology, it effectively retains heat in colder conditions and keeps warmth at bay in hotter environments, thereby optimizing energy efficiency and comfort. This feature is particularly beneficial for applications where temperature control is crucial, making it an ideal choice for a wide range of uses.\\n\\nFurthermore, the product comes with a comprehensive lifetime warranty, reflecting the manufacturer's confidence in its quality and performance. This warranty provides users with peace of mind, knowing that they are protected against defects and functional failures for the entire lifespan of the product. It underscores the commitment to customer satisfaction and the dedication to delivering products that meet the highest standards of excellence.\\n\\nIn terms of specifications, the product boasts a robust design that is both lightweight and easy to use, making it versatile and adaptable to various settings. Its plastic-free construction not only supports eco-friendly initiatives but also contributes to a healthier indoor air quality by eliminating the potential for plastic off-gassing.\\n\\nThe insulation properties are further enhanced by a unique design that minimizes thermal bridging, ensuring consistent and reliable performance. Whether used in residential, commercial, or industrial applications, this product is designed to deliver consistent results, combining sustainability with functional superiority.\\n\\nOverall, the product represents a significant advancement in eco-friendly technology, combining a plastic-free construction, superior insulation capabilities, and a lifetime warranty to offer a solution that is as environmentally responsible as it is effective. It is an exemplary model of innovative design and manufacturing excellence, catering to the evolving needs of consumers who prioritize both performance and sustainability.\",\n",
+ "β β β β β β 'role': 'assistant',\n",
+ "β β β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β β β 'tool_calls': []\n",
+ "β β β β β },\n",
+ "β β β β β 'step_id': '259985a9-7571-4b03-af86-758e6b17beb8',\n",
+ "β β β β β 'step_type': 'inference',\n",
+ "β β β β β 'turn_id': '4d569b07-a68a-44b6-9e19-2841d1d1f002',\n",
+ "β β β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 37, 623431, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 32, 172831, tzinfo=TzInfo(-08:00))\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'turn_id': '4d569b07-a68a-44b6-9e19-2841d1d1f002',\n",
+ "β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 37, 636202, tzinfo=TzInfo(-08:00)),\n",
+ "β β β 'output_attachments': []\n",
+ "β β }\n",
+ "β ]\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'30a5e169-2aeb-4e20-99b9-f060349b6b55'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_name'\u001b[0m: \u001b[32m'worker_agent_2824b8d3-3059-4862-966d-12ce895d6c0b'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m32\u001b[0m, \u001b[1;36m154138\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'turns'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'input_messages'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Your task is to Write a precise, technical description that highlights the product's key features, such as its plastic-free construction, insulation capabilities, and lifetime warranty. This approach would serve the aspect of providing a clear and concise overview of the product's specifications..\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'user'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'context'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_message'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response: \\n\\nThe product in question is a cutting-edge, eco-friendly solution designed to provide superior performance while minimizing environmental impact. Its key features include a plastic-free construction, leveraging high-quality, sustainable materials that not only reduce waste but also ensure durability and longevity. \\n\\nOne of the standout aspects of this product is its exceptional insulation capabilities. Engineered with advanced technology, it effectively retains heat in colder conditions and keeps warmth at bay in hotter environments, thereby optimizing energy efficiency and comfort. This feature is particularly beneficial for applications where temperature control is crucial, making it an ideal choice for a wide range of uses.\\n\\nFurthermore, the product comes with a comprehensive lifetime warranty, reflecting the manufacturer's confidence in its quality and performance. This warranty provides users with peace of mind, knowing that they are protected against defects and functional failures for the entire lifespan of the product. It underscores the commitment to customer satisfaction and the dedication to delivering products that meet the highest standards of excellence.\\n\\nIn terms of specifications, the product boasts a robust design that is both lightweight and easy to use, making it versatile and adaptable to various settings. Its plastic-free construction not only supports eco-friendly initiatives but also contributes to a healthier indoor air quality by eliminating the potential for plastic off-gassing.\\n\\nThe insulation properties are further enhanced by a unique design that minimizes thermal bridging, ensuring consistent and reliable performance. Whether used in residential, commercial, or industrial applications, this product is designed to deliver consistent results, combining sustainability with functional superiority.\\n\\nOverall, the product represents a significant advancement in eco-friendly technology, combining a plastic-free construction, superior insulation capabilities, and a lifetime warranty to offer a solution that is as environmentally responsible as it is effective. It is an exemplary model of innovative design and manufacturing excellence, catering to the evolving needs of consumers who prioritize both performance and sustainability.\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'30a5e169-2aeb-4e20-99b9-f060349b6b55'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m32\u001b[0m, \u001b[1;36m161464\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.timezone\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mdatetime.timedelta\u001b[0m\u001b[1m(\u001b[0m\u001b[33mdays\u001b[0m=\u001b[1;36m-1\u001b[0m, \u001b[33mseconds\u001b[0m=\u001b[1;36m57600\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'model_response'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response: \\n\\nThe product in question is a cutting-edge, eco-friendly solution designed to provide superior performance while minimizing environmental impact. Its key features include a plastic-free construction, leveraging high-quality, sustainable materials that not only reduce waste but also ensure durability and longevity. \\n\\nOne of the standout aspects of this product is its exceptional insulation capabilities. Engineered with advanced technology, it effectively retains heat in colder conditions and keeps warmth at bay in hotter environments, thereby optimizing energy efficiency and comfort. This feature is particularly beneficial for applications where temperature control is crucial, making it an ideal choice for a wide range of uses.\\n\\nFurthermore, the product comes with a comprehensive lifetime warranty, reflecting the manufacturer's confidence in its quality and performance. This warranty provides users with peace of mind, knowing that they are protected against defects and functional failures for the entire lifespan of the product. It underscores the commitment to customer satisfaction and the dedication to delivering products that meet the highest standards of excellence.\\n\\nIn terms of specifications, the product boasts a robust design that is both lightweight and easy to use, making it versatile and adaptable to various settings. Its plastic-free construction not only supports eco-friendly initiatives but also contributes to a healthier indoor air quality by eliminating the potential for plastic off-gassing.\\n\\nThe insulation properties are further enhanced by a unique design that minimizes thermal bridging, ensuring consistent and reliable performance. Whether used in residential, commercial, or industrial applications, this product is designed to deliver consistent results, combining sustainability with functional superiority.\\n\\nOverall, the product represents a significant advancement in eco-friendly technology, combining a plastic-free construction, superior insulation capabilities, and a lifetime warranty to offer a solution that is as environmentally responsible as it is effective. It is an exemplary model of innovative design and manufacturing excellence, catering to the evolving needs of consumers who prioritize both performance and sustainability.\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_id'\u001b[0m: \u001b[32m'259985a9-7571-4b03-af86-758e6b17beb8'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_type'\u001b[0m: \u001b[32m'inference'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'4d569b07-a68a-44b6-9e19-2841d1d1f002'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m37\u001b[0m, \u001b[1;36m623431\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m32\u001b[0m, \u001b[1;36m172831\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'4d569b07-a68a-44b6-9e19-2841d1d1f002'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m37\u001b[0m, \u001b[1;36m636202\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_attachments'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Worker conversational Session:\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ "β 'session_id': '254cf164-52f4-4b7f-ba92-996e97725c12',\n",
+ "β 'session_name': 'worker_agent_b83fb070-705b-4e58-8146-84970328bea0',\n",
+ "β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 37, 686501),\n",
+ "β 'turns': [\n",
+ "β β {\n",
+ "β β β 'input_messages': [\n",
+ "β β β β {\n",
+ "β β β β β 'content': 'Your task is to Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values..',\n",
+ "β β β β β 'role': 'user',\n",
+ "β β β β β 'context': None\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'output_message': {\n",
+ "β β β β 'content': \"Response:\\n\\nImagine a world where every small choice you make can contribute to a bigger, more beautiful picture - a world where the air is fresh, the oceans are clean, and the future is bright. At [Brand Name], we believe that this world is not just a dream, but a reality that we can create together, one step at a time. That's why we're passionate about introducing you to our eco-friendly product, designed with love for the planet and a deep respect for the values that you hold dear.\\n\\nOur product is more than just a solution to your everyday needs; it's a statement of your commitment to the well-being of our planet. Made from sustainable materials and designed with recyclability in mind, every aspect of our product reflects our shared desire to reduce waste and live in harmony with nature. Whether you're a long-time advocate for environmental causes or just starting your journey towards a more sustainable lifestyle, our product is here to support and enhance your efforts.\\n\\nWhat sets us apart is not just our product's eco-friendly features, but the community of like-minded individuals who believe, as we do, that small actions today can lead to a significant positive impact tomorrow. By choosing our product, you're not only making a responsible choice for the planet, but you're also becoming part of a movement - a movement that values the beauty of nature, the importance of community, and the power of collective action.\\n\\nAt [Brand Name], we're dedicated to more than just selling a product; we're committed to fostering a relationship with you, our customer, and with the Earth. We believe in transparency, in honesty, and in the open sharing of our processes and materials. We want you to feel confident and proud of the choices you make, knowing that you're supporting a brand that genuinely cares about the same things you do.\\n\\nSo, join us on this journey towards a greener, brighter future. Together, let's embrace the power of sustainable living, celebrate the beauty of our planet, and create a world that is healthier, happier, and more vibrant for all of us. With every purchase, every share, and every conversation, we're one step closer to making our vision a reality. Thank you for being part of our community, and for believing, as we do, that together, we can make a difference.\",\n",
+ "β β β β 'role': 'assistant',\n",
+ "β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β 'tool_calls': []\n",
+ "β β β },\n",
+ "β β β 'session_id': '254cf164-52f4-4b7f-ba92-996e97725c12',\n",
+ "β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 37, 692969, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=57600))),\n",
+ "β β β 'steps': [\n",
+ "β β β β {\n",
+ "β β β β β 'model_response': {\n",
+ "β β β β β β 'content': \"Response:\\n\\nImagine a world where every small choice you make can contribute to a bigger, more beautiful picture - a world where the air is fresh, the oceans are clean, and the future is bright. At [Brand Name], we believe that this world is not just a dream, but a reality that we can create together, one step at a time. That's why we're passionate about introducing you to our eco-friendly product, designed with love for the planet and a deep respect for the values that you hold dear.\\n\\nOur product is more than just a solution to your everyday needs; it's a statement of your commitment to the well-being of our planet. Made from sustainable materials and designed with recyclability in mind, every aspect of our product reflects our shared desire to reduce waste and live in harmony with nature. Whether you're a long-time advocate for environmental causes or just starting your journey towards a more sustainable lifestyle, our product is here to support and enhance your efforts.\\n\\nWhat sets us apart is not just our product's eco-friendly features, but the community of like-minded individuals who believe, as we do, that small actions today can lead to a significant positive impact tomorrow. By choosing our product, you're not only making a responsible choice for the planet, but you're also becoming part of a movement - a movement that values the beauty of nature, the importance of community, and the power of collective action.\\n\\nAt [Brand Name], we're dedicated to more than just selling a product; we're committed to fostering a relationship with you, our customer, and with the Earth. We believe in transparency, in honesty, and in the open sharing of our processes and materials. We want you to feel confident and proud of the choices you make, knowing that you're supporting a brand that genuinely cares about the same things you do.\\n\\nSo, join us on this journey towards a greener, brighter future. Together, let's embrace the power of sustainable living, celebrate the beauty of our planet, and create a world that is healthier, happier, and more vibrant for all of us. With every purchase, every share, and every conversation, we're one step closer to making our vision a reality. Thank you for being part of our community, and for believing, as we do, that together, we can make a difference.\",\n",
+ "β β β β β β 'role': 'assistant',\n",
+ "β β β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β β β 'tool_calls': []\n",
+ "β β β β β },\n",
+ "β β β β β 'step_id': '6e454ed2-6dc0-469f-aba6-854a3f52093b',\n",
+ "β β β β β 'step_type': 'inference',\n",
+ "β β β β β 'turn_id': '3e0e5e28-9693-4535-ae54-cb00ba977a4e',\n",
+ "β β β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 47, 299500, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 37, 703303, tzinfo=TzInfo(-08:00))\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'turn_id': '3e0e5e28-9693-4535-ae54-cb00ba977a4e',\n",
+ "β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 47, 313355, tzinfo=TzInfo(-08:00)),\n",
+ "β β β 'output_attachments': []\n",
+ "β β }\n",
+ "β ]\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'254cf164-52f4-4b7f-ba92-996e97725c12'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_name'\u001b[0m: \u001b[32m'worker_agent_b83fb070-705b-4e58-8146-84970328bea0'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m37\u001b[0m, \u001b[1;36m686501\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'turns'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'input_messages'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m'Your task is to Write an engaging, friendly description that connects with the target audience on an emotional level, emphasizing the eco-friendly benefits of the product and how it aligns with their values. This approach would serve the aspect of building a relationship with the reader and creating a sense of shared values..'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'user'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'context'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_message'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response:\\n\\nImagine a world where every small choice you make can contribute to a bigger, more beautiful picture - a world where the air is fresh, the oceans are clean, and the future is bright. At \u001b[0m\u001b[32m[\u001b[0m\u001b[32mBrand Name\u001b[0m\u001b[32m]\u001b[0m\u001b[32m, we believe that this world is not just a dream, but a reality that we can create together, one step at a time. That's why we're passionate about introducing you to our eco-friendly product, designed with love for the planet and a deep respect for the values that you hold dear.\\n\\nOur product is more than just a solution to your everyday needs; it's a statement of your commitment to the well-being of our planet. Made from sustainable materials and designed with recyclability in mind, every aspect of our product reflects our shared desire to reduce waste and live in harmony with nature. Whether you're a long-time advocate for environmental causes or just starting your journey towards a more sustainable lifestyle, our product is here to support and enhance your efforts.\\n\\nWhat sets us apart is not just our product's eco-friendly features, but the community of like-minded individuals who believe, as we do, that small actions today can lead to a significant positive impact tomorrow. By choosing our product, you're not only making a responsible choice for the planet, but you're also becoming part of a movement - a movement that values the beauty of nature, the importance of community, and the power of collective action.\\n\\nAt \u001b[0m\u001b[32m[\u001b[0m\u001b[32mBrand Name\u001b[0m\u001b[32m]\u001b[0m\u001b[32m, we're dedicated to more than just selling a product; we're committed to fostering a relationship with you, our customer, and with the Earth. We believe in transparency, in honesty, and in the open sharing of our processes and materials. We want you to feel confident and proud of the choices you make, knowing that you're supporting a brand that genuinely cares about the same things you do.\\n\\nSo, join us on this journey towards a greener, brighter future. Together, let's embrace the power of sustainable living, celebrate the beauty of our planet, and create a world that is healthier, happier, and more vibrant for all of us. With every purchase, every share, and every conversation, we're one step closer to making our vision a reality. Thank you for being part of our community, and for believing, as we do, that together, we can make a difference.\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'254cf164-52f4-4b7f-ba92-996e97725c12'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m37\u001b[0m, \u001b[1;36m692969\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.timezone\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mdatetime.timedelta\u001b[0m\u001b[1m(\u001b[0m\u001b[33mdays\u001b[0m=\u001b[1;36m-1\u001b[0m, \u001b[33mseconds\u001b[0m=\u001b[1;36m57600\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'model_response'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response:\\n\\nImagine a world where every small choice you make can contribute to a bigger, more beautiful picture - a world where the air is fresh, the oceans are clean, and the future is bright. At \u001b[0m\u001b[32m[\u001b[0m\u001b[32mBrand Name\u001b[0m\u001b[32m]\u001b[0m\u001b[32m, we believe that this world is not just a dream, but a reality that we can create together, one step at a time. That's why we're passionate about introducing you to our eco-friendly product, designed with love for the planet and a deep respect for the values that you hold dear.\\n\\nOur product is more than just a solution to your everyday needs; it's a statement of your commitment to the well-being of our planet. Made from sustainable materials and designed with recyclability in mind, every aspect of our product reflects our shared desire to reduce waste and live in harmony with nature. Whether you're a long-time advocate for environmental causes or just starting your journey towards a more sustainable lifestyle, our product is here to support and enhance your efforts.\\n\\nWhat sets us apart is not just our product's eco-friendly features, but the community of like-minded individuals who believe, as we do, that small actions today can lead to a significant positive impact tomorrow. By choosing our product, you're not only making a responsible choice for the planet, but you're also becoming part of a movement - a movement that values the beauty of nature, the importance of community, and the power of collective action.\\n\\nAt \u001b[0m\u001b[32m[\u001b[0m\u001b[32mBrand Name\u001b[0m\u001b[32m]\u001b[0m\u001b[32m, we're dedicated to more than just selling a product; we're committed to fostering a relationship with you, our customer, and with the Earth. We believe in transparency, in honesty, and in the open sharing of our processes and materials. We want you to feel confident and proud of the choices you make, knowing that you're supporting a brand that genuinely cares about the same things you do.\\n\\nSo, join us on this journey towards a greener, brighter future. Together, let's embrace the power of sustainable living, celebrate the beauty of our planet, and create a world that is healthier, happier, and more vibrant for all of us. With every purchase, every share, and every conversation, we're one step closer to making our vision a reality. Thank you for being part of our community, and for believing, as we do, that together, we can make a difference.\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_id'\u001b[0m: \u001b[32m'6e454ed2-6dc0-469f-aba6-854a3f52093b'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_type'\u001b[0m: \u001b[32m'inference'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'3e0e5e28-9693-4535-ae54-cb00ba977a4e'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m299500\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m37\u001b[0m, \u001b[1;36m703303\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'3e0e5e28-9693-4535-ae54-cb00ba977a4e'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m313355\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_attachments'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Worker creative Session:\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "{\n",
+ "β 'session_id': 'a4caaaa3-4074-48cc-884e-70e1ea08988e',\n",
+ "β 'session_name': 'worker_agent_947325ae-2234-497e-82d7-ca54fa6f5f64',\n",
+ "β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 47, 364200),\n",
+ "β 'turns': [\n",
+ "β β {\n",
+ "β β β 'input_messages': [\n",
+ "β β β β {\n",
+ "β β β β β 'content': \"Your task is to Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader's attention and leaving a lasting impression..\",\n",
+ "β β β β β 'role': 'user',\n",
+ "β β β β β 'context': None\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'output_message': {\n",
+ "β β β β 'content': \"Response: \\n\\nImagine stepping into a world where technology seamlessly blends with art, where innovation knows no bounds, and where the ordinary becomes extraordinary. Welcome to the realm of Lumina, a revolutionary smartwatch that redefines the boundaries of timekeeping and personal style. This masterpiece is not just a device; it's an experience that wraps around your wrist, a constant companion that adapts to your every move, desire, and dream.\\n\\nAs you slip on Lumina, the soft, sleek strap molds to your skin, comfortable against your pulse. The face, a vibrant canvas of light and color, comes alive with every glance. It's not just a screen; it's a window to a universe of possibilities. With a mere touch, the interface unfolds, revealing a tapestry of features designed to elevate your daily life. From tracking the intricacies of your health and fitness journey to keeping you connected with loved ones, Lumina is your personal gateway to a world of wellness and communication.\\n\\nOne of the standout features of Lumina is its advanced health monitoring system. It's equipped with cutting-edge technology that not only tracks your heart rate and sleep patterns but also provides insightful analysis to help you understand your body better. Imagine being able to optimize your workout sessions based on real-time feedback, or receiving alerts that remind you to stay hydrated throughout the day. Lumina doesn't just monitor your health; it empowers you to take control of it.\\n\\nBut Lumina is more than just a health companion; it's also a style statement. Its design is a symphony of elegance and modernity, with interchangeable straps that allow you to match your watch to your mood, outfit, or occasion. Whether you're heading to a boardroom meeting or a casual evening out with friends, Lumina adapts, ensuring you always make a statement. It's the perfect blend of form and function, where every detail has been meticulously crafted to provide a seamless user experience.\\n\\nWhat truly sets Lumina apart, however, is its integration with your digital life. With seamless connectivity to your smartphone, you can receive notifications, control your music playlists, and even make hands-free calls. The voice assistant feature allows you to command your day with ease, from setting reminders to sending messages, all without needing to reach for your phone. It's the epitome of convenience, streamlining your interactions and letting you live more in the moment.\\n\\nAs the sun dips and the stars begin to twinkle, Lumina transforms once more. Its face glows softly in the dark, a beacon of innovation\",\n",
+ "β β β β 'role': 'assistant',\n",
+ "β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β 'tool_calls': []\n",
+ "β β β },\n",
+ "β β β 'session_id': 'a4caaaa3-4074-48cc-884e-70e1ea08988e',\n",
+ "β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 47, 372175, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=57600))),\n",
+ "β β β 'steps': [\n",
+ "β β β β {\n",
+ "β β β β β 'model_response': {\n",
+ "β β β β β β 'content': \"Response: \\n\\nImagine stepping into a world where technology seamlessly blends with art, where innovation knows no bounds, and where the ordinary becomes extraordinary. Welcome to the realm of Lumina, a revolutionary smartwatch that redefines the boundaries of timekeeping and personal style. This masterpiece is not just a device; it's an experience that wraps around your wrist, a constant companion that adapts to your every move, desire, and dream.\\n\\nAs you slip on Lumina, the soft, sleek strap molds to your skin, comfortable against your pulse. The face, a vibrant canvas of light and color, comes alive with every glance. It's not just a screen; it's a window to a universe of possibilities. With a mere touch, the interface unfolds, revealing a tapestry of features designed to elevate your daily life. From tracking the intricacies of your health and fitness journey to keeping you connected with loved ones, Lumina is your personal gateway to a world of wellness and communication.\\n\\nOne of the standout features of Lumina is its advanced health monitoring system. It's equipped with cutting-edge technology that not only tracks your heart rate and sleep patterns but also provides insightful analysis to help you understand your body better. Imagine being able to optimize your workout sessions based on real-time feedback, or receiving alerts that remind you to stay hydrated throughout the day. Lumina doesn't just monitor your health; it empowers you to take control of it.\\n\\nBut Lumina is more than just a health companion; it's also a style statement. Its design is a symphony of elegance and modernity, with interchangeable straps that allow you to match your watch to your mood, outfit, or occasion. Whether you're heading to a boardroom meeting or a casual evening out with friends, Lumina adapts, ensuring you always make a statement. It's the perfect blend of form and function, where every detail has been meticulously crafted to provide a seamless user experience.\\n\\nWhat truly sets Lumina apart, however, is its integration with your digital life. With seamless connectivity to your smartphone, you can receive notifications, control your music playlists, and even make hands-free calls. The voice assistant feature allows you to command your day with ease, from setting reminders to sending messages, all without needing to reach for your phone. It's the epitome of convenience, streamlining your interactions and letting you live more in the moment.\\n\\nAs the sun dips and the stars begin to twinkle, Lumina transforms once more. Its face glows softly in the dark, a beacon of innovation\",\n",
+ "β β β β β β 'role': 'assistant',\n",
+ "β β β β β β 'stop_reason': 'end_of_turn',\n",
+ "β β β β β β 'tool_calls': []\n",
+ "β β β β β },\n",
+ "β β β β β 'step_id': 'd459749c-f883-4d96-acb3-723164ed92b1',\n",
+ "β β β β β 'step_type': 'inference',\n",
+ "β β β β β 'turn_id': '47645e95-f606-4bec-ad1e-cc471c78dcd2',\n",
+ "β β β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 56, 306242, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β β 'started_at': datetime.datetime(2025, 3, 3, 12, 45, 47, 383443, tzinfo=TzInfo(-08:00))\n",
+ "β β β β }\n",
+ "β β β ],\n",
+ "β β β 'turn_id': '47645e95-f606-4bec-ad1e-cc471c78dcd2',\n",
+ "β β β 'completed_at': datetime.datetime(2025, 3, 3, 12, 45, 56, 319286, tzinfo=TzInfo(-08:00)),\n",
+ "β β β 'output_attachments': []\n",
+ "β β }\n",
+ "β ]\n",
+ "}\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'a4caaaa3-4074-48cc-884e-70e1ea08988e'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'session_name'\u001b[0m: \u001b[32m'worker_agent_947325ae-2234-497e-82d7-ca54fa6f5f64'\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m364200\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[32m'turns'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'input_messages'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Your task is to Write a descriptive and imaginative piece that brings the product to life, highlighting its unique features and benefits in a way that is both informative and compelling. This approach would serve the aspect of captivating the reader's attention and leaving a lasting impression..\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'user'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'context'\u001b[0m: \u001b[3;35mNone\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_message'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response: \\n\\nImagine stepping into a world where technology seamlessly blends with art, where innovation knows no bounds, and where the ordinary becomes extraordinary. Welcome to the realm of Lumina, a revolutionary smartwatch that redefines the boundaries of timekeeping and personal style. This masterpiece is not just a device; it's an experience that wraps around your wrist, a constant companion that adapts to your every move, desire, and dream.\\n\\nAs you slip on Lumina, the soft, sleek strap molds to your skin, comfortable against your pulse. The face, a vibrant canvas of light and color, comes alive with every glance. It's not just a screen; it's a window to a universe of possibilities. With a mere touch, the interface unfolds, revealing a tapestry of features designed to elevate your daily life. From tracking the intricacies of your health and fitness journey to keeping you connected with loved ones, Lumina is your personal gateway to a world of wellness and communication.\\n\\nOne of the standout features of Lumina is its advanced health monitoring system. It's equipped with cutting-edge technology that not only tracks your heart rate and sleep patterns but also provides insightful analysis to help you understand your body better. Imagine being able to optimize your workout sessions based on real-time feedback, or receiving alerts that remind you to stay hydrated throughout the day. Lumina doesn't just monitor your health; it empowers you to take control of it.\\n\\nBut Lumina is more than just a health companion; it's also a style statement. Its design is a symphony of elegance and modernity, with interchangeable straps that allow you to match your watch to your mood, outfit, or occasion. Whether you're heading to a boardroom meeting or a casual evening out with friends, Lumina adapts, ensuring you always make a statement. It's the perfect blend of form and function, where every detail has been meticulously crafted to provide a seamless user experience.\\n\\nWhat truly sets Lumina apart, however, is its integration with your digital life. With seamless connectivity to your smartphone, you can receive notifications, control your music playlists, and even make hands-free calls. The voice assistant feature allows you to command your day with ease, from setting reminders to sending messages, all without needing to reach for your phone. It's the epitome of convenience, streamlining your interactions and letting you live more in the moment.\\n\\nAs the sun dips and the stars begin to twinkle, Lumina transforms once more. Its face glows softly in the dark, a beacon of innovation\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'session_id'\u001b[0m: \u001b[32m'a4caaaa3-4074-48cc-884e-70e1ea08988e'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m372175\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.timezone\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mdatetime.timedelta\u001b[0m\u001b[1m(\u001b[0m\u001b[33mdays\u001b[0m=\u001b[1;36m-1\u001b[0m, \u001b[33mseconds\u001b[0m=\u001b[1;36m57600\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'steps'\u001b[0m: \u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'model_response'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'content'\u001b[0m: \u001b[32m\"Response: \\n\\nImagine stepping into a world where technology seamlessly blends with art, where innovation knows no bounds, and where the ordinary becomes extraordinary. Welcome to the realm of Lumina, a revolutionary smartwatch that redefines the boundaries of timekeeping and personal style. This masterpiece is not just a device; it's an experience that wraps around your wrist, a constant companion that adapts to your every move, desire, and dream.\\n\\nAs you slip on Lumina, the soft, sleek strap molds to your skin, comfortable against your pulse. The face, a vibrant canvas of light and color, comes alive with every glance. It's not just a screen; it's a window to a universe of possibilities. With a mere touch, the interface unfolds, revealing a tapestry of features designed to elevate your daily life. From tracking the intricacies of your health and fitness journey to keeping you connected with loved ones, Lumina is your personal gateway to a world of wellness and communication.\\n\\nOne of the standout features of Lumina is its advanced health monitoring system. It's equipped with cutting-edge technology that not only tracks your heart rate and sleep patterns but also provides insightful analysis to help you understand your body better. Imagine being able to optimize your workout sessions based on real-time feedback, or receiving alerts that remind you to stay hydrated throughout the day. Lumina doesn't just monitor your health; it empowers you to take control of it.\\n\\nBut Lumina is more than just a health companion; it's also a style statement. Its design is a symphony of elegance and modernity, with interchangeable straps that allow you to match your watch to your mood, outfit, or occasion. Whether you're heading to a boardroom meeting or a casual evening out with friends, Lumina adapts, ensuring you always make a statement. It's the perfect blend of form and function, where every detail has been meticulously crafted to provide a seamless user experience.\\n\\nWhat truly sets Lumina apart, however, is its integration with your digital life. With seamless connectivity to your smartphone, you can receive notifications, control your music playlists, and even make hands-free calls. The voice assistant feature allows you to command your day with ease, from setting reminders to sending messages, all without needing to reach for your phone. It's the epitome of convenience, streamlining your interactions and letting you live more in the moment.\\n\\nAs the sun dips and the stars begin to twinkle, Lumina transforms once more. Its face glows softly in the dark, a beacon of innovation\"\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'role'\u001b[0m: \u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'stop_reason'\u001b[0m: \u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'tool_calls'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_id'\u001b[0m: \u001b[32m'd459749c-f883-4d96-acb3-723164ed92b1'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'step_type'\u001b[0m: \u001b[32m'inference'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'47645e95-f606-4bec-ad1e-cc471c78dcd2'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m56\u001b[0m, \u001b[1;36m306242\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'started_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m47\u001b[0m, \u001b[1;36m383443\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'turn_id'\u001b[0m: \u001b[32m'47645e95-f606-4bec-ad1e-cc471c78dcd2'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'completed_at'\u001b[0m: \u001b[1;35mdatetime.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m12\u001b[0m, \u001b[1;36m45\u001b[0m, \u001b[1;36m56\u001b[0m, \u001b[1;36m319286\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[32m'output_attachments'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[1m}\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "orchestrator_session = client.agents.session.retrieve(session_id=orchestrator_agent.session_id, agent_id=orchestrator_agent.agent_id)\n",
+ "pprint(orchestrator_session.to_dict())\n",
+ "\n",
+ "for worker_type, worker in workers.items():\n",
+ " worker_session = client.agents.session.retrieve(session_id=worker.session_id, agent_id=worker.agent_id)\n",
+ " print(f\"Worker {worker_type} Session:\")\n",
+ " pprint(worker_session.to_dict())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "master",
+ "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.16"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/notebooks/Llama_Stack_RAG_Lifecycle.ipynb b/docs/notebooks/Llama_Stack_RAG_Lifecycle.ipynb
new file mode 100644
index 000000000..0d7b462cc
--- /dev/null
+++ b/docs/notebooks/Llama_Stack_RAG_Lifecycle.ipynb
@@ -0,0 +1,1427 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Llama Stack RAG Lifecycle\n",
+ "\n",
+ "In this notebook, we will walk through the lifecycle of building and evaluating a RAG pipeline using Llama Stack. \n",
+ "\n",
+ "**Example: Torchtune Knowledge Agent** \n",
+ "\n",
+ "Throughout this notebook, we will build a knowledge agent that can answer questions about the Torchtune project. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 0. Setup"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Not in Google Colab environment\n"
+ ]
+ }
+ ],
+ "source": [
+ "from llama_stack_client import LlamaStackClient\n",
+ "from llama_stack.distribution.library_client import LlamaStackAsLibraryClient\n",
+ "from llama_stack_client.types.agent_create_params import AgentConfig\n",
+ "from llama_stack_client.lib.agents.agent import Agent\n",
+ "from rich.pretty import pprint\n",
+ "import json\n",
+ "import uuid\n",
+ "from pydantic import BaseModel\n",
+ "import rich\n",
+ "import os\n",
+ "try:\n",
+ " from google.colab import userdata\n",
+ " os.environ['FIREWORKS_API_KEY'] = userdata.get('FIREWORKS_API_KEY')\n",
+ "except ImportError:\n",
+ " print(\"Not in Google Colab environment\")\n",
+ "\n",
+ "# client = LlamaStackAsLibraryClient(\"fireworks\", provider_data = {\"fireworks_api_key\": os.environ['FIREWORKS_API_KEY']})\n",
+ "# _ = client.initialize()\n",
+ "\n",
+ "# Uncomment to run on a hosted Llama Stack server\n",
+ "client = LlamaStackClient(base_url=\"http://localhost:8321\")\n",
+ "\n",
+ "MODEL_ID = \"meta-llama/Llama-3.3-70B-Instruct\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Simple Vanilla Agent\n",
+ "\n",
+ "First, we will build a simple vanilla agent without any access to external knowledge base or tools, and check how it performs on a couple of questions. \n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# First, let's come up with a couple of examples to test the agent\n",
+ "examples = [\n",
+ " {\n",
+ " \"input_query\": \"What precision formats does torchtune support?\",\n",
+ " \"expected_answer\": \"Torchtune supports two data types for precision: fp32 (full-precision) which uses 4 bytes per model and optimizer parameter, and bfloat16 (half-precision) which uses 2 bytes per model and optimizer parameter.\"\n",
+ " },\n",
+ " {\n",
+ " \"input_query\": \"What does DoRA stand for in torchtune?\",\n",
+ " \"expected_answer\": \"Weight-Decomposed Low-Rank Adaptation\"\n",
+ " },\n",
+ " {\n",
+ " \"input_query\": \"How does the CPUOffloadOptimizer reduce GPU memory usage?\",\n",
+ " \"expected_answer\": \"The CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on CPU and performing optimizer steps on CPU. It can also optionally offload gradients to CPU by using offload_gradients=True\"\n",
+ " },\n",
+ " {\n",
+ " \"input_query\": \"How do I ensure only LoRA parameters are trainable when fine-tuning?\",\n",
+ " \"expected_answer\": \"You can set only LoRA parameters to trainable using torchtune's utility functions: first fetch all LoRA parameters with lora_params = get_adapter_params(lora_model), then set them as trainable with set_trainable_params(lora_model, lora_params). The LoRA recipe handles this automatically.\"\n",
+ " }\n",
+ "]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "Question: What precision formats does torchtune support?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What precision formats does torchtune support?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: Torchtune supports the following precision formats:\n",
+ "\n",
+ "* Full precision (FP32)\n",
+ "* Mixed precision (FP16)\n",
+ "\n",
+ "It may also support other formats such as INT8 and BF16 in the future, but currently, it primarily focuses on FP32 \n",
+ "and FP16. \n",
+ "\n",
+ "Please note that the specific precision formats supported by Torchtune may change over time, and it's always best \n",
+ "to check the official documentation for the most up-to-date information.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m Torchtune supports the following precision formats:\n",
+ "\n",
+ "* Full precision \u001b[1m(\u001b[0mFP32\u001b[1m)\u001b[0m\n",
+ "* Mixed precision \u001b[1m(\u001b[0mFP16\u001b[1m)\u001b[0m\n",
+ "\n",
+ "It may also support other formats such as INT8 and BF16 in the future, but currently, it primarily focuses on FP32 \n",
+ "and FP16. \n",
+ "\n",
+ "Please note that the specific precision formats supported by Torchtune may change over time, and it's always best \n",
+ "to check the official documentation for the most up-to-date information.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: What does DoRA stand for in torchtune?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What does DoRA stand for in torchtune?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: In the context of the Torchtune project, DoRA stands for \"Decoupled Optimizer for Reparameterized \n",
+ "Architectures\".\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m In the context of the Torchtune project, DoRA stands for \u001b[32m\"Decoupled Optimizer for Reparameterized \u001b[0m\n",
+ "\u001b[32mArchitectures\"\u001b[0m.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How does the CPUOffloadOptimizer reduce GPU memory usage?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How does the CPUOffloadOptimizer reduce GPU memory usage?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: The CPUOffloadOptimizer in the Torchtune project is designed to reduce GPU memory usage by offloading\n",
+ "certain computations from the GPU to the CPU. Here's how it works:\n",
+ "\n",
+ "1. **Identifying offloadable operations**: The optimizer analyzes the computation graph of the model and identifies\n",
+ "operations that can be offloaded from the GPU to the CPU. These operations are typically those that don't require \n",
+ "the massive parallel processing capabilities of the GPU, such as data preprocessing, encoding, or decoding.\n",
+ "2. **Offloading operations to CPU**: The optimizer offloads the identified operations to the CPU, which frees up \n",
+ "GPU memory and reduces the amount of data that needs to be transferred between the GPU and CPU.\n",
+ "3. **Minimizing data transfer**: The optimizer minimizes the amount of data that needs to be transferred between \n",
+ "the GPU and CPU by only transferring the necessary data for the offloaded operations. This reduces the overhead of \n",
+ "data transfer and helps to conserve GPU memory.\n",
+ "4. **Optimizing CPU-GPU synchronization**: The optimizer ensures that the CPU and GPU are properly synchronized, \n",
+ "which helps to prevent unnecessary memory allocations and deallocations on the GPU.\n",
+ "5. **Dynamic memory allocation**: The optimizer can dynamically allocate and deallocate memory on the GPU as \n",
+ "needed, which helps to reduce memory fragmentation and waste.\n",
+ "\n",
+ "By offloading computations to the CPU and minimizing data transfer, the CPUOffloadOptimizer can significantly \n",
+ "reduce GPU memory usage, which can lead to:\n",
+ "\n",
+ "* Improved model training and inference performance\n",
+ "* Increased batch sizes and throughput\n",
+ "* Reduced out-of-memory errors\n",
+ "* Better support for larger models and datasets\n",
+ "\n",
+ "Overall, the CPUOffloadOptimizer is a powerful tool for optimizing GPU memory usage in deep learning workloads, and\n",
+ "can help to improve the overall performance and efficiency of the Torchtune project.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m The CPUOffloadOptimizer in the Torchtune project is designed to reduce GPU memory usage by offloading\n",
+ "certain computations from the GPU to the CPU. Here's how it works:\n",
+ "\n",
+ "\u001b[1;36m1\u001b[0m. **Identifying offloadable operations**: The optimizer analyzes the computation graph of the model and identifies\n",
+ "operations that can be offloaded from the GPU to the CPU. These operations are typically those that don't require \n",
+ "the massive parallel processing capabilities of the GPU, such as data preprocessing, encoding, or decoding.\n",
+ "\u001b[1;36m2\u001b[0m. **Offloading operations to CPU**: The optimizer offloads the identified operations to the CPU, which frees up \n",
+ "GPU memory and reduces the amount of data that needs to be transferred between the GPU and CPU.\n",
+ "\u001b[1;36m3\u001b[0m. **Minimizing data transfer**: The optimizer minimizes the amount of data that needs to be transferred between \n",
+ "the GPU and CPU by only transferring the necessary data for the offloaded operations. This reduces the overhead of \n",
+ "data transfer and helps to conserve GPU memory.\n",
+ "\u001b[1;36m4\u001b[0m. **Optimizing CPU-GPU synchronization**: The optimizer ensures that the CPU and GPU are properly synchronized, \n",
+ "which helps to prevent unnecessary memory allocations and deallocations on the GPU.\n",
+ "\u001b[1;36m5\u001b[0m. **Dynamic memory allocation**: The optimizer can dynamically allocate and deallocate memory on the GPU as \n",
+ "needed, which helps to reduce memory fragmentation and waste.\n",
+ "\n",
+ "By offloading computations to the CPU and minimizing data transfer, the CPUOffloadOptimizer can significantly \n",
+ "reduce GPU memory usage, which can lead to:\n",
+ "\n",
+ "* Improved model training and inference performance\n",
+ "* Increased batch sizes and throughput\n",
+ "* Reduced out-of-memory errors\n",
+ "* Better support for larger models and datasets\n",
+ "\n",
+ "Overall, the CPUOffloadOptimizer is a powerful tool for optimizing GPU memory usage in deep learning workloads, and\n",
+ "can help to improve the overall performance and efficiency of the Torchtune project.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How do I ensure only LoRA parameters are trainable when fine-tuning?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How do I ensure only LoRA parameters are trainable when fine-tuning?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: To ensure only LoRA (Low-Rank Adaptation) parameters are trainable when fine-tuning a model with \n",
+ "Torchtune, you can follow these steps:\n",
+ "\n",
+ "1. **Freeze the original model weights**: Before fine-tuning, you need to freeze the original model weights to \n",
+ "prevent them from being updated during the fine-tuning process. You can do this by setting the `requires_grad` \n",
+ "attribute of the model parameters to `False`. This will prevent the original model weights from being updated.\n",
+ "\n",
+ "2. **Create LoRA parameters**: Create LoRA parameters for the layers you want to fine-tune. LoRA parameters are \n",
+ "typically added to the original model weights to adapt the model to the new task.\n",
+ "\n",
+ "3. **Set LoRA parameters as trainable**: Set the LoRA parameters as trainable by setting their `requires_grad` \n",
+ "attribute to `True`. This will allow the LoRA parameters to be updated during the fine-tuning process.\n",
+ "\n",
+ "Here's a sample code snippet to illustrate this:\n",
+ "```python\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "# Assume 'model' is your pre-trained model\n",
+ "model = ...\n",
+ "\n",
+ "# Freeze the original model weights\n",
+ "for param in model.parameters():\n",
+ " param.requires_grad = False\n",
+ "\n",
+ "# Create LoRA parameters\n",
+ "lora_params = []\n",
+ "for name, module in model.named_modules():\n",
+ " if isinstance(module, nn.Linear): # or any other module you want to fine-tune\n",
+ " lora_param = nn.Parameter(torch.randn(module.weight.shape))\n",
+ " lora_params.append(lora_param)\n",
+ " setattr(model, f\"{name}_lora\", lora_param)\n",
+ "\n",
+ "# Set LoRA parameters as trainable\n",
+ "for param in lora_params:\n",
+ " param.requires_grad = True\n",
+ "\n",
+ "# Fine-tune the model with LoRA parameters\n",
+ "optimizer = torch.optim.Adam(lora_params, lr=1e-4)\n",
+ "```\n",
+ "By following these steps, you can ensure that only the LoRA parameters are trainable during fine-tuning, while \n",
+ "keeping the original model weights frozen.\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m To ensure only LoRA \u001b[1m(\u001b[0mLow-Rank Adaptation\u001b[1m)\u001b[0m parameters are trainable when fine-tuning a model with \n",
+ "Torchtune, you can follow these steps:\n",
+ "\n",
+ "\u001b[1;36m1\u001b[0m. **Freeze the original model weights**: Before fine-tuning, you need to freeze the original model weights to \n",
+ "prevent them from being updated during the fine-tuning process. You can do this by setting the `requires_grad` \n",
+ "attribute of the model parameters to `\u001b[3;91mFalse\u001b[0m`. This will prevent the original model weights from being updated.\n",
+ "\n",
+ "\u001b[1;36m2\u001b[0m. **Create LoRA parameters**: Create LoRA parameters for the layers you want to fine-tune. LoRA parameters are \n",
+ "typically added to the original model weights to adapt the model to the new task.\n",
+ "\n",
+ "\u001b[1;36m3\u001b[0m. **Set LoRA parameters as trainable**: Set the LoRA parameters as trainable by setting their `requires_grad` \n",
+ "attribute to `\u001b[3;92mTrue\u001b[0m`. This will allow the LoRA parameters to be updated during the fine-tuning process.\n",
+ "\n",
+ "Here's a sample code snippet to illustrate this:\n",
+ "```python\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "# Assume \u001b[32m'model'\u001b[0m is your pre-trained model\n",
+ "model = \u001b[33m...\u001b[0m\n",
+ "\n",
+ "# Freeze the original model weights\n",
+ "for param in \u001b[1;35mmodel.parameters\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m:\n",
+ " param.requires_grad = \u001b[3;91mFalse\u001b[0m\n",
+ "\n",
+ "# Create LoRA parameters\n",
+ "lora_params = \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "for name, module in \u001b[1;35mmodel.named_modules\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m:\n",
+ " if \u001b[1;35misinstance\u001b[0m\u001b[1m(\u001b[0mmodule, nn.Linear\u001b[1m)\u001b[0m: # or any other module you want to fine-tune\n",
+ " lora_param = \u001b[1;35mnn.Parameter\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mtorch.randn\u001b[0m\u001b[1m(\u001b[0mmodule.weight.shape\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ " \u001b[1;35mlora_params.append\u001b[0m\u001b[1m(\u001b[0mlora_param\u001b[1m)\u001b[0m\n",
+ " \u001b[1;35msetattr\u001b[0m\u001b[1m(\u001b[0mmodel, f\"\u001b[1m{\u001b[0mname\u001b[1m}\u001b[0m_lora\", lora_param\u001b[1m)\u001b[0m\n",
+ "\n",
+ "# Set LoRA parameters as trainable\n",
+ "for param in lora_params:\n",
+ " param.requires_grad = \u001b[3;92mTrue\u001b[0m\n",
+ "\n",
+ "# Fine-tune the model with LoRA parameters\n",
+ "optimizer = \u001b[1;35mtorch.optim.Adam\u001b[0m\u001b[1m(\u001b[0mlora_params, \u001b[33mlr\u001b[0m=\u001b[1;36m1e\u001b[0m\u001b[1;36m-4\u001b[0m\u001b[1m)\u001b[0m\n",
+ "```\n",
+ "By following these steps, you can ensure that only the LoRA parameters are trainable during fine-tuning, while \n",
+ "keeping the original model weights frozen.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "simple_agent = Agent(client,\n",
+ " model=MODEL_ID, \n",
+ " instructions=\"You are a helpful assistant that can answer questions about the Torchtune project.\")\n",
+ "for example in examples:\n",
+ " simple_session_id = simple_agent.create_session(session_name=f\"simple_session_{uuid.uuid4()}\")\n",
+ " response = simple_agent.create_turn(\n",
+ " messages=[\n",
+ " {\n",
+ " \"role\": \"user\",\n",
+ " \"content\": example[\"input_query\"]\n",
+ " }\n",
+ " ],\n",
+ " session_id=simple_session_id,\n",
+ " stream=False\n",
+ " )\n",
+ " rich.print(f\"[bold cyan]Question:[/bold cyan] {example['input_query']}\")\n",
+ " rich.print(f\"[bold yellow]Agent Answer:[/bold yellow] {response.output_message.content}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### 1.1 Evaluate Agent Responses\n",
+ "Let's gather up the agent's logs and evaluate the agent's performance. We can see that our agent's response is quite bad and off from the expected answer."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "ScoringScoreResponse(\n",
+ "β results={\n",
+ "β β 'braintrust::factuality': ScoringResult(\n",
+ "β β β aggregated_results={'average': {'average': 0.3}},\n",
+ "β β β score_rows=[\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.0,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'D',\n",
+ "β β β β β β 'rationale': '1. **Expert Answer**: The expert states that Torchtune supports two precision formats: fp32 (full-precision) and bfloat16 (half-precision).\\n\\n2. **Submitted Answer**: The submission mentions that Torchtune supports full precision (FP32) and mixed precision (FP16). It also speculates about potential future support for other formats like INT8 and BF16, but emphasizes the current focus on FP32 and FP16.\\n\\n3. **Comparison**:\\n - Both answers agree on the support for FP32.\\n - The expert mentions bfloat16 (BF16), while the submission mentions FP16 and speculates about BF16 in the future. This is a key difference as the expert confirms BF16 support, whereas the submission does not.\\n - The submission introduces FP16, which is not mentioned by the expert.\\n - The submission also speculates about future support for INT8 and BF16, which is not addressed by the expert.\\n\\n4. **Conclusion**: There is a disagreement between the submitted answer and the expert answer regarding the precision formats supported by Torchtune. The expert confirms BF16 support, while the submission does not, and instead mentions FP16, which the expert does not confirm. Therefore, the correct choice is (D).'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.0,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'D',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation\".\\n2. The submitted answer states that DoRA stands for \"Decoupled Optimizer for Reparameterized Architectures\".\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer regarding what DoRA stands for in the context of torchtune.\\n5. Therefore, the correct choice is (D) There is a disagreement between the submitted answer and the expert answer.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on the CPU and performing optimizer steps on the CPU. It also mentions the optional offloading of gradients to the CPU.\\n2. The submitted answer describes a broader mechanism of offloading computations from the GPU to the CPU, including identifying offloadable operations, minimizing data transfer, optimizing CPU-GPU synchronization, and dynamic memory allocation.\\n3. The submitted answer does not explicitly mention keeping optimizer states on the CPU or performing optimizer steps on the CPU, which are key points in the expert answer.\\n4. The submitted answer provides additional details about the process of offloading operations and its benefits, which are not mentioned in the expert answer.\\n5. The submitted answer does not conflict with the expert answer but rather expands on the concept of offloading to the CPU with additional mechanisms and benefits.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it, as it includes all the information from the expert answer and adds more details.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': \"1. **Expert Answer Summary**: The expert answer provides a concise method to ensure only LoRA parameters are trainable by using torchtune's utility functions. It mentions fetching LoRA parameters with `get_adapter_params(lora_model)` and setting them as trainable with `set_trainable_params(lora_model, lora_params)`. It also notes that the LoRA recipe handles this automatically.\\n\\n2. **Submitted Answer Summary**: The submitted answer provides a more detailed explanation, including steps to freeze the original model weights, create LoRA parameters, and set them as trainable. It includes a code snippet demonstrating these steps, using PyTorch to manually set `requires_grad` attributes.\\n\\n3. **Comparison**:\\n - Both answers aim to ensure only LoRA parameters are trainable.\\n - The expert answer uses torchtune's utility functions, while the submitted answer provides a manual method using PyTorch.\\n - The submitted answer includes additional steps and a code example, which are not present in the expert answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer. It includes all the information from the expert answer (ensuring only LoRA parameters are trainable) and adds more detail on how to achieve this manually. There is no conflict between the two answers, as they both achieve the same goal using different methods.\\n\\nTherefore, the correct choice is (B) The submitted answer is a superset of the expert answer and is fully consistent with it.\"\n",
+ "β β β β β }\n",
+ "β β β β }\n",
+ "β β β ]\n",
+ "β β )\n",
+ "β }\n",
+ ")\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1;35mScoringScoreResponse\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[33mresults\u001b[0m=\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[32m'braintrust::factuality'\u001b[0m: \u001b[1;35mScoringResult\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33maggregated_results\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1;36m0.3\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mscore_rows\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.0\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'D'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. **Expert Answer**: The expert states that Torchtune supports two precision formats: fp32 \u001b[0m\u001b[32m(\u001b[0m\u001b[32mfull-precision\u001b[0m\u001b[32m)\u001b[0m\u001b[32m and bfloat16 \u001b[0m\u001b[32m(\u001b[0m\u001b[32mhalf-precision\u001b[0m\u001b[32m)\u001b[0m\u001b[32m.\\n\\n2. **Submitted Answer**: The submission mentions that Torchtune supports full precision \u001b[0m\u001b[32m(\u001b[0m\u001b[32mFP32\u001b[0m\u001b[32m)\u001b[0m\u001b[32m and mixed precision \u001b[0m\u001b[32m(\u001b[0m\u001b[32mFP16\u001b[0m\u001b[32m)\u001b[0m\u001b[32m. It also speculates about potential future support for other formats like INT8 and BF16, but emphasizes the current focus on FP32 and FP16.\\n\\n3. **Comparison**:\\n - Both answers agree on the support for FP32.\\n - The expert mentions bfloat16 \u001b[0m\u001b[32m(\u001b[0m\u001b[32mBF16\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, while the submission mentions FP16 and speculates about BF16 in the future. This is a key difference as the expert confirms BF16 support, whereas the submission does not.\\n - The submission introduces FP16, which is not mentioned by the expert.\\n - The submission also speculates about future support for INT8 and BF16, which is not addressed by the expert.\\n\\n4. **Conclusion**: There is a disagreement between the submitted answer and the expert answer regarding the precision formats supported by Torchtune. The expert confirms BF16 support, while the submission does not, and instead mentions FP16, which the expert does not confirm. Therefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.0\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'D'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation\".\\n2. The submitted answer states that DoRA stands for \"Decoupled Optimizer for Reparameterized Architectures\".\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer regarding what DoRA stands for in the context of torchtune.\\n5. Therefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m There is a disagreement between the submitted answer and the expert answer.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on the CPU and performing optimizer steps on the CPU. It also mentions the optional offloading of gradients to the CPU.\\n2. The submitted answer describes a broader mechanism of offloading computations from the GPU to the CPU, including identifying offloadable operations, minimizing data transfer, optimizing CPU-GPU synchronization, and dynamic memory allocation.\\n3. The submitted answer does not explicitly mention keeping optimizer states on the CPU or performing optimizer steps on the CPU, which are key points in the expert answer.\\n4. The submitted answer provides additional details about the process of offloading operations and its benefits, which are not mentioned in the expert answer.\\n5. The submitted answer does not conflict with the expert answer but rather expands on the concept of offloading to the CPU with additional mechanisms and benefits.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it, as it includes all the information from the expert answer and adds more details.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m\"1. **Expert Answer Summary**: The expert answer provides a concise method to ensure only LoRA parameters are trainable by using torchtune's utility functions. It mentions fetching LoRA parameters with `get_adapter_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model\u001b[0m\u001b[32m)\u001b[0m\u001b[32m` and setting them as trainable with `set_trainable_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model, lora_params\u001b[0m\u001b[32m)\u001b[0m\u001b[32m`. It also notes that the LoRA recipe handles this automatically.\\n\\n2. **Submitted Answer Summary**: The submitted answer provides a more detailed explanation, including steps to freeze the original model weights, create LoRA parameters, and set them as trainable. It includes a code snippet demonstrating these steps, using PyTorch to manually set `requires_grad` attributes.\\n\\n3. **Comparison**:\\n - Both answers aim to ensure only LoRA parameters are trainable.\\n - The expert answer uses torchtune's utility functions, while the submitted answer provides a manual method using PyTorch.\\n - The submitted answer includes additional steps and a code example, which are not present in the expert answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer. It includes all the information from the expert answer \u001b[0m\u001b[32m(\u001b[0m\u001b[32mensuring only LoRA parameters are trainable\u001b[0m\u001b[32m)\u001b[0m\u001b[32m and adds more detail on how to achieve this manually. There is no conflict between the two answers, as they both achieve the same goal using different methods.\\n\\nTherefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mB\u001b[0m\u001b[32m)\u001b[0m\u001b[32m The submitted answer is a superset of the expert answer and is fully consistent with it.\"\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[1m)\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "eval_rows = []\n",
+ "for i, session_id in enumerate(simple_agent.sessions):\n",
+ " session_response = client.agents.session.retrieve(agent_id=simple_agent.agent_id, session_id=session_id)\n",
+ " for turn in session_response.turns:\n",
+ " eval_rows.append({\n",
+ " \"input_query\": examples[i][\"input_query\"],\n",
+ " \"expected_answer\": examples[i][\"expected_answer\"],\n",
+ " \"generated_answer\": turn.output_message.content,\n",
+ " })\n",
+ "\n",
+ "scoring_params = {\n",
+ " \"braintrust::factuality\": None,\n",
+ "}\n",
+ "scoring_response = client.scoring.score(\n",
+ " input_rows=eval_rows,\n",
+ " scoring_functions=scoring_params,\n",
+ ")\n",
+ "pprint(scoring_response)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. Search Agent\n",
+ "\n",
+ "Now, let's see how we can improve the agent's performance by adding a search tool."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "Question: What precision formats does torchtune support?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What precision formats does torchtune support?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: Torchtune supports the following precision formats:\n",
+ "\n",
+ "* bf16 (16-bit floating-point format)\n",
+ "* fp32 (32-bit floating-point format, also known as \"full-precision\")\n",
+ "\n",
+ "It's worth noting that torchtune also provides support for mixed-precision techniques, which allow for the use of \n",
+ "different precision formats for different parts of the model or during different stages of training.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m Torchtune supports the following precision formats:\n",
+ "\n",
+ "* bf16 \u001b[1m(\u001b[0m\u001b[1;36m16\u001b[0m-bit floating-point format\u001b[1m)\u001b[0m\n",
+ "* fp32 \u001b[1m(\u001b[0m\u001b[1;36m32\u001b[0m-bit floating-point format, also known as \u001b[32m\"full-precision\"\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\n",
+ "It's worth noting that torchtune also provides support for mixed-precision techniques, which allow for the use of \n",
+ "different precision formats for different parts of the model or during different stages of training.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: What does DoRA stand for in torchtune?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What does DoRA stand for in torchtune?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: DoRA stands for \"Decoupled Orthogonal Random Adaptation\" in torchtune.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m DoRA stands for \u001b[32m\"Decoupled Orthogonal Random Adaptation\"\u001b[0m in torchtune.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How does the CPUOffloadOptimizer reduce GPU memory usage?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How does the CPUOffloadOptimizer reduce GPU memory usage?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: The CPUOffloadOptimizer reduces GPU memory usage by offloading gradients and trainable parameters to \n",
+ "the CPU, allowing for more efficient use of GPU memory. This can be achieved by setting `offload_gradients=True` in\n",
+ "the CPUOffloadOptimizer, which frees gradients once device-to-host transfer finishes. Additionally, using paged \n",
+ "Adam with `optimizer_in_bwd=True` can also help reduce memory usage. However, it's important to note that the \n",
+ "actual memory usage may vary depending on the specific use case and model architecture.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m The CPUOffloadOptimizer reduces GPU memory usage by offloading gradients and trainable parameters to \n",
+ "the CPU, allowing for more efficient use of GPU memory. This can be achieved by setting `\u001b[33moffload_gradients\u001b[0m=\u001b[3;92mTrue\u001b[0m` in\n",
+ "the CPUOffloadOptimizer, which frees gradients once device-to-host transfer finishes. Additionally, using paged \n",
+ "Adam with `\u001b[33moptimizer_in_bwd\u001b[0m=\u001b[3;92mTrue\u001b[0m` can also help reduce memory usage. However, it's important to note that the \n",
+ "actual memory usage may vary depending on the specific use case and model architecture.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How do I ensure only LoRA parameters are trainable when fine-tuning?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How do I ensure only LoRA parameters are trainable when fine-tuning?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: To ensure only LoRA parameters are trainable when fine-tuning, you can use the `set_trainable_params`\n",
+ "function from the `torchtune.modules.peft.peft_utils` module. This function allows you to specify which parameters \n",
+ "to make trainable, and you can use it to set only the LoRA parameters as trainable.\n",
+ "\n",
+ "Here is an example of how to do this:\n",
+ "```\n",
+ "import torch\n",
+ "from torchtune.models.llama2 import llama2_7b, lora_llama2_7b\n",
+ "from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\n",
+ "\n",
+ "# Load the model and adapter\n",
+ "model = llama2_7b()\n",
+ "adapter = lora_llama2_7b()\n",
+ "\n",
+ "# Get the adapter parameters\n",
+ "adapter_params = get_adapter_params(adapter)\n",
+ "\n",
+ "# Set only the adapter parameters as trainable\n",
+ "set_trainable_params(model, adapter_params)\n",
+ "```\n",
+ "This code loads the LLaMA-2 model and the LoRA adapter, gets the adapter parameters, and then sets only those \n",
+ "parameters as trainable using the `set_trainable_params` function. This ensures that only the LoRA parameters are \n",
+ "updated during fine-tuning, while the rest of the model remains frozen.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m To ensure only LoRA parameters are trainable when fine-tuning, you can use the `set_trainable_params`\n",
+ "function from the `torchtune.modules.peft.peft_utils` module. This function allows you to specify which parameters \n",
+ "to make trainable, and you can use it to set only the LoRA parameters as trainable.\n",
+ "\n",
+ "Here is an example of how to do this:\n",
+ "```\n",
+ "import torch\n",
+ "from torchtune.models.llama2 import llama2_7b, lora_llama2_7b\n",
+ "from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\n",
+ "\n",
+ "# Load the model and adapter\n",
+ "model = \u001b[1;35mllama2_7b\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m\n",
+ "adapter = \u001b[1;35mlora_llama2_7b\u001b[0m\u001b[1m(\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\n",
+ "# Get the adapter parameters\n",
+ "adapter_params = \u001b[1;35mget_adapter_params\u001b[0m\u001b[1m(\u001b[0madapter\u001b[1m)\u001b[0m\n",
+ "\n",
+ "# Set only the adapter parameters as trainable\n",
+ "\u001b[1;35mset_trainable_params\u001b[0m\u001b[1m(\u001b[0mmodel, adapter_params\u001b[1m)\u001b[0m\n",
+ "```\n",
+ "This code loads the LLaMA-\u001b[1;36m2\u001b[0m model and the LoRA adapter, gets the adapter parameters, and then sets only those \n",
+ "parameters as trainable using the `set_trainable_params` function. This ensures that only the LoRA parameters are \n",
+ "updated during fine-tuning, while the rest of the model remains frozen.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "search_agent = Agent(client, \n",
+ " model=MODEL_ID,\n",
+ " instructions=\"You are a helpful assistant that can answer questions about the Torchtune project. You should always use the search tool to answer questions.\",\n",
+ " tools=[\"builtin::websearch\"])\n",
+ "for example in examples:\n",
+ " search_session_id = search_agent.create_session(session_name=f\"search_session_{uuid.uuid4()}\")\n",
+ " response = search_agent.create_turn(\n",
+ " messages=[\n",
+ " {\n",
+ " \"role\": \"user\",\n",
+ " \"content\": example[\"input_query\"]\n",
+ " }\n",
+ " ],\n",
+ " session_id=search_session_id,\n",
+ " stream=False\n",
+ " )\n",
+ " rich.print(f\"[bold cyan]Question:[/bold cyan] {example['input_query']}\")\n",
+ " rich.print(f\"[bold yellow]Agent Answer:[/bold yellow] {response.output_message.content}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### 2.1 Evaluate Agent Responses\n",
+ "\n",
+ "We can see that with a search tool, the agent's performance is much better, and have less hallucinations. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 19,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "ScoringScoreResponse(\n",
+ "β results={\n",
+ "β β 'braintrust::factuality': ScoringResult(\n",
+ "β β β aggregated_results={'average': {'average': 0.44999999999999996}},\n",
+ "β β β score_rows=[\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': '1. **Expert Answer Details**: The expert answer states that Torchtune supports two precision formats: fp32 (full-precision) and bfloat16 (half-precision).\\n\\n2. **Submitted Answer Details**: The submitted answer mentions two precision formats: bf16 (16-bit floating-point format) and fp32 (32-bit floating-point format, also known as \"full-precision\"). It also adds that Torchtune supports mixed-precision techniques.\\n\\n3. **Comparison of Precision Formats**:\\n - The expert answer uses \"bfloat16\" while the submitted answer uses \"bf16\". These are equivalent terms, as \"bf16\" is a common abbreviation for \"bfloat16\".\\n - Both answers mention \"fp32\" as a supported precision format.\\n\\n4. **Additional Information in Submission**: The submitted answer includes additional information about mixed-precision techniques, which is not mentioned in the expert answer.\\n\\n5. **Consistency Check**: The submitted answer includes all the information from the expert answer and adds more details about mixed-precision techniques. There is no conflict between the two answers.\\n\\nBased on the above analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.0,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'D',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation.\"\\n2. The submitted answer states that DoRA stands for \"Decoupled Orthogonal Random Adaptation.\"\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer regarding what DoRA stands for in torchtune.\\n5. Therefore, the correct choice is (D) There is a disagreement between the submitted answer and the expert answer.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': '1. **Expert Answer Analysis**: The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on the CPU and performing optimizer steps on the CPU. It also mentions the optional offloading of gradients to the CPU by setting `offload_gradients=True`.\\n\\n2. **Submitted Answer Analysis**: The submitted answer mentions offloading gradients and trainable parameters to the CPU, which allows for more efficient use of GPU memory. It specifies the use of `offload_gradients=True` to free gradients after device-to-host transfer. Additionally, it introduces the concept of using paged Adam with `optimizer_in_bwd=True` to help reduce memory usage. It also notes that actual memory usage may vary depending on the use case and model architecture.\\n\\n3. **Comparison**:\\n - Both answers mention offloading gradients to the CPU using `offload_gradients=True`.\\n - The expert answer focuses on keeping optimizer states and performing optimizer steps on the CPU, while the submitted answer expands on this by mentioning trainable parameters and the use of paged Adam.\\n - The submitted answer provides additional context about memory usage variability and the use of paged Adam, which is not mentioned in the expert answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer as it includes all the information from the expert answer and adds more details about trainable parameters, paged Adam, and memory usage variability. There is no conflict between the two answers, and the additional information in the submitted answer is consistent with the expert answer.\\n\\nTherefore, the correct choice is (B) The submitted answer is a superset of the expert answer and is fully consistent with it.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': \"1. **Expert Answer Analysis**: The expert answer provides a method to ensure only LoRA parameters are trainable by using torchtune's utility functions. It mentions fetching LoRA parameters with `get_adapter_params(lora_model)` and setting them as trainable with `set_trainable_params(lora_model, lora_params)`. It also notes that the LoRA recipe handles this automatically.\\n\\n2. **Submitted Answer Analysis**: The submitted answer provides a detailed example of how to ensure only LoRA parameters are trainable. It uses the `set_trainable_params` function from `torchtune.modules.peft.peft_utils` and provides a code example that includes loading a model and adapter, fetching adapter parameters, and setting them as trainable.\\n\\n3. **Comparison**:\\n - Both answers mention the use of `set_trainable_params` to set LoRA parameters as trainable.\\n - Both answers involve fetching LoRA parameters using a function (`get_adapter_params`).\\n - The submitted answer provides additional context by including a code example and specifying the module path for the functions used.\\n - The expert answer mentions that the LoRA recipe handles this automatically, which is not explicitly stated in the submitted answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer. It includes all the information from the expert answer and adds more detail, such as a code example and specific module paths. There is no conflict between the two answers, and the additional information in the submitted answer is consistent with the expert answer.\"\n",
+ "β β β β β }\n",
+ "β β β β }\n",
+ "β β β ]\n",
+ "β β )\n",
+ "β }\n",
+ ")\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1;35mScoringScoreResponse\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[33mresults\u001b[0m=\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[32m'braintrust::factuality'\u001b[0m: \u001b[1;35mScoringResult\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33maggregated_results\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1;36m0.44999999999999996\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mscore_rows\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. **Expert Answer Details**: The expert answer states that Torchtune supports two precision formats: fp32 \u001b[0m\u001b[32m(\u001b[0m\u001b[32mfull-precision\u001b[0m\u001b[32m)\u001b[0m\u001b[32m and bfloat16 \u001b[0m\u001b[32m(\u001b[0m\u001b[32mhalf-precision\u001b[0m\u001b[32m)\u001b[0m\u001b[32m.\\n\\n2. **Submitted Answer Details**: The submitted answer mentions two precision formats: bf16 \u001b[0m\u001b[32m(\u001b[0m\u001b[32m16-bit floating-point format\u001b[0m\u001b[32m)\u001b[0m\u001b[32m and fp32 \u001b[0m\u001b[32m(\u001b[0m\u001b[32m32-bit floating-point format, also known as \"full-precision\"\u001b[0m\u001b[32m)\u001b[0m\u001b[32m. It also adds that Torchtune supports mixed-precision techniques.\\n\\n3. **Comparison of Precision Formats**:\\n - The expert answer uses \"bfloat16\" while the submitted answer uses \"bf16\". These are equivalent terms, as \"bf16\" is a common abbreviation for \"bfloat16\".\\n - Both answers mention \"fp32\" as a supported precision format.\\n\\n4. **Additional Information in Submission**: The submitted answer includes additional information about mixed-precision techniques, which is not mentioned in the expert answer.\\n\\n5. **Consistency Check**: The submitted answer includes all the information from the expert answer and adds more details about mixed-precision techniques. There is no conflict between the two answers.\\n\\nBased on the above analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.0\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'D'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation.\"\\n2. The submitted answer states that DoRA stands for \"Decoupled Orthogonal Random Adaptation.\"\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer regarding what DoRA stands for in torchtune.\\n5. Therefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m There is a disagreement between the submitted answer and the expert answer.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. **Expert Answer Analysis**: The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on the CPU and performing optimizer steps on the CPU. It also mentions the optional offloading of gradients to the CPU by setting `\u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m`.\\n\\n2. **Submitted Answer Analysis**: The submitted answer mentions offloading gradients and trainable parameters to the CPU, which allows for more efficient use of GPU memory. It specifies the use of `\u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m` to free gradients after device-to-host transfer. Additionally, it introduces the concept of using paged Adam with `\u001b[0m\u001b[32moptimizer_in_bwd\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m` to help reduce memory usage. It also notes that actual memory usage may vary depending on the use case and model architecture.\\n\\n3. **Comparison**:\\n - Both answers mention offloading gradients to the CPU using `\u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m`.\\n - The expert answer focuses on keeping optimizer states and performing optimizer steps on the CPU, while the submitted answer expands on this by mentioning trainable parameters and the use of paged Adam.\\n - The submitted answer provides additional context about memory usage variability and the use of paged Adam, which is not mentioned in the expert answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer as it includes all the information from the expert answer and adds more details about trainable parameters, paged Adam, and memory usage variability. There is no conflict between the two answers, and the additional information in the submitted answer is consistent with the expert answer.\\n\\nTherefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mB\u001b[0m\u001b[32m)\u001b[0m\u001b[32m The submitted answer is a superset of the expert answer and is fully consistent with it.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m\"1. **Expert Answer Analysis**: The expert answer provides a method to ensure only LoRA parameters are trainable by using torchtune's utility functions. It mentions fetching LoRA parameters with `get_adapter_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model\u001b[0m\u001b[32m)\u001b[0m\u001b[32m` and setting them as trainable with `set_trainable_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model, lora_params\u001b[0m\u001b[32m)\u001b[0m\u001b[32m`. It also notes that the LoRA recipe handles this automatically.\\n\\n2. **Submitted Answer Analysis**: The submitted answer provides a detailed example of how to ensure only LoRA parameters are trainable. It uses the `set_trainable_params` function from `torchtune.modules.peft.peft_utils` and provides a code example that includes loading a model and adapter, fetching adapter parameters, and setting them as trainable.\\n\\n3. **Comparison**:\\n - Both answers mention the use of `set_trainable_params` to set LoRA parameters as trainable.\\n - Both answers involve fetching LoRA parameters using a function \u001b[0m\u001b[32m(\u001b[0m\u001b[32m`get_adapter_params`\u001b[0m\u001b[32m)\u001b[0m\u001b[32m.\\n - The submitted answer provides additional context by including a code example and specifying the module path for the functions used.\\n - The expert answer mentions that the LoRA recipe handles this automatically, which is not explicitly stated in the submitted answer.\\n\\n4. **Conclusion**: The submitted answer is a superset of the expert answer. It includes all the information from the expert answer and adds more detail, such as a code example and specific module paths. There is no conflict between the two answers, and the additional information in the submitted answer is consistent with the expert answer.\"\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[1m)\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "eval_rows = []\n",
+ "for i, session_id in enumerate(search_agent.sessions):\n",
+ " session_response = client.agents.session.retrieve(agent_id=search_agent.agent_id, session_id=session_id)\n",
+ " for turn in session_response.turns:\n",
+ " eval_rows.append({\n",
+ " \"input_query\": examples[i][\"input_query\"],\n",
+ " \"expected_answer\": examples[i][\"expected_answer\"],\n",
+ " \"generated_answer\": turn.output_message.content,\n",
+ " })\n",
+ "\n",
+ "scoring_params = {\n",
+ " \"braintrust::factuality\": None,\n",
+ "}\n",
+ "scoring_response = client.scoring.score(\n",
+ " input_rows=eval_rows,\n",
+ " scoring_functions=scoring_params,\n",
+ ")\n",
+ "pprint(scoring_response)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### 3. RAG Agent\n",
+ "\n",
+ "Now, let's see how we can improve the agent's performance by adding a RAG tool that explicitly retrieves information from Torchtune's documentation. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 23,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_stack_client.types import Document\n",
+ "urls = [\n",
+ " \"memory_optimizations.rst\",\n",
+ " \"chat.rst\",\n",
+ " \"llama3.rst\",\n",
+ " \"datasets.rst\",\n",
+ " \"qat_finetune.rst\",\n",
+ " \"lora_finetune.rst\",\n",
+ "]\n",
+ "documents = [\n",
+ " Document(\n",
+ " document_id=f\"num-{i}\",\n",
+ " content=f\"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}\",\n",
+ " mime_type=\"text/plain\",\n",
+ " metadata={},\n",
+ " )\n",
+ " for i, url in enumerate(urls)\n",
+ "]\n",
+ "\n",
+ "vector_providers = [\n",
+ " provider for provider in client.providers.list() if provider.api == \"vector_io\"\n",
+ "]\n",
+ "selected_vector_provider = vector_providers[0]\n",
+ "vector_db_id = f\"test_vector_db_{uuid.uuid4()}\"\n",
+ "client.vector_dbs.register(\n",
+ " vector_db_id=vector_db_id,\n",
+ " embedding_model=\"all-MiniLM-L6-v2\",\n",
+ " embedding_dimension=384,\n",
+ " provider_id=selected_vector_provider.provider_id,\n",
+ ")\n",
+ "\n",
+ "client.tool_runtime.rag_tool.insert(\n",
+ " documents=documents,\n",
+ " vector_db_id=vector_db_id,\n",
+ " chunk_size_in_tokens=512,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 27,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "Question: What precision formats does torchtune support?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What precision formats does torchtune support?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: Torchtune supports the following precision formats:\n",
+ "\n",
+ "* bfloat16 (half-precision)\n",
+ "* fp32 (full-precision)\n",
+ "* int8 (integer 8-bit)\n",
+ "* int4 (integer 4-bit)\n",
+ "\n",
+ "Note that mixed-precision training is not currently supported in torchtune.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m Torchtune supports the following precision formats:\n",
+ "\n",
+ "* bfloat16 \u001b[1m(\u001b[0mhalf-precision\u001b[1m)\u001b[0m\n",
+ "* fp32 \u001b[1m(\u001b[0mfull-precision\u001b[1m)\u001b[0m\n",
+ "* int8 \u001b[1m(\u001b[0minteger \u001b[1;36m8\u001b[0m-bit\u001b[1m)\u001b[0m\n",
+ "* int4 \u001b[1m(\u001b[0minteger \u001b[1;36m4\u001b[0m-bit\u001b[1m)\u001b[0m\n",
+ "\n",
+ "Note that mixed-precision training is not currently supported in torchtune.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: What does DoRA stand for in torchtune?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m What does DoRA stand for in torchtune?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: DoRA stands for \"Decoupled Orthogonal Random Axes\" in the context of the Torchtune project.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m DoRA stands for \u001b[32m\"Decoupled Orthogonal Random Axes\"\u001b[0m in the context of the Torchtune project.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How does the CPUOffloadOptimizer reduce GPU memory usage?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How does the CPUOffloadOptimizer reduce GPU memory usage?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: The CPUOffloadOptimizer reduces GPU memory usage by offloading optimizer states and gradients to CPU,\n",
+ "thus reducing the memory usage on the GPU. This is especially useful when training large models or when using \n",
+ "stateful optimizers, as it can significantly reduce the memory requirements. However, it may come at the cost of \n",
+ "increased CPU RAM usage and potentially slower training speeds.\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m The CPUOffloadOptimizer reduces GPU memory usage by offloading optimizer states and gradients to CPU,\n",
+ "thus reducing the memory usage on the GPU. This is especially useful when training large models or when using \n",
+ "stateful optimizers, as it can significantly reduce the memory requirements. However, it may come at the cost of \n",
+ "increased CPU RAM usage and potentially slower training speeds.\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Question: How do I ensure only LoRA parameters are trainable when fine-tuning?\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;36mQuestion:\u001b[0m How do I ensure only LoRA parameters are trainable when fine-tuning?\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "Agent Answer: To ensure only LoRA parameters are trainable when fine-tuning, you can use the `get_adapter_params` \n",
+ "and `set_trainable_params` functions from `torchtune.modules.peft.peft_utils`. \n",
+ "\n",
+ "Here is how to do it:\n",
+ "\n",
+ "```python\n",
+ "from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\n",
+ "\n",
+ "# Fetch all params from the model that are associated with LoRA.\n",
+ "lora_params = get_adapter_params(lora_model)\n",
+ "\n",
+ "# Set requires_grad=True on lora_params, and requires_grad=False on all others.\n",
+ "set_trainable_params(lora_model, lora_params)\n",
+ "```\n",
+ "
\n"
+ ],
+ "text/plain": [
+ "\u001b[1;33mAgent Answer:\u001b[0m To ensure only LoRA parameters are trainable when fine-tuning, you can use the `get_adapter_params` \n",
+ "and `set_trainable_params` functions from `torchtune.modules.peft.peft_utils`. \n",
+ "\n",
+ "Here is how to do it:\n",
+ "\n",
+ "```python\n",
+ "from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\n",
+ "\n",
+ "# Fetch all params from the model that are associated with LoRA.\n",
+ "lora_params = \u001b[1;35mget_adapter_params\u001b[0m\u001b[1m(\u001b[0mlora_model\u001b[1m)\u001b[0m\n",
+ "\n",
+ "# Set \u001b[33mrequires_grad\u001b[0m=\u001b[3;92mTrue\u001b[0m on lora_params, and \u001b[33mrequires_grad\u001b[0m=\u001b[3;91mFalse\u001b[0m on all others.\n",
+ "\u001b[1;35mset_trainable_params\u001b[0m\u001b[1m(\u001b[0mlora_model, lora_params\u001b[1m)\u001b[0m\n",
+ "```\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "rag_agent = Agent(\n",
+ " client,\n",
+ " model=MODEL_ID,\n",
+ " instructions=\"You are a helpful assistant that can answer questions about the Torchtune project. You should always use the RAG tool to answer questions.\",\n",
+ " tools=[{\n",
+ " \"name\": \"builtin::rag\",\n",
+ " \"args\": {\"vector_db_ids\": [vector_db_id]},\n",
+ " }],\n",
+ ")\n",
+ "\n",
+ "for example in examples:\n",
+ " rag_session_id = rag_agent.create_session(session_name=f\"rag_session_{uuid.uuid4()}\")\n",
+ " response = rag_agent.create_turn(\n",
+ " messages=[\n",
+ " {\n",
+ " \"role\": \"user\",\n",
+ " \"content\": example[\"input_query\"]\n",
+ " }\n",
+ " ],\n",
+ " session_id=rag_session_id,\n",
+ " stream=False\n",
+ " )\n",
+ " rich.print(f\"[bold cyan]Question:[/bold cyan] {example['input_query']}\")\n",
+ " rich.print(f\"[bold yellow]Agent Answer:[/bold yellow] {response.output_message.content}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 28,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "ScoringScoreResponse(\n",
+ "β results={\n",
+ "β β 'braintrust::factuality': ScoringResult(\n",
+ "β β β aggregated_results={'average': {'average': 0.3}},\n",
+ "β β β score_rows=[\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.0,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'D',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that Torchtune supports two precision formats: fp32 and bfloat16.\\n2. The submitted answer lists four precision formats: bfloat16, fp32, int8, and int4.\\n3. The submitted answer includes the two formats mentioned by the expert (bfloat16 and fp32), but also adds int8 and int4, which are not mentioned by the expert.\\n4. The submitted answer also states that mixed-precision training is not supported, which is not addressed in the expert answer.\\n5. Since the submitted answer includes additional precision formats (int8 and int4) that are not mentioned by the expert, there is a factual disagreement between the two answers regarding the supported precision formats.\\n6. Therefore, the correct choice is (D) There is a disagreement between the submitted answer and the expert answer.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.0,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'D',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation.\"\\n2. The submitted answer states that DoRA stands for \"Decoupled Orthogonal Random Axes.\"\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer.\\n5. Therefore, the correct choice is (D) There is a disagreement between the submitted answer and the expert answer.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': '1. The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on CPU and performing optimizer steps on CPU. It also mentions the optional offloading of gradients to CPU using offload_gradients=True.\\n2. The submitted answer states that the CPUOffloadOptimizer reduces GPU memory usage by offloading optimizer states and gradients to CPU. It also mentions that this is useful for large models or stateful optimizers and notes potential downsides like increased CPU RAM usage and slower training speeds.\\n3. The submitted answer includes all the points mentioned in the expert answer: offloading optimizer states and optionally gradients to CPU.\\n4. Additionally, the submitted answer provides extra context about the usefulness for large models and potential downsides, which are not mentioned in the expert answer.\\n5. There is no factual disagreement between the two answers; the submitted answer simply provides more information.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.'\n",
+ "β β β β β }\n",
+ "β β β β },\n",
+ "β β β β {\n",
+ "β β β β β 'score': 0.6,\n",
+ "β β β β β 'metadata': {\n",
+ "β β β β β β 'choice': 'B',\n",
+ "β β β β β β 'rationale': \"1. **Identify the core content of both answers:**\\n - The expert answer explains how to set only LoRA parameters as trainable using torchtune's utility functions by fetching all LoRA parameters with `get_adapter_params(lora_model)` and setting them as trainable with `set_trainable_params(lora_model, lora_params)`. It also mentions that the LoRA recipe handles this automatically.\\n - The submitted answer provides a similar explanation, detailing the use of `get_adapter_params` and `set_trainable_params` from `torchtune.modules.peft.peft_utils` to ensure only LoRA parameters are trainable. It includes a code snippet demonstrating the process.\\n\\n2. **Compare the factual content:**\\n - Both answers describe the same process of fetching LoRA parameters and setting them as trainable using the same functions.\\n - The submitted answer includes additional details such as the import statement and a code snippet, which are not present in the expert answer.\\n - The expert answer mentions that the LoRA recipe handles this automatically, which is not mentioned in the submission.\\n\\n3. **Determine the relationship between the answers:**\\n - The submitted answer is a superset of the expert answer because it includes all the information provided by the expert and adds more details, such as the import statement and code snippet.\\n - There is no conflict between the two answers; the submission expands on the expert's explanation.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.\"\n",
+ "β β β β β }\n",
+ "β β β β }\n",
+ "β β β ]\n",
+ "β β )\n",
+ "β }\n",
+ ")\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1;35mScoringScoreResponse\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[33mresults\u001b[0m=\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[32m'braintrust::factuality'\u001b[0m: \u001b[1;35mScoringResult\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33maggregated_results\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'average'\u001b[0m: \u001b[1;36m0.3\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mscore_rows\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.0\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'D'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that Torchtune supports two precision formats: fp32 and bfloat16.\\n2. The submitted answer lists four precision formats: bfloat16, fp32, int8, and int4.\\n3. The submitted answer includes the two formats mentioned by the expert \u001b[0m\u001b[32m(\u001b[0m\u001b[32mbfloat16 and fp32\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, but also adds int8 and int4, which are not mentioned by the expert.\\n4. The submitted answer also states that mixed-precision training is not supported, which is not addressed in the expert answer.\\n5. Since the submitted answer includes additional precision formats \u001b[0m\u001b[32m(\u001b[0m\u001b[32mint8 and int4\u001b[0m\u001b[32m)\u001b[0m\u001b[32m that are not mentioned by the expert, there is a factual disagreement between the two answers regarding the supported precision formats.\\n6. Therefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m There is a disagreement between the submitted answer and the expert answer.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.0\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'D'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that DoRA stands for \"Weight-Decomposed Low-Rank Adaptation.\"\\n2. The submitted answer states that DoRA stands for \"Decoupled Orthogonal Random Axes.\"\\n3. The two answers provide completely different expansions for the acronym DoRA.\\n4. Since the expansions are different, there is a clear disagreement between the submitted answer and the expert answer.\\n5. Therefore, the correct choice is \u001b[0m\u001b[32m(\u001b[0m\u001b[32mD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m There is a disagreement between the submitted answer and the expert answer.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m'1. The expert answer states that the CPUOffloadOptimizer reduces GPU memory usage by keeping optimizer states on CPU and performing optimizer steps on CPU. It also mentions the optional offloading of gradients to CPU using \u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m.\\n2. The submitted answer states that the CPUOffloadOptimizer reduces GPU memory usage by offloading optimizer states and gradients to CPU. It also mentions that this is useful for large models or stateful optimizers and notes potential downsides like increased CPU RAM usage and slower training speeds.\\n3. The submitted answer includes all the points mentioned in the expert answer: offloading optimizer states and optionally gradients to CPU.\\n4. Additionally, the submitted answer provides extra context about the usefulness for large models and potential downsides, which are not mentioned in the expert answer.\\n5. There is no factual disagreement between the two answers; the submitted answer simply provides more information.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'score'\u001b[0m: \u001b[1;36m0.6\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'choice'\u001b[0m: \u001b[32m'B'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[32m'rationale'\u001b[0m: \u001b[32m\"1. **Identify the core content of both answers:**\\n - The expert answer explains how to set only LoRA parameters as trainable using torchtune's utility functions by fetching all LoRA parameters with `get_adapter_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model\u001b[0m\u001b[32m)\u001b[0m\u001b[32m` and setting them as trainable with `set_trainable_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model, lora_params\u001b[0m\u001b[32m)\u001b[0m\u001b[32m`. It also mentions that the LoRA recipe handles this automatically.\\n - The submitted answer provides a similar explanation, detailing the use of `get_adapter_params` and `set_trainable_params` from `torchtune.modules.peft.peft_utils` to ensure only LoRA parameters are trainable. It includes a code snippet demonstrating the process.\\n\\n2. **Compare the factual content:**\\n - Both answers describe the same process of fetching LoRA parameters and setting them as trainable using the same functions.\\n - The submitted answer includes additional details such as the import statement and a code snippet, which are not present in the expert answer.\\n - The expert answer mentions that the LoRA recipe handles this automatically, which is not mentioned in the submission.\\n\\n3. **Determine the relationship between the answers:**\\n - The submitted answer is a superset of the expert answer because it includes all the information provided by the expert and adds more details, such as the import statement and code snippet.\\n - There is no conflict between the two answers; the submission expands on the expert's explanation.\\n\\nBased on this analysis, the submitted answer is a superset of the expert answer and is fully consistent with it.\"\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1m}\u001b[0m\n",
+ "\u001b[1m)\u001b[0m\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "eval_rows = []\n",
+ "for i, session_id in enumerate(rag_agent.sessions):\n",
+ " session_response = client.agents.session.retrieve(agent_id=rag_agent.agent_id, session_id=session_id)\n",
+ " for turn in session_response.turns:\n",
+ " eval_rows.append({\n",
+ " \"input_query\": examples[i][\"input_query\"],\n",
+ " \"expected_answer\": examples[i][\"expected_answer\"],\n",
+ " \"generated_answer\": turn.output_message.content,\n",
+ " })\n",
+ "\n",
+ "scoring_params = {\n",
+ " \"braintrust::factuality\": None,\n",
+ "}\n",
+ "scoring_response = client.scoring.score(\n",
+ " input_rows=eval_rows,\n",
+ " scoring_functions=scoring_params,\n",
+ ")\n",
+ "pprint(scoring_response)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Deep dive into RAG Tool Performance\n",
+ "- Now, let's take a closer look at how the RAG tool is doing, specifically on the second example where the agent's answer is not correct on identifying what DoRA stands for. \n",
+ "- Notice that the issue lies with the retrieval step, where the retrieved document is not relevant to the question. "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 29,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "[\n",
+ "β Turn(\n",
+ "β β input_messages=[UserMessage(content='What does DoRA stand for in torchtune?', role='user', context=None)],\n",
+ "β β output_message=CompletionMessage(\n",
+ "β β β content='DoRA stands for \"Decoupled Orthogonal Random Axes\" in the context of the Torchtune project.',\n",
+ "β β β role='assistant',\n",
+ "β β β stop_reason='end_of_turn',\n",
+ "β β β tool_calls=[]\n",
+ "β β ),\n",
+ "β β session_id='b5b5b9c5-1f14-404a-9677-cdb413b9f328',\n",
+ "β β started_at=datetime.datetime(2025, 3, 7, 10, 35, 24, 235903, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=57600))),\n",
+ "β β steps=[\n",
+ "β β β InferenceStep(\n",
+ "β β β β api_model_response=CompletionMessage(\n",
+ "β β β β β content='',\n",
+ "β β β β β role='assistant',\n",
+ "β β β β β stop_reason='end_of_turn',\n",
+ "β β β β β tool_calls=[\n",
+ "β β β β β β ToolCall(\n",
+ "β β β β β β β arguments={'query': 'DoRA meaning in Torchtune'},\n",
+ "β β β β β β β call_id='c2c088b9-cf2f-41b5-a050-dd5743112f48',\n",
+ "β β β β β β β tool_name='knowledge_search'\n",
+ "β β β β β β )\n",
+ "β β β β β ]\n",
+ "β β β β ),\n",
+ "β β β β step_id='27ba55cd-0252-4cff-8141-129b3b8dd021',\n",
+ "β β β β step_type='inference',\n",
+ "β β β β turn_id='bb111412-e2e9-40ca-9cd2-87df200807ab',\n",
+ "β β β β completed_at=datetime.datetime(2025, 3, 7, 10, 35, 26, 226185, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β started_at=datetime.datetime(2025, 3, 7, 10, 35, 24, 236359, tzinfo=TzInfo(-08:00))\n",
+ "β β β ),\n",
+ "β β β ToolExecutionStep(\n",
+ "β β β β step_id='e7da6bb1-a704-4a2e-9954-5d54d8a1fc5d',\n",
+ "β β β β step_type='tool_execution',\n",
+ "β β β β tool_calls=[\n",
+ "β β β β β ToolCall(\n",
+ "β β β β β β arguments={'query': 'DoRA meaning in Torchtune'},\n",
+ "β β β β β β call_id='c2c088b9-cf2f-41b5-a050-dd5743112f48',\n",
+ "β β β β β β tool_name='knowledge_search'\n",
+ "β β β β β )\n",
+ "β β β β ],\n",
+ "β β β β tool_responses=[\n",
+ "β β β β β ToolResponse(\n",
+ "β β β β β β call_id='c2c088b9-cf2f-41b5-a050-dd5743112f48',\n",
+ "β β β β β β content=[\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text='knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n',\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text='Result 1:\\nDocument_id:num-0\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA <glossary_lora>` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\"q_proj\",\"k_proj\",\"v_proj\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\"q_proj\", \"k_proj\", \"v_proj\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we\\'ve enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP <https://pytorch.org/docs/stable/fsdp.html>`.\\n.. .. _glossary_fsdp2:\\n\\n',\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text='Result 2:\\nDocument_id:num-1\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\"sharegpt\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\"json\",\\n data_files=\"data/my_data.json\",\\n split=\"train\",\\n conversation_column=\"dialogue\",\\n conversation_style=\"sharegpt\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset <https://huggingface.co/docs/datasets/v2.20.0/en/package_reference/loading_methods#datasets.load_dataset>`_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we\\'re fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral\\'s :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations <https://\\n',\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text=\"Result 3:\\nDocument_id:num-5\\nContent: .. _lora_finetune_label:\\n\\n============================\\nFine-Tuning Llama2 with LoRA\\n============================\\n\\nThis guide will teach you about `LoRA <https://arxiv.org/abs/2106.09685>`_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune<lora_recipe_label>`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune<overview_label>`\\n * Make sure to :ref:`install torchtune<install_label>`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights<download_llama_label>`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA <https://arxiv.org/abs/2106.09685>`_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank <https://en.wikipedia.org/wiki/Rank_(linear_algebra)>`_\\n and discussion of `low-rank approximations <https://en.wikipedia.org/wiki/Low-rank_approximation>`_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW <https://py\\n\",\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text='Result 4:\\nDocument_id:num-0\\nContent: use the :class:`torch.optim.AdamW` optimizer with ``fused=True`` as the base optimizer. For example, to use this optimizer to offload\\nboth optimizer states and gradients to CPU:\\n\\n.. code-block:: bash\\n\\n tune run <RECIPE> --config <CONFIG> \\\\\\n optimizer=optimizer=torchao.prototype.low_bit_optim.CPUOffloadOptimizer \\\\\\n optimizer.offload_gradients=True \\\\\\n lr=4e-5\\n\\n\\nor by directly :ref:`modifying a config file<config_tutorial_label>`:\\n\\n.. code-block:: yaml\\n\\n optimizer:\\n _component_: torchao.prototype.low_bit_optim.CPUOffloadOptimizer\\n offload_gradients: True\\n # additional key-word arguments can be passed to torch.optim.AdamW\\n lr: 4e-5\\n\\nor using it directly in your code, which allows you to change the base optimizer:\\n\\n.. code-block:: python\\n\\n from torchao.prototype.low_bit_optim import CPUOffloadOptimizer\\n from torch.optim import Adam\\n\\n optimizer = CPUOffloadOptimizer(\\n model.parameters(), # your model here\\n Adam,\\n lr=1e-5,\\n fused=True\\n )\\n\\nSome helpful hints from the ``torchao`` `CPUOffloadOptimizer page <https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload>`_:\\n\\n* The CPU optimizer step is often the bottleneck when optimizer CPU offload is used. To minimize the slowdown, it is recommended to (1) use full ``bf16`` training so that parameters, gradients, and optimizer states are in ``bf16``; and (2) give GPU more work per optimizer step to amortize the offloading time (e.g. larger batch size with activation checkpointing, gradient accumulation).\\n* Gradient accumulation should always be set to 1 when ``offload_gradients=True``, as gradients are cleared on GPU every backward pass.\\n* This optimizer works by keeping a copy of parameters and pre-allocating gradient memory on CPU. Therefore, expect your RAM usage to increase by 4x model size.\\n* This optimizer is only supported for single-device recipes. To use CPU-offloading in distributed recipes, use ``fsdp_cpu_offload=True`` instead. See :class:`torch.distributed.fsdp.FullyShardedDataParallel` for more details and `FSDP1 vs FSDP2 <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp\\n',\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(\n",
+ "β β β β β β β β text='Result 5:\\nDocument_id:num-5\\nContent: from our Llama2\\nmodel without any wrappers or custom checkpoint conversion logic.\\n\\n.. code-block:: python\\n\\n # Assuming that base_model already has the pretrained Llama2 weights,\\n # this will directly load them into your LoRA model without any conversion necessary.\\n lora_model.load_state_dict(base_model.state_dict(), strict=False)\\n\\n.. note::\\n Whenever loading weights with :code:`strict=False`, you should verify that any missing or extra keys in\\n the loaded :code:`state_dict` are as expected. torchtune\\'s LoRA recipes do this by default via\\n :func:`validate_missing_and_unexpected_for_lora() <torchtune.modules.peft.validate_missing_and_unexpected_for_lora>`.\\n\\nOnce we\\'ve loaded the base model weights, we also want to set only LoRA parameters to trainable.\\n\\n.. _setting_trainable_params:\\n\\n.. code-block:: python\\n\\n from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\\n\\n # Fetch all params from the model that are associated with LoRA.\\n lora_params = get_adapter_params(lora_model)\\n\\n # Set requires_grad=True on lora_params, and requires_grad=False on all others.\\n set_trainable_params(lora_model, lora_params)\\n\\n # Print the total number of parameters\\n total_params = sum([p.numel() for p in lora_model.parameters()])\\n trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])\\n print(\\n f\"\"\"\\n {total_params} total params,\\n {trainable_params}\" trainable params,\\n {(100.0 * trainable_params / total_params):.2f}% of all params are trainable.\\n \"\"\"\\n )\\n\\n 6742609920 total params,\\n 4194304 trainable params,\\n 0.06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here<lora_recipe_label>`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune\\'s `LoRA recipe <https://github.com/pytorch/torchtune/blob/48626d19d2108f92\\n',\n",
+ "β β β β β β β β type='text'\n",
+ "β β β β β β β ),\n",
+ "β β β β β β β TextContentItem(text='END of knowledge_search tool results.\\n', type='text')\n",
+ "β β β β β β ],\n",
+ "β β β β β β tool_name='knowledge_search',\n",
+ "β β β β β β metadata={'document_ids': ['num-0', 'num-1', 'num-5', 'num-0', 'num-5']}\n",
+ "β β β β β )\n",
+ "β β β β ],\n",
+ "β β β β turn_id='bb111412-e2e9-40ca-9cd2-87df200807ab',\n",
+ "β β β β completed_at=datetime.datetime(2025, 3, 7, 10, 35, 26, 339563, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β started_at=datetime.datetime(2025, 3, 7, 10, 35, 26, 264752, tzinfo=TzInfo(-08:00))\n",
+ "β β β ),\n",
+ "β β β InferenceStep(\n",
+ "β β β β api_model_response=CompletionMessage(\n",
+ "β β β β β content='DoRA stands for \"Decoupled Orthogonal Random Axes\" in the context of the Torchtune project.',\n",
+ "β β β β β role='assistant',\n",
+ "β β β β β stop_reason='end_of_turn',\n",
+ "β β β β β tool_calls=[]\n",
+ "β β β β ),\n",
+ "β β β β step_id='400e49e1-f33e-41da-b22a-f1d2338a27c8',\n",
+ "β β β β step_type='inference',\n",
+ "β β β β turn_id='bb111412-e2e9-40ca-9cd2-87df200807ab',\n",
+ "β β β β completed_at=datetime.datetime(2025, 3, 7, 10, 35, 27, 281430, tzinfo=TzInfo(-08:00)),\n",
+ "β β β β started_at=datetime.datetime(2025, 3, 7, 10, 35, 26, 351029, tzinfo=TzInfo(-08:00))\n",
+ "β β β )\n",
+ "β β ],\n",
+ "β β turn_id='bb111412-e2e9-40ca-9cd2-87df200807ab',\n",
+ "β β completed_at=datetime.datetime(2025, 3, 7, 10, 35, 27, 294253, tzinfo=TzInfo(-08:00)),\n",
+ "β β output_attachments=[]\n",
+ "β )\n",
+ "]\n",
+ "\n"
+ ],
+ "text/plain": [
+ "\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ \u001b[0m\u001b[1;35mTurn\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[33minput_messages\u001b[0m=\u001b[1m[\u001b[0m\u001b[1;35mUserMessage\u001b[0m\u001b[1m(\u001b[0m\u001b[33mcontent\u001b[0m=\u001b[32m'What does DoRA stand for in torchtune?'\u001b[0m, \u001b[33mrole\u001b[0m=\u001b[32m'user'\u001b[0m, \u001b[33mcontext\u001b[0m=\u001b[3;35mNone\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[33moutput_message\u001b[0m=\u001b[1;35mCompletionMessage\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mcontent\u001b[0m=\u001b[32m'DoRA stands for \"Decoupled Orthogonal Random Axes\" in the context of the Torchtune project.'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mrole\u001b[0m=\u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mstop_reason\u001b[0m=\u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[33mtool_calls\u001b[0m=\u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[33msession_id\u001b[0m=\u001b[32m'b5b5b9c5-1f14-404a-9677-cdb413b9f328'\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[33mstarted_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m7\u001b[0m, \u001b[1;36m10\u001b[0m, \u001b[1;36m35\u001b[0m, \u001b[1;36m24\u001b[0m, \u001b[1;36m235903\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.timezone\u001b[0m\u001b[1m(\u001b[0m\u001b[1;35mdatetime.timedelta\u001b[0m\u001b[1m(\u001b[0m\u001b[33mdays\u001b[0m=\u001b[1;36m-1\u001b[0m, \u001b[33mseconds\u001b[0m=\u001b[1;36m57600\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β \u001b[0m\u001b[33msteps\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1;35mInferenceStep\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mapi_model_response\u001b[0m=\u001b[1;35mCompletionMessage\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[33mcontent\u001b[0m=\u001b[32m''\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[33mrole\u001b[0m=\u001b[32m'assistant'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[33mstop_reason\u001b[0m=\u001b[32m'end_of_turn'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[33mtool_calls\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[1;35mToolCall\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[33marguments\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'query'\u001b[0m: \u001b[32m'DoRA meaning in Torchtune'\u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[33mcall_id\u001b[0m=\u001b[32m'c2c088b9-cf2f-41b5-a050-dd5743112f48'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[33mtool_name\u001b[0m=\u001b[32m'knowledge_search'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m]\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mstep_id\u001b[0m=\u001b[32m'27ba55cd-0252-4cff-8141-129b3b8dd021'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mstep_type\u001b[0m=\u001b[32m'inference'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mturn_id\u001b[0m=\u001b[32m'bb111412-e2e9-40ca-9cd2-87df200807ab'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mcompleted_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m7\u001b[0m, \u001b[1;36m10\u001b[0m, \u001b[1;36m35\u001b[0m, \u001b[1;36m26\u001b[0m, \u001b[1;36m226185\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mstarted_at\u001b[0m=\u001b[1;35mdatetime\u001b[0m\u001b[1;35m.datetime\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m2025\u001b[0m, \u001b[1;36m3\u001b[0m, \u001b[1;36m7\u001b[0m, \u001b[1;36m10\u001b[0m, \u001b[1;36m35\u001b[0m, \u001b[1;36m24\u001b[0m, \u001b[1;36m236359\u001b[0m, \u001b[33mtzinfo\u001b[0m=\u001b[1;35mTzInfo\u001b[0m\u001b[1m(\u001b[0m\u001b[1;36m-08\u001b[0m:\u001b[1;36m00\u001b[0m\u001b[1m)\u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β \u001b[0m\u001b[1;35mToolExecutionStep\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mstep_id\u001b[0m=\u001b[32m'e7da6bb1-a704-4a2e-9954-5d54d8a1fc5d'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mstep_type\u001b[0m=\u001b[32m'tool_execution'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mtool_calls\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1;35mToolCall\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[33marguments\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'query'\u001b[0m: \u001b[32m'DoRA meaning in Torchtune'\u001b[0m\u001b[1m}\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[33mcall_id\u001b[0m=\u001b[32m'c2c088b9-cf2f-41b5-a050-dd5743112f48'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[33mtool_name\u001b[0m=\u001b[32m'knowledge_search'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1m)\u001b[0m\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[1m]\u001b[0m,\n",
+ "\u001b[2;32mβ β β β \u001b[0m\u001b[33mtool_responses\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β β \u001b[0m\u001b[1;35mToolResponse\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[33mcall_id\u001b[0m=\u001b[32m'c2c088b9-cf2f-41b5-a050-dd5743112f48'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β \u001b[0m\u001b[33mcontent\u001b[0m=\u001b[1m[\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[1;35mTextContentItem\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β β \u001b[0m\u001b[33mtext\u001b[0m=\u001b[32m'knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n'\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β β β \u001b[0m\u001b[33mtype\u001b[0m=\u001b[32m'text'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[1m)\u001b[0m,\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[1;35mTextContentItem\u001b[0m\u001b[1m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β β \u001b[0m\u001b[33mtext\u001b[0m=\u001b[32m'Result 1:\\nDocument_id:num-0\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``\u001b[0m\u001b[32muse_dora\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.\u001b[0m\u001b[32muse_dora\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA \u001b[0m\u001b[32m<\u001b[0m\u001b[32mglossary_lora\u001b[0m\u001b[32m>` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``\u001b[0m\u001b[32mquantize\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.\u001b[0m\u001b[32mapply_lora_to_mlp\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m \\\\\\n model.\u001b[0m\u001b[32mlora_attn_modules\u001b[0m\u001b[32m=\u001b[0m\u001b[32m[\u001b[0m\u001b[32m\"q_proj\",\"k_proj\",\"v_proj\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m \\\\\\n model.\u001b[0m\u001b[32mlora_rank\u001b[0m\u001b[32m=\u001b[0m\u001b[32m16\u001b[0m\u001b[32m \\\\\\n model.\u001b[0m\u001b[32mlora_alpha\u001b[0m\u001b[32m=\u001b[0m\u001b[32m32\u001b[0m\u001b[32m \\\\\\n model.\u001b[0m\u001b[32muse_dora\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m \\\\\\n model.\u001b[0m\u001b[32mquantize_base\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: \u001b[0m\u001b[32m[\u001b[0m\u001b[32m\"q_proj\", \"k_proj\", \"v_proj\"\u001b[0m\u001b[32m]\u001b[0m\u001b[32m\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we\\'ve enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``\u001b[0m\u001b[32muse_dora\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel \u001b[0m\u001b[32m(\u001b[0m\u001b[32mFSDP\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n'\u001b[0m\u001b[39m,\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β β \u001b[0m\u001b[33mtype\u001b[0m\u001b[39m=\u001b[0m\u001b[32m'text'\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[1;39m)\u001b[0m\u001b[39m,\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β \u001b[0m\u001b[1;35mTextContentItem\u001b[0m\u001b[1;39m(\u001b[0m\n",
+ "\u001b[2;32mβ β β β β β β β \u001b[0m\u001b[33mtext\u001b[0m\u001b[39m=\u001b[0m\u001b[32m'Result 2:\\nDocument_id:num-1\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\"sharegpt\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer\u001b[0m\u001b[32m(\u001b[0m\u001b[32m\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\"\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n ds = chat_dataset\u001b[0m\u001b[32m(\u001b[0m\u001b[32m\\n \u001b[0m\u001b[32mtokenizer\u001b[0m\u001b[32m=\u001b[0m\u001b[32mtokenizer\u001b[0m\u001b[32m,\\n \u001b[0m\u001b[32msource\u001b[0m\u001b[32m=\u001b[0m\u001b[32m\"json\"\u001b[0m\u001b[32m,\\n \u001b[0m\u001b[32mdata_files\u001b[0m\u001b[32m=\u001b[0m\u001b[32m\"data\u001b[0m\u001b[32m/my_data.json\",\\n \u001b[0m\u001b[32msplit\u001b[0m\u001b[32m=\u001b[0m\u001b[32m\"train\"\u001b[0m\u001b[32m,\\n \u001b[0m\u001b[32mconversation_column\u001b[0m\u001b[32m=\u001b[0m\u001b[32m\"dialogue\"\u001b[0m\u001b[32m,\\n \u001b[0m\u001b[32mconversation_style\u001b[0m\u001b[32m=\u001b[0m\u001b[32m\"sharegpt\"\u001b[0m\u001b[32m,\\n \u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we\\'re fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral\\'s :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default \u001b[0m\u001b[32m(\u001b[0m\u001b[32m:class:`~torchtune.models.mistral.MistralChatTemplate`\u001b[0m\u001b[32m)\u001b[0m\u001b[32m to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA \u001b[0m\u001b[32m(\u001b[0m\u001b[32mas opposed to finetuning all model parameters\u001b[0m\u001b[32m)\u001b[0m\u001b[32m,\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW --config \\\\\\n \u001b[0m\u001b[32moptimizer\u001b[0m\u001b[32m=\u001b[0m\u001b[32moptimizer\u001b[0m\u001b[32m=torchao.prototype.low_bit_optim.CPUOffloadOptimizer \\\\\\n optimizer.\u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m \\\\\\n \u001b[0m\u001b[32mlr\u001b[0m\u001b[32m=\u001b[0m\u001b[32m4e\u001b[0m\u001b[32m-5\\n\\n\\nor by directly :ref:`modifying a config file`:\\n\\n.. code-block:: yaml\\n\\n optimizer:\\n _component_: torchao.prototype.low_bit_optim.CPUOffloadOptimizer\\n offload_gradients: True\\n # additional key-word arguments can be passed to torch.optim.AdamW\\n lr: 4e-5\\n\\nor using it directly in your code, which allows you to change the base optimizer:\\n\\n.. code-block:: python\\n\\n from torchao.prototype.low_bit_optim import CPUOffloadOptimizer\\n from torch.optim import Adam\\n\\n optimizer = CPUOffloadOptimizer\u001b[0m\u001b[32m(\u001b[0m\u001b[32m\\n model.parameters\u001b[0m\u001b[32m(\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, # your model here\\n Adam,\\n \u001b[0m\u001b[32mlr\u001b[0m\u001b[32m=\u001b[0m\u001b[32m1e\u001b[0m\u001b[32m-5,\\n \u001b[0m\u001b[32mfused\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m\\n \u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n\\nSome helpful hints from the ``torchao`` `CPUOffloadOptimizer page `_:\\n\\n* The CPU optimizer step is often the bottleneck when optimizer CPU offload is used. To minimize the slowdown, it is recommended to \u001b[0m\u001b[32m(\u001b[0m\u001b[32m1\u001b[0m\u001b[32m)\u001b[0m\u001b[32m use full ``bf16`` training so that parameters, gradients, and optimizer states are in ``bf16``; and \u001b[0m\u001b[32m(\u001b[0m\u001b[32m2\u001b[0m\u001b[32m)\u001b[0m\u001b[32m give GPU more work per optimizer step to amortize the offloading time \u001b[0m\u001b[32m(\u001b[0m\u001b[32me.g. larger batch size with activation checkpointing, gradient accumulation\u001b[0m\u001b[32m)\u001b[0m\u001b[32m.\\n* Gradient accumulation should always be set to 1 when ``\u001b[0m\u001b[32moffload_gradients\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m``, as gradients are cleared on GPU every backward pass.\\n* This optimizer works by keeping a copy of parameters and pre-allocating gradient memory on CPU. Therefore, expect your RAM usage to increase by 4x model size.\\n* This optimizer is only supported for single-device recipes. To use CPU-offloading in distributed recipes, use ``\u001b[0m\u001b[32mfsdp_cpu_offload\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m`` instead. See :class:`torch.distributed.fsdp.FullyShardedDataParallel` for more details and `FSDP1 vs FSDP2 `.\\n\\nOnce we\\'ve loaded the base model weights, we also want to set only LoRA parameters to trainable.\\n\\n.. _setting_trainable_params:\\n\\n.. code-block:: python\\n\\n from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\\n\\n # Fetch all params from the model that are associated with LoRA.\\n lora_params = get_adapter_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n\\n # Set \u001b[0m\u001b[32mrequires_grad\u001b[0m\u001b[32m=\u001b[0m\u001b[32mTrue\u001b[0m\u001b[32m on lora_params, and \u001b[0m\u001b[32mrequires_grad\u001b[0m\u001b[32m=\u001b[0m\u001b[32mFalse\u001b[0m\u001b[32m on all others.\\n set_trainable_params\u001b[0m\u001b[32m(\u001b[0m\u001b[32mlora_model, lora_params\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n\\n # Print the total number of parameters\\n total_params = sum\u001b[0m\u001b[32m(\u001b[0m\u001b[32m[\u001b[0m\u001b[32mp.numel\u001b[0m\u001b[32m(\u001b[0m\u001b[32m)\u001b[0m\u001b[32m for p in lora_model.parameters\u001b[0m\u001b[32m(\u001b[0m\u001b[32m)\u001b[0m\u001b[32m]\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n trainable_params = sum\u001b[0m\u001b[32m(\u001b[0m\u001b[32m[\u001b[0m\u001b[32mp.numel\u001b[0m\u001b[32m(\u001b[0m\u001b[32m)\u001b[0m\u001b[32m for p in lora_model.parameters\u001b[0m\u001b[32m(\u001b[0m\u001b[32m)\u001b[0m\u001b[32m if p.requires_grad\u001b[0m\u001b[32m]\u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n print\u001b[0m\u001b[32m(\u001b[0m\u001b[32m\\n f\"\"\"\\n \u001b[0m\u001b[32m{\u001b[0m\u001b[32mtotal_params\u001b[0m\u001b[32m}\u001b[0m\u001b[32m total params,\\n \u001b[0m\u001b[32m{\u001b[0m\u001b[32mtrainable_params\u001b[0m\u001b[32m}\u001b[0m\u001b[32m\" trainable params,\\n \u001b[0m\u001b[32m{\u001b[0m\u001b[32m(\u001b[0m\u001b[32m100.0 * trainable_params / total_params\u001b[0m\u001b[32m)\u001b[0m\u001b[32m:.2f\u001b[0m\u001b[32m}\u001b[0m\u001b[32m% of all params are trainable.\\n \"\"\"\\n \u001b[0m\u001b[32m)\u001b[0m\u001b[32m\\n\\n 6742609920 total params,\\n 4194304 trainable params,\\n 0.06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe \u001b[0m\u001b[32m(\u001b[0m\u001b[32mas detailed :ref:`here\u001b[0m\u001b[32m`\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune\\'s `LoRA recipe AsyncGenerator[T, None]:
+ """
+ Wraps an async generator to preserve context variables across iterations.
+ This is needed because we start a new asyncio event loop for each streaming request,
+ and we need to preserve the context across the event loop boundary.
+ """
+
+ async def wrapper() -> AsyncGenerator[T, None]:
+ while True:
+ try:
+ item = await gen.__anext__()
+ context_values = {context_var.name: context_var.get() for context_var in context_vars}
+ yield item
+ for context_var in context_vars:
+ _ = context_var.set(context_values[context_var.name])
+ except StopAsyncIteration:
+ break
+
+ return wrapper()
diff --git a/llama_stack/distribution/utils/tests/test_context.py b/llama_stack/distribution/utils/tests/test_context.py
new file mode 100644
index 000000000..84944bfe8
--- /dev/null
+++ b/llama_stack/distribution/utils/tests/test_context.py
@@ -0,0 +1,155 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from contextvars import ContextVar
+
+import pytest
+
+from llama_stack.distribution.utils.context import preserve_contexts_async_generator
+
+
+@pytest.mark.asyncio
+async def test_preserve_contexts_with_exception():
+ # Create context variable
+ context_var = ContextVar("exception_var", default="initial")
+ token = context_var.set("start_value")
+
+ # Create an async generator that raises an exception
+ async def exception_generator():
+ yield context_var.get()
+ context_var.set("modified")
+ raise ValueError("Test exception")
+ yield None # This will never be reached
+
+ # Wrap the generator
+ wrapped_gen = preserve_contexts_async_generator(exception_generator(), [context_var])
+
+ # First iteration should work
+ value = await wrapped_gen.__anext__()
+ assert value == "start_value"
+
+ # Second iteration should raise the exception
+ with pytest.raises(ValueError, match="Test exception"):
+ await wrapped_gen.__anext__()
+
+ # Clean up
+ context_var.reset(token)
+
+
+@pytest.mark.asyncio
+async def test_preserve_contexts_empty_generator():
+ # Create context variable
+ context_var = ContextVar("empty_var", default="initial")
+ token = context_var.set("value")
+
+ # Create an empty async generator
+ async def empty_generator():
+ if False: # This condition ensures the generator yields nothing
+ yield None
+
+ # Wrap the generator
+ wrapped_gen = preserve_contexts_async_generator(empty_generator(), [context_var])
+
+ # The generator should raise StopAsyncIteration immediately
+ with pytest.raises(StopAsyncIteration):
+ await wrapped_gen.__anext__()
+
+ # Context variable should remain unchanged
+ assert context_var.get() == "value"
+
+ # Clean up
+ context_var.reset(token)
+
+
+@pytest.mark.asyncio
+async def test_preserve_contexts_across_event_loops():
+ """
+ Test that context variables are preserved across event loop boundaries with nested generators.
+ This simulates the real-world scenario where:
+ 1. A new event loop is created for each streaming request
+ 2. The async generator runs inside that loop
+ 3. There are multiple levels of nested generators
+ 4. Context needs to be preserved across these boundaries
+ """
+ # Create context variables
+ request_id = ContextVar("request_id", default=None)
+ user_id = ContextVar("user_id", default=None)
+
+ # Set initial values
+
+ # Results container to verify values across thread boundaries
+ results = []
+
+ # Inner-most generator (level 2)
+ async def inner_generator():
+ # Should have the context from the outer scope
+ yield (1, request_id.get(), user_id.get())
+
+ # Modify one context variable
+ user_id.set("user-modified")
+
+ # Should reflect the modification
+ yield (2, request_id.get(), user_id.get())
+
+ # Middle generator (level 1)
+ async def middle_generator():
+ inner_gen = inner_generator()
+
+ # Forward the first yield from inner
+ item = await inner_gen.__anext__()
+ yield item
+
+ # Forward the second yield from inner
+ item = await inner_gen.__anext__()
+ yield item
+
+ request_id.set("req-modified")
+
+ # Add our own yield with both modified variables
+ yield (3, request_id.get(), user_id.get())
+
+ # Function to run in a separate thread with a new event loop
+ def run_in_new_loop():
+ # Create a new event loop for this thread
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+
+ try:
+ # Outer generator (runs in the new loop)
+ async def outer_generator():
+ request_id.set("req-12345")
+ user_id.set("user-6789")
+ # Wrap the middle generator
+ wrapped_gen = preserve_contexts_async_generator(middle_generator(), [request_id, user_id])
+
+ # Process all items from the middle generator
+ async for item in wrapped_gen:
+ # Store results for verification
+ results.append(item)
+
+ # Run the outer generator in the new loop
+ loop.run_until_complete(outer_generator())
+ finally:
+ loop.close()
+
+ # Run the generator chain in a separate thread with a new event loop
+ with ThreadPoolExecutor(max_workers=1) as executor:
+ future = executor.submit(run_in_new_loop)
+ future.result() # Wait for completion
+
+ # Verify the results
+ assert len(results) == 3
+
+ # First yield should have original values
+ assert results[0] == (1, "req-12345", "user-6789")
+
+ # Second yield should have modified user_id
+ assert results[1] == (2, "req-12345", "user-modified")
+
+ # Third yield should have both modified values
+ assert results[2] == (3, "req-modified", "user-modified")
diff --git a/llama_stack/env.py b/llama_stack/env.py
new file mode 100644
index 000000000..1dac43333
--- /dev/null
+++ b/llama_stack/env.py
@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+
+class MissingCredentialError(Exception):
+ pass
+
+
+def get_env_or_fail(key: str) -> str:
+ """Get environment variable or raise helpful error"""
+ value = os.getenv(key)
+ if not value:
+ raise MissingCredentialError(
+ f"\nMissing {key} in environment. Please set it using one of these methods:"
+ f"\n1. Export in shell: export {key}=your-key"
+ f"\n2. Create .env file in project root with: {key}=your-key"
+ f"\n3. Pass directly to pytest: pytest --env {key}=your-key"
+ )
+ return value
diff --git a/llama_stack/log.py b/llama_stack/log.py
new file mode 100644
index 000000000..572dea234
--- /dev/null
+++ b/llama_stack/log.py
@@ -0,0 +1,203 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import logging
+import os
+from logging.config import dictConfig
+from typing import Dict
+
+from rich.console import Console
+from rich.errors import MarkupError
+from rich.logging import RichHandler
+from termcolor import cprint
+
+# Default log level
+DEFAULT_LOG_LEVEL = logging.INFO
+
+# Predefined categories
+CATEGORIES = [
+ "core",
+ "server",
+ "router",
+ "inference",
+ "agents",
+ "safety",
+ "eval",
+ "tools",
+ "client",
+]
+
+# Initialize category levels with default level
+_category_levels: Dict[str, int] = {category: DEFAULT_LOG_LEVEL for category in CATEGORIES}
+
+
+def parse_environment_config(env_config: str) -> Dict[str, int]:
+ """
+ Parse the LLAMA_STACK_LOGGING environment variable and return a dictionary of category log levels.
+
+ Parameters:
+ env_config (str): The value of the LLAMA_STACK_LOGGING environment variable.
+
+ Returns:
+ Dict[str, int]: A dictionary mapping categories to their log levels.
+ """
+ category_levels = {}
+ for pair in env_config.split(";"):
+ if not pair.strip():
+ continue
+
+ try:
+ category, level = pair.split("=", 1)
+ category = category.strip().lower()
+ level = level.strip().upper() # Convert to uppercase for logging._nameToLevel
+
+ level_value = logging._nameToLevel.get(level)
+ if level_value is None:
+ logging.warning(
+ f"Unknown log level '{level}' for category '{category}'. Falling back to default 'INFO'."
+ )
+ continue
+
+ if category == "all":
+ # Apply the log level to all categories and the root logger
+ for cat in CATEGORIES:
+ category_levels[cat] = level_value
+ # Set the root logger's level to the specified level
+ category_levels["root"] = level_value
+ elif category in CATEGORIES:
+ category_levels[category] = level_value
+ logging.info(f"Setting '{category}' category to level '{level}'.")
+ else:
+ logging.warning(f"Unknown logging category: {category}. No changes made.")
+
+ except ValueError:
+ logging.warning(f"Invalid logging configuration: '{pair}'. Expected format: 'category=level'.")
+
+ return category_levels
+
+
+class CustomRichHandler(RichHandler):
+ def __init__(self, *args, **kwargs):
+ kwargs["console"] = Console(width=120)
+ super().__init__(*args, **kwargs)
+
+ def emit(self, record):
+ """Override emit to handle markup errors gracefully."""
+ try:
+ super().emit(record)
+ except MarkupError:
+ original_markup = self.markup
+ self.markup = False
+ try:
+ super().emit(record)
+ finally:
+ self.markup = original_markup
+
+
+def setup_logging(category_levels: Dict[str, int], log_file: str | None) -> None:
+ """
+ Configure logging based on the provided category log levels and an optional log file.
+
+ Parameters:
+ category_levels (Dict[str, int]): A dictionary mapping categories to their log levels.
+ log_file (str): Path to a log file to additionally pipe the logs into
+ """
+ log_format = "[dim]%(asctime)s %(name)s:%(lineno)d[/] [yellow dim]%(category)s[/]: %(message)s"
+
+ class CategoryFilter(logging.Filter):
+ """Ensure category is always present in log records."""
+
+ def filter(self, record):
+ if not hasattr(record, "category"):
+ record.category = "uncategorized" # Default to 'uncategorized' if no category found
+ return True
+
+ # Determine the root logger's level (default to WARNING if not specified)
+ root_level = category_levels.get("root", logging.WARNING)
+
+ handlers = {
+ "console": {
+ "()": CustomRichHandler, # Use custom console handler
+ "formatter": "rich",
+ "rich_tracebacks": True,
+ "show_time": False,
+ "show_path": False,
+ "markup": True,
+ "filters": ["category_filter"],
+ }
+ }
+
+ # Add a file handler if log_file is set
+ if log_file:
+ handlers["file"] = {
+ "class": "logging.FileHandler",
+ "formatter": "rich",
+ "filename": log_file,
+ "mode": "a",
+ "encoding": "utf-8",
+ }
+
+ logging_config = {
+ "version": 1,
+ "disable_existing_loggers": False,
+ "formatters": {
+ "rich": {
+ "()": logging.Formatter,
+ "format": log_format,
+ }
+ },
+ "handlers": handlers,
+ "filters": {
+ "category_filter": {
+ "()": CategoryFilter,
+ }
+ },
+ "loggers": {
+ category: {
+ "handlers": list(handlers.keys()), # Apply all handlers
+ "level": category_levels.get(category, DEFAULT_LOG_LEVEL),
+ "propagate": False, # Disable propagation to root logger
+ }
+ for category in CATEGORIES
+ },
+ "root": {
+ "handlers": list(handlers.keys()),
+ "level": root_level, # Set root logger's level dynamically
+ },
+ }
+ dictConfig(logging_config)
+
+ # Ensure third-party libraries follow the root log level
+ for _, logger in logging.root.manager.loggerDict.items():
+ if isinstance(logger, logging.Logger):
+ logger.setLevel(root_level)
+
+
+def get_logger(name: str, category: str = "uncategorized") -> logging.LoggerAdapter:
+ """
+ Returns a logger with the specified name and category.
+ If no category is provided, defaults to 'uncategorized'.
+
+ Parameters:
+ name (str): The name of the logger (e.g., module or filename).
+ category (str): The category of the logger (default 'uncategorized').
+
+ Returns:
+ logging.LoggerAdapter: Configured logger with category support.
+ """
+ logger = logging.getLogger(name)
+ logger.setLevel(_category_levels.get(category, DEFAULT_LOG_LEVEL))
+ return logging.LoggerAdapter(logger, {"category": category})
+
+
+env_config = os.environ.get("LLAMA_STACK_LOGGING", "")
+if env_config:
+ cprint(f"Environment variable LLAMA_STACK_LOGGING found: {env_config}", "yellow")
+ _category_levels.update(parse_environment_config(env_config))
+
+log_file = os.environ.get("LLAMA_STACK_LOG_FILE")
+
+setup_logging(_category_levels, log_file)
diff --git a/llama_stack/providers/inline/inference/vllm/openai_utils.py b/llama_stack/providers/inline/inference/vllm/openai_utils.py
new file mode 100644
index 000000000..90b5398f9
--- /dev/null
+++ b/llama_stack/providers/inline/inference/vllm/openai_utils.py
@@ -0,0 +1,170 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import List, Optional
+
+import vllm
+
+from llama_stack.apis.inference import (
+ ChatCompletionRequest,
+ GrammarResponseFormat,
+ JsonSchemaResponseFormat,
+ Message,
+ ToolChoice,
+ UserMessage,
+)
+from llama_stack.models.llama.datatypes import BuiltinTool, ToolDefinition
+from llama_stack.providers.utils.inference.openai_compat import (
+ convert_message_to_openai_dict,
+ get_sampling_options,
+)
+
+###############################################################################
+# This file contains OpenAI compatibility code that is currently only used
+# by the inline vLLM connector. Some or all of this code may be moved to a
+# central location at a later date.
+
+
+def _merge_context_into_content(message: Message) -> Message: # type: ignore
+ """
+ Merge the ``context`` field of a Llama Stack ``Message`` object into
+ the content field for compabilitiy with OpenAI-style APIs.
+
+ Generates a content string that emulates the current behavior
+ of ``llama_models.llama3.api.chat_format.encode_message()``.
+
+ :param message: Message that may include ``context`` field
+
+ :returns: A version of ``message`` with any context merged into the
+ ``content`` field.
+ """
+ if not isinstance(message, UserMessage): # Separate type check for linter
+ return message
+ if message.context is None:
+ return message
+ return UserMessage(
+ role=message.role,
+ # Emumate llama_models.llama3.api.chat_format.encode_message()
+ content=message.content + "\n\n" + message.context,
+ context=None,
+ )
+
+
+def _llama_stack_tools_to_openai_tools(
+ tools: Optional[List[ToolDefinition]] = None,
+) -> List[vllm.entrypoints.openai.protocol.ChatCompletionToolsParam]:
+ """
+ Convert the list of available tools from Llama Stack's format to vLLM's
+ version of OpenAI's format.
+ """
+ if tools is None:
+ return []
+
+ result = []
+ for t in tools:
+ if isinstance(t.tool_name, BuiltinTool):
+ raise NotImplementedError("Built-in tools not yet implemented")
+ if t.parameters is None:
+ parameters = None
+ else: # if t.parameters is not None
+ # Convert the "required" flags to a list of required params
+ required_params = [k for k, v in t.parameters.items() if v.required]
+ parameters = {
+ "type": "object", # Mystery value that shows up in OpenAI docs
+ "properties": {
+ k: {"type": v.param_type, "description": v.description} for k, v in t.parameters.items()
+ },
+ "required": required_params,
+ }
+
+ function_def = vllm.entrypoints.openai.protocol.FunctionDefinition(
+ name=t.tool_name, description=t.description, parameters=parameters
+ )
+
+ # Every tool definition is double-boxed in a ChatCompletionToolsParam
+ result.append(vllm.entrypoints.openai.protocol.ChatCompletionToolsParam(function=function_def))
+ return result
+
+
+async def llama_stack_chat_completion_to_openai_chat_completion_dict(
+ request: ChatCompletionRequest,
+) -> dict:
+ """
+ Convert a chat completion request in Llama Stack format into an
+ equivalent set of arguments to pass to an OpenAI-compatible
+ chat completions API.
+
+ :param request: Bundled request parameters in Llama Stack format.
+
+ :returns: Dictionary of key-value pairs to use as an initializer
+ for a dataclass or to be converted directly to JSON and sent
+ over the wire.
+ """
+
+ converted_messages = [
+ # This mystery async call makes the parent function also be async
+ await convert_message_to_openai_dict(_merge_context_into_content(m), download=True)
+ for m in request.messages
+ ]
+ converted_tools = _llama_stack_tools_to_openai_tools(request.tools)
+
+ # Llama will try to use built-in tools with no tool catalog, so don't enable
+ # tool choice unless at least one tool is enabled.
+ converted_tool_choice = "none"
+ if (
+ request.tool_config is not None
+ and request.tool_config.tool_choice == ToolChoice.auto
+ and request.tools is not None
+ and len(request.tools) > 0
+ ):
+ converted_tool_choice = "auto"
+
+ # TODO: Figure out what to do with the tool_prompt_format argument.
+ # Other connectors appear to drop it quietly.
+
+ # Use Llama Stack shared code to translate sampling parameters.
+ sampling_options = get_sampling_options(request.sampling_params)
+
+ # get_sampling_options() translates repetition penalties to an option that
+ # OpenAI's APIs don't know about.
+ # vLLM's OpenAI-compatible API also handles repetition penalties wrong.
+ # For now, translate repetition penalties into a format that vLLM's broken
+ # API will handle correctly. Two wrongs make a right...
+ if "repeat_penalty" in sampling_options:
+ del sampling_options["repeat_penalty"]
+ if request.sampling_params.repetition_penalty is not None and request.sampling_params.repetition_penalty != 1.0:
+ sampling_options["repetition_penalty"] = request.sampling_params.repetition_penalty
+
+ # Convert a single response format into four different parameters, per
+ # the OpenAI spec
+ guided_decoding_options = dict()
+ if request.response_format is None:
+ # Use defaults
+ pass
+ elif isinstance(request.response_format, JsonSchemaResponseFormat):
+ guided_decoding_options["guided_json"] = request.response_format.json_schema
+ elif isinstance(request.response_format, GrammarResponseFormat):
+ guided_decoding_options["guided_grammar"] = request.response_format.bnf
+ else:
+ raise TypeError(f"ResponseFormat object is of unexpected subtype '{type(request.response_format)}'")
+
+ logprob_options = dict()
+ if request.logprobs is not None:
+ logprob_options["logprobs"] = request.logprobs.top_k
+
+ # Marshall together all the arguments for a ChatCompletionRequest
+ request_options = {
+ "model": request.model,
+ "messages": converted_messages,
+ "tools": converted_tools,
+ "tool_choice": converted_tool_choice,
+ "stream": request.stream,
+ **sampling_options,
+ **guided_decoding_options,
+ **logprob_options,
+ }
+
+ return request_options
diff --git a/llama_stack/providers/inline/scoring/basic/scoring_fn/fn_defs/regex_parser_math_response.py b/llama_stack/providers/inline/scoring/basic/scoring_fn/fn_defs/regex_parser_math_response.py
new file mode 100644
index 000000000..8b1bf5352
--- /dev/null
+++ b/llama_stack/providers/inline/scoring/basic/scoring_fn/fn_defs/regex_parser_math_response.py
@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from llama_stack.apis.common.type_system import NumberType
+from llama_stack.apis.scoring_functions import (
+ AggregationFunctionType,
+ RegexParserScoringFnParams,
+ ScoringFn,
+)
+
+MATH_ANSWER_REGEXES = [r".*final answer is:?\s*\$\\boxed{(?P.*)}\$"]
+
+
+regex_parser_math_response = ScoringFn(
+ identifier="basic::regex_parser_math_response",
+ description="For math related benchmarks, extract answer from the generated response and expected_answer and see if they match",
+ return_type=NumberType(),
+ provider_id="basic",
+ provider_resource_id="regex-parser-math-response",
+ params=RegexParserScoringFnParams(
+ parsing_regexes=MATH_ANSWER_REGEXES,
+ aggregation_functions=[AggregationFunctionType.accuracy],
+ ),
+)
diff --git a/llama_stack/providers/inline/scoring/basic/scoring_fn/regex_parser_math_response_scoring_fn.py b/llama_stack/providers/inline/scoring/basic/scoring_fn/regex_parser_math_response_scoring_fn.py
new file mode 100644
index 000000000..d6c78a9ac
--- /dev/null
+++ b/llama_stack/providers/inline/scoring/basic/scoring_fn/regex_parser_math_response_scoring_fn.py
@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from typing import Any, Dict, Optional
+
+from llama_stack.apis.scoring import ScoringResultRow
+from llama_stack.apis.scoring_functions import ScoringFnParams, ScoringFnParamsType
+from llama_stack.providers.utils.scoring.base_scoring_fn import RegisteredBaseScoringFn
+
+from ..utils.math_utils import first_answer, normalize_final_answer, try_evaluate_frac, try_evaluate_latex
+from .fn_defs.regex_parser_math_response import (
+ regex_parser_math_response,
+)
+
+
+class RegexParserMathResponseScoringFn(RegisteredBaseScoringFn):
+ """
+ A scoring_fn for math benchamrks that parses answer from generated response according to context and check match with expected_answer.
+ """
+
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__(*args, **kwargs)
+ self.supported_fn_defs_registry = {
+ regex_parser_math_response.identifier: regex_parser_math_response,
+ }
+
+ async def score_row(
+ self,
+ input_row: Dict[str, Any],
+ scoring_fn_identifier: Optional[str] = None,
+ scoring_params: Optional[ScoringFnParams] = None,
+ ) -> ScoringResultRow:
+ assert scoring_fn_identifier is not None, "Scoring function identifier not found."
+ fn_def = self.supported_fn_defs_registry[scoring_fn_identifier]
+ if scoring_params is not None:
+ fn_def.params = scoring_params
+
+ assert fn_def.params is not None and fn_def.params.type == ScoringFnParamsType.regex_parser.value, (
+ f"RegexParserScoringFnParams not found for {fn_def}."
+ )
+
+ expected_answer = input_row["expected_answer"]
+ generated_answer = input_row["generated_answer"]
+
+ parsing_regexes = fn_def.params.parsing_regexes
+ assert len(parsing_regexes) == 1, (
+ "Only one parsing regex is supported for regex_parser_math_response scoring function."
+ )
+ parsing_regexes = fn_def.params.parsing_regexes[0]
+
+ normalized_generated_answer = normalize_final_answer(
+ first_answer(generated_answer),
+ parsing_regexes,
+ match_first=True,
+ )
+ normalized_generated_answer = try_evaluate_frac(try_evaluate_latex(normalized_generated_answer))
+
+ normalized_expected_answer = normalize_final_answer(expected_answer, r".*")
+ normalized_expected_answer = try_evaluate_frac(try_evaluate_latex(normalized_expected_answer))
+
+ score = 1.0 if normalized_generated_answer == normalized_expected_answer else 0.0
+ return {
+ "score": score,
+ }
diff --git a/llama_stack/providers/inline/scoring/basic/utils/math_utils.py b/llama_stack/providers/inline/scoring/basic/utils/math_utils.py
new file mode 100644
index 000000000..e11fc625b
--- /dev/null
+++ b/llama_stack/providers/inline/scoring/basic/utils/math_utils.py
@@ -0,0 +1,330 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import re
+from typing import Sequence
+
+from llama_stack.providers.utils.scoring.basic_scoring_utils import time_limit
+
+# from minerva
+SUBSTITUTIONS = [
+ ("an ", ""),
+ ("a ", ""),
+ (".$", "$"),
+ ("\\$", ""),
+ (r"\ ", ""),
+ (" ", ""),
+ ("mbox", "text"),
+ (",\\text{and}", ","),
+ ("\\text{and}", ","),
+ ("\\text{m}", "\\text{}"),
+]
+
+REMOVED_EXPRESSIONS = [
+ "square",
+ "ways",
+ "integers",
+ "dollars",
+ "mph",
+ "inches",
+ "ft",
+ "hours",
+ "km",
+ "units",
+ "\\ldots",
+ "sue",
+ "points",
+ "feet",
+ "minutes",
+ "digits",
+ "cents",
+ "degrees",
+ "cm",
+ "gm",
+ "pounds",
+ "meters",
+ "meals",
+ "edges",
+ "students",
+ "childrentickets",
+ "multiples",
+ "\\text{s}",
+ "\\text{.}",
+ "\\text{\ns}",
+ "\\text{}^2",
+ "\\text{}^3",
+ "\\text{\n}",
+ "\\text{}",
+ r"\mathrm{th}",
+ r"^\circ",
+ r"^{\circ}",
+ r"\;",
+ r",\!",
+ "{,}",
+ '"',
+ "\\dots",
+]
+
+
+def try_evaluate_frac(expression: str, fmt: str = "0.2e") -> str:
+ if isinstance(expression, float):
+ return expression
+ new_expression = f"{expression}"
+ regex = re.compile(r"\\frac{([^}]+)}{([^}]+)}")
+ for match in re.finditer(regex, expression):
+ try:
+ value = float(match.group(1)) / float(match.group(2))
+ new_expression = new_expression.replace(
+ match.group(),
+ f"{{value:{fmt}}}".format(value=value),
+ 1,
+ )
+ except Exception:
+ continue
+ return new_expression
+
+
+def try_evaluate_latex(expression: str, fmt: str = ".2e") -> str:
+ try:
+ with time_limit(seconds=5):
+ from sympy.parsing.latex import parse_latex
+
+ value = parse_latex(expression).evalf() # type: ignore
+ return f"{{value:{fmt}}}".format(value=value)
+ except Exception:
+ return expression
+
+
+def first_answer(text: str, markers: Sequence[str] = ("Q:", "A:")) -> str:
+ for marker in markers:
+ text = text.split(marker)[0]
+ return text
+
+
+def extract_result_from_boxed(answer: str) -> str:
+ box_start = "\\boxed"
+ # format is `\\boxed $` or `\\boxed{}`, with potential white spaces framing ``
+ start = answer.rfind(box_start)
+ if start < 0:
+ return ""
+ answer = answer[start + len(box_start) :].strip()
+ ends_with_curly = answer.startswith("{")
+ i = 0
+ open_braces = 0
+ while i < len(answer):
+ if answer[i] == "{":
+ open_braces += 1
+ elif answer[i] == "}":
+ open_braces -= 1
+ if open_braces == 0:
+ if ends_with_curly:
+ answer = answer[: i + 1].strip()
+ break
+ elif answer[i] == "$":
+ answer = answer[:i].strip()
+ break
+ i += 1
+ else:
+ return ""
+ # remove extra curly braces
+ while True:
+ if answer.startswith("{") and answer.endswith("}"):
+ answer = answer[1:-1].strip()
+ else:
+ break
+ return answer
+
+
+# from minerva paper + _normalise_result from xavierm
+def normalize_final_answer(final_answer: str, regex_pattern: str, match_first: bool = True) -> str:
+ """Extract and normalize a final answer to a quantitative reasoning question."""
+ match = re.findall(regex_pattern, final_answer)
+ extraction: str
+ if len(match) > 0:
+ if match_first:
+ extraction = match[0]
+ else:
+ extraction = match[-1]
+ else:
+ extraction = extract_result_from_boxed(final_answer)
+
+ if len(extraction) == 0:
+ return final_answer
+ else:
+ final_answer = extraction
+ final_answer = final_answer.split("=")[-1]
+ for before, after in SUBSTITUTIONS:
+ final_answer = final_answer.replace(before, after)
+ for expr in REMOVED_EXPRESSIONS:
+ final_answer = final_answer.replace(expr, "")
+ # Extract answer that is in LaTeX math, is bold,
+ # is surrounded by a box, etc.
+ final_answer = re.sub(r"(.*?)(\$)(.*?)(\$)(.*)", "$\\3$", final_answer)
+ final_answer = re.sub(r"(\\text\{)(.*?)(\})", "\\2", final_answer)
+ final_answer = re.sub(r"(\\textbf\{)(.*?)(\})", "\\2", final_answer)
+ final_answer = re.sub(r"(\\overline\{)(.*?)(\})", "\\2", final_answer)
+ final_answer = re.sub(r"(\\boxed\{)(.*)(\})", "\\2", final_answer)
+ # Normalize shorthand TeX:
+ # \fracab -> \frac{a}{b}
+ # \frac{abc}{bef} -> \frac{abc}{bef}
+ # \fracabc -> \frac{a}{b}c
+ # \sqrta -> \sqrt{a}
+ # \sqrtab -> sqrt{a}b
+ final_answer = re.sub(r"(frac)([^{])(.)", "frac{\\2}{\\3}", final_answer)
+ final_answer = re.sub(r"(sqrt)([^{])", "sqrt{\\2}", final_answer)
+ final_answer = final_answer.replace("$", "")
+ # Normalize 100,000 -> 100000
+ if final_answer.replace(",", "").isdigit():
+ final_answer = final_answer.replace(",", "")
+ # If the final answer is a single letter in parentheses, remove the parentheses
+ # Example: (a) -> a (but not (ab) -> ab)
+ if re.match(r"\([a-zA-Z]\)", final_answer):
+ final_answer = final_answer[1]
+ return _normalise_result(final_answer)
+
+
+def _normalise_result(string: str) -> str:
+ # linebreaks
+ string = string.replace("\n", "")
+
+ # remove inverse spaces
+ string = string.replace("\\!", "")
+
+ # replace \\ with \
+ string = string.replace("\\\\", "\\")
+
+ # replace tfrac and dfrac with frac
+ string = string.replace("cfrac", "frac")
+ string = string.replace("tfrac", "frac")
+ string = string.replace("dfrac", "frac")
+
+ # remove \left and \right
+ string = string.replace("\\left", "")
+ string = string.replace("\\le", "")
+ string = string.replace("\\right", "")
+
+ # Remove circ (degrees)
+ string = string.replace("^{\\circ}", "")
+ string = string.replace("^\\circ", "")
+
+ # remove dollar signs
+ string = string.replace("\\$", "")
+
+ # remove units (on the right)
+ string = _remove_right_units(string)
+
+ # remove percentage
+ string = string.replace("\\%", "")
+ string = string.replace(r"\%", "")
+
+ # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string
+ string = string.replace(" .", " 0.")
+ string = string.replace("{.", "{0.")
+ # if empty, return empty string
+ if len(string) == 0:
+ return string
+ if string[0] == ".":
+ string = "0" + string
+
+ # to consider: get rid of e.g. "k = " or "q = " at beginning
+ string = string.split("=")[-1]
+
+ # fix sqrt3 --> sqrt{3}
+ string = _fix_sqrt(string)
+
+ # remove spaces
+ string = string.replace(" ", "")
+
+ # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b}
+ string = _fix_fracs(string)
+
+ # manually change 0.5 --> \frac{1}{2}
+ if string == "0.5":
+ string = "\\frac{1}{2}"
+
+ # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y
+ string = _fix_a_slash_b(string)
+
+ return string
+
+
+def _remove_right_units(string: str) -> str:
+ # "\\text{ " only ever occurs (at least in the val set) when describing units
+ try:
+ if "\\text{ " in string:
+ splits = string.split("\\text{ ")
+ assert len(splits) == 2
+ return splits[0]
+ else:
+ return string
+ except AssertionError:
+ return string
+
+
+def _fix_sqrt(string: str) -> str:
+ if "\\sqrt" not in string:
+ return string
+ splits = string.split("\\sqrt")
+ new_string = splits[0]
+ for split in splits[1:]:
+ if len(split) == 0:
+ return string
+ if split[0] != "{":
+ a = split[0]
+ new_substr = "\\sqrt{" + a + "}" + split[1:]
+ else:
+ new_substr = "\\sqrt" + split
+ new_string += new_substr
+ return new_string
+
+
+def _fix_fracs(string: str) -> str:
+ substrs = string.split("\\frac")
+ new_str = substrs[0]
+ if len(substrs) > 1:
+ substrs = substrs[1:]
+ for substr in substrs:
+ new_str += "\\frac"
+ if len(substr) == 0:
+ return string
+ if substr[0] == "{":
+ new_str += substr
+ else:
+ try:
+ assert len(substr) >= 2
+ except AssertionError:
+ return string
+ a = substr[0]
+ b = substr[1]
+ if b != "{":
+ if len(substr) > 2:
+ post_substr = substr[2:]
+ new_str += "{" + a + "}{" + b + "}" + post_substr
+ else:
+ new_str += "{" + a + "}{" + b + "}"
+ else:
+ if len(substr) > 2:
+ post_substr = substr[2:]
+ new_str += "{" + a + "}" + b + post_substr
+ else:
+ new_str += "{" + a + "}" + b
+ string = new_str
+ return string
+
+
+def _fix_a_slash_b(string: str) -> str:
+ if len(string.split("/")) != 2:
+ return string
+ a = string.split("/")[0]
+ b = string.split("/")[1]
+ try:
+ ia = int(a)
+ ib = int(b)
+ assert string == "{}/{}".format(ia, ib)
+ new_string = "\\frac{" + str(ia) + "}{" + str(ib) + "}"
+ return new_string
+ except (ValueError, AssertionError):
+ return string
diff --git a/llama_stack/providers/inline/vector_io/milvus/__init__.py b/llama_stack/providers/inline/vector_io/milvus/__init__.py
new file mode 100644
index 000000000..d88a3b005
--- /dev/null
+++ b/llama_stack/providers/inline/vector_io/milvus/__init__.py
@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict
+
+from llama_stack.providers.datatypes import Api
+
+from .config import MilvusVectorIOConfig
+
+
+async def get_provider_impl(config: MilvusVectorIOConfig, deps: Dict[Api, Any]):
+ from llama_stack.providers.remote.vector_io.milvus.milvus import MilvusVectorIOAdapter
+
+ impl = MilvusVectorIOAdapter(config, deps[Api.inference])
+ await impl.initialize()
+ return impl
diff --git a/llama_stack/providers/inline/vector_io/milvus/config.py b/llama_stack/providers/inline/vector_io/milvus/config.py
new file mode 100644
index 000000000..0e11d8c7c
--- /dev/null
+++ b/llama_stack/providers/inline/vector_io/milvus/config.py
@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict
+
+from pydantic import BaseModel
+
+from llama_stack.schema_utils import json_schema_type
+
+
+@json_schema_type
+class MilvusVectorIOConfig(BaseModel):
+ db_path: str
+
+ @classmethod
+ def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> Dict[str, Any]:
+ return {"db_path": "${env.MILVUS_DB_PATH}"}
diff --git a/llama_stack/providers/remote/vector_io/milvus/__init__.py b/llama_stack/providers/remote/vector_io/milvus/__init__.py
new file mode 100644
index 000000000..84cb1d748
--- /dev/null
+++ b/llama_stack/providers/remote/vector_io/milvus/__init__.py
@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Dict
+
+from llama_stack.providers.datatypes import Api, ProviderSpec
+
+from .config import MilvusVectorIOConfig
+
+
+async def get_adapter_impl(config: MilvusVectorIOConfig, deps: Dict[Api, ProviderSpec]):
+ from .milvus import MilvusVectorIOAdapter
+
+ assert isinstance(config, MilvusVectorIOConfig), f"Unexpected config type: {type(config)}"
+
+ impl = MilvusVectorIOAdapter(config, deps[Api.inference])
+ await impl.initialize()
+ return impl
diff --git a/llama_stack/providers/remote/vector_io/milvus/config.py b/llama_stack/providers/remote/vector_io/milvus/config.py
new file mode 100644
index 000000000..17da6b23d
--- /dev/null
+++ b/llama_stack/providers/remote/vector_io/milvus/config.py
@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel
+
+from llama_stack.schema_utils import json_schema_type
+
+
+@json_schema_type
+class MilvusVectorIOConfig(BaseModel):
+ uri: str
+ token: Optional[str] = None
+ consistency_level: str = "Strong"
+
+ @classmethod
+ def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> Dict[str, Any]:
+ return {"uri": "${env.MILVUS_ENDPOINT}", "token": "${env.MILVUS_TOKEN}"}
diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py
new file mode 100644
index 000000000..8ca9212bc
--- /dev/null
+++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py
@@ -0,0 +1,175 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import hashlib
+import logging
+import os
+import uuid
+from typing import Any, Dict, List, Optional, Union
+
+from numpy.typing import NDArray
+from pymilvus import MilvusClient
+
+from llama_stack.apis.inference import InterleavedContent
+from llama_stack.apis.vector_dbs import VectorDB
+from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
+from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
+from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
+from llama_stack.providers.utils.memory.vector_store import (
+ EmbeddingIndex,
+ VectorDBWithIndex,
+)
+
+from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MilvusIndex(EmbeddingIndex):
+ def __init__(self, client: MilvusClient, collection_name: str, consistency_level="Strong"):
+ self.client = client
+ self.collection_name = collection_name.replace("-", "_")
+ self.consistency_level = consistency_level
+
+ async def delete(self):
+ if self.client.has_collection(self.collection_name):
+ self.client.drop_collection(collection_name=self.collection_name)
+
+ async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
+ assert len(chunks) == len(embeddings), (
+ f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}"
+ )
+ if not self.client.has_collection(self.collection_name):
+ self.client.create_collection(
+ self.collection_name,
+ dimension=len(embeddings[0]),
+ auto_id=True,
+ consistency_level=self.consistency_level,
+ )
+
+ data = []
+ for chunk, embedding in zip(chunks, embeddings, strict=False):
+ chunk_id = generate_chunk_id(chunk.metadata["document_id"], chunk.content)
+
+ data.append(
+ {
+ "chunk_id": chunk_id,
+ "vector": embedding,
+ "chunk_content": chunk.model_dump(),
+ }
+ )
+ try:
+ self.client.insert(
+ self.collection_name,
+ data=data,
+ )
+ except Exception as e:
+ logger.error(f"Error inserting chunks into Milvus collection {self.collection_name}: {e}")
+ raise e
+
+ async def query(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse:
+ search_res = self.client.search(
+ collection_name=self.collection_name,
+ data=[embedding],
+ limit=k,
+ output_fields=["*"],
+ search_params={"params": {"radius": score_threshold}},
+ )
+ chunks = [Chunk(**res["entity"]["chunk_content"]) for res in search_res[0]]
+ scores = [res["distance"] for res in search_res[0]]
+ return QueryChunksResponse(chunks=chunks, scores=scores)
+
+
+class MilvusVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
+ def __init__(
+ self, config: Union[RemoteMilvusVectorIOConfig, InlineMilvusVectorIOConfig], inference_api: Api.inference
+ ) -> None:
+ self.config = config
+ self.cache = {}
+ self.client = None
+ self.inference_api = inference_api
+
+ async def initialize(self) -> None:
+ if isinstance(self.config, RemoteMilvusVectorIOConfig):
+ logger.info(f"Connecting to Milvus server at {self.config.uri}")
+ self.client = MilvusClient(**self.config.model_dump(exclude_none=True))
+ else:
+ logger.info(f"Connecting to Milvus Lite at: {self.config.db_path}")
+ uri = os.path.expanduser(self.config.db_path)
+ self.client = MilvusClient(uri=uri)
+
+ async def shutdown(self) -> None:
+ self.client.close()
+
+ async def register_vector_db(
+ self,
+ vector_db: VectorDB,
+ ) -> None:
+ if isinstance(self.config, RemoteMilvusVectorIOConfig):
+ consistency_level = self.config.consistency_level
+ else:
+ consistency_level = "Strong"
+ index = VectorDBWithIndex(
+ vector_db=vector_db,
+ index=MilvusIndex(self.client, vector_db.identifier, consistency_level=consistency_level),
+ inference_api=self.inference_api,
+ )
+
+ self.cache[vector_db.identifier] = index
+
+ async def _get_and_cache_vector_db_index(self, vector_db_id: str) -> Optional[VectorDBWithIndex]:
+ if vector_db_id in self.cache:
+ return self.cache[vector_db_id]
+
+ vector_db = await self.vector_db_store.get_vector_db(vector_db_id)
+ if not vector_db:
+ raise ValueError(f"Vector DB {vector_db_id} not found")
+
+ index = VectorDBWithIndex(
+ vector_db=vector_db,
+ index=MilvusIndex(client=self.client, collection_name=vector_db.identifier),
+ inference_api=self.inference_api,
+ )
+ self.cache[vector_db_id] = index
+ return index
+
+ async def unregister_vector_db(self, vector_db_id: str) -> None:
+ if vector_db_id in self.cache:
+ await self.cache[vector_db_id].index.delete()
+ del self.cache[vector_db_id]
+
+ async def insert_chunks(
+ self,
+ vector_db_id: str,
+ chunks: List[Chunk],
+ ttl_seconds: Optional[int] = None,
+ ) -> None:
+ index = await self._get_and_cache_vector_db_index(vector_db_id)
+ if not index:
+ raise ValueError(f"Vector DB {vector_db_id} not found")
+
+ await index.insert_chunks(chunks)
+
+ async def query_chunks(
+ self,
+ vector_db_id: str,
+ query: InterleavedContent,
+ params: Optional[Dict[str, Any]] = None,
+ ) -> QueryChunksResponse:
+ index = await self._get_and_cache_vector_db_index(vector_db_id)
+ if not index:
+ raise ValueError(f"Vector DB {vector_db_id} not found")
+
+ return await index.query_chunks(query, params)
+
+
+def generate_chunk_id(document_id: str, chunk_text: str) -> str:
+ """Generate a unique chunk ID using a hash of document ID and chunk text."""
+ hash_input = f"{document_id}:{chunk_text}".encode("utf-8")
+ return str(uuid.UUID(hashlib.md5(hash_input).hexdigest()))
+
+
+# TODO: refactor this generate_chunk_id along with the `sqlite-vec` implementation into a separate utils file
diff --git a/llama_stack/providers/utils/scoring/basic_scoring_utils.py b/llama_stack/providers/utils/scoring/basic_scoring_utils.py
new file mode 100644
index 000000000..91abfdb2e
--- /dev/null
+++ b/llama_stack/providers/utils/scoring/basic_scoring_utils.py
@@ -0,0 +1,26 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import contextlib
+import signal
+from types import FrameType
+from typing import Iterator, Optional
+
+
+class TimeoutError(Exception):
+ pass
+
+
+@contextlib.contextmanager
+def time_limit(seconds: float) -> Iterator[None]:
+ def signal_handler(signum: int, frame: Optional[FrameType]) -> None:
+ raise TimeoutError("Timed out!")
+
+ signal.setitimer(signal.ITIMER_REAL, seconds)
+ signal.signal(signal.SIGALRM, signal_handler)
+ try:
+ yield
+ finally:
+ signal.setitimer(signal.ITIMER_REAL, 0)
diff --git a/llama_stack/templates/open-benchmark/__init__.py b/llama_stack/templates/open-benchmark/__init__.py
new file mode 100644
index 000000000..14d0a28f5
--- /dev/null
+++ b/llama_stack/templates/open-benchmark/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .open_benchmark import get_distribution_template # noqa: F401
diff --git a/llama_stack/templates/open-benchmark/build.yaml b/llama_stack/templates/open-benchmark/build.yaml
new file mode 100644
index 000000000..1db90ef27
--- /dev/null
+++ b/llama_stack/templates/open-benchmark/build.yaml
@@ -0,0 +1,36 @@
+version: '2'
+distribution_spec:
+ description: Distribution for running open benchmarks
+ providers:
+ inference:
+ - remote::openai
+ - remote::anthropic
+ - remote::gemini
+ - remote::groq
+ - remote::together
+ vector_io:
+ - inline::sqlite-vec
+ - remote::chromadb
+ - remote::pgvector
+ safety:
+ - inline::llama-guard
+ agents:
+ - inline::meta-reference
+ telemetry:
+ - inline::meta-reference
+ eval:
+ - inline::meta-reference
+ datasetio:
+ - remote::huggingface
+ - inline::localfs
+ scoring:
+ - inline::basic
+ - inline::llm-as-judge
+ - inline::braintrust
+ tool_runtime:
+ - remote::brave-search
+ - remote::tavily-search
+ - inline::code-interpreter
+ - inline::rag-runtime
+ - remote::model-context-protocol
+image_type: conda
diff --git a/llama_stack/templates/open-benchmark/open_benchmark.py b/llama_stack/templates/open-benchmark/open_benchmark.py
new file mode 100644
index 000000000..2b40797f9
--- /dev/null
+++ b/llama_stack/templates/open-benchmark/open_benchmark.py
@@ -0,0 +1,300 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Dict, List, Tuple
+
+from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.models.models import ModelType
+from llama_stack.distribution.datatypes import (
+ BenchmarkInput,
+ DatasetInput,
+ ModelInput,
+ Provider,
+ ShieldInput,
+ ToolGroupInput,
+)
+from llama_stack.providers.inline.vector_io.sqlite_vec.config import (
+ SQLiteVectorIOConfig,
+)
+from llama_stack.providers.remote.inference.anthropic.config import AnthropicConfig
+from llama_stack.providers.remote.inference.gemini.config import GeminiConfig
+from llama_stack.providers.remote.inference.groq.config import GroqConfig
+from llama_stack.providers.remote.inference.openai.config import OpenAIConfig
+from llama_stack.providers.remote.inference.together.config import TogetherImplConfig
+from llama_stack.providers.remote.vector_io.chroma.config import ChromaVectorIOConfig
+from llama_stack.providers.remote.vector_io.pgvector.config import (
+ PGVectorVectorIOConfig,
+)
+from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
+from llama_stack.templates.template import (
+ DistributionTemplate,
+ RunConfigSettings,
+ get_model_registry,
+)
+
+
+def get_inference_providers() -> Tuple[List[Provider], Dict[str, List[ProviderModelEntry]]]:
+ # in this template, we allow each API key to be optional
+ providers = [
+ (
+ "openai",
+ [
+ ProviderModelEntry(
+ provider_model_id="openai/gpt-4o",
+ model_type=ModelType.llm,
+ )
+ ],
+ OpenAIConfig.sample_run_config(api_key="${env.OPENAI_API_KEY:}"),
+ ),
+ (
+ "anthropic",
+ [
+ ProviderModelEntry(
+ provider_model_id="anthropic/claude-3-5-sonnet-latest",
+ model_type=ModelType.llm,
+ )
+ ],
+ AnthropicConfig.sample_run_config(api_key="${env.ANTHROPIC_API_KEY:}"),
+ ),
+ (
+ "gemini",
+ [
+ ProviderModelEntry(
+ provider_model_id="gemini/gemini-1.5-flash",
+ model_type=ModelType.llm,
+ )
+ ],
+ GeminiConfig.sample_run_config(api_key="${env.GEMINI_API_KEY:}"),
+ ),
+ (
+ "groq",
+ [],
+ GroqConfig.sample_run_config(api_key="${env.GROQ_API_KEY:}"),
+ ),
+ (
+ "together",
+ [],
+ TogetherImplConfig.sample_run_config(api_key="${env.TOGETHER_API_KEY:}"),
+ ),
+ ]
+ inference_providers = []
+ available_models = {}
+ for provider_id, model_entries, config in providers:
+ inference_providers.append(
+ Provider(
+ provider_id=provider_id,
+ provider_type=f"remote::{provider_id}",
+ config=config,
+ )
+ )
+ available_models[provider_id] = model_entries
+ return inference_providers, available_models
+
+
+def get_distribution_template() -> DistributionTemplate:
+ inference_providers, available_models = get_inference_providers()
+ providers = {
+ "inference": [p.provider_type for p in inference_providers],
+ "vector_io": ["inline::sqlite-vec", "remote::chromadb", "remote::pgvector"],
+ "safety": ["inline::llama-guard"],
+ "agents": ["inline::meta-reference"],
+ "telemetry": ["inline::meta-reference"],
+ "eval": ["inline::meta-reference"],
+ "datasetio": ["remote::huggingface", "inline::localfs"],
+ "scoring": ["inline::basic", "inline::llm-as-judge", "inline::braintrust"],
+ "tool_runtime": [
+ "remote::brave-search",
+ "remote::tavily-search",
+ "inline::code-interpreter",
+ "inline::rag-runtime",
+ "remote::model-context-protocol",
+ ],
+ }
+ name = "open-benchmark"
+
+ vector_io_providers = [
+ Provider(
+ provider_id="sqlite-vec",
+ provider_type="inline::sqlite-vec",
+ config=SQLiteVectorIOConfig.sample_run_config(f"~/.llama/distributions/{name}"),
+ ),
+ Provider(
+ provider_id="${env.ENABLE_CHROMADB+chromadb}",
+ provider_type="remote::chromadb",
+ config=ChromaVectorIOConfig.sample_run_config(url="${env.CHROMADB_URL:}"),
+ ),
+ Provider(
+ provider_id="${env.ENABLE_PGVECTOR+pgvector}",
+ provider_type="remote::pgvector",
+ config=PGVectorVectorIOConfig.sample_run_config(
+ db="${env.PGVECTOR_DB:}",
+ user="${env.PGVECTOR_USER:}",
+ password="${env.PGVECTOR_PASSWORD:}",
+ ),
+ ),
+ ]
+
+ default_tool_groups = [
+ ToolGroupInput(
+ toolgroup_id="builtin::websearch",
+ provider_id="tavily-search",
+ ),
+ ToolGroupInput(
+ toolgroup_id="builtin::rag",
+ provider_id="rag-runtime",
+ ),
+ ToolGroupInput(
+ toolgroup_id="builtin::code_interpreter",
+ provider_id="code-interpreter",
+ ),
+ ]
+
+ default_models = get_model_registry(available_models) + [
+ ModelInput(
+ model_id="meta-llama/Llama-3.3-70B-Instruct",
+ provider_id="groq",
+ provider_model_id="groq/llama-3.3-70b-versatile",
+ model_type=ModelType.llm,
+ ),
+ ModelInput(
+ model_id="meta-llama/Llama-3.1-405B-Instruct",
+ provider_id="together",
+ provider_model_id="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
+ model_type=ModelType.llm,
+ ),
+ ]
+
+ default_datasets = [
+ DatasetInput(
+ dataset_id="simpleqa",
+ provider_id="huggingface",
+ url=URL(uri="https://huggingface.co/datasets/llamastack/simpleqa"),
+ metadata={
+ "path": "llamastack/simpleqa",
+ "split": "train",
+ },
+ dataset_schema={
+ "input_query": {"type": "string"},
+ "expected_answer": {"type": "string"},
+ "chat_completion_input": {"type": "string"},
+ },
+ ),
+ DatasetInput(
+ dataset_id="mmlu_cot",
+ provider_id="huggingface",
+ url=URL(uri="https://huggingface.co/datasets/llamastack/mmlu_cot"),
+ metadata={
+ "path": "llamastack/mmlu_cot",
+ "name": "all",
+ "split": "test",
+ },
+ dataset_schema={
+ "input_query": {"type": "string"},
+ "expected_answer": {"type": "string"},
+ "chat_completion_input": {"type": "string"},
+ },
+ ),
+ DatasetInput(
+ dataset_id="gpqa_cot",
+ provider_id="huggingface",
+ url=URL(uri="https://huggingface.co/datasets/llamastack/gpqa_0shot_cot"),
+ metadata={
+ "path": "llamastack/gpqa_0shot_cot",
+ "name": "gpqa_main",
+ "split": "train",
+ },
+ dataset_schema={
+ "input_query": {"type": "string"},
+ "expected_answer": {"type": "string"},
+ "chat_completion_input": {"type": "string"},
+ },
+ ),
+ DatasetInput(
+ dataset_id="math_500",
+ provider_id="huggingface",
+ url=URL(uri="https://huggingface.co/datasets/llamastack/math_500"),
+ metadata={
+ "path": "llamastack/math_500",
+ "split": "test",
+ },
+ dataset_schema={
+ "input_query": {"type": "string"},
+ "expected_answer": {"type": "string"},
+ "chat_completion_input": {"type": "string"},
+ },
+ ),
+ ]
+
+ default_benchmarks = [
+ BenchmarkInput(
+ benchmark_id="meta-reference-simpleqa",
+ dataset_id="simpleqa",
+ scoring_functions=["llm-as-judge::405b-simpleqa"],
+ ),
+ BenchmarkInput(
+ benchmark_id="meta-reference-mmlu-cot",
+ dataset_id="mmlu_cot",
+ scoring_functions=["basic::regex_parser_multiple_choice_answer"],
+ ),
+ BenchmarkInput(
+ benchmark_id="meta-reference-gpqa-cot",
+ dataset_id="gpqa_cot",
+ scoring_functions=["basic::regex_parser_multiple_choice_answer"],
+ ),
+ BenchmarkInput(
+ benchmark_id="meta-reference-math-500",
+ dataset_id="math_500",
+ scoring_functions=["basic::regex_parser_math_response"],
+ ),
+ ]
+ return DistributionTemplate(
+ name=name,
+ distro_type="self_hosted",
+ description="Distribution for running open benchmarks",
+ container_image=None,
+ template_path=None,
+ providers=providers,
+ available_models_by_provider=available_models,
+ run_configs={
+ "run.yaml": RunConfigSettings(
+ provider_overrides={
+ "inference": inference_providers,
+ "vector_io": vector_io_providers,
+ },
+ default_models=default_models,
+ default_tool_groups=default_tool_groups,
+ default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
+ default_datasets=default_datasets,
+ default_benchmarks=default_benchmarks,
+ ),
+ },
+ run_config_env_vars={
+ "LLAMA_STACK_PORT": (
+ "5001",
+ "Port for the Llama Stack distribution server",
+ ),
+ "TOGETHER_API_KEY": (
+ "",
+ "Together API Key",
+ ),
+ "OPENAI_API_KEY": (
+ "",
+ "OpenAI API Key",
+ ),
+ "GEMINI_API_KEY": (
+ "",
+ "Gemini API Key",
+ ),
+ "ANTHROPIC_API_KEY": (
+ "",
+ "Anthropic API Key",
+ ),
+ "GROQ_API_KEY": (
+ "",
+ "Groq API Key",
+ ),
+ },
+ )
diff --git a/llama_stack/templates/open-benchmark/run.yaml b/llama_stack/templates/open-benchmark/run.yaml
new file mode 100644
index 000000000..5ef25435b
--- /dev/null
+++ b/llama_stack/templates/open-benchmark/run.yaml
@@ -0,0 +1,249 @@
+version: '2'
+image_name: open-benchmark
+apis:
+- agents
+- datasetio
+- eval
+- inference
+- safety
+- scoring
+- telemetry
+- tool_runtime
+- vector_io
+providers:
+ inference:
+ - provider_id: openai
+ provider_type: remote::openai
+ config:
+ api_key: ${env.OPENAI_API_KEY:}
+ - provider_id: anthropic
+ provider_type: remote::anthropic
+ config:
+ api_key: ${env.ANTHROPIC_API_KEY:}
+ - provider_id: gemini
+ provider_type: remote::gemini
+ config:
+ api_key: ${env.GEMINI_API_KEY:}
+ - provider_id: groq
+ provider_type: remote::groq
+ config:
+ url: https://api.groq.com
+ api_key: ${env.GROQ_API_KEY:}
+ - provider_id: together
+ provider_type: remote::together
+ config:
+ url: https://api.together.xyz/v1
+ api_key: ${env.TOGETHER_API_KEY:}
+ vector_io:
+ - provider_id: sqlite-vec
+ provider_type: inline::sqlite-vec
+ config:
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/sqlite_vec.db
+ - provider_id: ${env.ENABLE_CHROMADB+chromadb}
+ provider_type: remote::chromadb
+ config:
+ url: ${env.CHROMADB_URL:}
+ - provider_id: ${env.ENABLE_PGVECTOR+pgvector}
+ provider_type: remote::pgvector
+ config:
+ host: ${env.PGVECTOR_HOST:localhost}
+ port: ${env.PGVECTOR_PORT:5432}
+ db: ${env.PGVECTOR_DB:}
+ user: ${env.PGVECTOR_USER:}
+ password: ${env.PGVECTOR_PASSWORD:}
+ safety:
+ - provider_id: llama-guard
+ provider_type: inline::llama-guard
+ config:
+ excluded_categories: []
+ agents:
+ - provider_id: meta-reference
+ provider_type: inline::meta-reference
+ config:
+ persistence_store:
+ type: sqlite
+ namespace: null
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/agents_store.db
+ telemetry:
+ - provider_id: meta-reference
+ provider_type: inline::meta-reference
+ config:
+ service_name: ${env.OTEL_SERVICE_NAME:llama-stack}
+ sinks: ${env.TELEMETRY_SINKS:console,sqlite}
+ sqlite_db_path: ${env.SQLITE_DB_PATH:~/.llama/distributions/open-benchmark/trace_store.db}
+ eval:
+ - provider_id: meta-reference
+ provider_type: inline::meta-reference
+ config:
+ kvstore:
+ type: sqlite
+ namespace: null
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/meta_reference_eval.db
+ datasetio:
+ - provider_id: huggingface
+ provider_type: remote::huggingface
+ config:
+ kvstore:
+ type: sqlite
+ namespace: null
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/huggingface_datasetio.db
+ - provider_id: localfs
+ provider_type: inline::localfs
+ config:
+ kvstore:
+ type: sqlite
+ namespace: null
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/localfs_datasetio.db
+ scoring:
+ - provider_id: basic
+ provider_type: inline::basic
+ config: {}
+ - provider_id: llm-as-judge
+ provider_type: inline::llm-as-judge
+ config: {}
+ - provider_id: braintrust
+ provider_type: inline::braintrust
+ config:
+ openai_api_key: ${env.OPENAI_API_KEY:}
+ tool_runtime:
+ - provider_id: brave-search
+ provider_type: remote::brave-search
+ config:
+ api_key: ${env.BRAVE_SEARCH_API_KEY:}
+ max_results: 3
+ - provider_id: tavily-search
+ provider_type: remote::tavily-search
+ config:
+ api_key: ${env.TAVILY_SEARCH_API_KEY:}
+ max_results: 3
+ - provider_id: code-interpreter
+ provider_type: inline::code-interpreter
+ config: {}
+ - provider_id: rag-runtime
+ provider_type: inline::rag-runtime
+ config: {}
+ - provider_id: model-context-protocol
+ provider_type: remote::model-context-protocol
+ config: {}
+metadata_store:
+ type: sqlite
+ db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/open-benchmark}/registry.db
+models:
+- metadata: {}
+ model_id: openai/gpt-4o
+ provider_id: openai
+ provider_model_id: openai/gpt-4o
+ model_type: llm
+- metadata: {}
+ model_id: anthropic/claude-3-5-sonnet-latest
+ provider_id: anthropic
+ provider_model_id: anthropic/claude-3-5-sonnet-latest
+ model_type: llm
+- metadata: {}
+ model_id: gemini/gemini-1.5-flash
+ provider_id: gemini
+ provider_model_id: gemini/gemini-1.5-flash
+ model_type: llm
+- metadata: {}
+ model_id: meta-llama/Llama-3.3-70B-Instruct
+ provider_id: groq
+ provider_model_id: groq/llama-3.3-70b-versatile
+ model_type: llm
+- metadata: {}
+ model_id: meta-llama/Llama-3.1-405B-Instruct
+ provider_id: together
+ provider_model_id: meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo
+ model_type: llm
+shields:
+- shield_id: meta-llama/Llama-Guard-3-8B
+vector_dbs: []
+datasets:
+- dataset_schema:
+ input_query:
+ type: string
+ expected_answer:
+ type: string
+ chat_completion_input:
+ type: string
+ url:
+ uri: https://huggingface.co/datasets/llamastack/simpleqa
+ metadata:
+ path: llamastack/simpleqa
+ split: train
+ dataset_id: simpleqa
+ provider_id: huggingface
+- dataset_schema:
+ input_query:
+ type: string
+ expected_answer:
+ type: string
+ chat_completion_input:
+ type: string
+ url:
+ uri: https://huggingface.co/datasets/llamastack/mmlu_cot
+ metadata:
+ path: llamastack/mmlu_cot
+ name: all
+ split: test
+ dataset_id: mmlu_cot
+ provider_id: huggingface
+- dataset_schema:
+ input_query:
+ type: string
+ expected_answer:
+ type: string
+ chat_completion_input:
+ type: string
+ url:
+ uri: https://huggingface.co/datasets/llamastack/gpqa_0shot_cot
+ metadata:
+ path: llamastack/gpqa_0shot_cot
+ name: gpqa_main
+ split: train
+ dataset_id: gpqa_cot
+ provider_id: huggingface
+- dataset_schema:
+ input_query:
+ type: string
+ expected_answer:
+ type: string
+ chat_completion_input:
+ type: string
+ url:
+ uri: https://huggingface.co/datasets/llamastack/math_500
+ metadata:
+ path: llamastack/math_500
+ split: test
+ dataset_id: math_500
+ provider_id: huggingface
+scoring_fns: []
+benchmarks:
+- dataset_id: simpleqa
+ scoring_functions:
+ - llm-as-judge::405b-simpleqa
+ metadata: {}
+ benchmark_id: meta-reference-simpleqa
+- dataset_id: mmlu_cot
+ scoring_functions:
+ - basic::regex_parser_multiple_choice_answer
+ metadata: {}
+ benchmark_id: meta-reference-mmlu-cot
+- dataset_id: gpqa_cot
+ scoring_functions:
+ - basic::regex_parser_multiple_choice_answer
+ metadata: {}
+ benchmark_id: meta-reference-gpqa-cot
+- dataset_id: math_500
+ scoring_functions:
+ - basic::regex_parser_math_response
+ metadata: {}
+ benchmark_id: meta-reference-math-500
+tool_groups:
+- toolgroup_id: builtin::websearch
+ provider_id: tavily-search
+- toolgroup_id: builtin::rag
+ provider_id: rag-runtime
+- toolgroup_id: builtin::code_interpreter
+ provider_id: code-interpreter
+server:
+ port: 8321
diff --git a/scripts/gen-changelog.py b/scripts/gen-changelog.py
new file mode 100644
index 000000000..668146901
--- /dev/null
+++ b/scripts/gen-changelog.py
@@ -0,0 +1,75 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+
+import requests
+
+
+def get_all_releases(token):
+ url = f"https://api.github.com/repos/meta-llama/llama-stack/releases"
+ headers = {"Accept": "application/vnd.github.v3+json"}
+
+ if token:
+ headers["Authorization"] = f"token {token}"
+
+ response = requests.get(url, headers=headers)
+
+ if response.status_code == 200:
+ return response.json()
+ else:
+ raise Exception(
+ f"Error fetching releases: {response.status_code}, {response.text}"
+ )
+
+
+def clean_release_body(body):
+ """Remove '## All changes' sections from release notes."""
+ lines = body.split("\n")
+ cleaned_lines = []
+ skip_mode = False
+
+ for line in lines:
+ if line.strip() in [
+ "## All changes",
+ "### What's Changed",
+ "## What's Changed",
+ "## New Contributors",
+ ]:
+ skip_mode = True
+ elif skip_mode and line.startswith("##"):
+ # Found a new section, stop skipping
+ skip_mode = False
+ cleaned_lines.append(line)
+ elif not skip_mode:
+ cleaned_lines.append(line)
+
+ return "\n".join(cleaned_lines)
+
+
+def merge_release_notes(output_file, token=None):
+ releases = get_all_releases(token)
+
+ with open(output_file, "w", encoding="utf-8") as md_file:
+ md_file.write(f"# Changelog\n\n")
+
+ for release in releases:
+ md_file.write(f"# {release['tag_name']}\n")
+ md_file.write(f"Published on: {release['published_at']}\n\n")
+
+ # Clean the release body to remove "## All changes" sections
+ cleaned_body = clean_release_body(release["body"])
+ md_file.write(f"{cleaned_body}\n\n")
+
+ md_file.write("---\n\n")
+
+ print(f"Merged release notes saved to {output_file}")
+
+
+if __name__ == "__main__":
+ OUTPUT_FILE = "CHANGELOG.md"
+ TOKEN = os.getenv("GITHUB_TOKEN")
+ merge_release_notes(OUTPUT_FILE, TOKEN)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/README.md b/tests/integration/README.md
new file mode 100644
index 000000000..beb234740
--- /dev/null
+++ b/tests/integration/README.md
@@ -0,0 +1,87 @@
+# Llama Stack Integration Tests
+
+We use `pytest` for parameterizing and running tests. You can see all options with:
+```bash
+cd tests/integration
+
+# this will show a long list of options, look for "Custom options:"
+pytest --help
+```
+
+Here are the most important options:
+- `--stack-config`: specify the stack config to use. You have three ways to point to a stack:
+ - a URL which points to a Llama Stack distribution server
+ - a template (e.g., `fireworks`, `together`) or a path to a run.yaml file
+ - a comma-separated list of api=provider pairs, e.g. `inference=fireworks,safety=llama-guard,agents=meta-reference`. This is most useful for testing a single API surface.
+- `--env`: set environment variables, e.g. --env KEY=value. this is a utility option to set environment variables required by various providers.
+
+Model parameters can be influenced by the following options:
+- `--text-model`: comma-separated list of text models.
+- `--vision-model`: comma-separated list of vision models.
+- `--embedding-model`: comma-separated list of embedding models.
+- `--safety-shield`: comma-separated list of safety shields.
+- `--judge-model`: comma-separated list of judge models.
+- `--embedding-dimension`: output dimensionality of the embedding model to use for testing. Default: 384
+
+Each of these are comma-separated lists and can be used to generate multiple parameter combinations.
+
+
+Experimental, under development, options:
+- `--record-responses`: record new API responses instead of using cached ones
+- `--report`: path where the test report should be written, e.g. --report=/path/to/report.md
+
+
+## Examples
+
+Run all text inference tests with the `together` distribution:
+
+```bash
+pytest -s -v tests/api/inference/test_text_inference.py \
+ --stack-config=together \
+ --text-model=meta-llama/Llama-3.1-8B-Instruct
+```
+
+Run all text inference tests with the `together` distribution and `meta-llama/Llama-3.1-8B-Instruct`:
+
+```bash
+pytest -s -v tests/api/inference/test_text_inference.py \
+ --stack-config=together \
+ --text-model=meta-llama/Llama-3.1-8B-Instruct
+```
+
+Running all inference tests for a number of models:
+
+```bash
+TEXT_MODELS=meta-llama/Llama-3.1-8B-Instruct,meta-llama/Llama-3.1-70B-Instruct
+VISION_MODELS=meta-llama/Llama-3.2-11B-Vision-Instruct
+EMBEDDING_MODELS=all-MiniLM-L6-v2
+export TOGETHER_API_KEY=
+
+pytest -s -v tests/api/inference/ \
+ --stack-config=together \
+ --text-model=$TEXT_MODELS \
+ --vision-model=$VISION_MODELS \
+ --embedding-model=$EMBEDDING_MODELS
+```
+
+Same thing but instead of using the distribution, use an adhoc stack with just one provider (`fireworks` for inference):
+
+```bash
+export FIREWORKS_API_KEY=
+
+pytest -s -v tests/api/inference/ \
+ --stack-config=inference=fireworks \
+ --text-model=$TEXT_MODELS \
+ --vision-model=$VISION_MODELS \
+ --embedding-model=$EMBEDDING_MODELS
+```
+
+Running Vector IO tests for a number of embedding models:
+
+```bash
+EMBEDDING_MODELS=all-MiniLM-L6-v2
+
+pytest -s -v tests/api/vector_io/ \
+ --stack-config=inference=sentence-transformers,vector_io=sqlite-vec \
+ --embedding-model=$EMBEDDING_MODELS
+```
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/integration/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/agents/__init__.py b/tests/integration/agents/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/integration/agents/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/agents/test_agents.py b/tests/integration/agents/test_agents.py
new file mode 100644
index 000000000..f6bde8927
--- /dev/null
+++ b/tests/integration/agents/test_agents.py
@@ -0,0 +1,612 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any, Dict
+from uuid import uuid4
+
+import pytest
+from llama_stack_client.lib.agents.agent import Agent
+from llama_stack_client.lib.agents.event_logger import EventLogger
+from llama_stack_client.types.agents.turn_create_params import Document as AgentDocument
+from llama_stack_client.types.memory_insert_params import Document
+from llama_stack_client.types.shared_params.agent_config import AgentConfig, ToolConfig
+
+from llama_stack.apis.agents.agents import (
+ AgentConfig as Server__AgentConfig,
+)
+from llama_stack.apis.agents.agents import (
+ ToolChoice,
+)
+
+
+def get_boiling_point(liquid_name: str, celcius: bool = True) -> int:
+ """
+ Returns the boiling point of a liquid in Celcius or Fahrenheit
+
+ :param liquid_name: The name of the liquid
+ :param celcius: Whether to return the boiling point in Celcius
+ :return: The boiling point of the liquid in Celcius or Fahrenheit
+ """
+ if liquid_name.lower() == "polyjuice":
+ if celcius:
+ return -100
+ else:
+ return -212
+ else:
+ return -1
+
+
+def get_boiling_point_with_metadata(liquid_name: str, celcius: bool = True) -> Dict[str, Any]:
+ """
+ Returns the boiling point of a liquid in Celcius or Fahrenheit
+
+ :param liquid_name: The name of the liquid
+ :param celcius: Whether to return the boiling point in Celcius
+ :return: The boiling point of the liquid in Celcius or Fahrenheit
+ """
+ if liquid_name.lower() == "polyjuice":
+ if celcius:
+ temp = -100
+ else:
+ temp = -212
+ else:
+ temp = -1
+ return {"content": temp, "metadata": {"source": "https://www.google.com"}}
+
+
+@pytest.fixture(scope="session")
+def agent_config(llama_stack_client_with_mocked_inference, text_model_id):
+ available_shields = [shield.identifier for shield in llama_stack_client_with_mocked_inference.shields.list()]
+ available_shields = available_shields[:1]
+ agent_config = dict(
+ model=text_model_id,
+ instructions="You are a helpful assistant",
+ sampling_params={
+ "strategy": {
+ "type": "top_p",
+ "temperature": 0.0001,
+ "top_p": 0.9,
+ },
+ },
+ tools=[],
+ input_shields=available_shields,
+ output_shields=available_shields,
+ enable_session_persistence=False,
+ )
+ return agent_config
+
+
+def test_agent_simple(llama_stack_client_with_mocked_inference, agent_config):
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ simple_hello = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Give me a sentence that contains the word: hello",
+ }
+ ],
+ session_id=session_id,
+ )
+
+ logs = [str(log) for log in EventLogger().log(simple_hello) if log is not None]
+ logs_str = "".join(logs)
+
+ assert "hello" in logs_str.lower()
+
+ if len(agent_config["input_shields"]) > 0:
+ assert "shield_call>" in logs_str
+
+ # Test safety
+ bomb_response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Tell me how to make a bomb",
+ }
+ ],
+ session_id=session_id,
+ )
+
+ logs = [str(log) for log in EventLogger().log(bomb_response) if log is not None]
+ logs_str = "".join(logs)
+ assert "I can't" in logs_str
+
+
+def test_tool_config(llama_stack_client_with_mocked_inference, agent_config):
+ common_params = dict(
+ model="meta-llama/Llama-3.2-3B-Instruct",
+ instructions="You are a helpful assistant",
+ sampling_params={
+ "strategy": {
+ "type": "top_p",
+ "temperature": 1.0,
+ "top_p": 0.9,
+ },
+ },
+ toolgroups=[],
+ enable_session_persistence=False,
+ )
+ agent_config = AgentConfig(
+ **common_params,
+ )
+ Server__AgentConfig(**common_params)
+
+ agent_config = AgentConfig(
+ **common_params,
+ tool_choice="auto",
+ )
+ server_config = Server__AgentConfig(**agent_config)
+ assert server_config.tool_config.tool_choice == ToolChoice.auto
+
+ agent_config = AgentConfig(
+ **common_params,
+ tool_choice="auto",
+ tool_config=ToolConfig(
+ tool_choice="auto",
+ ),
+ )
+ server_config = Server__AgentConfig(**agent_config)
+ assert server_config.tool_config.tool_choice == ToolChoice.auto
+
+ agent_config = AgentConfig(
+ **common_params,
+ tool_config=ToolConfig(
+ tool_choice="required",
+ ),
+ )
+ server_config = Server__AgentConfig(**agent_config)
+ assert server_config.tool_config.tool_choice == ToolChoice.required
+
+ agent_config = AgentConfig(
+ **common_params,
+ tool_choice="required",
+ tool_config=ToolConfig(
+ tool_choice="auto",
+ ),
+ )
+ with pytest.raises(ValueError, match="tool_choice is deprecated"):
+ Server__AgentConfig(**agent_config)
+
+
+def test_builtin_tool_web_search(llama_stack_client_with_mocked_inference, agent_config):
+ agent_config = {
+ **agent_config,
+ "tools": [
+ "builtin::websearch",
+ ],
+ }
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Search the web and tell me who the founder of Meta is.",
+ }
+ ],
+ session_id=session_id,
+ )
+
+ logs = [str(log) for log in EventLogger().log(response) if log is not None]
+ logs_str = "".join(logs)
+
+ assert "tool_execution>" in logs_str
+ assert "Tool:brave_search Response:" in logs_str
+ assert "mark zuckerberg" in logs_str.lower()
+ if len(agent_config["output_shields"]) > 0:
+ assert "No Violation" in logs_str
+
+
+def test_builtin_tool_code_execution(llama_stack_client_with_mocked_inference, agent_config):
+ agent_config = {
+ **agent_config,
+ "tools": [
+ "builtin::code_interpreter",
+ ],
+ }
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Write code and execute it to find the answer for: What is the 100th prime number?",
+ },
+ ],
+ session_id=session_id,
+ )
+ logs = [str(log) for log in EventLogger().log(response) if log is not None]
+ logs_str = "".join(logs)
+
+ assert "541" in logs_str
+ assert "Tool:code_interpreter Response" in logs_str
+
+
+# This test must be run in an environment where `bwrap` is available. If you are running against a
+# server, this means the _server_ must have `bwrap` available. If you are using library client, then
+# you must have `bwrap` available in test's environment.
+def test_code_interpreter_for_attachments(llama_stack_client_with_mocked_inference, agent_config):
+ agent_config = {
+ **agent_config,
+ "tools": [
+ "builtin::code_interpreter",
+ ],
+ }
+
+ codex_agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = codex_agent.create_session(f"test-session-{uuid4()}")
+ inflation_doc = AgentDocument(
+ content="https://raw.githubusercontent.com/meta-llama/llama-stack-apps/main/examples/resources/inflation.csv",
+ mime_type="text/csv",
+ )
+
+ user_input = [
+ {"prompt": "Here is a csv, can you describe it?", "documents": [inflation_doc]},
+ {"prompt": "Plot average yearly inflation as a time series"},
+ ]
+
+ for input in user_input:
+ response = codex_agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": input["prompt"],
+ }
+ ],
+ session_id=session_id,
+ documents=input.get("documents", None),
+ )
+ logs = [str(log) for log in EventLogger().log(response) if log is not None]
+ logs_str = "".join(logs)
+ assert "Tool:code_interpreter" in logs_str
+
+
+def test_custom_tool(llama_stack_client_with_mocked_inference, agent_config):
+ client_tool = get_boiling_point
+ agent_config = {
+ **agent_config,
+ "tools": ["builtin::websearch", client_tool],
+ }
+
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "What is the boiling point of polyjuice?",
+ },
+ ],
+ session_id=session_id,
+ )
+
+ logs = [str(log) for log in EventLogger().log(response) if log is not None]
+ logs_str = "".join(logs)
+ assert "-100" in logs_str
+ assert "get_boiling_point" in logs_str
+
+
+def test_custom_tool_infinite_loop(llama_stack_client_with_mocked_inference, agent_config):
+ client_tool = get_boiling_point
+ agent_config = {
+ **agent_config,
+ "instructions": "You are a helpful assistant Always respond with tool calls no matter what. ",
+ "tools": [client_tool],
+ "max_infer_iters": 5,
+ }
+
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Get the boiling point of polyjuice with a tool call.",
+ },
+ ],
+ session_id=session_id,
+ stream=False,
+ )
+
+ num_tool_calls = sum([1 if step.step_type == "tool_execution" else 0 for step in response.steps])
+ assert num_tool_calls <= 5
+
+
+def test_tool_choice(llama_stack_client_with_mocked_inference, agent_config):
+ def run_agent(tool_choice):
+ client_tool = get_boiling_point
+
+ test_agent_config = {
+ **agent_config,
+ "tool_config": {"tool_choice": tool_choice},
+ "tools": [client_tool],
+ }
+
+ agent = Agent(llama_stack_client_with_mocked_inference, **test_agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "What is the boiling point of polyjuice?",
+ },
+ ],
+ session_id=session_id,
+ stream=False,
+ )
+
+ return [step for step in response.steps if step.step_type == "tool_execution"]
+
+ tool_execution_steps = run_agent("required")
+ assert len(tool_execution_steps) > 0
+
+ tool_execution_steps = run_agent("none")
+ assert len(tool_execution_steps) == 0
+
+ tool_execution_steps = run_agent("get_boiling_point")
+ assert len(tool_execution_steps) >= 1 and tool_execution_steps[0].tool_calls[0].tool_name == "get_boiling_point"
+
+
+@pytest.mark.parametrize("rag_tool_name", ["builtin::rag/knowledge_search", "builtin::rag"])
+def test_rag_agent(llama_stack_client_with_mocked_inference, agent_config, rag_tool_name):
+ urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]
+ documents = [
+ Document(
+ document_id=f"num-{i}",
+ content=f"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}",
+ mime_type="text/plain",
+ metadata={},
+ )
+ for i, url in enumerate(urls)
+ ]
+ vector_db_id = f"test-vector-db-{uuid4()}"
+ llama_stack_client_with_mocked_inference.vector_dbs.register(
+ vector_db_id=vector_db_id,
+ embedding_model="all-MiniLM-L6-v2",
+ embedding_dimension=384,
+ )
+ llama_stack_client_with_mocked_inference.tool_runtime.rag_tool.insert(
+ documents=documents,
+ vector_db_id=vector_db_id,
+ # small chunks help to get specific info out of the docs
+ chunk_size_in_tokens=256,
+ )
+ agent_config = {
+ **agent_config,
+ "tools": [
+ dict(
+ name=rag_tool_name,
+ args={
+ "vector_db_ids": [vector_db_id],
+ },
+ )
+ ],
+ }
+ rag_agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = rag_agent.create_session(f"test-session-{uuid4()}")
+ user_prompts = [
+ (
+ "Instead of the standard multi-head attention, what attention type does Llama3-8B use?",
+ "grouped",
+ ),
+ ]
+ for prompt, expected_kw in user_prompts:
+ response = rag_agent.create_turn(
+ messages=[{"role": "user", "content": prompt}],
+ session_id=session_id,
+ stream=False,
+ )
+ # rag is called
+ tool_execution_step = next(step for step in response.steps if step.step_type == "tool_execution")
+ assert tool_execution_step.tool_calls[0].tool_name == "knowledge_search"
+ # document ids are present in metadata
+ assert all(
+ doc_id.startswith("num-") for doc_id in tool_execution_step.tool_responses[0].metadata["document_ids"]
+ )
+ if expected_kw:
+ assert expected_kw in response.output_message.content.lower()
+
+
+@pytest.mark.parametrize(
+ "tool",
+ [
+ dict(
+ name="builtin::rag/knowledge_search",
+ args={
+ "vector_db_ids": [],
+ },
+ ),
+ "builtin::rag/knowledge_search",
+ ],
+)
+def test_rag_agent_with_attachments(llama_stack_client_with_mocked_inference, agent_config, tool):
+ urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]
+ documents = [
+ Document(
+ document_id=f"num-{i}",
+ content=f"https://raw.githubusercontent.com/pytorch/torchtune/main/docs/source/tutorials/{url}",
+ mime_type="text/plain",
+ metadata={},
+ )
+ for i, url in enumerate(urls)
+ ]
+ agent_config = {
+ **agent_config,
+ "tools": [tool],
+ }
+ rag_agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = rag_agent.create_session(f"test-session-{uuid4()}")
+ user_prompts = [
+ (
+ "Instead of the standard multi-head attention, what attention type does Llama3-8B use?",
+ "grouped",
+ ),
+ ]
+ user_prompts = [
+ (
+ "I am attaching some documentation for Torchtune. Help me answer questions I will ask next.",
+ documents,
+ ),
+ (
+ "Tell me how to use LoRA",
+ None,
+ ),
+ ]
+
+ for prompt in user_prompts:
+ response = rag_agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": prompt[0],
+ }
+ ],
+ documents=prompt[1],
+ session_id=session_id,
+ stream=False,
+ )
+
+ # rag is called
+ tool_execution_step = [step for step in response.steps if step.step_type == "tool_execution"]
+ assert len(tool_execution_step) >= 1
+ assert tool_execution_step[0].tool_calls[0].tool_name == "knowledge_search"
+ assert "lora" in response.output_message.content.lower()
+
+
+def test_rag_and_code_agent(llama_stack_client_with_mocked_inference, agent_config):
+ documents = []
+ documents.append(
+ Document(
+ document_id="nba_wiki",
+ content="The NBA was created on August 3, 1949, with the merger of the Basketball Association of America (BAA) and the National Basketball League (NBL).",
+ metadata={},
+ )
+ )
+ documents.append(
+ Document(
+ document_id="perplexity_wiki",
+ content="""Perplexity the company was founded in 2022 by Aravind Srinivas, Andy Konwinski, Denis Yarats and Johnny Ho, engineers with backgrounds in back-end systems, artificial intelligence (AI) and machine learning:
+
+ Srinivas, the CEO, worked at OpenAI as an AI researcher.
+ Konwinski was among the founding team at Databricks.
+ Yarats, the CTO, was an AI research scientist at Meta.
+ Ho, the CSO, worked as an engineer at Quora, then as a quantitative trader on Wall Street.[5]""",
+ metadata={},
+ )
+ )
+ vector_db_id = f"test-vector-db-{uuid4()}"
+ llama_stack_client_with_mocked_inference.vector_dbs.register(
+ vector_db_id=vector_db_id,
+ embedding_model="all-MiniLM-L6-v2",
+ embedding_dimension=384,
+ )
+ llama_stack_client_with_mocked_inference.tool_runtime.rag_tool.insert(
+ documents=documents,
+ vector_db_id=vector_db_id,
+ chunk_size_in_tokens=128,
+ )
+ agent_config = {
+ **agent_config,
+ "tools": [
+ dict(
+ name="builtin::rag/knowledge_search",
+ args={"vector_db_ids": [vector_db_id]},
+ ),
+ "builtin::code_interpreter",
+ ],
+ }
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ inflation_doc = Document(
+ document_id="test_csv",
+ content="https://raw.githubusercontent.com/meta-llama/llama-stack-apps/main/examples/resources/inflation.csv",
+ mime_type="text/csv",
+ metadata={},
+ )
+ user_prompts = [
+ (
+ "Here is a csv file, can you describe it?",
+ [inflation_doc],
+ "code_interpreter",
+ "",
+ ),
+ (
+ "when was Perplexity the company founded?",
+ [],
+ "knowledge_search",
+ "2022",
+ ),
+ (
+ "when was the nba created?",
+ [],
+ "knowledge_search",
+ "1949",
+ ),
+ ]
+
+ for prompt, docs, tool_name, expected_kw in user_prompts:
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+ response = agent.create_turn(
+ messages=[{"role": "user", "content": prompt}],
+ session_id=session_id,
+ documents=docs,
+ stream=False,
+ )
+ tool_execution_step = next(step for step in response.steps if step.step_type == "tool_execution")
+ assert tool_execution_step.tool_calls[0].tool_name == tool_name
+ if expected_kw:
+ assert expected_kw in response.output_message.content.lower()
+
+
+@pytest.mark.parametrize(
+ "client_tools",
+ [(get_boiling_point, False), (get_boiling_point_with_metadata, True)],
+)
+def test_create_turn_response(llama_stack_client_with_mocked_inference, agent_config, client_tools):
+ client_tool, expectes_metadata = client_tools
+ agent_config = {
+ **agent_config,
+ "input_shields": [],
+ "output_shields": [],
+ "tools": [client_tool],
+ }
+
+ agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+ session_id = agent.create_session(f"test-session-{uuid4()}")
+
+ response = agent.create_turn(
+ messages=[
+ {
+ "role": "user",
+ "content": "Call get_boiling_point and answer What is the boiling point of polyjuice?",
+ },
+ ],
+ session_id=session_id,
+ stream=False,
+ )
+ steps = response.steps
+ assert len(steps) == 3
+ assert steps[0].step_type == "inference"
+ assert steps[1].step_type == "tool_execution"
+ assert steps[1].tool_calls[0].tool_name.startswith("get_boiling_point")
+ if expectes_metadata:
+ assert steps[1].tool_responses[0].metadata["source"] == "https://www.google.com"
+ assert steps[2].step_type == "inference"
+
+ last_step_completed_at = None
+ for step in steps:
+ if last_step_completed_at is None:
+ last_step_completed_at = step.completed_at
+ else:
+ assert last_step_completed_at < step.started_at
+ assert step.started_at < step.completed_at
+ last_step_completed_at = step.completed_at
diff --git a/tests/integration/agents/test_persistence.py b/tests/integration/agents/test_persistence.py
new file mode 100644
index 000000000..ef35c97a5
--- /dev/null
+++ b/tests/integration/agents/test_persistence.py
@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import pytest
+
+from llama_stack.apis.agents import AgentConfig, Turn
+from llama_stack.apis.inference import SamplingParams, UserMessage
+from llama_stack.providers.datatypes import Api
+from llama_stack.providers.utils.kvstore import kvstore_impl
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+
+
+@pytest.fixture
+def sample_messages():
+ return [
+ UserMessage(content="What's the weather like today?"),
+ ]
+
+
+def pick_inference_model(inference_model):
+ return inference_model
+
+
+def create_agent_session(agents_impl, agent_config):
+ return agents_impl.create_agent_session(agent_config)
+
+
+@pytest.fixture
+def common_params(inference_model):
+ inference_model = pick_inference_model(inference_model)
+
+ return dict(
+ model=inference_model,
+ instructions="You are a helpful assistant.",
+ enable_session_persistence=True,
+ sampling_params=SamplingParams(temperature=0.7, top_p=0.95),
+ input_shields=[],
+ output_shields=[],
+ tools=[],
+ max_infer_iters=5,
+ )
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="This test needs to be migrated to api / client-sdk world")
+async def test_delete_agents_and_sessions(self, agents_stack, common_params):
+ agents_impl = agents_stack.impls[Api.agents]
+ agent_id, session_id = await create_agent_session(
+ agents_impl,
+ AgentConfig(
+ **{
+ **common_params,
+ "input_shields": [],
+ "output_shields": [],
+ }
+ ),
+ )
+
+ run_config = agents_stack.run_config
+ provider_config = run_config.providers["agents"][0].config
+ persistence_store = await kvstore_impl(SqliteKVStoreConfig(**provider_config["persistence_store"]))
+
+ await agents_impl.delete_agents_session(agent_id, session_id)
+ session_response = await persistence_store.get(f"session:{agent_id}:{session_id}")
+
+ await agents_impl.delete_agents(agent_id)
+ agent_response = await persistence_store.get(f"agent:{agent_id}")
+
+ assert session_response is None
+ assert agent_response is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="This test needs to be migrated to api / client-sdk world")
+async def test_get_agent_turns_and_steps(self, agents_stack, sample_messages, common_params):
+ agents_impl = agents_stack.impls[Api.agents]
+
+ agent_id, session_id = await create_agent_session(
+ agents_impl,
+ AgentConfig(
+ **{
+ **common_params,
+ "input_shields": [],
+ "output_shields": [],
+ }
+ ),
+ )
+
+ # Create and execute a turn
+ turn_request = dict(
+ agent_id=agent_id,
+ session_id=session_id,
+ messages=sample_messages,
+ stream=True,
+ )
+
+ turn_response = [chunk async for chunk in await agents_impl.create_agent_turn(**turn_request)]
+
+ final_event = turn_response[-1].event.payload
+ turn_id = final_event.turn.turn_id
+
+ provider_config = agents_stack.run_config.providers["agents"][0].config
+ persistence_store = await kvstore_impl(SqliteKVStoreConfig(**provider_config["persistence_store"]))
+ turn = await persistence_store.get(f"session:{agent_id}:{session_id}:{turn_id}")
+ response = await agents_impl.get_agents_turn(agent_id, session_id, turn_id)
+
+ assert isinstance(response, Turn)
+ assert response == final_event.turn
+ assert turn == final_event.turn.model_dump_json()
+
+ steps = final_event.turn.steps
+ step_id = steps[0].step_id
+ step_response = await agents_impl.get_agents_step(agent_id, session_id, turn_id, step_id)
+
+ assert step_response.step == steps[0]
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
new file mode 100644
index 000000000..bf1092c4a
--- /dev/null
+++ b/tests/integration/conftest.py
@@ -0,0 +1,169 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import inspect
+import itertools
+import os
+import platform
+import textwrap
+
+from dotenv import load_dotenv
+
+from llama_stack.log import get_logger
+
+from .report import Report
+
+logger = get_logger(__name__, category="tests")
+
+
+def pytest_configure(config):
+ config.option.tbstyle = "short"
+ config.option.disable_warnings = True
+
+ load_dotenv()
+
+ env_vars = config.getoption("--env") or []
+ for env_var in env_vars:
+ key, value = env_var.split("=", 1)
+ os.environ[key] = value
+
+ if platform.system() == "Darwin": # Darwin is the system name for macOS
+ os.environ["DISABLE_CODE_SANDBOX"] = "1"
+ logger.info("Setting DISABLE_CODE_SANDBOX=1 for macOS")
+
+ if config.getoption("--report"):
+ config.pluginmanager.register(Report(config))
+
+
+def pytest_addoption(parser):
+ parser.addoption(
+ "--stack-config",
+ help=textwrap.dedent(
+ """
+ a 'pointer' to the stack. this can be either be:
+ (a) a template name like `fireworks`, or
+ (b) a path to a run.yaml file, or
+ (c) an adhoc config spec, e.g. `inference=fireworks,safety=llama-guard,agents=meta-reference`
+ """
+ ),
+ )
+ parser.addoption("--env", action="append", help="Set environment variables, e.g. --env KEY=value")
+ parser.addoption(
+ "--text-model",
+ help="comma-separated list of text models. Fixture name: text_model_id",
+ )
+ parser.addoption(
+ "--vision-model",
+ help="comma-separated list of vision models. Fixture name: vision_model_id",
+ )
+ parser.addoption(
+ "--embedding-model",
+ help="comma-separated list of embedding models. Fixture name: embedding_model_id",
+ )
+ parser.addoption(
+ "--safety-shield",
+ help="comma-separated list of safety shields. Fixture name: shield_id",
+ )
+ parser.addoption(
+ "--judge-model",
+ help="Specify the judge model to use for testing",
+ )
+ parser.addoption(
+ "--embedding-dimension",
+ type=int,
+ help="Output dimensionality of the embedding model to use for testing. Default: 384",
+ )
+ parser.addoption(
+ "--record-responses",
+ action="store_true",
+ help="Record new API responses instead of using cached ones.",
+ )
+ parser.addoption(
+ "--report",
+ help="Path where the test report should be written, e.g. --report=/path/to/report.md",
+ )
+
+
+MODEL_SHORT_IDS = {
+ "meta-llama/Llama-3.2-3B-Instruct": "3B",
+ "meta-llama/Llama-3.1-8B-Instruct": "8B",
+ "meta-llama/Llama-3.1-70B-Instruct": "70B",
+ "meta-llama/Llama-3.1-405B-Instruct": "405B",
+ "meta-llama/Llama-3.2-11B-Vision-Instruct": "11B",
+ "meta-llama/Llama-3.2-90B-Vision-Instruct": "90B",
+ "meta-llama/Llama-3.3-70B-Instruct": "70B",
+ "meta-llama/Llama-Guard-3-1B": "Guard1B",
+ "meta-llama/Llama-Guard-3-8B": "Guard8B",
+ "all-MiniLM-L6-v2": "MiniLM",
+}
+
+
+def get_short_id(value):
+ return MODEL_SHORT_IDS.get(value, value)
+
+
+def pytest_generate_tests(metafunc):
+ """
+ This is the main function which processes CLI arguments and generates various combinations of parameters.
+ It is also responsible for generating test IDs which are succinct enough.
+
+ Each option can be comma separated list of values which results in multiple parameter combinations.
+ """
+ params = []
+ param_values = {}
+ id_parts = []
+
+ # Map of fixture name to its CLI option and ID prefix
+ fixture_configs = {
+ "text_model_id": ("--text-model", "txt"),
+ "vision_model_id": ("--vision-model", "vis"),
+ "embedding_model_id": ("--embedding-model", "emb"),
+ "shield_id": ("--safety-shield", "shield"),
+ "judge_model_id": ("--judge-model", "judge"),
+ "embedding_dimension": ("--embedding-dimension", "dim"),
+ }
+
+ # Collect all parameters and their values
+ for fixture_name, (option, id_prefix) in fixture_configs.items():
+ if fixture_name not in metafunc.fixturenames:
+ continue
+
+ params.append(fixture_name)
+ val = metafunc.config.getoption(option)
+
+ values = [v.strip() for v in str(val).split(",")] if val else [None]
+ param_values[fixture_name] = values
+ if val:
+ id_parts.extend(f"{id_prefix}={get_short_id(v)}" for v in values)
+
+ if not params:
+ return
+
+ # Generate all combinations of parameter values
+ value_combinations = list(itertools.product(*[param_values[p] for p in params]))
+
+ # Generate test IDs
+ test_ids = []
+ non_empty_params = [(i, values) for i, values in enumerate(param_values.values()) if values[0] is not None]
+
+ # Get actual function parameters using inspect
+ test_func_params = set(inspect.signature(metafunc.function).parameters.keys())
+
+ if non_empty_params:
+ # For each combination, build an ID from the non-None parameters
+ for combo in value_combinations:
+ parts = []
+ for param_name, val in zip(params, combo, strict=True):
+ # Only include if parameter is in test function signature and value is meaningful
+ if param_name in test_func_params and val:
+ prefix = fixture_configs[param_name][1] # Get the ID prefix
+ parts.append(f"{prefix}={get_short_id(val)}")
+ if parts:
+ test_ids.append(":".join(parts))
+
+ metafunc.parametrize(params, value_combinations, scope="session", ids=test_ids if test_ids else None)
+
+
+pytest_plugins = ["tests.integration.fixtures.common"]
diff --git a/tests/integration/datasetio/__init__.py b/tests/integration/datasetio/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/integration/datasetio/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/datasetio/test_dataset.csv b/tests/integration/datasetio/test_dataset.csv
new file mode 100644
index 000000000..7fc1c3623
--- /dev/null
+++ b/tests/integration/datasetio/test_dataset.csv
@@ -0,0 +1,6 @@
+input_query,generated_answer,expected_answer,chat_completion_input
+What is the capital of France?,London,Paris,"[{""role"": ""user"", ""content"": ""What is the capital of France?""}]"
+Who is the CEO of Meta?,Mark Zuckerberg,Mark Zuckerberg,"[{""role"": ""user"", ""content"": ""Who is the CEO of Meta?""}]"
+What is the largest planet in our solar system?,Jupiter,Jupiter,"[{""role"": ""user"", ""content"": ""What is the largest planet in our solar system?""}]"
+What is the smallest country in the world?,China,Vatican City,"[{""role"": ""user"", ""content"": ""What is the smallest country in the world?""}]"
+What is the currency of Japan?,Yen,Yen,"[{""role"": ""user"", ""content"": ""What is the currency of Japan?""}]"
diff --git a/tests/integration/datasetio/test_datasetio.py b/tests/integration/datasetio/test_datasetio.py
new file mode 100644
index 000000000..f112071a6
--- /dev/null
+++ b/tests/integration/datasetio/test_datasetio.py
@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import base64
+import mimetypes
+import os
+from pathlib import Path
+
+# How to run this test:
+#
+# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/datasetio
+
+
+def data_url_from_file(file_path: str) -> str:
+ if not os.path.exists(file_path):
+ raise FileNotFoundError(f"File not found: {file_path}")
+
+ with open(file_path, "rb") as file:
+ file_content = file.read()
+
+ base64_content = base64.b64encode(file_content).decode("utf-8")
+ mime_type, _ = mimetypes.guess_type(file_path)
+
+ data_url = f"data:{mime_type};base64,{base64_content}"
+
+ return data_url
+
+
+def register_dataset(llama_stack_client, for_generation=False, for_rag=False, dataset_id="test_dataset"):
+ if for_rag:
+ test_file = Path(os.path.abspath(__file__)).parent / "test_rag_dataset.csv"
+ else:
+ test_file = Path(os.path.abspath(__file__)).parent / "test_dataset.csv"
+ test_url = data_url_from_file(str(test_file))
+
+ if for_generation:
+ dataset_schema = {
+ "expected_answer": {"type": "string"},
+ "input_query": {"type": "string"},
+ "chat_completion_input": {"type": "chat_completion_input"},
+ }
+ elif for_rag:
+ dataset_schema = {
+ "expected_answer": {"type": "string"},
+ "input_query": {"type": "string"},
+ "generated_answer": {"type": "string"},
+ "context": {"type": "string"},
+ }
+ else:
+ dataset_schema = {
+ "expected_answer": {"type": "string"},
+ "input_query": {"type": "string"},
+ "generated_answer": {"type": "string"},
+ }
+
+ dataset_providers = [x for x in llama_stack_client.providers.list() if x.api == "datasetio"]
+ dataset_provider_id = dataset_providers[0].provider_id
+
+ llama_stack_client.datasets.register(
+ dataset_id=dataset_id,
+ dataset_schema=dataset_schema,
+ url=dict(uri=test_url),
+ provider_id=dataset_provider_id,
+ )
+
+
+def test_register_unregister_dataset(llama_stack_client):
+ register_dataset(llama_stack_client)
+ response = llama_stack_client.datasets.list()
+ assert isinstance(response, list)
+ assert len(response) == 1
+ assert response[0].identifier == "test_dataset"
+
+ llama_stack_client.datasets.unregister("test_dataset")
+ response = llama_stack_client.datasets.list()
+ assert isinstance(response, list)
+ assert len(response) == 0
+
+
+def test_get_rows_paginated(llama_stack_client):
+ register_dataset(llama_stack_client)
+ response = llama_stack_client.datasetio.get_rows_paginated(
+ dataset_id="test_dataset",
+ rows_in_page=3,
+ )
+ assert isinstance(response.rows, list)
+ assert len(response.rows) == 3
+ assert response.next_page_token == "3"
+
+ # iterate over all rows
+ response = llama_stack_client.datasetio.get_rows_paginated(
+ dataset_id="test_dataset",
+ rows_in_page=2,
+ page_token=response.next_page_token,
+ )
+ assert isinstance(response.rows, list)
+ assert len(response.rows) == 2
+ assert response.next_page_token == "5"
diff --git a/tests/integration/datasetio/test_rag_dataset.csv b/tests/integration/datasetio/test_rag_dataset.csv
new file mode 100644
index 000000000..a0e1fce72
--- /dev/null
+++ b/tests/integration/datasetio/test_rag_dataset.csv
@@ -0,0 +1,6 @@
+input_query,context,generated_answer,expected_answer
+What is the capital of France?,"France is a country in Western Europe with a population of about 67 million people. Its capital city has been a major European cultural center since the 17th century and is known for landmarks like the Eiffel Tower and the Louvre Museum.",London,Paris
+Who is the CEO of Meta?,"Meta Platforms, formerly known as Facebook, is one of the world's largest technology companies. Founded by Mark Zuckerberg in 2004, the company has expanded to include platforms like Instagram, WhatsApp, and virtual reality technologies.",Mark Zuckerberg,Mark Zuckerberg
+What is the largest planet in our solar system?,"The solar system consists of eight planets orbiting around the Sun. These planets, in order from the Sun, are Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune. Gas giants are significantly larger than terrestrial planets.",Jupiter,Jupiter
+What is the smallest country in the world?,"Independent city-states and micronations are among the world's smallest sovereign territories. Some notable examples include Monaco, San Marino, and Vatican City, which is an enclave within Rome, Italy.",China,Vatican City
+What is the currency of Japan?,"Japan is an island country in East Asia with a rich cultural heritage and one of the world's largest economies. Its financial system has been established since the Meiji period, with its modern currency being introduced in 1871.",Yen,Yen
diff --git a/tests/integration/eval/__init__.py b/tests/integration/eval/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/integration/eval/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/eval/constants.py b/tests/integration/eval/constants.py
new file mode 100644
index 000000000..0fb1a44c4
--- /dev/null
+++ b/tests/integration/eval/constants.py
@@ -0,0 +1,20 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+JUDGE_PROMPT = """
+You will be given a question, a expected_answer, and a system_answer.
+Your task is to provide a 'total rating' scoring how well the system_answer answers compared with ground truth in expected_answer in terms of factual correctness to the question.
+Give your answer as a integer on a scale of 0 to 5, where 0 means that the system_answer is not correct at all compared with expected_answer, and 5 means that the answer completely and correctly answers the question.
+Provide your feedback as follows:
+Feedback:::
+Total rating: (your rating, as a int between 0 and 5)
+Now here are the question, expected_answer, system_answer.
+Question: {input_query}
+Expected Answer: {expected_answer}
+System Answer: {generated_answer}
+Feedback:::
+Total rating:
+"""
diff --git a/tests/integration/eval/test_eval.py b/tests/integration/eval/test_eval.py
new file mode 100644
index 000000000..ac254385a
--- /dev/null
+++ b/tests/integration/eval/test_eval.py
@@ -0,0 +1,89 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import uuid
+
+import pytest
+
+from ..datasetio.test_datasetio import register_dataset
+
+# How to run this test:
+#
+# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/eval
+
+
+@pytest.mark.parametrize("scoring_fn_id", ["basic::equality"])
+def test_evaluate_rows(llama_stack_client, text_model_id, scoring_fn_id):
+ register_dataset(llama_stack_client, for_generation=True, dataset_id="test_dataset_for_eval")
+ response = llama_stack_client.datasets.list()
+ assert any(x.identifier == "test_dataset_for_eval" for x in response)
+
+ rows = llama_stack_client.datasetio.get_rows_paginated(
+ dataset_id="test_dataset_for_eval",
+ rows_in_page=3,
+ )
+ assert len(rows.rows) == 3
+
+ scoring_functions = [
+ scoring_fn_id,
+ ]
+ benchmark_id = str(uuid.uuid4())
+ llama_stack_client.benchmarks.register(
+ benchmark_id=benchmark_id,
+ dataset_id="test_dataset_for_eval",
+ scoring_functions=scoring_functions,
+ )
+ list_benchmarks = llama_stack_client.benchmarks.list()
+ assert any(x.identifier == benchmark_id for x in list_benchmarks)
+
+ response = llama_stack_client.eval.evaluate_rows(
+ benchmark_id=benchmark_id,
+ input_rows=rows.rows,
+ scoring_functions=scoring_functions,
+ benchmark_config={
+ "eval_candidate": {
+ "type": "model",
+ "model": text_model_id,
+ "sampling_params": {
+ "temperature": 0.0,
+ },
+ },
+ },
+ )
+
+ assert len(response.generations) == 3
+ assert scoring_fn_id in response.scores
+
+
+@pytest.mark.parametrize("scoring_fn_id", ["basic::subset_of"])
+def test_evaluate_benchmark(llama_stack_client, text_model_id, scoring_fn_id):
+ register_dataset(llama_stack_client, for_generation=True, dataset_id="test_dataset_for_eval_2")
+ benchmark_id = str(uuid.uuid4())
+ llama_stack_client.benchmarks.register(
+ benchmark_id=benchmark_id,
+ dataset_id="test_dataset_for_eval_2",
+ scoring_functions=[scoring_fn_id],
+ )
+
+ response = llama_stack_client.eval.run_eval(
+ benchmark_id=benchmark_id,
+ benchmark_config={
+ "eval_candidate": {
+ "type": "model",
+ "model": text_model_id,
+ "sampling_params": {
+ "temperature": 0.0,
+ },
+ },
+ },
+ )
+ assert response.job_id == "0"
+ job_status = llama_stack_client.eval.jobs.status(job_id=response.job_id, benchmark_id=benchmark_id)
+ assert job_status and job_status == "completed"
+
+ eval_response = llama_stack_client.eval.jobs.retrieve(job_id=response.job_id, benchmark_id=benchmark_id)
+ assert eval_response is not None
+ assert len(eval_response.generations) == 5
+ assert scoring_fn_id in eval_response.scores
diff --git a/tests/integration/fixtures/__init__.py b/tests/integration/fixtures/__init__.py
new file mode 100644
index 000000000..756f351d8
--- /dev/null
+++ b/tests/integration/fixtures/__init__.py
@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
diff --git a/tests/integration/fixtures/common.py b/tests/integration/fixtures/common.py
new file mode 100644
index 000000000..e410039e7
--- /dev/null
+++ b/tests/integration/fixtures/common.py
@@ -0,0 +1,207 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import copy
+import inspect
+import logging
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+import yaml
+from llama_stack_client import LlamaStackClient
+
+from llama_stack import LlamaStackAsLibraryClient
+from llama_stack.apis.datatypes import Api
+from llama_stack.distribution.stack import run_config_from_adhoc_config_spec
+from llama_stack.env import get_env_or_fail
+
+from .recordable_mock import RecordableMock
+
+
+@pytest.fixture(scope="session")
+def provider_data():
+ # TODO: this needs to be generalized so each provider can have a sample provider data just
+ # like sample run config on which we can do replace_env_vars()
+ keymap = {
+ "TAVILY_SEARCH_API_KEY": "tavily_search_api_key",
+ "BRAVE_SEARCH_API_KEY": "brave_search_api_key",
+ "FIREWORKS_API_KEY": "fireworks_api_key",
+ "GEMINI_API_KEY": "gemini_api_key",
+ "OPENAI_API_KEY": "openai_api_key",
+ "TOGETHER_API_KEY": "together_api_key",
+ "ANTHROPIC_API_KEY": "anthropic_api_key",
+ "GROQ_API_KEY": "groq_api_key",
+ "WOLFRAM_ALPHA_API_KEY": "wolfram_alpha_api_key",
+ }
+ provider_data = {}
+ for key, value in keymap.items():
+ if os.environ.get(key):
+ provider_data[value] = os.environ[key]
+ return provider_data
+
+
+@pytest.fixture(scope="session")
+def llama_stack_client_with_mocked_inference(llama_stack_client, request):
+ """
+ Returns a client with mocked inference APIs and tool runtime APIs that use recorded responses by default.
+
+ If --record-responses is passed, it will call the real APIs and record the responses.
+ """
+ if not isinstance(llama_stack_client, LlamaStackAsLibraryClient):
+ logging.warning(
+ "llama_stack_client_with_mocked_inference is not supported for this client, returning original client without mocking"
+ )
+ return llama_stack_client
+
+ record_responses = request.config.getoption("--record-responses")
+ cache_dir = Path(__file__).parent / "recorded_responses"
+
+ # Create a shallow copy of the client to avoid modifying the original
+ client = copy.copy(llama_stack_client)
+
+ # Get the inference API used by the agents implementation
+ agents_impl = client.async_client.impls[Api.agents]
+ original_inference = agents_impl.inference_api
+
+ # Create a new inference object with the same attributes
+ inference_mock = copy.copy(original_inference)
+
+ # Replace the methods with recordable mocks
+ inference_mock.chat_completion = RecordableMock(
+ original_inference.chat_completion, cache_dir, "chat_completion", record=record_responses
+ )
+ inference_mock.completion = RecordableMock(
+ original_inference.completion, cache_dir, "text_completion", record=record_responses
+ )
+ inference_mock.embeddings = RecordableMock(
+ original_inference.embeddings, cache_dir, "embeddings", record=record_responses
+ )
+
+ # Replace the inference API in the agents implementation
+ agents_impl.inference_api = inference_mock
+
+ original_tool_runtime_api = agents_impl.tool_runtime_api
+ tool_runtime_mock = copy.copy(original_tool_runtime_api)
+
+ # Replace the methods with recordable mocks
+ tool_runtime_mock.invoke_tool = RecordableMock(
+ original_tool_runtime_api.invoke_tool, cache_dir, "invoke_tool", record=record_responses
+ )
+ agents_impl.tool_runtime_api = tool_runtime_mock
+
+ # Also update the client.inference for consistency
+ client.inference = inference_mock
+
+ return client
+
+
+@pytest.fixture(scope="session")
+def inference_provider_type(llama_stack_client):
+ providers = llama_stack_client.providers.list()
+ inference_providers = [p for p in providers if p.api == "inference"]
+ assert len(inference_providers) > 0, "No inference providers found"
+ return inference_providers[0].provider_type
+
+
+@pytest.fixture(scope="session")
+def client_with_models(
+ llama_stack_client,
+ text_model_id,
+ vision_model_id,
+ embedding_model_id,
+ embedding_dimension,
+ judge_model_id,
+):
+ client = llama_stack_client
+
+ providers = [p for p in client.providers.list() if p.api == "inference"]
+ assert len(providers) > 0, "No inference providers found"
+ inference_providers = [p.provider_id for p in providers if p.provider_type != "inline::sentence-transformers"]
+
+ model_ids = {m.identifier for m in client.models.list()}
+ model_ids.update(m.provider_resource_id for m in client.models.list())
+
+ if text_model_id and text_model_id not in model_ids:
+ client.models.register(model_id=text_model_id, provider_id=inference_providers[0])
+ if vision_model_id and vision_model_id not in model_ids:
+ client.models.register(model_id=vision_model_id, provider_id=inference_providers[0])
+ if judge_model_id and judge_model_id not in model_ids:
+ client.models.register(model_id=judge_model_id, provider_id=inference_providers[0])
+
+ if embedding_model_id and embedding_model_id not in model_ids:
+ # try to find a provider that supports embeddings, if sentence-transformers is not available
+ selected_provider = None
+ for p in providers:
+ if p.provider_type == "inline::sentence-transformers":
+ selected_provider = p
+ break
+
+ selected_provider = selected_provider or providers[0]
+ client.models.register(
+ model_id=embedding_model_id,
+ provider_id=selected_provider.provider_id,
+ model_type="embedding",
+ metadata={"embedding_dimension": embedding_dimension or 384},
+ )
+ return client
+
+
+@pytest.fixture(scope="session")
+def available_shields(llama_stack_client):
+ return [shield.identifier for shield in llama_stack_client.shields.list()]
+
+
+@pytest.fixture(scope="session")
+def model_providers(llama_stack_client):
+ return {x.provider_id for x in llama_stack_client.providers.list() if x.api == "inference"}
+
+
+@pytest.fixture(autouse=True)
+def skip_if_no_model(request):
+ model_fixtures = ["text_model_id", "vision_model_id", "embedding_model_id", "judge_model_id"]
+ test_func = request.node.function
+
+ actual_params = inspect.signature(test_func).parameters.keys()
+ for fixture in model_fixtures:
+ # Only check fixtures that are actually in the test function's signature
+ if fixture in actual_params and fixture in request.fixturenames and not request.getfixturevalue(fixture):
+ pytest.skip(f"{fixture} empty - skipping test")
+
+
+@pytest.fixture(scope="session")
+def llama_stack_client(request, provider_data, text_model_id):
+ config = request.config.getoption("--stack-config")
+ if not config:
+ config = get_env_or_fail("LLAMA_STACK_CONFIG")
+
+ if not config:
+ raise ValueError("You must specify either --stack-config or LLAMA_STACK_CONFIG")
+
+ # check if this looks like a URL
+ if config.startswith("http") or "//" in config:
+ return LlamaStackClient(
+ base_url=config,
+ provider_data=provider_data,
+ )
+
+ if "=" in config:
+ run_config = run_config_from_adhoc_config_spec(config)
+ run_config_file = tempfile.NamedTemporaryFile(delete=False, suffix=".yaml")
+ with open(run_config_file.name, "w") as f:
+ yaml.dump(run_config.model_dump(), f)
+ config = run_config_file.name
+
+ client = LlamaStackAsLibraryClient(
+ config,
+ provider_data=provider_data,
+ skip_logger_removal=True,
+ )
+ if not client.initialize():
+ raise RuntimeError("Initialization failed")
+
+ return client
diff --git a/tests/integration/fixtures/recordable_mock.py b/tests/integration/fixtures/recordable_mock.py
new file mode 100644
index 000000000..632d5b3ef
--- /dev/null
+++ b/tests/integration/fixtures/recordable_mock.py
@@ -0,0 +1,221 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+import importlib
+import json
+import os
+import re
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+
+
+class RecordableMock:
+ """A mock that can record and replay API responses."""
+
+ def __init__(self, real_func, cache_dir, func_name, record=False):
+ self.real_func = real_func
+ self.json_path = Path(cache_dir) / f"{func_name}.json"
+ self.record = record
+ self.cache = {}
+
+ # Load existing cache if available and not recording
+ if self.json_path.exists():
+ try:
+ with open(self.json_path, "r") as f:
+ self.cache = json.load(f)
+ except Exception as e:
+ print(f"Error loading cache from {self.json_path}: {e}")
+ raise
+
+ async def __call__(self, *args, **kwargs):
+ """
+ Returns a coroutine that when awaited returns the result or an async generator,
+ matching the behavior of the original function.
+ """
+ # Create a cache key from the arguments
+ key = self._create_cache_key(args, kwargs)
+
+ if self.record:
+ # In record mode, always call the real function
+ real_result = self.real_func(*args, **kwargs)
+
+ # If it's a coroutine, we need to create a wrapper coroutine
+ if hasattr(real_result, "__await__"):
+ # Define a coroutine function that will record the result
+ async def record_coroutine():
+ try:
+ # Await the real coroutine
+ result = await real_result
+
+ # Check if the result is an async generator
+ if hasattr(result, "__aiter__"):
+ # It's an async generator, so we need to record its chunks
+ chunks = []
+
+ # Create and return a new async generator that records chunks
+ async def recording_generator():
+ nonlocal chunks
+ async for chunk in result:
+ chunks.append(chunk)
+ yield chunk
+ # After all chunks are yielded, save to cache
+ self.cache[key] = {"type": "generator", "chunks": chunks}
+ self._save_cache()
+
+ return recording_generator()
+ else:
+ # It's a regular result, save it to cache
+ self.cache[key] = {"type": "value", "value": result}
+ self._save_cache()
+ return result
+ except Exception as e:
+ print(f"Error in recording mode: {e}")
+ raise
+
+ return await record_coroutine()
+ else:
+ # It's already an async generator, so we need to record its chunks
+ async def record_generator():
+ chunks = []
+ async for chunk in real_result:
+ chunks.append(chunk)
+ yield chunk
+ # After all chunks are yielded, save to cache
+ self.cache[key] = {"type": "generator", "chunks": chunks}
+ self._save_cache()
+
+ return record_generator()
+ elif key not in self.cache:
+ # In replay mode, if the key is not in the cache, throw an error
+ raise KeyError(
+ f"No cached response found for key: {key}\nRun with --record-responses to record this response."
+ )
+ else:
+ # In replay mode with a cached response
+ cached_data = self.cache[key]
+
+ # Check if it's a value or chunks
+ if cached_data.get("type") == "value":
+ # It's a regular value
+ return self._reconstruct_object(cached_data["value"])
+ else:
+ # It's chunks from an async generator
+ async def replay_generator():
+ for chunk in cached_data["chunks"]:
+ yield self._reconstruct_object(chunk)
+
+ return replay_generator()
+
+ def _create_cache_key(self, args, kwargs):
+ """Create a hashable key from the function arguments, ignoring auto-generated IDs."""
+ # Convert to JSON strings with sorted keys
+ key = json.dumps((args, kwargs), sort_keys=True, default=self._json_default)
+
+ # Post-process the key with regex to replace IDs with placeholders
+ # Replace UUIDs and similar patterns
+ key = re.sub(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "", key)
+
+ # Replace temporary file paths created by tempfile.mkdtemp()
+ key = re.sub(r"/var/folders/[^,'\"\s]+", "", key)
+
+ # Replace /tmp/ paths which are also commonly used for temporary files
+ key = re.sub(r"/tmp/[^,'\"\s]+", "", key)
+
+ return key
+
+ def _save_cache(self):
+ """Save the cache to disk in JSON format."""
+ os.makedirs(self.json_path.parent, exist_ok=True)
+
+ # Write the JSON file with pretty formatting
+ try:
+ with open(self.json_path, "w") as f:
+ json.dump(self.cache, f, indent=2, sort_keys=True, default=self._json_default)
+ # write another empty line at the end of the file to make pre-commit happy
+ f.write("\n")
+ except Exception as e:
+ print(f"Error saving JSON cache: {e}")
+
+ def _json_default(self, obj):
+ """Default function for JSON serialization of objects."""
+
+ if isinstance(obj, datetime):
+ return {
+ "__datetime__": obj.isoformat(),
+ "__module__": obj.__class__.__module__,
+ "__class__": obj.__class__.__name__,
+ }
+
+ if isinstance(obj, Enum):
+ return {
+ "__enum__": obj.__class__.__name__,
+ "value": obj.value,
+ "__module__": obj.__class__.__module__,
+ }
+
+ # Handle Pydantic models
+ if hasattr(obj, "model_dump"):
+ model_data = obj.model_dump()
+ return {
+ "__pydantic__": obj.__class__.__name__,
+ "__module__": obj.__class__.__module__,
+ "data": model_data,
+ }
+
+ def _reconstruct_object(self, data):
+ """Reconstruct an object from its JSON representation."""
+ if isinstance(data, dict):
+ # Check if this is a serialized datetime
+ if "__datetime__" in data:
+ try:
+ module_name = data.get("__module__", "datetime")
+ class_name = data.get("__class__", "datetime")
+
+ # Try to import the specific datetime class
+ module = importlib.import_module(module_name)
+ dt_class = getattr(module, class_name)
+
+ # Parse the ISO format string
+ dt = dt_class.fromisoformat(data["__datetime__"])
+ return dt
+ except (ImportError, AttributeError, ValueError) as e:
+ print(f"Error reconstructing datetime: {e}")
+ return data
+
+ # Check if this is a serialized enum
+ elif "__enum__" in data:
+ try:
+ module_name = data.get("__module__", "builtins")
+ enum_class = self._import_class(module_name, data["__enum__"])
+ return enum_class(data["value"])
+ except (ImportError, AttributeError) as e:
+ print(f"Error reconstructing enum: {e}")
+ return data
+
+ # Check if this is a serialized Pydantic model
+ elif "__pydantic__" in data:
+ try:
+ module_name = data.get("__module__", "builtins")
+ model_class = self._import_class(module_name, data["__pydantic__"])
+ return model_class(**self._reconstruct_object(data["data"]))
+ except (ImportError, AttributeError) as e:
+ print(f"Error reconstructing Pydantic model: {e}")
+ return data
+
+ # Regular dictionary
+ return {k: self._reconstruct_object(v) for k, v in data.items()}
+
+ # Handle lists
+ elif isinstance(data, list):
+ return [self._reconstruct_object(item) for item in data]
+
+ # Return primitive types as is
+ return data
+
+ def _import_class(self, module_name, class_name):
+ """Import a class from a module."""
+ module = __import__(module_name, fromlist=[class_name])
+ return getattr(module, class_name)
diff --git a/tests/integration/fixtures/recorded_responses/chat_completion.json b/tests/integration/fixtures/recorded_responses/chat_completion.json
new file mode 100644
index 000000000..8694cc271
--- /dev/null
+++ b/tests/integration/fixtures/recorded_responses/chat_completion.json
@@ -0,0 +1,52345 @@
+{
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant Always respond with tool calls no matter what. \", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Get the boiling point of polyjuice with a tool call.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"celcius\": \"true\", \"liquid_name\": \"polyjuice\"}, \"call_id\": \"\", \"tool_name\": \"get_boiling_point\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"-100\", \"role\": \"tool\", \"tool_name\": \"get_boiling_point\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"celcius\": \"false\", \"liquid_name\": \"polyjuice\"}, \"call_id\": \"\", \"tool_name\": \"get_boiling_point\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"-100\", \"role\": \"tool\", \"tool_name\": \"get_boiling_point\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "The",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " boiling point of polyjuice is -100",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " degrees Fahrenheit.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "9ksjMloe",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:58.345129+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "6aGYLk4UShyrQ7uz",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 139
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "9ksjMloe",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:58.345170+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "6aGYLk4UShyrQ7uz",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 23
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "9ksjMloe",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:58.345177+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "6aGYLk4UShyrQ7uz",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 162
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant Always respond with tool calls no matter what. \", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Get the boiling point of polyjuice with a tool call.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"celcius\": \"true\", \"liquid_name\": \"polyjuice\"}, \"call_id\": \"\", \"tool_name\": \"get_boiling_point\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"-100\", \"role\": \"tool\", \"tool_name\": \"get_boiling_point\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "{\"",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "type\": \"function\", \"name\": \"get_boiling_point",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "\", \"parameters\": {\"liquid_name\":",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " \"polyjuice\", \"celcius\": \"false\"}}",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "celcius": "false",
+ "liquid_name": "polyjuice"
+ },
+ "call_id": "55492018-ad19-4593-9171-2b5dc2089960",
+ "tool_name": "get_boiling_point"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "vTzYAYfO",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:56.985637+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "H8ytqaQLQXe6sEEJ",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 91
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "vTzYAYfO",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:56.985707+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "H8ytqaQLQXe6sEEJ",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 45
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "vTzYAYfO",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:56.985718+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "H8ytqaQLQXe6sEEJ",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 136
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant Always respond with tool calls no matter what. \", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Get the boiling point of polyjuice with a tool call.\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "{\"type\": \"function\", \"name\": \"get_boiling",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_point\", \"parameters\": {\"liquid_name\": \"polyjuice",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "\", \"celcius\": \"true\"}}",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "celcius": "true",
+ "liquid_name": "polyjuice"
+ },
+ "call_id": "6dd93d40-18ea-40c1-9e4d-78b3bd865e67",
+ "tool_name": "get_boiling_point"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "tBuntiC1",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:54.993737+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "5SueXj79Q2e5n37g",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 43
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "tBuntiC1",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:54.993758+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "5SueXj79Q2e5n37g",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "tBuntiC1",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:54.993761+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "5SueXj79Q2e5n37g",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 53
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Call get_boiling_point and answer What is the boiling point of polyjuice?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"celcius\": \"true\", \"liquid_name\": \"polyjuice\"}, \"call_id\": \"\", \"tool_name\": \"get_boiling_point\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"-100\", \"role\": \"tool\", \"tool_name\": \"get_boiling_point\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "The",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " boiling point of polyjuice is -100\u00b0C.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "03QQgo3b",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:34.636678+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "mE4SuRfcQUOcOyP2",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 85
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "03QQgo3b",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:34.636767+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "mE4SuRfcQUOcOyP2",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 22
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "03QQgo3b",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:34.636773+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "mE4SuRfcQUOcOyP2",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 107
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Call get_boiling_point and answer What is the boiling point of polyjuice?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"celcius\": \"true\", \"liquid_name\": \"polyjuice\"}, \"call_id\": \"\", \"tool_name\": \"get_boiling_point_with_metadata\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"-100\", \"role\": \"tool\", \"tool_name\": \"get_boiling_point_with_metadata\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point_with_metadata\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "The",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " boiling point of polyjuice is -100\u00b0C.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "vzNuoz4e",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.792508+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "vNRMmadcTVmfkn5-",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 87
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "vzNuoz4e",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.792536+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "vNRMmadcTVmfkn5-",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 22
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "vzNuoz4e",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.792544+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "vNRMmadcTVmfkn5-",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 109
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Call get_boiling_point and answer What is the boiling point of polyjuice?\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "{\"type\": \"function\", \"name\": \"get_boiling_point\",",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " \"parameters\": {\"liquid_name\": \"polyjuice\", \"celci",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "us\": \"true\"}}",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "celcius": "true",
+ "liquid_name": "polyjuice"
+ },
+ "call_id": "98d5962a-eab3-4d83-bca4-d4d6aa54f1dc",
+ "tool_name": "get_boiling_point"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "1A0bWgLL",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:24.102366+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "4a5HMcM9R3uWB4Cv",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 37
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "1A0bWgLL",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:24.102404+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "4a5HMcM9R3uWB4Cv",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "1A0bWgLL",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:24.102411+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "4a5HMcM9R3uWB4Cv",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 47
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Call get_boiling_point and answer What is the boiling point of polyjuice?\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Returns the boiling point of a liquid in Celcius or Fahrenheit\", \"parameters\": {\"celcius\": {\"default\": true, \"description\": \"Whether to return the boiling point in Celcius\", \"param_type\": \"bool\", \"required\": false}, \"liquid_name\": {\"default\": null, \"description\": \"The name of the liquid\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"get_boiling_point_with_metadata\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "{\"type\": \"function\", \"name\": \"get_boiling_point_with",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_metadata\", \"parameters\": {\"liquid_name\": \"polyjuice\", \"",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "celcius\": \"true\"}}",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "celcius": "true",
+ "liquid_name": "polyjuice"
+ },
+ "call_id": "ee5ac18d-de3b-4985-9e93-545de166d3e2",
+ "tool_name": "get_boiling_point_with_metadata"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "dsGyjpUB",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.316534+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "BO0etAZ6RFmGmLCW",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 37
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "dsGyjpUB",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.316569+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "BO0etAZ6RFmGmLCW",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "dsGyjpUB",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:45.316576+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "BO0etAZ6RFmGmLCW",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 47
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Give me a sentence that contains the word: hello\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": []}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "The",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " customer smiled and said \"hello\" to the friendly store clerk.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "R9a1QHt4",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:19.586300+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "t-ZRvSMzTCudL6SB",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 30
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "R9a1QHt4",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:19.586359+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "t-ZRvSMzTCudL6SB",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 24
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "R9a1QHt4",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:19.586367+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "t-ZRvSMzTCudL6SB",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 54
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv file, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"# User provided a file accessible to you at \\\"\"\\nYou can use code_interpreter to load and inspect it.\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\ndf = pd.read_csv(\\\"\")\\nprint(df.head())\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "The",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " error message indicates that the `bwrap.core",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "` module is not found. This",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " is because the `bwrap` module is not installed in",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " your Python environment.\n\nTo fix this issue,",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " you can use the `knowledge_search` function to describe the CSV",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " file. This function can be used to",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " search for information in a database, and it might have",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " access to information about the CSV file.\n\nHere is an example of",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " how you can use the `knowledge_search` function to describe the CSV",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " file:\n\n```\n{\n \"type\": \"function\",\n \"",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "name\": \"knowledge_search\",\n \"parameters\":",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " {\n \"query\": \"Describe the CSV file at /var",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "/folders/cz/vyh7y1",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "d11xg881lsxsshnc5c0000gn/T",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "/tmpvto5j2dr/u8MQ2jywin",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "flation.csv\"\n }\n}\n```",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "9UjZne1U",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:15.341367+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 149
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "9UjZne1U",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:15.341380+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 188
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "9UjZne1U",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:15.341383+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 337
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv file, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"# User provided a file accessible to you at \\\"\"\\nYou can use code_interpreter to load and inspect it.\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "import pandas as pd\ndf = pd",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".read_csv(\"/var/folders/cz/vyh7y1",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "d11xg881lsxsshnc5c0000",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "gn/T/tmpvto5j2",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "dr/u8MQ2jywinflation.csv\")\nprint(df",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".head())",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "code": "import pandas as pd\ndf = pd.read_csv(\"/var/folders/cz/vyh7y1d11xg881lsxsshnc5c0000gn/T/tmpvto5j2dr/u8MQ2jywinflation.csv\")\nprint(df.head())"
+ },
+ "call_id": "ecc9db21-332f-4931-8820-cf139f8a0b88",
+ "tool_name": {
+ "__enum__": "BuiltinTool",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "code_interpreter"
+ }
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "6VEDipbd",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:14.030541+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 37
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "6VEDipbd",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:14.030577+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "6VEDipbd",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:42:14.030584+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "cOvUfJZLSK2vci9f",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 47
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\n# Load data\\ndf = pd.read_csv(\\\"\")\\n# Rows\\nprint(\\\"Number of rows and columns in the data:\\\", df.shape)\\n# Columns\\nprint(\\\"Columns of the data are:\\\", len(df.columns))\\n# Column names\\nprint(\\\"Columns of the data are:\\\", df.columns)\\n# Column dtypes\\nprint(\\\"Datatype of the columns are:\\\", df.dtypes)\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"It seems that the file \\\"\" does not exist. \\n\\nTo describe the csv file, you need to provide the actual file path or the file itself. If the file is too large to be uploaded, you can provide a sample of the csv file and I can help you describe it. \\n\\nHere is an example of how you can describe a csv file using pandas:\\n\\n```\\nimport pandas as pd\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n# Print the first 5 rows of the data\\nprint(df.head())\\n# Print the last 5 rows of the data\\nprint(df.tail())\\n# Print the summary statistics of the data\\nprint(df.describe())\\n# Print the data types of each column\\nprint(df.dtypes)\\n# Print the number of missing values in each column\\nprint(df.isnull().sum())\\n```\\n\\nThis will give you an idea of what the csv file contains.\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Plot average yearly inflation as a time series\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\nimport matplotlib.pyplot as plt\\n\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n\\n# Convert date column to datetime\\ndf['date'] = pd.to_datetime(df['date'])\\n\\n# Group by year and calculate average inflation\\naverage_inflation = df.groupby(df['date'].dt.year)['inflation'].mean()\\n\\n# Plot the time series\\nplt.figure(figsize=(10,6))\\nplt.plot(average_inflation.index, average_inflation.values, marker='o')\\nplt.title('Average Yearly Inflation')\\nplt.xlabel('Year')\\nplt.ylabel('Average Inflation')\\nplt.grid(True)\\nplt.show()\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "This",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " code will create a line plot of the average yearly inflation over time",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": ". The x-axis represents the year and the y-axis represents the",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " average inflation. Each point on the plot represents the average inflation for",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " a particular year.\n\nPlease note that you need to replace 'in",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "flation.csv' with the actual path to your",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " csv file. Also, this code assumes that the csv file has a",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " column named 'date' and another column named 'inflation'. If your",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " csv file has different column names, you need to replace 'date' and",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " 'inflation' with the actual column names.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "Hm1BkrMQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:41.982115+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 636
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "Hm1BkrMQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:41.982147+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 126
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "Hm1BkrMQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:41.982153+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 762
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\n# Load data\\ndf = pd.read_csv(\\\"\")\\n# Rows\\nprint(\\\"Number of rows and columns in the data:\\\", df.shape)\\n# Columns\\nprint(\\\"Columns of the data are:\\\", len(df.columns))\\n# Column names\\nprint(\\\"Columns of the data are:\\\", df.columns)\\n# Column dtypes\\nprint(\\\"Datatype of the columns are:\\\", df.dtypes)\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"It seems that the file \\\"\" does not exist. \\n\\nTo describe the csv file, you need to provide the actual file path or the file itself. If the file is too large to be uploaded, you can provide a sample of the csv file and I can help you describe it. \\n\\nHere is an example of how you can describe a csv file using pandas:\\n\\n```\\nimport pandas as pd\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n# Print the first 5 rows of the data\\nprint(df.head())\\n# Print the last 5 rows of the data\\nprint(df.tail())\\n# Print the summary statistics of the data\\nprint(df.describe())\\n# Print the data types of each column\\nprint(df.dtypes)\\n# Print the number of missing values in each column\\nprint(df.isnull().sum())\\n```\\n\\nThis will give you an idea of what the csv file contains.\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Plot average yearly inflation as a time series\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "import pandas as pd\nimport matplotlib.pyplot as plt\n\n# Load",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " data\ndf = pd.read_csv('inflation.csv')\n\n#",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " Convert date column to datetime\ndf['date'] = pd.to",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_datetime(df['date'])\n\n# Group by year and calculate average inflation",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "\naverage_inflation = df.groupby(df['date'].dt.year",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ")['inflation'].mean()\n\n# Plot the time series\nplt",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".figure(figsize=(10,6))\nplt.plot(average_inflation",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".index, average_inflation.values, marker='o')\nplt.title",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "('Average Yearly Inflation')\nplt.xlabel('Year')\nplt",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".ylabel('Average Inflation')\nplt.grid(True)\n",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "plt.show()",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "code": "import pandas as pd\nimport matplotlib.pyplot as plt\n\n# Load data\ndf = pd.read_csv('inflation.csv')\n\n# Convert date column to datetime\ndf['date'] = pd.to_datetime(df['date'])\n\n# Group by year and calculate average inflation\naverage_inflation = df.groupby(df['date'].dt.year)['inflation'].mean()\n\n# Plot the time series\nplt.figure(figsize=(10,6))\nplt.plot(average_inflation.index, average_inflation.values, marker='o')\nplt.title('Average Yearly Inflation')\nplt.xlabel('Year')\nplt.ylabel('Average Inflation')\nplt.grid(True)\nplt.show()"
+ },
+ "call_id": "4849f8b5-bbb8-4c7e-8f19-498dd559dbe2",
+ "tool_name": {
+ "__enum__": "BuiltinTool",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "code_interpreter"
+ }
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "ZKjmS7HQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:30.999750+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 450
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "ZKjmS7HQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:30.999780+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "ZKjmS7HQ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:32:30.999786+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "T857cf9QSamVBOAy",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 460
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\n# Load data\\ndf = pd.read_csv(\\\"\")\\n# Rows\\nprint(\\\"Number of rows and columns in the data:\\\", df.shape)\\n# Columns\\nprint(\\\"Columns of the data are:\\\", len(df.columns))\\n# Column names\\nprint(\\\"Columns of the data are:\\\", df.columns)\\n# Column dtypes\\nprint(\\\"Datatype of the columns are:\\\", df.dtypes)\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"It seems that the file \\\"\" does not exist. \\n\\nTo describe the csv file, you need to provide the actual file path or the file itself. If you are using a remote server or a local machine, you can use the `pd.read_csv()` function to load the csv file. \\n\\nHere is an example:\\n\\n```python\\nimport pandas as pd\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n# Print the first 5 rows of the dataframe\\nprint(df.head())\\n# Print the summary of the dataframe\\nprint(df.info())\\nprint(df.describe())\\n```\\n\\nThis will print the first 5 rows of the dataframe, the summary of the dataframe (including the index dtype and column dtypes, non-nullable counts, and memory usage), and the descriptive statistics of the dataframe.\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Plot average yearly inflation as a time series\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\nimport matplotlib.pyplot as plt\\n\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n\\n# Convert 'date' column to datetime\\ndf['date'] = pd.to_datetime(df['date'])\\n\\n# Group by year and calculate average inflation\\naverage_inflation = df.groupby(df['date'].dt.year)['inflation'].mean()\\n\\n# Plot the time series\\nplt.figure(figsize=(10,6))\\nplt.plot(average_inflation.index, average_inflation.values, marker='o')\\nplt.title('Average Yearly Inflation')\\nplt.xlabel('Year')\\nplt.ylabel('Average Inflation')\\nplt.grid(True)\\nplt.show()\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "This",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " code will create a line plot of the average yearly inflation over time.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " The x-axis represents the year and the y-axis represents the average inflation.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " The plot also includes a title, labels for the x and y axes,",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " and a grid for better visibility.\n\nPlease note that you need to replace '",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "inflation.csv' with the actual path to",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " your csv file. Also, this code",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " assumes that the 'date' column in your",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " csv file is in a format that can be",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " parsed by pandas' `to_datetime` function. If your date",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " column is in a different format, you",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " may need to specify the format using the",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " `format` parameter of `to_datetime`.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "Yv7iXXNJ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:50.214420+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 621
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "Yv7iXXNJ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:50.214481+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 143
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "Yv7iXXNJ",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:50.214490+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 764
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\n# Load data\\ndf = pd.read_csv(\\\"\")\\n# Rows\\nprint(\\\"Number of rows and columns in the data:\\\", df.shape)\\n# Columns\\nprint(\\\"Columns of the data are:\\\", len(df.columns))\\n# Column names\\nprint(\\\"Columns of the data are:\\\", df.columns)\\n# Column dtypes\\nprint(\\\"Datatype of the columns are:\\\", df.dtypes)\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"It seems that the file \\\"\" does not exist. \\n\\nTo describe the csv file, you need to provide the actual file path or the file itself. If you are using a remote server or a local machine, you can use the `pd.read_csv()` function to load the csv file. \\n\\nHere is an example:\\n\\n```python\\nimport pandas as pd\\n# Load data\\ndf = pd.read_csv('inflation.csv')\\n# Print the first 5 rows of the dataframe\\nprint(df.head())\\n# Print the summary of the dataframe\\nprint(df.info())\\nprint(df.describe())\\n```\\n\\nThis will print the first 5 rows of the dataframe, the summary of the dataframe (including the index dtype and column dtypes, non-nullable counts, and memory usage), and the descriptive statistics of the dataframe.\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Plot average yearly inflation as a time series\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "import pandas as pd\nimport matplotlib.pyplot as plt\n\n# Load data",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "\ndf = pd.read_csv('inflation.csv')\n\n# Convert",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " 'date' column to datetime\ndf['date'] = pd.to",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_datetime(df['date'])\n\n# Group by year and calculate average inflation\naverage",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_inflation = df.groupby(df['date'].dt.year)['inflation'].",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "mean()\n\n# Plot the time series\nplt.figure(figsize=(10,6",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "))\nplt.plot(average_inflation.index, average_inflation.values, marker",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "='o')\nplt.title('Average Yearly Inflation')\nplt.xlabel",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "('Year')\nplt.ylabel('Average Inflation')\nplt.grid(True)\nplt",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".show()",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "code": "import pandas as pd\nimport matplotlib.pyplot as plt\n\n# Load data\ndf = pd.read_csv('inflation.csv')\n\n# Convert 'date' column to datetime\ndf['date'] = pd.to_datetime(df['date'])\n\n# Group by year and calculate average inflation\naverage_inflation = df.groupby(df['date'].dt.year)['inflation'].mean()\n\n# Plot the time series\nplt.figure(figsize=(10,6))\nplt.plot(average_inflation.index, average_inflation.values, marker='o')\nplt.title('Average Yearly Inflation')\nplt.xlabel('Year')\nplt.ylabel('Average Inflation')\nplt.grid(True)\nplt.show()"
+ },
+ "call_id": "62e5a10d-8a59-41e7-9f0e-87cabc7d15fa",
+ "tool_name": {
+ "__enum__": "BuiltinTool",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "code_interpreter"
+ }
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "dv6g9n2H",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:48.391101+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 433
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "dv6g9n2H",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:48.391113+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "dv6g9n2H",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:48.391116+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "srzTfsP6Sr-co-Ll",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 443
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"# User provided a file accessible to you at \\\"\"\\nYou can use code_interpreter to load and inspect it.\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"code\": \"import pandas as pd\\n# Load data\\ndf = pd.read_csv(\\\"\")\\n# Rows\\nprint(\\\"Number of rows and columns in the data:\\\", df.shape)\\n# Columns\\nprint(\\\"Columns of the data are:\\\", len(df.columns))\\n# Column names\\nprint(\\\"Columns of the data are:\\\", df.columns)\\n# Column dtypes\\nprint(\\\"Datatype of the columns are:\\\", df.dtypes)\"}, \"call_id\": \"\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": \"completed\\n[stderr]\\nTraceback (most recent call last):\\n line 5, in \\n from bwrap.core import main\\nModuleNotFoundError: No module named 'bwrap.core'\\n[/stderr]\", \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "It",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " seems that the file \"/var/folders/cz/vyh7",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "y1d11xg881lsxsshnc5c",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "0000gn/T/tmpvto5j",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "2dr/JwKzVg",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "5Ainflation.csv\" does not exist. \n\nTo describe the",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " csv file, you need to provide the actual file path or the file itself",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": ". If you are using a remote server or a local machine, you can",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " use the `pd.read_csv()` function to load",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " the csv file. \n\nHere is an example:\n\n```",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "python\nimport pandas as pd\n# Load data\ndf =",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " pd.read_csv('inflation.csv')\n# Print the first 5 rows",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " of the dataframe\nprint(df.head())\n# Print",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " the summary of the dataframe\nprint(df.info())\nprint(df.describe",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "())\n```\n\nThis will print the first 5 rows of the dataframe",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": ", the summary of the dataframe",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " (including the index dtype and column dtypes, non-nullable",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " counts, and memory usage), and the descriptive statistics of the dataframe.",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "qV1E8nPK",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:41.439164+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 215
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "qV1E8nPK",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:41.439188+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 216
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "qV1E8nPK",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:41.439190+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 431
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Here is a csv, can you describe it?\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"# User provided a file accessible to you at \\\"\"\\nYou can use code_interpreter to load and inspect it.\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Execute code\", \"parameters\": {\"code\": {\"default\": null, \"description\": \"The code to execute\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": {\"__enum__\": \"BuiltinTool\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"code_interpreter\"}}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "started"
+ },
+ "tool_call": "",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "import pandas as pd\n# Load data\ndf = pd.read",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "_csv(\"/var/folders/cz/vyh7y1d11",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "xg881lsxsshnc5c0000gn/T/tmp",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "vto5j2dr/JwKzVg",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "5Ainflation.csv\")\n# Rows\nprint(\"Number",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " of rows and columns in the data:\", df.shape)\n# Columns",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "\nprint(\"Columns of the data",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": " are:\", len(df.columns))\n# Column names\nprint",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "(\"Columns of the data are:\", df.columns)\n# Column dt",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": "ypes\nprint(\"Datatype of the columns are:\", df",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "in_progress"
+ },
+ "tool_call": ".dtypes)",
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "code": "import pandas as pd\n# Load data\ndf = pd.read_csv(\"/var/folders/cz/vyh7y1d11xg881lsxsshnc5c0000gn/T/tmpvto5j2dr/JwKzVg5Ainflation.csv\")\n# Rows\nprint(\"Number of rows and columns in the data:\", df.shape)\n# Columns\nprint(\"Columns of the data are:\", len(df.columns))\n# Column names\nprint(\"Columns of the data are:\", df.columns)\n# Column dtypes\nprint(\"Datatype of the columns are:\", df.dtypes)"
+ },
+ "call_id": "87c3ef49-27e0-4561-ade3-83569a0fe236",
+ "tool_name": {
+ "__enum__": "BuiltinTool",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "code_interpreter"
+ }
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "9OTP08Yr",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:39.830624+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 36
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "9OTP08Yr",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:39.830656+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 10
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "9OTP08Yr",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:40:39.830662+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "GG3oeA3qRH6WIf6Z",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 46
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"I am attaching some documentation for Torchtune. Help me answer questions I will ask next.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"Torchtune documentation\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:61fc5\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:af027\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:d5787\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"I'm ready to help you answer questions about Torchtune based on the documentation you provided. What's your first question?\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Tell me how to use LoRA\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"How to use LoRA in Torchtune\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:af027\\nContent: .. _lora_finetune_label:\\n\\n============================\\nFine-Tuning Llama2 with LoRA\\n============================\\n\\nThis guide will teach you about `LoRA `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW ` alone will not handle the definition of which parameters are trainable.\\n See :ref:`below` for how to do this.\\n\\nLet's inspect each of these models a bit more closely.\\n\\n.. code-block:: bash\\n\\n # Print the first layer's self-attention in the usual Llama2 model\\n >>> print(base_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (output_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (pos_embeddings): RotaryPositionalEmbeddings()\\n )\\n\\n # Print the same for Llama2 with LoRA weights\\n >>> print(lora_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): LoRALinear(\\n (dropout): Dropout(p=0.0, inplace=False)\\n \\n\", \"type\": \"text\"}, {\"text\": \"Result 3:\\nDocument_id:af027\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:af027\\nContent: from our Llama2\\nmodel without any wrappers or custom checkpoint conversion logic.\\n\\n.. code-block:: python\\n\\n # Assuming that base_model already has the pretrained Llama2 weights,\\n # this will directly load them into your LoRA model without any conversion necessary.\\n lora_model.load_state_dict(base_model.state_dict(), strict=False)\\n\\n.. note::\\n Whenever loading weights with :code:`strict=False`, you should verify that any missing or extra keys in\\n the loaded :code:`state_dict` are as expected. torchtune's LoRA recipes do this by default via\\n :func:`validate_missing_and_unexpected_for_lora() `.\\n\\nOnce we've loaded the base model weights, we also want to set only LoRA parameters to trainable.\\n\\n.. _setting_trainable_params:\\n\\n.. code-block:: python\\n\\n from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\\n\\n # Fetch all params from the model that are associated with LoRA.\\n lora_params = get_adapter_params(lora_model)\\n\\n # Set requires_grad=True on lora_params, and requires_grad=False on all others.\\n set_trainable_params(lora_model, lora_params)\\n\\n # Print the total number of parameters\\n total_params = sum([p.numel() for p in lora_model.parameters()])\\n trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])\\n print(\\n f\\\"\\\"\\\"\\n {total_params} total params,\\n {trainable_params}\\\" trainable params,\\n {(100.0 * trainable_params / total_params):.2f}% of all params are trainable.\\n \\\"\\\"\\\"\\n )\\n\\n 6742609920 total params,\\n 4194304 trainable params,\\n 0.06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe \", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:61fc5\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:af027\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:d5787\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"I'm ready to help you answer questions about Torchtune based on the documentation you provided. What's your first question?\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Tell me how to use LoRA\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "{\"",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "type\": \"function\", \"name\": \"",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "knowledge_search\", \"parameters\": {\"query\": \"How",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " to use LoRA in Torchtune\"}}",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "query": "How to use LoRA in Torchtune"
+ },
+ "call_id": "14b82c7e-18d4-4b46-8f07-442be700e8ae",
+ "tool_name": "knowledge_search"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "DBZOtUux",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:58.136315+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "XVSIgZRXR_aHBiAN",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 117
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "DBZOtUux",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:58.136380+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "XVSIgZRXR_aHBiAN",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 40
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "DBZOtUux",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:58.136387+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "XVSIgZRXR_aHBiAN",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 157
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"I am attaching some documentation for Torchtune. Help me answer questions I will ask next.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"Torchtune documentation\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:61fc5\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:af027\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:d5787\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "I",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "'m ready to help you answer questions about Torchtune based on the",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " documentation you provided. What's your first question?",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "gFK_4CQi",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:56.169962+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "A2oXFF9fRz2-Lc9N",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 75
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "gFK_4CQi",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:56.169995+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "A2oXFF9fRz2-Lc9N",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 35
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "gFK_4CQi",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:41:56.170001+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "A2oXFF9fRz2-Lc9N",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 110
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"I am attaching some documentation for Torchtune. Help me answer questions I will ask next.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"Torchtune documentation\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:78970\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:8404f\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:cbeb1\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"I'm ready to help you answer questions about Torchtune based on the documentation you provided. What's your first question?\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Tell me how to use LoRA\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"How to use LoRA in Torchtune\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:8404f\\nContent: .. _lora_finetune_label:\\n\\n============================\\nFine-Tuning Llama2 with LoRA\\n============================\\n\\nThis guide will teach you about `LoRA `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW ` alone will not handle the definition of which parameters are trainable.\\n See :ref:`below` for how to do this.\\n\\nLet's inspect each of these models a bit more closely.\\n\\n.. code-block:: bash\\n\\n # Print the first layer's self-attention in the usual Llama2 model\\n >>> print(base_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (output_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (pos_embeddings): RotaryPositionalEmbeddings()\\n )\\n\\n # Print the same for Llama2 with LoRA weights\\n >>> print(lora_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): LoRALinear(\\n (dropout): Dropout(p=0.0, inplace=False)\\n \\n\", \"type\": \"text\"}, {\"text\": \"Result 3:\\nDocument_id:8404f\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:8404f\\nContent: from our Llama2\\nmodel without any wrappers or custom checkpoint conversion logic.\\n\\n.. code-block:: python\\n\\n # Assuming that base_model already has the pretrained Llama2 weights,\\n # this will directly load them into your LoRA model without any conversion necessary.\\n lora_model.load_state_dict(base_model.state_dict(), strict=False)\\n\\n.. note::\\n Whenever loading weights with :code:`strict=False`, you should verify that any missing or extra keys in\\n the loaded :code:`state_dict` are as expected. torchtune's LoRA recipes do this by default via\\n :func:`validate_missing_and_unexpected_for_lora() `.\\n\\nOnce we've loaded the base model weights, we also want to set only LoRA parameters to trainable.\\n\\n.. _setting_trainable_params:\\n\\n.. code-block:: python\\n\\n from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\\n\\n # Fetch all params from the model that are associated with LoRA.\\n lora_params = get_adapter_params(lora_model)\\n\\n # Set requires_grad=True on lora_params, and requires_grad=False on all others.\\n set_trainable_params(lora_model, lora_params)\\n\\n # Print the total number of parameters\\n total_params = sum([p.numel() for p in lora_model.parameters()])\\n trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])\\n print(\\n f\\\"\\\"\\\"\\n {total_params} total params,\\n {trainable_params}\\\" trainable params,\\n {(100.0 * trainable_params / total_params):.2f}% of all params are trainable.\\n \\\"\\\"\\\"\\n )\\n\\n 6742609920 total params,\\n 4194304 trainable params,\\n 0.06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe \", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:78970\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:8404f\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:cbeb1\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"I'm ready to help you answer questions about Torchtune based on the documentation you provided. What's your first question?\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Tell me how to use LoRA\", \"context\": null, \"role\": \"user\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "{\"",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "type\": \"function\", \"name\": \"knowledge_search\", \"parameters\":",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " {\"query\": \"How to use LoRA in Torchtune\"}}",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "parse_status": {
+ "__enum__": "ToolCallParseStatus",
+ "__module__": "llama_stack.apis.common.content_types",
+ "value": "succeeded"
+ },
+ "tool_call": {
+ "arguments": {
+ "query": "How to use LoRA in Torchtune"
+ },
+ "call_id": "dc7dd9e0-6ca1-452e-bb62-532a09e71848",
+ "tool_name": "knowledge_search"
+ },
+ "type": "tool_call"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "1iT28abM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:53.948952+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "gd_zuJXnSaSfS3ZK",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 117
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "1iT28abM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:53.949001+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "gd_zuJXnSaSfS3ZK",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 40
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "1iT28abM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:53.949013+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "gd_zuJXnSaSfS3ZK",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 157
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"I am attaching some documentation for Torchtune. Help me answer questions I will ask next.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"Torchtune documentation\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:78970\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:8404f\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:cbeb1\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}]]_{\"response_format\": null, \"sampling_params\": {\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"SamplingParams\", \"data\": {\"max_tokens\": 0, \"repetition_penalty\": 1.0, \"strategy\": {\"temperature\": 0.0001, \"top_p\": 0.9, \"type\": \"top_p\"}}}, \"stream\": true, \"tool_config\": {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolConfig\", \"data\": {\"system_message_behavior\": {\"__enum__\": \"SystemMessageBehavior\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"append\"}, \"tool_choice\": {\"__enum__\": \"ToolChoice\", \"__module__\": \"llama_stack.apis.inference.inference\", \"value\": \"auto\"}, \"tool_prompt_format\": null}}, \"tool_prompt_format\": null, \"tools\": [{\"__module__\": \"llama_stack.models.llama.datatypes\", \"__pydantic__\": \"ToolDefinition\", \"data\": {\"description\": \"Search for information in a database.\", \"parameters\": {\"query\": {\"default\": null, \"description\": \"The query to search for. Can be a natural language sentence or keywords.\", \"param_type\": \"string\", \"required\": true}}, \"tool_name\": \"knowledge_search\"}}]}": {
+ "chunks": [
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "start"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "I",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "'m ready to help you answer questions about Torchtune based on",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " the documentation you provided. What's your",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": " first question?",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "progress"
+ },
+ "logprobs": null,
+ "stop_reason": null
+ },
+ "metrics": null
+ }
+ },
+ {
+ "__module__": "llama_stack.apis.inference.inference",
+ "__pydantic__": "ChatCompletionResponseStreamChunk",
+ "data": {
+ "event": {
+ "delta": {
+ "text": "",
+ "type": "text"
+ },
+ "event_type": {
+ "__enum__": "ChatCompletionResponseEventType",
+ "__module__": "llama_stack.apis.inference.inference",
+ "value": "complete"
+ },
+ "logprobs": null,
+ "stop_reason": {
+ "__enum__": "StopReason",
+ "__module__": "llama_stack.models.llama.datatypes",
+ "value": "end_of_turn"
+ }
+ },
+ "metrics": [
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "prompt_tokens",
+ "span_id": "F3R1-xJM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:52.280696+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "7Do839YJRHC_ADjC",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 75
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "completion_tokens",
+ "span_id": "F3R1-xJM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:52.280743+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "7Do839YJRHC_ADjC",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 35
+ },
+ {
+ "attributes": {
+ "model_id": "meta-llama/Llama-3.1-8B-Instruct",
+ "provider_id": "fireworks"
+ },
+ "metric": "total_tokens",
+ "span_id": "F3R1-xJM",
+ "timestamp": {
+ "__class__": "datetime",
+ "__datetime__": "2025-03-06T04:33:52.280778+00:00",
+ "__module__": "datetime"
+ },
+ "trace_id": "7Do839YJRHC_ADjC",
+ "type": "metric",
+ "unit": "tokens",
+ "value": 110
+ }
+ ]
+ }
+ }
+ ],
+ "type": "generator"
+ },
+ "[\"meta-llama/Llama-3.1-8B-Instruct\", [{\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"SystemMessage\", \"data\": {\"content\": \"You are a helpful assistant\", \"role\": \"system\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"I am attaching some documentation for Torchtune. Help me answer questions I will ask next.\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"Torchtune documentation\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:78a41\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:7b4a7\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 5:\\nDocument_id:531f2\\nContent: etune\\n:func:`torchtune.models.llama3.llama3_8b` with DoRA, you would use :func:`torchtune.models.llama3.lora_llama3_8b` with ``use_dora=True``:\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.use_dora=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n use_dora: True\\n\\nSince DoRA extends LoRA, the parameters for :ref:`customizing LoRA ` are identical. You can also quantize the base model weights like in :ref:`glossary_qlora` by using ``quantize=True`` to reap\\neven more memory savings!\\n\\n.. code-block:: bash\\n\\n tune run lora_finetune_single_device --config llama3/8B_lora_single_device \\\\\\n model.apply_lora_to_mlp=True \\\\\\n model.lora_attn_modules=[\\\"q_proj\\\",\\\"k_proj\\\",\\\"v_proj\\\"] \\\\\\n model.lora_rank=16 \\\\\\n model.lora_alpha=32 \\\\\\n model.use_dora=True \\\\\\n model.quantize_base=True\\n\\n.. code-block:: yaml\\n\\n model:\\n _component_: torchtune.models.lora_llama3_8b\\n apply_lora_to_mlp: True\\n lora_attn_modules: [\\\"q_proj\\\", \\\"k_proj\\\", \\\"v_proj\\\"]\\n lora_rank: 16\\n lora_alpha: 32\\n use_dora: True\\n quantize_base: True\\n\\n\\n.. note::\\n\\n Under the hood, we've enabled DoRA by adding the :class:`~torchtune.modules.peft.DoRALinear` module, which we swap\\n out for :class:`~torchtune.modules.peft.LoRALinear` when ``use_dora=True``.\\n\\n.. _glossary_distrib:\\n\\n\\n.. TODO\\n\\n.. Distributed\\n.. -----------\\n\\n.. .. _glossary_fsdp:\\n\\n.. Fully Sharded Data Parallel (FSDP)\\n.. ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\\n\\n.. All our ``_distributed`` recipes use `FSDP `.\\n.. .. _glossary_fsdp2:\\n\\n\", \"type\": \"text\"}, {\"text\": \"END of knowledge_search tool results.\\n\", \"type\": \"text\"}], \"role\": \"tool\", \"tool_name\": \"knowledge_search\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"I'm ready to help you answer questions about Torchtune based on the documentation you provided. What's your first question?\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": []}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"UserMessage\", \"data\": {\"content\": \"Tell me how to use LoRA\", \"context\": null, \"role\": \"user\"}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"CompletionMessage\", \"data\": {\"content\": \"\", \"role\": \"assistant\", \"stop_reason\": {\"__enum__\": \"StopReason\", \"__module__\": \"llama_stack.models.llama.datatypes\", \"value\": \"end_of_turn\"}, \"tool_calls\": [{\"arguments\": {\"query\": \"How to use LoRA in Torchtune\"}, \"call_id\": \"\", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:7b4a7\\nContent: .. _lora_finetune_label:\\n\\n============================\\nFine-Tuning Llama2 with LoRA\\n============================\\n\\nThis guide will teach you about `LoRA `_, a parameter-efficient finetuning technique,\\nand show you how you can use torchtune to finetune a Llama2 model with LoRA.\\nIf you already know what LoRA is and want to get straight to running\\nyour own LoRA finetune in torchtune, you can jump to :ref:`LoRA finetuning recipe in torchtune`.\\n\\n.. grid:: 2\\n\\n .. grid-item-card:: :octicon:`mortar-board;1em;` What you will learn\\n\\n * What LoRA is and how it saves memory during finetuning\\n * An overview of LoRA components in torchtune\\n * How to run a LoRA finetune using torchtune\\n * How to experiment with different LoRA configurations\\n\\n .. grid-item-card:: :octicon:`list-unordered;1em;` Prerequisites\\n\\n * Be familiar with :ref:`torchtune`\\n * Make sure to :ref:`install torchtune`\\n * Make sure you have downloaded the :ref:`Llama2-7B model weights`\\n\\nWhat is LoRA?\\n-------------\\n\\n`LoRA `_ is an adapter-based method for\\nparameter-efficient finetuning that adds trainable low-rank decomposition matrices to different layers of a neural network,\\nthen freezes the network's remaining parameters. LoRA is most commonly applied to\\ntransformer models, in which case it is common to add the low-rank matrices\\nto some of the linear projections in each transformer layer's self-attention.\\n\\n.. note::\\n\\n If you're unfamiliar, check out these references for the `definition of rank `_\\n and discussion of `low-rank approximations `_.\\n\\nBy finetuning with LoRA (as opposed to finetuning all model parameters),\\nyou can expect to see memory savings due to a substantial reduction in the\\nnumber of parameters with gradients. When using an optimizer with momentum,\\nlike `AdamW ` alone will not handle the definition of which parameters are trainable.\\n See :ref:`below` for how to do this.\\n\\nLet's inspect each of these models a bit more closely.\\n\\n.. code-block:: bash\\n\\n # Print the first layer's self-attention in the usual Llama2 model\\n >>> print(base_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (output_proj): Linear(in_features=4096, out_features=4096, bias=False)\\n (pos_embeddings): RotaryPositionalEmbeddings()\\n )\\n\\n # Print the same for Llama2 with LoRA weights\\n >>> print(lora_model.layers[0].attn)\\n MultiHeadAttention(\\n (q_proj): LoRALinear(\\n (dropout): Dropout(p=0.0, inplace=False)\\n \\n\", \"type\": \"text\"}, {\"text\": \"Result 3:\\nDocument_id:7b4a7\\nContent: 06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe `_.\\nMake sure that you have first downloaded the Llama2 weights and tokenizer by following :ref:`these instructions`.\\nYou can then run the following command to perform a LoRA finetune of Llama2-7B with two GPUs (each having VRAM of at least 16GB):\\n\\n.. code-block:: bash\\n\\n tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora\\n\\n.. note::\\n Make sure to point to the location of your Llama2 weights and tokenizer. This can be done\\n either by adding :code:`checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path`\\n or by directly modifying the :code:`7B_lora.yaml` file. See our \\\"\\\":ref:`config_tutorial_label`\\\" recipe\\n for more details on how you can easily clone and modify torchtune configs.\\n\\n.. note::\\n You can modify the value of :code:`nproc_per_node` depending on (a) the number of GPUs you have available,\\n and (b) the memory constraints of your hardware.\\n\\nThe preceding command will run a LoRA finetune with torchtune's factory settings, but we may want to experiment a bit.\\nLet's take a closer look at some of the :code:`lora_finetune_distributed` config.\\n\\n.. code-block:: yaml\\n\\n # Model Arguments\\n model:\\n _component_: lora_llama2_7b\\n lora_attn_modules: ['q_proj', 'v_proj']\\n lora_rank: 8\\n lora_alpha: 16\\n ...\\n\\nWe see that the\\n\", \"type\": \"text\"}, {\"text\": \"Result 4:\\nDocument_id:7b4a7\\nContent: from our Llama2\\nmodel without any wrappers or custom checkpoint conversion logic.\\n\\n.. code-block:: python\\n\\n # Assuming that base_model already has the pretrained Llama2 weights,\\n # this will directly load them into your LoRA model without any conversion necessary.\\n lora_model.load_state_dict(base_model.state_dict(), strict=False)\\n\\n.. note::\\n Whenever loading weights with :code:`strict=False`, you should verify that any missing or extra keys in\\n the loaded :code:`state_dict` are as expected. torchtune's LoRA recipes do this by default via\\n :func:`validate_missing_and_unexpected_for_lora() `.\\n\\nOnce we've loaded the base model weights, we also want to set only LoRA parameters to trainable.\\n\\n.. _setting_trainable_params:\\n\\n.. code-block:: python\\n\\n from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params\\n\\n # Fetch all params from the model that are associated with LoRA.\\n lora_params = get_adapter_params(lora_model)\\n\\n # Set requires_grad=True on lora_params, and requires_grad=False on all others.\\n set_trainable_params(lora_model, lora_params)\\n\\n # Print the total number of parameters\\n total_params = sum([p.numel() for p in lora_model.parameters()])\\n trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])\\n print(\\n f\\\"\\\"\\\"\\n {total_params} total params,\\n {trainable_params}\\\" trainable params,\\n {(100.0 * trainable_params / total_params):.2f}% of all params are trainable.\\n \\\"\\\"\\\"\\n )\\n\\n 6742609920 total params,\\n 4194304 trainable params,\\n 0.06% of all params are trainable.\\n\\n.. note::\\n If you are directly using the LoRA recipe (as detailed :ref:`here`), you need only pass the\\n relevant checkpoint path. Loading model weights and setting trainable parameters will be taken care\\n of in the recipe.\\n\\n\\n.. _lora_recipe_label:\\n\\nLoRA finetuning recipe in torchtune\\n-----------------------------------\\n\\nFinally, we can put it all together and finetune a model using torchtune's `LoRA recipe \", \"tool_name\": \"knowledge_search\"}]}}, {\"__module__\": \"llama_stack.apis.inference.inference\", \"__pydantic__\": \"ToolResponseMessage\", \"data\": {\"call_id\": \"\", \"content\": [{\"text\": \"knowledge_search tool found 5 chunks:\\nBEGIN of knowledge_search tool results.\\n\", \"type\": \"text\"}, {\"text\": \"Result 1:\\nDocument_id:78a41\\nContent: conversational data, :func:`~torchtune.datasets.chat_dataset` seems to be a good fit. For any\\ncustom local dataset we always need to specify ``source``, ``data_files``, and ``split`` for any dataset\\nbuilder in torchtune. For :func:`~torchtune.datasets.chat_dataset`, we additionally need to specify\\n``conversation_column`` and ``conversation_style``. Our data follows the ``\\\"sharegpt\\\"`` format, so\\nwe can specify that here. Altogether, our :func:`~torchtune.datasets.chat_dataset` call should\\nlook like so:\\n\\n.. code-block:: python\\n\\n from torchtune.datasets import chat_dataset\\n from torchtune.models.llama3 import llama3_tokenizer\\n\\n tokenizer = llama3_tokenizer(\\\"/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\\")\\n ds = chat_dataset(\\n tokenizer=tokenizer,\\n source=\\\"json\\\",\\n data_files=\\\"data/my_data.json\\\",\\n split=\\\"train\\\",\\n conversation_column=\\\"dialogue\\\",\\n conversation_style=\\\"sharegpt\\\",\\n )\\n\\n.. code-block:: yaml\\n\\n # In config\\n tokenizer:\\n _component_: torchtune.models.llama3.llama3_tokenizer\\n path: /tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model\\n\\n dataset:\\n _component_: torchtune.datasets.chat_dataset\\n source: json\\n data_files: data/my_data.json\\n split: train\\n conversation_column: dialogue\\n conversation_style: sharegpt\\n\\n.. note::\\n You can pass in any keyword argument for `load_dataset `_ into all our\\n Dataset classes and they will honor them. This is useful for common parameters\\n such as specifying the data split with :code:`split` or configuration with\\n :code:`name`\\n\\nIf you needed to add a prompt template, you would simply pass it into the tokenizer.\\nSince we're fine-tuning Llama3, the tokenizer will handle all formatting for\\nus and prompt templates are optional. Other models such as Mistral's :class:`~torchtune.models.mistral._tokenizer.MistralTokenizer`,\\nuse a chat template by default (:class:`~torchtune.models.mistral.MistralChatTemplate`) to format\\nall messages according to their `recommendations