diff --git a/examples/tnt-llm/tnt-llm.ipynb b/examples/tnt-llm/tnt-llm.ipynb
new file mode 100644
index 000000000..c8a9e2e84
--- /dev/null
+++ b/examples/tnt-llm/tnt-llm.ipynb
@@ -0,0 +1,757 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "86b44172-7131-44a3-a825-ac6a7347b7a7",
+ "metadata": {},
+ "source": [
+ "## Phase 1: Taxonomy Generation\n",
+ "\n",
+ "\n",
+ "Scratch: get runs"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 79,
+ "id": "a7e4f593-d3e5-4483-a2c7-b7e29c83a261",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# \"Data\"\n",
+ "\n",
+ "from datetime import datetime, timedelta\n",
+ "\n",
+ "from langsmith import Client\n",
+ "\n",
+ "client = Client(timeout_ms=30_000)\n",
+ "\n",
+ "yday = datetime.now() - timedelta(days=3)\n",
+ "runs = list(\n",
+ " client.list_runs(\n",
+ " # project_id=\"f53ccf51-57c1-4c97-afdf-7ca5569945cd\",\n",
+ " project_name=\"chat-langchain\",\n",
+ " filter=\"eq(is_root, true)\",\n",
+ " start_time=yday,\n",
+ " select=[\"inputs\", \"outputs\", \"feedback_stats\"],\n",
+ " )\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 111,
+ "id": "759c7ee9-ce44-42c8-8c6f-6a1e6bfe4f25",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import random\n",
+ "\n",
+ "\n",
+ "def run_to_doc(run) -> Doc:\n",
+ " turns = []\n",
+ " idx = 0\n",
+ " for turn in run.inputs[\"chat_history\"] or []:\n",
+ " key, value = next(iter(turn.items()))\n",
+ " turns.append(f\"<{key} idx={idx}>\\n{value}\\n\")\n",
+ " idx += 1\n",
+ " turns.append(\n",
+ " f\"\"\"\n",
+ "\n",
+ "{run.inputs['question']}\n",
+ "\"\"\"\n",
+ " )\n",
+ " if run.outputs and run.outputs[\"output\"]:\n",
+ " turns.append(\n",
+ " f\"\"\"\n",
+ "{run.outputs['output']}\n",
+ "\"\"\"\n",
+ " )\n",
+ " return {\n",
+ " \"id\": str(run.id),\n",
+ " \"content\": (\"\\n\".join(turns)),\n",
+ " }\n",
+ "\n",
+ "\n",
+ "docs = [run_to_doc(run) for run in runs if run.inputs]\n",
+ "docs = random.sample(docs, min(len(docs), 1000))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 261,
+ "id": "cd9eda3e-3943-426a-bba2-8b9b81c00ec9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "use_case = (\n",
+ " \"Generate the taxonomy that can be used both to label the user intent\"\n",
+ " \" as well as to identify any required documentation (references, how-tos, etc.)\"\n",
+ " \" that would benefit the user.\"\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a07c8526-614a-4154-abfe-91e768a58a46",
+ "metadata": {},
+ "source": [
+ "#### 1.a Summarize Docs"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 272,
+ "id": "ff02c2a1-18b5-4848-96bb-27ff00978570",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import operator\n",
+ "import re\n",
+ "from typing import Annotated, List, Optional, Sequence, TypedDict\n",
+ "\n",
+ "from langchain import hub\n",
+ "from langchain_anthropic import ChatAnthropic\n",
+ "from langchain_core.messages import AIMessage, BaseMessage, HumanMessage\n",
+ "from langchain_core.output_parsers import StrOutputParser\n",
+ "from langchain_core.prompts import ChatPromptTemplate\n",
+ "from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough\n",
+ "\n",
+ "\n",
+ "class Doc(TypedDict):\n",
+ " id: str\n",
+ " content: str\n",
+ " summary: Optional[str]\n",
+ " explanation: Optional[str]\n",
+ "\n",
+ "\n",
+ "class TaxonomyGenerationState(TypedDict):\n",
+ " documents: List[Doc]\n",
+ " # Indices\n",
+ " minibatches: List[List[int]]\n",
+ " # Candidate Taxonomies\n",
+ " clusters: Annotated[List[List[dict]], operator.add]\n",
+ "\n",
+ "\n",
+ "# Phase 1.a: Generate summaries + explanations\n",
+ "\n",
+ "\n",
+ "summary_prompt = hub.pull(\"wfh/tnt-llm-summary-generation\").partial(\n",
+ " summary_length=20, explanation_length=30\n",
+ ")\n",
+ "\n",
+ "\n",
+ "def parse_summary(xml_string: str) -> dict:\n",
+ " summary_pattern = r\"(.*?)\"\n",
+ " explanation_pattern = r\"(.*?)\"\n",
+ "\n",
+ " summary_match = re.search(summary_pattern, xml_string, re.DOTALL)\n",
+ " explanation_match = re.search(explanation_pattern, xml_string, re.DOTALL)\n",
+ "\n",
+ " summary = summary_match.group(1).strip() if summary_match else \"\"\n",
+ " explanation = explanation_match.group(1).strip() if explanation_match else \"\"\n",
+ "\n",
+ " return {\"summary\": summary, \"explanation\": explanation}\n",
+ "\n",
+ "\n",
+ "summary_llm_chain = (\n",
+ " summary_prompt\n",
+ " | ChatAnthropic(model=\"claude-3-haiku-20240307\")\n",
+ " | StrOutputParser()\n",
+ " # Customize the tracing name for easier organization\n",
+ ").with_config(run_name=\"GenerateSummary\")\n",
+ "summary_chain = summary_llm_chain | parse_xml\n",
+ "\n",
+ "\n",
+ "# Now combine as a \"map\" operation in a map-reduce chain\n",
+ "# Input: state\n",
+ "# Output: state U summaries\n",
+ "# Processes docs in parallel\n",
+ "def get_content(state: TaxonomyGenerationState):\n",
+ " docs = state[\"documents\"]\n",
+ " return [{\"content\": doc[\"content\"]} for doc in docs]\n",
+ "\n",
+ "\n",
+ "map_step = RunnablePassthrough.assign(\n",
+ " summaries=get_content\n",
+ " | RunnableLambda(func=summary_chain.batch, afunc=summary_chain.abatch)\n",
+ ")\n",
+ "\n",
+ "\n",
+ "def reduce_summaries(combined: dict) -> TaxonomyGenerationState:\n",
+ " summaries = combined[\"summaries\"]\n",
+ " documents = combined[\"documents\"]\n",
+ " return {\n",
+ " \"documents\": [\n",
+ " {\n",
+ " \"id\": doc[\"id\"],\n",
+ " \"content\": doc[\"content\"],\n",
+ " \"summary\": summ_info[\"summary\"],\n",
+ " \"explanation\": summ_info[\"explanation\"],\n",
+ " }\n",
+ " for doc, summ_info in zip(documents, summaries)\n",
+ " ]\n",
+ " }\n",
+ "\n",
+ "\n",
+ "# This is actually the node itself!\n",
+ "map_reduce_chain = map_step | reduce_summaries"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 114,
+ "id": "69c28ede-7a6c-485b-9606-f39e8ec215c3",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "Failed to batch ingest runs: LangSmithConnectionError('Connection error caused failure to post https://api.smith.langchain.com/runs/batch in LangSmith API. Please confirm your internet connection.. SSLError(MaxRetryError(\"HTTPSConnectionPool(host=\\'api.smith.langchain.com\\', port=443): Max retries exceeded with url: /runs/batch (Caused by SSLError(SSLEOFError(8, \\'EOF occurred in violation of protocol (_ssl.c:2393)\\')))\"))')\n",
+ "Failed to batch ingest runs: LangSmithConnectionError('Connection error caused failure to post https://api.smith.langchain.com/runs/batch in LangSmith API. Please confirm your internet connection.. SSLError(MaxRetryError(\"HTTPSConnectionPool(host=\\'api.smith.langchain.com\\', port=443): Max retries exceeded with url: /runs/batch (Caused by SSLError(SSLEOFError(8, \\'EOF occurred in violation of protocol (_ssl.c:2393)\\')))\"))')\n"
+ ]
+ }
+ ],
+ "source": [
+ "summarized_docs = map_reduce_chain.invoke({\"documents\": docs}, {\"max_concurrency\": 5})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 121,
+ "id": "df55b666-0123-4fce-b15c-7126ff84acdb",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import json\n",
+ "\n",
+ "with open(\"docs.json\", \"w\") as f:\n",
+ " json.dump(summarized_docs, f)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "654cbad5-3af9-4c16-9d43-cbb903957944",
+ "metadata": {},
+ "source": [
+ "#### 1.b Split into Minibatches"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 236,
+ "id": "3e0139c3-b5ba-42b9-9367-33533d66eb58",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_minibatches(state: TaxonomyGenerationState, config: RunnableConfig):\n",
+ " batch_size = config[\"configurable\"].get(\"batch_size\", 200)\n",
+ " original = state[\"documents\"]\n",
+ " indices = list(range(len(original)))\n",
+ " random.shuffle(indices)\n",
+ " if len(indices) < batch_size:\n",
+ " # Don't pad needlessly if we can't fill a single batch\n",
+ " return [indices]\n",
+ "\n",
+ " num_full_batches = len(indices) // batch_size\n",
+ "\n",
+ " batches = [\n",
+ " indices[i * batch_size : (i + 1) * batch_size] for i in range(num_full_batches)\n",
+ " ]\n",
+ "\n",
+ " leftovers = len(indices) % batch_size\n",
+ " if leftovers:\n",
+ " last_batch = indices[num_full_batches * batch_size :]\n",
+ " elements_to_add = batch_size - leftovers\n",
+ " last_batch += random.sample(indices, elements_to_add)\n",
+ " batches.append(last_batch)\n",
+ "\n",
+ " return {\n",
+ " \"minibatches\": batches,\n",
+ " }"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 243,
+ "id": "870071b4-cba9-4711-8af3-70a6f7ecd225",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "batched_state = get_minibatches(summarized_docs, {\"configurable\": {}})\n",
+ "batched_state = {**summarized_docs, **batched_state}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1ca9ca18-f43a-46bd-bd28-7bec90932dd9",
+ "metadata": {},
+ "source": [
+ "#### 1.c Generate Taxonomy"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 273,
+ "id": "224ed013-2963-489c-b734-315cad701d59",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import random\n",
+ "import re\n",
+ "from typing import Any, Dict\n",
+ "\n",
+ "# We instruct our LLMs to generate 10\n",
+ "# intent categories and 25 domain categories for taxonomy generation.\n",
+ "\n",
+ "\n",
+ "def parse_taxa(output_text: str) -> Dict:\n",
+ " cluster_matches = re.findall(\n",
+ " r\"\\s*(.*?)\\s*(.*?)\\s*(.*?)\\s*\",\n",
+ " output_text,\n",
+ " re.DOTALL,\n",
+ " )\n",
+ " clusters = [\n",
+ " {\"id\": id.strip(), \"name\": name.strip(), \"description\": description.strip()}\n",
+ " for id, name, description in cluster_matches\n",
+ " ]\n",
+ " # We don't parse the explanation since it isn't used downstream\n",
+ " return {\"clusters\": clusters}\n",
+ "\n",
+ "\n",
+ "def format_docs(docs: List[Doc]) -> str:\n",
+ " xml_table = \"\\n\"\n",
+ " for doc in docs:\n",
+ " xml_table += f'{doc[\"summary\"]}\\n'\n",
+ " xml_table += \"\"\n",
+ " return xml_table\n",
+ "\n",
+ "\n",
+ "def format_taxonomy(clusters):\n",
+ " xml = \"\\n\"\n",
+ " for label in clusters:\n",
+ " xml += \" \\n\"\n",
+ " xml += f' {label[\"id\"]}\\n'\n",
+ " xml += f' {label[\"name\"]}\\n'\n",
+ " xml += f' {label[\"description\"]}\\n'\n",
+ " xml += \" \\n\"\n",
+ " xml += \"\"\n",
+ " return xml\n",
+ "\n",
+ "\n",
+ "# def generate_taxonomy(\n",
+ "# state: TaxonomyGenerationState, config: RunnableConfig\n",
+ "# ) -> TaxonomyGenerationState:\n",
+ "# \"\"\"Prompt an LLM to generate an initial taxonomy.\"\"\"\n",
+ "# configurable = config[\"configurable\"]\n",
+ "# docs = state[\"documents\"]\n",
+ "# mb_indices = state[\"minibatches\"][0]\n",
+ "# first_minibatch = [docs[idx] for idx in mb_indices]\n",
+ "# data_table_xml = format_docs(first_minibatch)\n",
+ "# initial_taxonomy = generate_taxonomy_chain.invoke(\n",
+ "# {\n",
+ "# \"data_xml\": data_table_xml,\n",
+ "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n",
+ "# \"cluster_description_length\": configurable.get(\n",
+ "# \"cluster_description_length\", 30\n",
+ "# ),\n",
+ "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n",
+ "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n",
+ "# }\n",
+ "# )\n",
+ "\n",
+ "# return {\n",
+ "# \"clusters\": [initial_taxonomy[\"clusters\"]],\n",
+ "# }"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c2f11188-3d1d-4d44-a5cf-886d46884b48",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "taxonomy_generation_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-generation\").partial(\n",
+ " use_case=\"Generate the taxonomy that can be used to label the user intent in the conversation.\",\n",
+ ")\n",
+ "\n",
+ "taxonomy_generation_llm = ChatAnthropic(\n",
+ " model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000\n",
+ ")\n",
+ "taxa_gen_llm_chain = (\n",
+ " taxonomy_generation_prompt | taxonomy_generation_llm | StrOutputParser()\n",
+ ").with_config(run_name=\"GenerateTaxonomy\")\n",
+ "\n",
+ "\n",
+ "generate_taxonomy_chain = taxa_gen_llm_chain | parse_generation_output"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 258,
+ "id": "f4b2820e-e002-4d51-8efe-feb32ab83bd1",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "taxonomies = generate_taxonomy(batched_state, {\"configurable\": {\"max_concurrency\": 5}})\n",
+ "taxonomies = {**batched_state, **taxonomies}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "6df7a089-b268-47b6-b8a3-91606cfe1bec",
+ "metadata": {},
+ "source": [
+ "#### 1.c Update taxonomy\n",
+ "\n",
+ "Mostly the same."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 274,
+ "id": "ce10047e-7062-4935-8bd7-ffacba4bd099",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "taxonomy_update_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-update\")\n",
+ "\n",
+ "taxa_update_llm_chain = (\n",
+ " taxonomy_update_prompt | taxonomy_update_llm | StrOutputParser()\n",
+ ").with_config(run_name=\"UpdateTaxonomy\")\n",
+ "\n",
+ "\n",
+ "update_taxonomy_chain = taxa_update_llm_chain | parse_taxa\n",
+ "\n",
+ "# def update_taxonomy(\n",
+ "# state: TaxonomyGenerationState, config: RunnableConfig\n",
+ "# ) -> TaxonomyGenerationState:\n",
+ "# \"\"\"Prompt an LLM to update the taxonomy based on the current state.\"\"\"\n",
+ "# configurable = config[\"configurable\"]\n",
+ "# docs = state[\"documents\"]\n",
+ "# minibatches = state[\"minibatches\"]\n",
+ "# previous_taxonomy = state[\"clusters\"][-1]\n",
+ "# which_mb = len(state[\"clusters\"]) % len(minibatches)\n",
+ "# mb_indices = minibatches[which_mb]\n",
+ "# minibatch = [docs[idx] for idx in mb_indices]\n",
+ "# # The new data we will be using to\n",
+ "# data_table_xml = format_docs(minibatch)\n",
+ "# cluster_table_xml = format_taxonomy(previous_taxonomy)\n",
+ "# updated_taxonomy = update_taxonomy_chain.invoke(\n",
+ "# {\n",
+ "# \"data_xml\": data_table_xml,\n",
+ "# \"use_case\": configurable[\"use_case\"],\n",
+ "# \"cluster_table_xml\": cluster_table_xml,\n",
+ "# \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n",
+ "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n",
+ "# \"cluster_description_length\": configurable.get(\n",
+ "# \"cluster_description_length\", 30\n",
+ "# ),\n",
+ "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n",
+ "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n",
+ "# }\n",
+ "# )\n",
+ "# return {\n",
+ "# \"clusters\": [updated_taxonomy[\"clusters\"]],\n",
+ "# }"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 270,
+ "id": "6bf6b09f-5777-49df-bf47-5ae5ff3d2bce",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "updated_taxonomies = update_taxonomy(\n",
+ " taxonomies, {\"configurable\": {\"use_case\": use_case}}\n",
+ ")\n",
+ "updated_taxonomies = {**taxonomies, **updated_taxonomies}"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 271,
+ "id": "e6c532c4-5468-4e35-a557-35ea7f0b5651",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "[{'id': '1',\n",
+ " 'name': 'Manage file operations',\n",
+ " 'description': 'Load, read, write, and list files in various formats like JSON, CSV, and Parquet.'},\n",
+ " {'id': '2',\n",
+ " 'name': 'Resolve code and network issues',\n",
+ " 'description': 'Debug directory errors, network firewall problems, and create AI-powered models.'},\n",
+ " {'id': '3',\n",
+ " 'name': 'Implement conversational AI',\n",
+ " 'description': 'Handle user queries, provide responses, and manage user-AI interactions.'},\n",
+ " {'id': '4',\n",
+ " 'name': 'Build chatbots with LangChain',\n",
+ " 'description': 'Develop chatbots that use LangChain for knowledge retrieval, conversation history, and response generation.'},\n",
+ " {'id': '5',\n",
+ " 'name': 'Integrate LangChain with external services',\n",
+ " 'description': 'Combine LangChain with tools like Ollama, Mistral, and Azure OpenAI for advanced capabilities.'},\n",
+ " {'id': '6',\n",
+ " 'name': 'Process web content with LangChain',\n",
+ " 'description': 'Load web pages, extract text, and handle various file formats from online sources.'},\n",
+ " {'id': '7',\n",
+ " 'name': 'Manage LangChain vector stores',\n",
+ " 'description': 'Ingest, store, and retrieve text data using vector embeddings and database solutions.'}]"
+ ]
+ },
+ "execution_count": 271,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "updated_taxonomies[\"clusters\"][-1]"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c2abc558-507f-44b7-9812-54fb1e9b3c64",
+ "metadata": {},
+ "source": [
+ "#### 1.d Review Taxonomy"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 275,
+ "id": "fe3ae6e6-80ac-4e71-8a43-5d8c0d64fa2c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "taxonomy_review_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-review\")\n",
+ "\n",
+ "taxa_review_llm_chain = (\n",
+ " taxonomy_generation_llm | taxonomy_review_llm | StrOutputParser()\n",
+ ").with_config(run_name=\"ReviewTaxonomy\")\n",
+ "\n",
+ "\n",
+ "review_taxonomy_chain = taxa_review_llm_chain | parse_taxa\n",
+ "\n",
+ "\n",
+ "# def review_taxonomy(\n",
+ "# state: TaxonomyGenerationState, config: RunnableConfig\n",
+ "# ) -> TaxonomyGenerationState:\n",
+ "# \"\"\"Prompt an LLM to update the taxonomy based on the current state.\"\"\"\n",
+ "# configurable = config[\"configurable\"]\n",
+ "# docs = state[\"documents\"]\n",
+ "# minibatches = state[\"minibatches\"]\n",
+ "# previous_taxonomy = state[\"clusters\"][-1]\n",
+ "# mb_indices = random.sample\n",
+ "# which_mb = len(state[\"clusters\"]) % len(minibatches)\n",
+ "# mb_indices = minibatches[which_mb]\n",
+ "# minibatch = [docs[idx] for idx in mb_indices]\n",
+ "# # The new data we will be using to\n",
+ "# data_table_xml = format_docs(minibatch)\n",
+ "# cluster_table_xml = format_taxonomy(previous_taxonomy)\n",
+ "# reviewed_taxonomy = review_taxonomy_chain.invoke(\n",
+ "# {\n",
+ "# \"data_xml\": data_table_xml,\n",
+ "# \"use_case\": configurable[\"use_case\"],\n",
+ "# \"cluster_table_xml\": cluster_table_xml,\n",
+ "# \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n",
+ "# \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n",
+ "# \"cluster_description_length\": configurable.get(\n",
+ "# \"cluster_description_length\", 30\n",
+ "# ),\n",
+ "# \"explanation_length\": configurable.get(\"explanation_length\", 20),\n",
+ "# \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n",
+ "# }\n",
+ "# )\n",
+ "# return {\n",
+ "# \"clusters\": [reviewed_taxonomy[\"clusters\"]],\n",
+ "# }"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0039cf1c-54d5-4e9e-8dd6-a5cebfaec92d",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def invoke_taxonomy_chain(\n",
+ " chain: Callable,\n",
+ " state: TaxonomyGenerationState,\n",
+ " config: RunnableConfig,\n",
+ " mb_indices: List[int],\n",
+ ") -> TaxonomyGenerationState:\n",
+ " configurable = config[\"configurable\"]\n",
+ " docs = state[\"documents\"]\n",
+ " minibatch = [docs[idx] for idx in mb_indices]\n",
+ " data_table_xml = format_docs(minibatch)\n",
+ "\n",
+ " previous_taxonomy = state[\"clusters\"][-1] if state[\"clusters\"] else []\n",
+ " cluster_table_xml = format_taxonomy(previous_taxonomy)\n",
+ "\n",
+ " updated_taxonomy = chain.invoke(\n",
+ " {\n",
+ " \"data_xml\": data_table_xml,\n",
+ " \"use_case\": configurable[\"use_case\"],\n",
+ " \"cluster_table_xml\": cluster_table_xml,\n",
+ " \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n",
+ " \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n",
+ " \"cluster_description_length\": configurable.get(\n",
+ " \"cluster_description_length\", 30\n",
+ " ),\n",
+ " \"explanation_length\": configurable.get(\"explanation_length\", 20),\n",
+ " \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n",
+ " }\n",
+ " )\n",
+ "\n",
+ " return {\n",
+ " **state,\n",
+ " \"clusters\": state[\"clusters\"] + [updated_taxonomy[\"clusters\"]],\n",
+ " }\n",
+ "\n",
+ "\n",
+ "def generate_taxonomy(\n",
+ " state: TaxonomyGenerationState, config: RunnableConfig\n",
+ ") -> TaxonomyGenerationState:\n",
+ " return invoke_taxonomy_chain(\n",
+ " generate_taxonomy_chain, state, config, state[\"minibatches\"][0]\n",
+ " )\n",
+ "\n",
+ "\n",
+ "def update_taxonomy(\n",
+ " state: TaxonomyGenerationState, config: RunnableConfig\n",
+ ") -> TaxonomyGenerationState:\n",
+ " which_mb = len(state[\"clusters\"]) % len(state[\"minibatches\"])\n",
+ " return invoke_taxonomy_chain(\n",
+ " update_taxonomy_chain, state, config, state[\"minibatches\"][which_mb]\n",
+ " )\n",
+ "\n",
+ "\n",
+ "def review_taxonomy(\n",
+ " state: TaxonomyGenerationState, config: RunnableConfig\n",
+ ") -> TaxonomyGenerationState:\n",
+ " return invoke_taxonomy_chain(\n",
+ " review_taxonomy_chain, state, config, random.choice(state[\"minibatches\"])\n",
+ " )"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 276,
+ "id": "e407115c-85bd-415f-8012-2ea061b1215c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "reviewed_taxonomies = review_taxonomy(\n",
+ " updated_taxonomies, {\"configurable\": {\"use_case\": use_case}}\n",
+ ")\n",
+ "reviewed_taxonomies = {**updated_taxonomies, **reviewed_taxonomies}"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f1f97ea4-53e5-4f55-8d73-b5b2234a47d9",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Define the graph\n",
+ "graph = StateGraph(TaxonomyGenerationState)\n",
+ "graph.add_node(\"summarize\", map_reduce_chain)\n",
+ "graph.add_node(\"get_minibatches\", get_minibatches)\n",
+ "graph.add_node(\"generate_taxonomy\", generate_taxonomy)\n",
+ "graph.add_node(\"update_taxonomy\", update_taxonomy)\n",
+ "graph.add_node(\"review_taxonomy\", review_taxonomy)\n",
+ "\n",
+ "graph.add_edge(\"summarize\", \"add_minibatches\")\n",
+ "graph.add_edge(\"add_minibatches\", \"generate_taxonomy\")\n",
+ "graph.add_edge(\"generate_taxonomy\", \"update_taxonomy\")\n",
+ "def should_review(\n",
+ "graph.add_conditional_edge(should_review)\n",
+ "graph.set_finish_point(\"review_taxonomy\")\n",
+ "\n",
+ "graph.set_entry_point(\"summarize\")\n",
+ "app = graph.compile()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9e1da39e-4c8f-407b-8011-f54931df7c7c",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Phase 2: Text Classification\n",
+ "\n",
+ "\n",
+ "# Define the state for the text classification graph\n",
+ "class TextClassificationState(TypedDict):\n",
+ " messages: Annotated[Sequence[BaseMessage], operator.add]\n",
+ " labels: Annotated[Sequence[str], operator.add]\n",
+ "\n",
+ "\n",
+ "def assign_labels(state: TextClassificationState) -> TextClassificationState:\n",
+ " \"\"\"Prompt an LLM to assign labels to the given text.\"\"\"\n",
+ " llm = ChatOpenAI(temperature=0)\n",
+ " messages = [HumanMessage(content=text) for text in state[\"messages\"]]\n",
+ " responses = llm.batch(messages)\n",
+ " labels = [response.additional_kwargs[\"labels\"] for response in responses]\n",
+ " return {\"messages\": messages, \"labels\": labels}\n",
+ "\n",
+ "\n",
+ "# Define the TnT-LLM workflow using LangGraph\n",
+ "\n",
+ "# Phase 1: Taxonomy Generation\n",
+ "taxonomy_graph = StateGraph(TaxonomyGenerationState)\n",
+ "taxonomy_graph.add_node(\"summarize\", summarize_text)\n",
+ "taxonomy_graph.add_node(\"generate\", generate_taxonomy)\n",
+ "taxonomy_graph.add_node(\"update\", update_taxonomy)\n",
+ "taxonomy_graph.add_node(\"review\", review_taxonomy)\n",
+ "\n",
+ "taxonomy_graph.add_edge(\"summarize\", \"generate\")\n",
+ "taxonomy_graph.add_edge(\"generate\", \"update\")\n",
+ "taxonomy_graph.add_conditional_edges(\n",
+ " \"update\",\n",
+ " lambda state: \"update\" if len(state[\"summaries\"]) < 10 else \"review\",\n",
+ ")\n",
+ "taxonomy_graph.add_edge(\"review\", END)\n",
+ "\n",
+ "taxonomy_graph.set_entry_point(\"summarize\")\n",
+ "taxonomy_generator = taxonomy_graph.compile()\n",
+ "\n",
+ "# Phase 2: Text Classification\n",
+ "classification_graph = StateGraph(TextClassificationState)\n",
+ "classification_graph.add_node(\"assign_labels\", assign_labels)\n",
+ "\n",
+ "classification_graph.add_edge(\"assign_labels\", END)\n",
+ "\n",
+ "classification_graph.set_entry_point(\"assign_labels\")\n",
+ "text_classifier = classification_graph.compile()"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.2"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}