2023-04-30 18:14:09 +00:00
|
|
|
{
|
|
|
|
"cells": [
|
|
|
|
{
|
|
|
|
"cell_type": "markdown",
|
|
|
|
"id": "593f7553-7038-498e-96d4-8255e5ce34f0",
|
|
|
|
"metadata": {},
|
|
|
|
"source": [
|
2023-06-16 18:52:56 +00:00
|
|
|
"# Custom chain\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
"\n",
|
|
|
|
"To implement your own custom chain you can subclass `Chain` and implement the following methods:"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
|
|
|
"execution_count": 11,
|
|
|
|
"id": "c19c736e-ca74-4726-bb77-0a849bcc2960",
|
|
|
|
"metadata": {
|
2023-07-29 15:44:32 +00:00
|
|
|
"tags": []
|
2023-04-30 18:14:09 +00:00
|
|
|
},
|
|
|
|
"outputs": [],
|
|
|
|
"source": [
|
|
|
|
"from __future__ import annotations\n",
|
|
|
|
"\n",
|
|
|
|
"from typing import Any, Dict, List, Optional\n",
|
|
|
|
"\n",
|
|
|
|
"from pydantic import Extra\n",
|
|
|
|
"\n",
|
2023-07-29 15:44:32 +00:00
|
|
|
"from langchain.schema.language_model import BaseLanguageModel\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
"from langchain.callbacks.manager import (\n",
|
|
|
|
" AsyncCallbackManagerForChainRun,\n",
|
|
|
|
" CallbackManagerForChainRun,\n",
|
|
|
|
")\n",
|
|
|
|
"from langchain.chains.base import Chain\n",
|
|
|
|
"from langchain.prompts.base import BasePromptTemplate\n",
|
|
|
|
"\n",
|
|
|
|
"\n",
|
|
|
|
"class MyCustomChain(Chain):\n",
|
|
|
|
" \"\"\"\n",
|
|
|
|
" An example of a custom chain.\n",
|
|
|
|
" \"\"\"\n",
|
|
|
|
"\n",
|
|
|
|
" prompt: BasePromptTemplate\n",
|
|
|
|
" \"\"\"Prompt object to use.\"\"\"\n",
|
|
|
|
" llm: BaseLanguageModel\n",
|
|
|
|
" output_key: str = \"text\" #: :meta private:\n",
|
|
|
|
"\n",
|
|
|
|
" class Config:\n",
|
|
|
|
" \"\"\"Configuration for this pydantic object.\"\"\"\n",
|
|
|
|
"\n",
|
|
|
|
" extra = Extra.forbid\n",
|
|
|
|
" arbitrary_types_allowed = True\n",
|
|
|
|
"\n",
|
|
|
|
" @property\n",
|
|
|
|
" def input_keys(self) -> List[str]:\n",
|
|
|
|
" \"\"\"Will be whatever keys the prompt expects.\n",
|
|
|
|
"\n",
|
|
|
|
" :meta private:\n",
|
|
|
|
" \"\"\"\n",
|
|
|
|
" return self.prompt.input_variables\n",
|
|
|
|
"\n",
|
|
|
|
" @property\n",
|
|
|
|
" def output_keys(self) -> List[str]:\n",
|
|
|
|
" \"\"\"Will always return text key.\n",
|
|
|
|
"\n",
|
|
|
|
" :meta private:\n",
|
|
|
|
" \"\"\"\n",
|
|
|
|
" return [self.output_key]\n",
|
|
|
|
"\n",
|
|
|
|
" def _call(\n",
|
|
|
|
" self,\n",
|
|
|
|
" inputs: Dict[str, Any],\n",
|
|
|
|
" run_manager: Optional[CallbackManagerForChainRun] = None,\n",
|
|
|
|
" ) -> Dict[str, str]:\n",
|
|
|
|
" # Your custom chain logic goes here\n",
|
|
|
|
" # This is just an example that mimics LLMChain\n",
|
|
|
|
" prompt_value = self.prompt.format_prompt(**inputs)\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
"\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" # Whenever you call a language model, or another chain, you should pass\n",
|
|
|
|
" # a callback manager to it. This allows the inner run to be tracked by\n",
|
|
|
|
" # any callbacks that are registered on the outer run.\n",
|
|
|
|
" # You can always obtain a callback manager for this by calling\n",
|
|
|
|
" # `run_manager.get_child()` as shown below.\n",
|
|
|
|
" response = self.llm.generate_prompt(\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
" [prompt_value], callbacks=run_manager.get_child() if run_manager else None\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" )\n",
|
|
|
|
"\n",
|
|
|
|
" # If you want to log something about this run, you can do so by calling\n",
|
|
|
|
" # methods on the `run_manager`, as shown below. This will trigger any\n",
|
|
|
|
" # callbacks that are registered for that event.\n",
|
|
|
|
" if run_manager:\n",
|
|
|
|
" run_manager.on_text(\"Log something about this run\")\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
"\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" return {self.output_key: response.generations[0][0].text}\n",
|
|
|
|
"\n",
|
|
|
|
" async def _acall(\n",
|
|
|
|
" self,\n",
|
|
|
|
" inputs: Dict[str, Any],\n",
|
|
|
|
" run_manager: Optional[AsyncCallbackManagerForChainRun] = None,\n",
|
|
|
|
" ) -> Dict[str, str]:\n",
|
|
|
|
" # Your custom chain logic goes here\n",
|
|
|
|
" # This is just an example that mimics LLMChain\n",
|
|
|
|
" prompt_value = self.prompt.format_prompt(**inputs)\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
"\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" # Whenever you call a language model, or another chain, you should pass\n",
|
|
|
|
" # a callback manager to it. This allows the inner run to be tracked by\n",
|
|
|
|
" # any callbacks that are registered on the outer run.\n",
|
|
|
|
" # You can always obtain a callback manager for this by calling\n",
|
|
|
|
" # `run_manager.get_child()` as shown below.\n",
|
|
|
|
" response = await self.llm.agenerate_prompt(\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
" [prompt_value], callbacks=run_manager.get_child() if run_manager else None\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" )\n",
|
|
|
|
"\n",
|
|
|
|
" # If you want to log something about this run, you can do so by calling\n",
|
|
|
|
" # methods on the `run_manager`, as shown below. This will trigger any\n",
|
|
|
|
" # callbacks that are registered for that event.\n",
|
|
|
|
" if run_manager:\n",
|
|
|
|
" await run_manager.on_text(\"Log something about this run\")\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
"\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
" return {self.output_key: response.generations[0][0].text}\n",
|
|
|
|
"\n",
|
|
|
|
" @property\n",
|
|
|
|
" def _chain_type(self) -> str:\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
" return \"my_custom_chain\""
|
2023-04-30 18:14:09 +00:00
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
|
|
|
"execution_count": 12,
|
|
|
|
"id": "18361f89",
|
2023-07-29 15:44:32 +00:00
|
|
|
"metadata": {},
|
2023-04-30 18:14:09 +00:00
|
|
|
"outputs": [
|
|
|
|
{
|
|
|
|
"name": "stdout",
|
|
|
|
"output_type": "stream",
|
|
|
|
"text": [
|
|
|
|
"\n",
|
|
|
|
"\n",
|
|
|
|
"\u001b[1m> Entering new MyCustomChain chain...\u001b[0m\n",
|
|
|
|
"Log something about this run\n",
|
|
|
|
"\u001b[1m> Finished chain.\u001b[0m\n"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"data": {
|
|
|
|
"text/plain": [
|
|
|
|
"'Why did the callback function feel lonely? Because it was always waiting for someone to call it back!'"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
"execution_count": 12,
|
|
|
|
"metadata": {},
|
|
|
|
"output_type": "execute_result"
|
|
|
|
}
|
|
|
|
],
|
|
|
|
"source": [
|
|
|
|
"from langchain.callbacks.stdout import StdOutCallbackHandler\n",
|
|
|
|
"from langchain.chat_models.openai import ChatOpenAI\n",
|
|
|
|
"from langchain.prompts.prompt import PromptTemplate\n",
|
|
|
|
"\n",
|
|
|
|
"\n",
|
|
|
|
"chain = MyCustomChain(\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
" prompt=PromptTemplate.from_template(\"tell us a joke about {topic}\"),\n",
|
|
|
|
" llm=ChatOpenAI(),\n",
|
2023-04-30 18:14:09 +00:00
|
|
|
")\n",
|
|
|
|
"\n",
|
2023-06-16 18:52:56 +00:00
|
|
|
"chain.run({\"topic\": \"callbacks\"}, callbacks=[StdOutCallbackHandler()])"
|
2023-04-30 18:14:09 +00:00
|
|
|
]
|
|
|
|
}
|
|
|
|
],
|
|
|
|
"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",
|
2023-06-16 18:52:56 +00:00
|
|
|
"version": "3.11.3"
|
2023-04-30 18:14:09 +00:00
|
|
|
}
|
|
|
|
},
|
|
|
|
"nbformat": 4,
|
|
|
|
"nbformat_minor": 5
|
|
|
|
}
|