diff --git a/docs/use_cases/evaluation/evaluating_traced_examples.ipynb b/docs/use_cases/evaluation/evaluating_traced_examples.ipynb index 581c4122d28..392cffc8d0b 100644 --- a/docs/use_cases/evaluation/evaluating_traced_examples.ipynb +++ b/docs/use_cases/evaluation/evaluating_traced_examples.ipynb @@ -5,7 +5,7 @@ "id": "1a4596ea-a631-416d-a2a4-3577c140493d", "metadata": {}, "source": [ - "## Running Chains on Traced Datasets\n", + "# Running Chains on Traced Datasets\n", "\n", "Developing applications with language models can be uniquely challenging. To manage this complexity and ensure reliable performance, LangChain provides tracing and evaluation functionality through . This notebook demonstrates how to run Chains, which are language model functions, on previously captured datasets or traces. Some common use cases for this approach include:\n", "\n", @@ -27,24 +27,30 @@ "source": [ "from langchain.client import LangChainClient\n", "\n", - "client = LangChainClient()" + "client = LangChainClient(\n", + " api_url=\"http://localhost:8000\",\n", + " api_key=None,\n", + ")" ] }, { "cell_type": "markdown", "id": "db79dea2-fbaa-4c12-9083-f6154b51e2d3", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "## Seed an example dataset\n", "\n", - "If you have been using LangChainPlus already, you may have datasets available. You can generate these from `Run`'s captured through the LangChain tracing API.\n", + "If you have been using LangChainPlus already, you may have datasets available. To view all saved datasets, run:\n", "\n", "```\n", "datasets = client.list_datasets()\n", "datasets\n", "```\n", + "Datasets can be created in a number of ways, most often by collecting `Run`'s captured through the LangChain tracing API.\n", "\n", - "Assuming you're running locally for the first time, you may not have runs stored, so we'll first upload an example evaluation dataset." + "However, this notebook assumes you're running locally for the first time, so we'll start by uploading an example evaluation dataset." ] }, { @@ -62,12 +68,34 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "60d14593-c61f-449f-a38f-772ca43707c2", "metadata": { "tags": [] }, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Found cached dataset json (/Users/wfh/.cache/huggingface/datasets/LangChainDatasets___json/LangChainDatasets--agent-search-calculator-8a025c0ce5fb99d2/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "76dc2a0a3d3442d997569ad1b08a82d4", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + " 0%| | 0/1 [00:00, line 1). Please try again with a valid numerical expression\n", - "Chain failed for example b925eaa0-6b63-45c1-a3c3-c11ad53ca0eb. Error: 'VariableNode' object is not callable. Please try again with a valid numerical expression\n" + "Chain failed for example d015179a-f4a7-40e1-ba2f-6ca90ed45fed. Error: unknown format from LLM: Assuming we don't have any information about the actual number of points scored in the 2023 super bowl, we cannot provide a mathematical expression to solve this problem.\n", + "Chain failed for example d151c47c-fe80-4b58-8a28-5d5523d42b27. Error: 'age'. Please try again with a valid numerical expression\n", + "Chain failed for example 0d63e5f0-83ce-49b2-826b-ced0ab4ca701. Error: invalid syntax. Perhaps you forgot a comma? (, line 1). Please try again with a valid numerical expression\n", + "Chain failed for example 2dc584d8-dbe8-4146-8872-f92440f1048e. Error: 'VariableNode' object is not callable. Please try again with a valid numerical expression\n" ] } ], @@ -263,6 +291,17 @@ ")" ] }, + { + "cell_type": "markdown", + "id": "d2737458-b20c-4288-8790-1f4a8d237b2a", + "metadata": {}, + "source": [ + "## Reviewing the Chain Results\n", + "\n", + "The method called above returns a dictionary mapping Example IDs to the output of the chain.\n", + "You can directly inspect the results below." + ] + }, { "cell_type": "code", "execution_count": 9, @@ -274,16 +313,16 @@ { "data": { "text/plain": [ - "['The current population of Canada as of May 3, 2023 is 38,677,281.',\n", - " \"Anwar Hadid's age raised to the 0.43 power is approximately 3.87.\",\n", - " {'Error': \"'age'. Please try again with a valid numerical expression\"},\n", - " 'The distance between Paris and Boston is approximately 3448 miles.',\n", - " {'Error': \"unknown format from LLM: Assuming we don't have any information about the actual number of points scored in the 2023 super bowl, we cannot provide a mathematical expression to solve this problem.\"},\n", - " {'Error': 'invalid syntax. Perhaps you forgot a comma? (, line 1). Please try again with a valid numerical expression'},\n", - " '0',\n", - " '1.9347796717823205',\n", - " \"Devin Booker is Kendall Jenner's boyfriend and his height (in inches) raised to .13 power is approximately 1.7589.\",\n", - " {'Error': \"'VariableNode' object is not callable. Please try again with a valid numerical expression\"}]" + "{'1781b1db-45cb-428f-8000-00482ccf7142': 'The distance between Paris and Boston is 3448 miles.',\n", + " 'd015179a-f4a7-40e1-ba2f-6ca90ed45fed': {'Error': \"unknown format from LLM: Assuming we don't have any information about the actual number of points scored in the 2023 super bowl, we cannot provide a mathematical expression to solve this problem.\"},\n", + " 'd151c47c-fe80-4b58-8a28-5d5523d42b27': {'Error': \"'age'. Please try again with a valid numerical expression\"},\n", + " 'd437c17f-85b6-4d30-a4ff-8f2b7456ae43': 'The current population of Canada as of 2023 is 38,677,281.',\n", + " 'a4d4ddbd-eb82-4300-9473-97715baf8f26': \"Anwar Hadid's age raised to the .43 power is approximately 2.68.\",\n", + " '0d63e5f0-83ce-49b2-826b-ced0ab4ca701': {'Error': 'invalid syntax. Perhaps you forgot a comma? (, line 1). Please try again with a valid numerical expression'},\n", + " '7702d8f5-96d1-46ad-9919-cb40efed224c': '1.9347796717823205',\n", + " 'dcf59228-f3d4-4359-83d4-9032ec00a950': '0',\n", + " 'd93f5863-cdb1-4529-ae33-71013c8e5854': \"Devin Booker is Kendall Jenner's boyfriend and his height (in inches) raised to .13 power is approximately 1.7589.\",\n", + " '2dc584d8-dbe8-4146-8872-f92440f1048e': {'Error': \"'VariableNode' object is not callable. Please try again with a valid numerical expression\"}}" ] }, "execution_count": 9, @@ -324,11 +363,207 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 11, "id": "b05f1aab-6bc9-4963-9059-c5cd9c656992", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "# You can review all the chain runs over a given example as follows:\n", + "example_id = next(iter(chain_results))\n", + "example = client.read_example(example_id)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "d28b7752-7abb-4b4d-8575-1e98d6baf1b1", + "metadata": { + "tags": [] + }, + "outputs": [ + { + "data": { + "text/plain": [ + "[{'child_runs': [{'child_runs': [{'start_time': '2023-05-04T01:40:58.004292',\n", + " 'end_time': '2023-05-04T01:41:00.124338',\n", + " 'extra': {},\n", + " 'execution_order': 3,\n", + " 'serialized': {'name': 'ChatOpenAI'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 2,\n", + " 'parent_tool_run_id': None,\n", + " 'prompts': ['Human: Answer the following questions as best you can. You have access to the following tools:\\n\\nSearch: A search engine. Useful for when you need to answer questions about current events. Input should be a search query.\\nCalculator: Useful for when you need to answer questions about math.\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [Search, Calculator]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: how far is it from paris to boston in miles\\nThought:'],\n", + " 'response': {'generations': [[{'text': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"'}]],\n", + " 'llm_output': {'token_usage': {'prompt_tokens': 178,\n", + " 'completion_tokens': 27,\n", + " 'total_tokens': 205},\n", + " 'model_name': 'gpt-3.5-turbo'}},\n", + " 'id': 1,\n", + " 'type': 'llm'}],\n", + " 'start_time': '2023-05-04T01:40:58.003567',\n", + " 'end_time': '2023-05-04T01:41:00.124705',\n", + " 'extra': {},\n", + " 'execution_order': 2,\n", + " 'serialized': {'name': 'LLMChain'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 1,\n", + " 'parent_tool_run_id': None,\n", + " 'inputs': {'input': 'how far is it from paris to boston in miles',\n", + " 'agent_scratchpad': '',\n", + " 'stop': ['\\nObservation:', '\\n\\tObservation:']},\n", + " 'outputs': {'text': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"'},\n", + " 'id': 2,\n", + " 'type': 'chain'},\n", + " {'child_runs': [],\n", + " 'start_time': '2023-05-04T01:41:00.125929',\n", + " 'end_time': '2023-05-04T01:41:00.228395',\n", + " 'extra': {'color': 'green',\n", + " 'llm_prefix': 'Thought:',\n", + " 'observation_prefix': 'Observation: '},\n", + " 'execution_order': 4,\n", + " 'serialized': {'name': 'Search',\n", + " 'description': 'A search engine. Useful for when you need to answer questions about current events. Input should be a search query.'},\n", + " 'session_id': 2,\n", + " 'example_id': None,\n", + " 'error': None,\n", + " 'parent_chain_run_id': 1,\n", + " 'parent_tool_run_id': None,\n", + " 'tool_input': 'distance between Paris and Boston in miles',\n", + " 'output': 'The distance between Paris (Paris Charles de Gaulle Airport) and Boston (Logan International Airport) is 3448 miles / 5550 kilometers / 2997 nautical miles.',\n", + " 'action': \"{'name': 'Search', 'description': 'A search engine. Useful for when you need to answer questions about current events. Input should be a search query.'}\",\n", + " 'id': 1,\n", + " 'type': 'tool'},\n", + " {'child_runs': [{'start_time': '2023-05-04T01:41:00.229748',\n", + " 'end_time': '2023-05-04T01:41:01.915701',\n", + " 'extra': {},\n", + " 'execution_order': 6,\n", + " 'serialized': {'name': 'ChatOpenAI'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 3,\n", + " 'parent_tool_run_id': None,\n", + " 'prompts': ['Human: Answer the following questions as best you can. You have access to the following tools:\\n\\nSearch: A search engine. Useful for when you need to answer questions about current events. Input should be a search query.\\nCalculator: Useful for when you need to answer questions about math.\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [Search, Calculator]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: how far is it from paris to boston in miles\\nThought:I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"\\nObservation: The distance between Paris (Paris Charles de Gaulle Airport) and Boston (Logan International Airport) is 3448 miles / 5550 kilometers / 2997 nautical miles.\\nThought:'],\n", + " 'response': {'generations': [[{'text': 'I have the answer to the question\\nFinal Answer: The distance between Paris and Boston is 3448 miles.'}]],\n", + " 'llm_output': {'token_usage': {'prompt_tokens': 247,\n", + " 'completion_tokens': 23,\n", + " 'total_tokens': 270},\n", + " 'model_name': 'gpt-3.5-turbo'}},\n", + " 'id': 2,\n", + " 'type': 'llm'}],\n", + " 'start_time': '2023-05-04T01:41:00.229351',\n", + " 'end_time': '2023-05-04T01:41:01.916173',\n", + " 'extra': {},\n", + " 'execution_order': 5,\n", + " 'serialized': {'name': 'LLMChain'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 1,\n", + " 'parent_tool_run_id': None,\n", + " 'inputs': {'input': 'how far is it from paris to boston in miles',\n", + " 'agent_scratchpad': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"\\nObservation: The distance between Paris (Paris Charles de Gaulle Airport) and Boston (Logan International Airport) is 3448 miles / 5550 kilometers / 2997 nautical miles.\\nThought:',\n", + " 'stop': ['\\nObservation:', '\\n\\tObservation:']},\n", + " 'outputs': {'text': 'I have the answer to the question\\nFinal Answer: The distance between Paris and Boston is 3448 miles.'},\n", + " 'id': 3,\n", + " 'type': 'chain'}],\n", + " 'start_time': '2023-05-04T01:40:58.002934',\n", + " 'end_time': '2023-05-04T01:41:01.917236',\n", + " 'extra': {},\n", + " 'execution_order': 1,\n", + " 'serialized': {'name': 'AgentExecutor'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': None,\n", + " 'parent_tool_run_id': None,\n", + " 'inputs': {'input': 'how far is it from paris to boston in miles'},\n", + " 'outputs': {'output': 'The distance between Paris and Boston is 3448 miles.'},\n", + " 'id': 1,\n", + " 'type': 'chain'},\n", + " {'child_runs': [{'start_time': '2023-05-04T01:40:58.004292',\n", + " 'end_time': '2023-05-04T01:41:00.124338',\n", + " 'extra': {},\n", + " 'execution_order': 3,\n", + " 'serialized': {'name': 'ChatOpenAI'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 2,\n", + " 'parent_tool_run_id': None,\n", + " 'prompts': ['Human: Answer the following questions as best you can. You have access to the following tools:\\n\\nSearch: A search engine. Useful for when you need to answer questions about current events. Input should be a search query.\\nCalculator: Useful for when you need to answer questions about math.\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [Search, Calculator]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: how far is it from paris to boston in miles\\nThought:'],\n", + " 'response': {'generations': [[{'text': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"'}]],\n", + " 'llm_output': {'token_usage': {'prompt_tokens': 178,\n", + " 'completion_tokens': 27,\n", + " 'total_tokens': 205},\n", + " 'model_name': 'gpt-3.5-turbo'}},\n", + " 'id': 1,\n", + " 'type': 'llm'}],\n", + " 'start_time': '2023-05-04T01:40:58.003567',\n", + " 'end_time': '2023-05-04T01:41:00.124705',\n", + " 'extra': {},\n", + " 'execution_order': 2,\n", + " 'serialized': {'name': 'LLMChain'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 1,\n", + " 'parent_tool_run_id': None,\n", + " 'inputs': {'input': 'how far is it from paris to boston in miles',\n", + " 'agent_scratchpad': '',\n", + " 'stop': ['\\nObservation:', '\\n\\tObservation:']},\n", + " 'outputs': {'text': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"'},\n", + " 'id': 2,\n", + " 'type': 'chain'},\n", + " {'child_runs': [{'start_time': '2023-05-04T01:41:00.229748',\n", + " 'end_time': '2023-05-04T01:41:01.915701',\n", + " 'extra': {},\n", + " 'execution_order': 6,\n", + " 'serialized': {'name': 'ChatOpenAI'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 3,\n", + " 'parent_tool_run_id': None,\n", + " 'prompts': ['Human: Answer the following questions as best you can. You have access to the following tools:\\n\\nSearch: A search engine. Useful for when you need to answer questions about current events. Input should be a search query.\\nCalculator: Useful for when you need to answer questions about math.\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [Search, Calculator]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: how far is it from paris to boston in miles\\nThought:I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"\\nObservation: The distance between Paris (Paris Charles de Gaulle Airport) and Boston (Logan International Airport) is 3448 miles / 5550 kilometers / 2997 nautical miles.\\nThought:'],\n", + " 'response': {'generations': [[{'text': 'I have the answer to the question\\nFinal Answer: The distance between Paris and Boston is 3448 miles.'}]],\n", + " 'llm_output': {'token_usage': {'prompt_tokens': 247,\n", + " 'completion_tokens': 23,\n", + " 'total_tokens': 270},\n", + " 'model_name': 'gpt-3.5-turbo'}},\n", + " 'id': 2,\n", + " 'type': 'llm'}],\n", + " 'start_time': '2023-05-04T01:41:00.229351',\n", + " 'end_time': '2023-05-04T01:41:01.916173',\n", + " 'extra': {},\n", + " 'execution_order': 5,\n", + " 'serialized': {'name': 'LLMChain'},\n", + " 'session_id': 2,\n", + " 'example_id': '1781b1db-45cb-428f-8000-00482ccf7142',\n", + " 'error': None,\n", + " 'parent_chain_run_id': 1,\n", + " 'parent_tool_run_id': None,\n", + " 'inputs': {'input': 'how far is it from paris to boston in miles',\n", + " 'agent_scratchpad': 'I need to find the distance between Paris and Boston\\nAction: Search\\nAction Input: \"distance between Paris and Boston in miles\"\\nObservation: The distance between Paris (Paris Charles de Gaulle Airport) and Boston (Logan International Airport) is 3448 miles / 5550 kilometers / 2997 nautical miles.\\nThought:',\n", + " 'stop': ['\\nObservation:', '\\n\\tObservation:']},\n", + " 'outputs': {'text': 'I have the answer to the question\\nFinal Answer: The distance between Paris and Boston is 3448 miles.'},\n", + " 'id': 3,\n", + " 'type': 'chain'}]" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# For example, view the chain runs on this example\n", + "example.chain_runs" + ] } ], "metadata": { diff --git a/langchain/client/langchain.py b/langchain/client/langchain.py index 27b5571f016..14d4b15ffed 100644 --- a/langchain/client/langchain.py +++ b/langchain/client/langchain.py @@ -92,35 +92,6 @@ class LangChainClient(BaseSettings): headers["authorization"] = f"Bearer {self.api_key}" return headers - @staticmethod - async def _arun_chain( - example: Example, langchain_tracer: LangChainTracer, chain: Chain - ) -> Union[dict, str]: - """Run the chain asynchronously.""" - previous_example_id = langchain_tracer.example_id - langchain_tracer.example_id = example.id - try: - chain_output = await chain.arun( - example.inputs, callbacks=[langchain_tracer] - ) - except Exception as e: - logger.warning(f"Chain failed for example {example.id}. Error: {e}") - return {"Error": str(e)} - finally: - langchain_tracer.example_id = previous_example_id - return chain_output - - @staticmethod - async def _arun_chain_over_buffer( - buffer: Sequence[Example], tracers: Sequence[LangChainTracer], chain: Chain - ) -> List[Union[str, dict]]: - """Run the chain asynchronously over a buffer of examples.""" - batch_results = [ - LangChainClient._arun_chain(example, tracer, chain) - for example, tracer in zip(buffer, tracers) - ] - return await asyncio.gather(*batch_results) - @xor_args(("session_name", "session_id")) def read_session( self, *, session_name: Optional[str] = None, session_id: Optional[int] = None @@ -306,6 +277,15 @@ class LangChainClient(BaseSettings): return [Dataset(**dataset) for dataset in results] # Examples APIs. + + def read_example(self, example_id: str) -> Example: + """Read an example from the LangChain+ API.""" + response = requests.get( + self.api_url + f"/examples/{example_id}", headers=self._headers + ) + _raise_rich_error(response) + return Example(**response.json()) + def list_examples(self, dataset_id: Optional[str] = None) -> Iterable[Example]: """List the datasets on the LangChain+ API.""" params = {} if dataset_id is None else {"dataset": dataset_id} @@ -328,32 +308,100 @@ class LangChainClient(BaseSettings): results = await response.json() return [Example(**dataset) for dataset in results] + @staticmethod + async def _arun_chain( + example: Example, langchain_tracer: LangChainTracer, chain: Chain + ) -> Union[dict, str]: + """Run the chain asynchronously.""" + previous_example_id = langchain_tracer.example_id + langchain_tracer.example_id = example.id + try: + chain_output = await chain.arun( + example.inputs, callbacks=[langchain_tracer] + ) + except Exception as e: + logger.warning(f"Chain failed for example {example.id}. Error: {e}") + return {"Error": str(e)} + finally: + langchain_tracer.example_id = previous_example_id + return chain_output + + @staticmethod + async def _worker( + queue: asyncio.Queue, + tracers: List[LangChainTracer], + chain: Chain, + results: Dict[str, Any], + ) -> None: + """Worker for running the chain on examples.""" + while True: + example: Optional[Example] = await queue.get() + if example is None: + break + + tracer = tracers.pop() + result = await LangChainClient._arun_chain(example, tracer, chain) + results[example.id] = result + tracers.append(tracer) + queue.task_done() + + @staticmethod + async def _arun_chain_over_buffer( + buffer: Sequence[Example], tracers: Sequence[LangChainTracer], chain: Chain + ) -> List[Union[str, dict]]: + """Run the chain asynchronously over a buffer of examples.""" + batch_results = [ + LangChainClient._arun_chain(example, tracer, chain) + for example, tracer in zip(buffer, tracers) + ] + return await asyncio.gather(*batch_results) + async def arun_chain_on_dataset( self, dataset_name: str, chain: Chain, batch_size: int = 5, session_name: Optional[str] = None, - ) -> List[Any]: - """Run the chain on the specified dataset""" + ) -> Dict[str, Any]: + """Run the chain on a dataset and store traces to the specified session name. + + Args: + dataset_name: Name of the dataset to run the chain on. + chain: Chain to run on the dataset. + batch_size: Number of async workers to run in parallel. + session_name: Name of the session to store the traces in. + + Returns: + A dictionary mapping example ids to the chain outputs. + """ if session_name is not None: self.create_session(session_name) + dataset = await self.aread_dataset(dataset_name=dataset_name) - tracers = [LangChainTracer() for _ in range(batch_size)] - for tracer in tracers: + + tracers = [] + for _ in range(batch_size): + tracer = LangChainTracer() tracer.load_session(session_name or "default") - graded_outputs = [] + tracers.append(tracer) + + results: Dict[str, Any] = {} examples = await self.alist_examples(dataset_id=dataset.id) - buffer = [] + + queue = asyncio.Queue() + workers = [ + asyncio.create_task(LangChainClient._worker(queue, tracers, chain, results)) + for _ in range(batch_size) + ] + for example in examples: - buffer.append(example) - if len(buffer) == batch_size: - batch_results = await self._arun_chain_over_buffer( - buffer, tracers, chain - ) - graded_outputs.extend(batch_results) - buffer = [] - if buffer: - batch_results = await self._arun_chain_over_buffer(buffer, tracers, chain) - graded_outputs.extend(batch_results) - return graded_outputs + await queue.put(example) + + await queue.join() # Wait for all tasks to complete + + for _ in workers: + await queue.put(None) # Signal the workers to exit + + await asyncio.gather(*workers) + + return results