Add a callback handler for Context (https://getcontext.ai) (#7151)

### Description

Adding a callback handler for Context. Context is a product analytics
platform for AI chat experiences to help you understand how users are
interacting with your product.

I've added the callback library + an example notebook showing its use.

### Dependencies

Requires the user to install the `context-python` library. The library
is lazily-loaded when the callback is instantiated.

### Announcing the feature

We spoke with Harrison a few weeks ago about also doing a blog post
announcing our integration, so will coordinate this with him. Our
Twitter handle for the company is @getcontextai, and the founders are
@_agamble and @HenrySG.

Thanks in advance!

docs/extras/modules/callbacks/integrations/context.ipynb

@@ -0,0 +1,220 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Context\n",
+    "\n",
+    "![Context - Product Analytics for AI Chatbots](https://go.getcontext.ai/langchain.png)\n",
+    "\n",
+    "[Context](https://getcontext.ai/) provides product analytics for AI chatbots.\n",
+    "\n",
+    "Context helps you understand how users are interacting with your AI chat products.\n",
+    "Gain critical insights, optimise poor experiences, and minimise brand risks.\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In this guide we will show you how to integrate with Context."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Installation and Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "vscode": {
+     "languageId": "shellscript"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "$ pip install context-python --upgrade"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Getting API Credentials\n",
+    "\n",
+    "To get your Context API token:\n",
+    "\n",
+    "1. Go to the settings page within your Context account (https://go.getcontext.ai/settings).\n",
+    "2. Generate a new API Token.\n",
+    "3. Store this token somewhere secure."
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Setup Context\n",
+    "\n",
+    "To use the `ContextCallbackHandler`, import the handler from Langchain and instantiate it with your Context API token.\n",
+    "\n",
+    "Ensure you have installed the `context-python` package before using the handler."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from langchain.callbacks import ContextCallbackHandler\n",
+    "\n",
+    "token = os.environ[\"CONTEXT_API_TOKEN\"]\n",
+    "\n",
+    "context_callback = ContextCallbackHandler(token)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Usage\n",
+    "### Using the Context callback within a Chat Model\n",
+    "\n",
+    "The Context callback handler can be used to directly record transcripts between users and AI assistants.\n",
+    "\n",
+    "#### Example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain.schema import (\n",
+    "    SystemMessage,\n",
+    "    HumanMessage,\n",
+    ")\n",
+    "from langchain.callbacks import ContextCallbackHandler\n",
+    "\n",
+    "token = os.environ[\"CONTEXT_API_TOKEN\"]\n",
+    "\n",
+    "chat = ChatOpenAI(\n",
+    "    headers={\"user_id\": \"123\"}, temperature=0, callbacks=[ContextCallbackHandler(token)]\n",
+    ")\n",
+    "\n",
+    "messages = [\n",
+    "    SystemMessage(\n",
+    "        content=\"You are a helpful assistant that translates English to French.\"\n",
+    "    ),\n",
+    "    HumanMessage(content=\"I love programming.\"),\n",
+    "]\n",
+    "\n",
+    "print(chat(messages))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Using the Context callback within Chains\n",
+    "\n",
+    "The Context callback handler can also be used to record the inputs and outputs of chains. Note that intermediate steps of the chain are not recorded - only the starting inputs and final outputs.\n",
+    "\n",
+    "__Note:__ Ensure that you pass the same context object to the chat model and the chain.\n",
+    "\n",
+    "Wrong:\n",
+    "> ```python\n",
+    "> chat = ChatOpenAI(temperature=0.9, callbacks=[ContextCallbackHandler(token)])\n",
+    "> chain = LLMChain(llm=chat, prompt=chat_prompt_template, callbacks=[ContextCallbackHandler(token)])\n",
+    "> ```\n",
+    "\n",
+    "Correct:\n",
+    ">```python\n",
+    ">handler = ContextCallbackHandler(token)\n",
+    ">chat = ChatOpenAI(temperature=0.9, callbacks=[callback])\n",
+    ">chain = LLMChain(llm=chat, prompt=chat_prompt_template, callbacks=[callback])\n",
+    ">```\n",
+    "\n",
+    "#### Example"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain import LLMChain\n",
+    "from langchain.prompts import PromptTemplate\n",
+    "from langchain.prompts.chat import (\n",
+    "    ChatPromptTemplate,\n",
+    "    HumanMessagePromptTemplate,\n",
+    ")\n",
+    "from langchain.callbacks import ContextCallbackHandler\n",
+    "\n",
+    "token = os.environ[\"CONTEXT_API_TOKEN\"]\n",
+    "\n",
+    "human_message_prompt = HumanMessagePromptTemplate(\n",
+    "    prompt=PromptTemplate(\n",
+    "        template=\"What is a good name for a company that makes {product}?\",\n",
+    "        input_variables=[\"product\"],\n",
+    "    )\n",
+    ")\n",
+    "chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])\n",
+    "callback = ContextCallbackHandler(token)\n",
+    "chat = ChatOpenAI(temperature=0.9, callbacks=[callback])\n",
+    "chain = LLMChain(llm=chat, prompt=chat_prompt_template, callbacks=[callback])\n",
+    "print(chain.run(\"colorful socks\"))"
+   ]
+  }
+ ],
+ "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.3"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "a53ebf4a859167383b364e7e7521d0add3c2dbbdecce4edf676e8c4634ff3fbb"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

langchain/callbacks/__init__.py

@@ -6,6 +6,7 @@ from langchain.callbacks.arize_callback import ArizeCallbackHandler
 from langchain.callbacks.arthur_callback import ArthurCallbackHandler
 from langchain.callbacks.clearml_callback import ClearMLCallbackHandler
 from langchain.callbacks.comet_ml_callback import CometCallbackHandler
+from langchain.callbacks.context_callback import ContextCallbackHandler
 from langchain.callbacks.file import FileCallbackHandler
 from langchain.callbacks.flyte_callback import FlyteCallbackHandler
 from langchain.callbacks.human import HumanApprovalCallbackHandler
@@ -36,6 +37,7 @@ __all__ = [
     "ArthurCallbackHandler",
     "ClearMLCallbackHandler",
     "CometCallbackHandler",
+    "ContextCallbackHandler",
     "FileCallbackHandler",
     "HumanApprovalCallbackHandler",
     "InfinoCallbackHandler",

langchain/callbacks/context_callback.py

@@ -0,0 +1,193 @@
+"""Callback handler for Context AI"""
+import os
+from typing import Any, Dict, List
+from uuid import UUID
+
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import (
+    BaseMessage,
+    LLMResult,
+)
+
+
+def import_context() -> Any:
+    try:
+        import getcontext  # noqa: F401
+        from getcontext.generated.models import (
+            Conversation,
+            Message,
+            MessageRole,
+            Rating,
+        )
+        from getcontext.token import Credential  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "To use the context callback manager you need to have the "
+            "`getcontext` python package installed (version >=0.3.0). "
+            "Please install it with `pip install --upgrade python-context`"
+        )
+    return getcontext, Credential, Conversation, Message, MessageRole, Rating
+
+
+class ContextCallbackHandler(BaseCallbackHandler):
+    """Callback Handler that records transcripts to Context (https://getcontext.ai).
+
+    Keyword Args:
+        token (optional): The token with which to authenticate requests to Context.
+            Visit https://go.getcontext.ai/settings to generate a token.
+            If not provided, the value of the `CONTEXT_TOKEN` environment
+            variable will be used.
+
+    Raises:
+        ImportError: if the `context-python` package is not installed.
+
+    Chat Example:
+        >>> from langchain.llms import ChatOpenAI
+        >>> from langchain.callbacks import ContextCallbackHandler
+        >>> context_callback = ContextCallbackHandler(
+        ...     token="<CONTEXT_TOKEN_HERE>",
+        ... )
+        >>> chat = ChatOpenAI(
+        ...     temperature=0,
+        ...     headers={"user_id": "123"},
+        ...     callbacks=[context_callback],
+        ...     openai_api_key="API_KEY_HERE",
+        ... )
+        >>> messages = [
+        ...     SystemMessage(content="You translate English to French."),
+        ...     HumanMessage(content="I love programming with LangChain."),
+        ... ]
+        >>> chat(messages)
+
+    Chain Example:
+        >>> from langchain import LLMChain
+        >>> from langchain.llms import ChatOpenAI
+        >>> from langchain.callbacks import ContextCallbackHandler
+        >>> context_callback = ContextCallbackHandler(
+        ...     token="<CONTEXT_TOKEN_HERE>",
+        ... )
+        >>> human_message_prompt = HumanMessagePromptTemplate(
+        ...     prompt=PromptTemplate(
+        ...         template="What is a good name for a company that makes {product}?",
+        ...         input_variables=["product"],
+        ...    ),
+        ... )
+        >>> chat_prompt_template = ChatPromptTemplate.from_messages(
+        ...   [human_message_prompt]
+        ... )
+        >>> callback = ContextCallbackHandler(token)
+        >>> # Note: the same callback object must be shared between the
+        ...   LLM and the chain.
+        >>> chat = ChatOpenAI(temperature=0.9, callbacks=[callback])
+        >>> chain = LLMChain(
+        ...   llm=chat,
+        ...   prompt=chat_prompt_template,
+        ...   callbacks=[callback]
+        ... )
+        >>> chain.run("colorful socks")
+    """
+
+    def __init__(self, token: str = "", verbose: bool = False, **kwargs: Any) -> None:
+        (
+            self.context,
+            self.credential,
+            self.conversation_model,
+            self.message_model,
+            self.message_role_model,
+            self.rating_model,
+        ) = import_context()
+
+        token = token or os.environ.get("CONTEXT_TOKEN") or ""
+
+        self.client = self.context.ContextAPI(credential=self.credential(token))
+
+        self.chain_run_id = None
+
+        self.llm_model = None
+
+        self.messages: List[Any] = []
+        self.metadata: Dict[str, str] = {}
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> Any:
+        """Run when the chat model is started."""
+        llm_model = kwargs.get("invocation_params", {}).get("model", None)
+        if llm_model is not None:
+            self.metadata["llm_model"] = llm_model
+
+        if len(messages) == 0:
+            return
+
+        for message in messages[0]:
+            role = self.message_role_model.SYSTEM
+            if message.type == "human":
+                role = self.message_role_model.USER
+            elif message.type == "system":
+                role = self.message_role_model.SYSTEM
+            elif message.type == "ai":
+                role = self.message_role_model.ASSISTANT
+
+            self.messages.append(
+                self.message_model(
+                    message=message.content,
+                    role=role,
+                )
+            )
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Run when LLM ends."""
+        if len(response.generations) == 0 or len(response.generations[0]) == 0:
+            return
+
+        if not self.chain_run_id:
+            generation = response.generations[0][0]
+            self.messages.append(
+                self.message_model(
+                    message=generation.text,
+                    role=self.message_role_model.ASSISTANT,
+                )
+            )
+
+            self._log_conversation()
+
+    def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> None:
+        """Run when chain starts."""
+        self.chain_run_id = kwargs.get("run_id", None)
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        """Run when chain ends."""
+        self.messages.append(
+            self.message_model(
+                message=outputs["text"],
+                role=self.message_role_model.ASSISTANT,
+            )
+        )
+
+        self._log_conversation()
+
+        self.chain_run_id = None
+
+    def _log_conversation(self) -> None:
+        """Log the conversation to the context API."""
+        if len(self.messages) == 0:
+            return
+
+        self.client.log.conversation_upsert(
+            body={
+                "conversation": self.conversation_model(
+                    messages=self.messages,
+                    metadata=self.metadata,
+                )
+            }
+        )
+
+        self.messages = []
+        self.metadata = {}