Upgraded LangChain, fixed prompts for Bedrock

Author: 3coins

docs/source/users/index.md

@@ -116,6 +116,7 @@ Jupyter AI supports the following model providers:
 |---------------------|----------------------|----------------------------|---------------------------------|
 | AI21                | `ai21`               | `AI21_API_KEY`             | `ai21`                          |
 | Anthropic           | `anthropic`          | `ANTHROPIC_API_KEY`        | `anthropic`                     |
+| Anthropic (chat)    | `anthropic-chat`     | `ANTHROPIC_API_KEY`        | `anthropic`                     |
 | Bedrock             | `amazon-bedrock`     | N/A                        | `boto3`                         |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
@@ -464,6 +465,7 @@ We currently support the following language model providers:
 
 - `ai21`
 - `anthropic`
+- `anthropic-chat`
 - `cohere`
 - `huggingface_hub`
 - `openai`

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -15,7 +15,9 @@
     AnthropicProvider,
     AzureChatOpenAIProvider,
     BaseProvider,
+    BedrockChatProvider,
     BedrockProvider,
+    ChatAnthropicProvider,
     ChatOpenAINewProvider,
     ChatOpenAIProvider,
     CohereProvider,

packages/jupyter-ai-magics/jupyter_ai_magics/magics.py

@@ -13,6 +13,7 @@
 from IPython.display import HTML, JSON, Markdown, Math
 from jupyter_ai_magics.utils import decompose_model_id, get_lm_providers
 from langchain.chains import LLMChain
+from langchain.schema import HumanMessage
 
 from .parsers import (
     CellArgs,
@@ -125,6 +126,12 @@ def __init__(self, shell):
             "no longer supported. Instead, please use: "
             "`from langchain.chat_models import ChatOpenAI`",
         )
+        # suppress warning when using old Anthropic provider
+        warnings.filterwarnings(
+            "ignore",
+            message="This Anthropic LLM is deprecated. Please use "
+            "`from langchain.chat_models import ChatAnthropic` instead",
+        )
 
         self.providers = get_lm_providers()
 
@@ -410,6 +417,9 @@ def _get_provider(self, provider_id: Optional[str]) -> BaseProvider:
 
         return self.providers[provider_id]
 
+    def _is_chat_model(self, provider_id: str) -> bool:
+        return provider_id in ["anthropic-chat", "bedrock-chat"]
+
     def display_output(self, output, display_format, md):
         # build output display
         DisplayClass = DISPLAYS_BY_FORMAT[display_format]
@@ -529,8 +539,12 @@ def run_ai_cell(self, args: CellArgs, prompt: str):
         ip = get_ipython()
         prompt = prompt.format_map(FormatDict(ip.user_ns))
 
-        # generate output from model via provider
-        result = provider.generate([prompt])
+        if self._is_chat_model(provider.id):
+            result = provider.generate([[HumanMessage(content=prompt)]])
+        else:
+            # generate output from model via provider
+            result = provider.generate([prompt])
+
         output = result.generations[0][0].text
 
         # if openai-chat, append exchange to transcript

packages/jupyter-ai-magics/jupyter_ai_magics/providers.py

@@ -8,8 +8,12 @@
 from typing import Any, ClassVar, Coroutine, Dict, List, Literal, Optional, Union
 
 from jsonpath_ng import parse
-from langchain import PromptTemplate
-from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
+from langchain.chat_models import (
+    AzureChatOpenAI,
+    BedrockChat,
+    ChatAnthropic,
+    ChatOpenAI,
+)
 from langchain.llms import (
     AI21,
     Anthropic,
@@ -23,6 +27,8 @@
 )
 from langchain.llms.sagemaker_endpoint import LLMContentHandler
 from langchain.llms.utils import enforce_stop_tokens
+from langchain.prompts import PromptTemplate
+from langchain.schema import LLMResult
 from langchain.utils import get_from_dict_or_env
 from pydantic import BaseModel, Extra, root_validator
 
@@ -187,6 +193,18 @@ async def _call_in_executor(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
         _call_with_args = functools.partial(self._call, *args, **kwargs)
         return await loop.run_in_executor(executor, _call_with_args)
 
+    async def _generate_in_executor(
+        self, *args, **kwargs
+    ) -> Coroutine[Any, Any, LLMResult]:
+        """
+        Calls self._call() asynchronously in a separate thread for providers
+        without an async implementation. Requires the event loop to be running.
+        """
+        executor = ThreadPoolExecutor(max_workers=1)
+        loop = asyncio.get_running_loop()
+        _call_with_args = functools.partial(self._generate, *args, **kwargs)
+        return await loop.run_in_executor(executor, _call_with_args)
+
     def update_prompt_template(self, format: str, template: str):
         """
         Changes the class-level prompt template for a given format.
@@ -235,8 +253,28 @@ class AnthropicProvider(BaseProvider, Anthropic):
         "claude-v1.0",
         "claude-v1.2",
         "claude-2",
+        "claude-2.0",
+        "claude-instant-v1",
+        "claude-instant-v1.0",
+        "claude-instant-v1.2",
+    ]
+    model_id_key = "model"
+    pypi_package_deps = ["anthropic"]
+    auth_strategy = EnvAuthStrategy(name="ANTHROPIC_API_KEY")
+
+
+class ChatAnthropicProvider(BaseProvider, ChatAnthropic):
+    id = "anthropic-chat"
+    name = "ChatAnthropic"
+    models = [
+        "claude-v1",
+        "claude-v1.0",
+        "claude-v1.2",
+        "claude-2",
+        "claude-2.0",
         "claude-instant-v1",
         "claude-instant-v1.0",
+        "claude-instant-v1.2",
     ]
     model_id_key = "model"
     pypi_package_deps = ["anthropic"]
@@ -576,16 +614,56 @@ class BedrockProvider(BaseProvider, Bedrock):
     id = "bedrock"
     name = "Amazon Bedrock"
     models = [
-        "amazon.titan-tg1-large",
+        "amazon.titan-text-express-v1",
         "anthropic.claude-v1",
+        "anthropic.claude-v2",
         "anthropic.claude-instant-v1",
+        "ai21.j2-ultra-v1",
+        "ai21.j2-mid-v1",
+        "cohere.command-text-v14",
+    ]
+    model_id_key = "model_id"
+    pypi_package_deps = ["boto3"]
+    auth_strategy = AwsAuthStrategy()
+    fields = [
+        TextField(
+            key="credentials_profile_name",
+            label="AWS profile (optional)",
+            format="text",
+        ),
+        TextField(key="region_name", label="Region name (optional)", format="text"),
+    ]
+
+    async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
+        return await self._call_in_executor(*args, **kwargs)
+
+
+class BedrockChatProvider(BaseProvider, BedrockChat):
+    id = "bedrock-chat"
+    name = "Amazon Bedrock Chat"
+    models = [
+        "amazon.titan-text-express-v1",
+        "anthropic.claude-v1",
         "anthropic.claude-v2",
-        "ai21.j2-jumbo-instruct",
-        "ai21.j2-grande-instruct",
+        "anthropic.claude-instant-v1",
+        "ai21.j2-ultra-v1",
+        "ai21.j2-mid-v1",
+        "cohere.command-text-v14",
     ]
     model_id_key = "model_id"
     pypi_package_deps = ["boto3"]
     auth_strategy = AwsAuthStrategy()
+    fields = [
+        TextField(
+            key="credentials_profile_name",
+            label="AWS profile (optional)",
+            format="text",
+        ),
+        TextField(key="region_name", label="Region name (optional)", format="text"),
+    ]
 
     async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]:
         return await self._call_in_executor(*args, **kwargs)
+
+    async def _agenerate(self, *args, **kwargs) -> Coroutine[Any, Any, LLMResult]:
+        return await self._generate_in_executor(*args, **kwargs)

packages/jupyter-ai-magics/pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "ipython",
     "pydantic~=1.0",
     "importlib_metadata>=5.2.0",
-    "langchain==0.0.277",
+    "langchain==0.0.306",
     "typing_extensions>=4.5.0",
     "click~=8.0",
     "jsonpath-ng>=1.5.3,<2",
@@ -44,7 +44,7 @@ test = [
 
 all = [
     "ai21",
-    "anthropic~=0.2.10",
+    "anthropic~=0.3.0",
     "cohere",
     "gpt4all",
     "huggingface_hub",
@@ -66,6 +66,8 @@ openai-chat-new = "jupyter_ai_magics:ChatOpenAINewProvider"
 azure-chat-openai = "jupyter_ai_magics:AzureChatOpenAIProvider"
 sagemaker-endpoint = "jupyter_ai_magics:SmEndpointProvider"
 amazon-bedrock = "jupyter_ai_magics:BedrockProvider"
+anthropic-chat = "jupyter_ai_magics:ChatAnthropicProvider"
+amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"

packages/jupyter-ai/jupyter_ai/chat_handlers/ask.py

@@ -11,7 +11,7 @@
 class AskChatHandler(BaseChatHandler):
     """Processes messages prefixed with /ask. This actor will
     send the message as input to a RetrieverQA chain, that
-    follows the Retrieval and Generation (RAG) tehnique to
+    follows the Retrieval and Generation (RAG) technique to
     query the documents from the index, and sends this context
     to the LLM to generate the final reply.
     """
@@ -29,7 +29,7 @@ def create_llm_chain(
         self.llm = provider(**provider_params)
         self.chat_history = []
         self.llm_chain = ConversationalRetrievalChain.from_llm(
-            self.llm, self._retriever
+            self.llm, self._retriever, verbose=True
         )
 
     async def _process_message(self, message: HumanChatMessage):

packages/jupyter-ai/jupyter_ai/chat_handlers/default.py

@@ -1,16 +1,21 @@
-from typing import Dict, List, Type
+from typing import Any, Dict, List, Type
 
 from jupyter_ai.models import ChatMessage, ClearMessage, HumanChatMessage
-from jupyter_ai_magics.providers import BaseProvider
-from langchain import ConversationChain
+from jupyter_ai_magics.providers import (
+    BaseProvider,
+    BedrockChatProvider,
+    BedrockProvider,
+)
+from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferWindowMemory
 from langchain.prompts import (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
     MessagesPlaceholder,
     SystemMessagePromptTemplate,
 )
-from langchain.schema import AIMessage
+from langchain.schema import AIMessage, ChatMessage
+from langchain.schema.messages import BaseMessage
 
 from .base import BaseChatHandler
 
@@ -26,6 +31,20 @@
 """.strip()
 
 
+class HistoryPlaceholderTemplate(MessagesPlaceholder):
+    def format_messages(self, **kwargs: Any) -> List[BaseMessage]:
+        values = super().format_messages(**kwargs)
+        corrected_values = []
+        for v in values:
+            if isinstance(v, AIMessage):
+                corrected_values.append(
+                    ChatMessage(role="Assistant", content=v.content)
+                )
+            else:
+                corrected_values.append(v)
+        return corrected_values
+
+
 class DefaultChatHandler(BaseChatHandler):
     def __init__(self, chat_history: List[ChatMessage], *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -36,16 +55,32 @@ def create_llm_chain(
         self, provider: Type[BaseProvider], provider_params: Dict[str, str]
     ):
         llm = provider(**provider_params)
-        prompt_template = ChatPromptTemplate.from_messages(
-            [
-                SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT).format(
-                    provider_name=llm.name, local_model_id=llm.model_id
-                ),
-                MessagesPlaceholder(variable_name="history"),
-                HumanMessagePromptTemplate.from_template("{input}"),
-                AIMessage(content=""),
-            ]
-        )
+        if provider == BedrockChatProvider or provider == BedrockProvider:
+            prompt_template = ChatPromptTemplate.from_messages(
+                [
+                    ChatMessage(
+                        role="Instructions",
+                        content=SYSTEM_PROMPT.format(
+                            provider_name=llm.name, local_model_id=llm.model_id
+                        ),
+                    ),
+                    HistoryPlaceholderTemplate(variable_name="history"),
+                    HumanMessagePromptTemplate.from_template("{input}"),
+                    ChatMessage(role="Assistant", content=""),
+                ]
+            )
+        else:
+            prompt_template = ChatPromptTemplate.from_messages(
+                [
+                    SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT).format(
+                        provider_name=llm.name, local_model_id=llm.model_id
+                    ),
+                    MessagesPlaceholder(variable_name="history"),
+                    HumanMessagePromptTemplate.from_template("{input}"),
+                    AIMessage(content=""),
+                ]
+            )
+
         self.llm = llm
         self.llm_chain = ConversationChain(
             llm=llm, prompt=prompt_template, verbose=True, memory=self.memory

packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py

@@ -15,14 +15,14 @@
     IndexMetadata,
 )
 from jupyter_core.paths import jupyter_data_dir
-from langchain import FAISS
 from langchain.schema import BaseRetriever, Document
 from langchain.text_splitter import (
     LatexTextSplitter,
     MarkdownTextSplitter,
     PythonCodeTextSplitter,
     RecursiveCharacterTextSplitter,
 )
+from langchain.vectorstores import FAISS
 
 from .base import BaseChatHandler
 

packages/jupyter-ai/pyproject.toml

@@ -28,7 +28,7 @@ dependencies = [
     "openai~=0.26",
     "aiosqlite>=0.18",
     "importlib_metadata>=5.2.0",
-    "langchain==0.0.277",
+    "langchain==0.0.306",
     "tiktoken", # required for OpenAIEmbeddings
     "jupyter_ai_magics",
     "dask[distributed]",