Backport to 1.x

docs/source/users/index.md

@@ -39,7 +39,9 @@ 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`                         |
+| Bedrock (chat)      | `amazon-bedrock-chat`| N/A                        | `boto3`                         |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `openai`                        |
@@ -437,6 +439,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

@@ -2,6 +2,7 @@
 
 # expose embedding model providers on the package root
 from .embedding_providers import (
+    BedrockEmbeddingsProvider,
     CohereEmbeddingsProvider,
     HfHubEmbeddingsProvider,
     OpenAIEmbeddingsProvider,
@@ -15,7 +16,9 @@
     AnthropicProvider,
     AzureChatOpenAIProvider,
     BaseProvider,
+    BedrockChatProvider,
     BedrockProvider,
+    ChatAnthropicProvider,
     ChatOpenAINewProvider,
     ChatOpenAIProvider,
     CohereProvider,

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

@@ -1,7 +1,13 @@
 from typing import ClassVar, List, Type
 
-from jupyter_ai_magics.providers import AuthStrategy, EnvAuthStrategy, Field
+from jupyter_ai_magics.providers import (
+    AuthStrategy,
+    AwsAuthStrategy,
+    EnvAuthStrategy,
+    Field,
+)
 from langchain.embeddings import (
+    BedrockEmbeddings,
     CohereEmbeddings,
     HuggingFaceHubEmbeddings,
     OpenAIEmbeddings,
@@ -54,7 +60,8 @@ def __init__(self, *args, **kwargs):
             )
 
         model_kwargs = {}
-        model_kwargs[self.__class__.model_id_key] = kwargs["model_id"]
+        if self.__class__.model_id_key != "model_id":
+            model_kwargs[self.__class__.model_id_key] = kwargs["model_id"]
 
         super().__init__(*args, **kwargs, **model_kwargs)
 
@@ -88,3 +95,12 @@ class HfHubEmbeddingsProvider(BaseEmbeddingsProvider, HuggingFaceHubEmbeddings):
     pypi_package_deps = ["huggingface_hub", "ipywidgets"]
     auth_strategy = EnvAuthStrategy(name="HUGGINGFACEHUB_API_TOKEN")
     registry = True
+
+
+class BedrockEmbeddingsProvider(BaseEmbeddingsProvider, BedrockEmbeddings):
+    id = "bedrock"
+    name = "Bedrock"
+    models = ["amazon.titan-embed-text-v1"]
+    model_id_key = "model_id"
+    pypi_package_deps = ["boto3"]
+    auth_strategy = AwsAuthStrategy()

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,
@@ -138,6 +139,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()
 
@@ -542,8 +549,19 @@ def run_ai_cell(self, args: CellArgs, prompt: str):
 
         provider = Provider(**provider_params)
 
-        # generate output from model via provider
-        result = provider.generate([prompt])
+        # Apply a prompt template.
+        prompt = provider.get_prompt_template(args.format).format(prompt=prompt)
+
+        # interpolate user namespace into prompt
+        ip = get_ipython()
+        prompt = prompt.format_map(FormatDict(ip.user_ns))
+
+        if provider.is_chat_provider:
+            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,7 +8,13 @@
 from typing import Any, ClassVar, Coroutine, Dict, List, Literal, Optional, Union
 
 from jsonpath_ng import parse
-from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
+from langchain.chat_models import (
+    AzureChatOpenAI,
+    BedrockChat,
+    ChatAnthropic,
+    ChatOpenAI,
+)
+from langchain.chat_models.base import BaseChatModel
 from langchain.llms import (
     AI21,
     Anthropic,
@@ -22,6 +28,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
 
@@ -152,6 +160,39 @@ 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._generate() 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.
+        """
+        self.prompt_templates[format] = PromptTemplate.from_template(template)
+
+    def get_prompt_template(self, format) -> PromptTemplate:
+        """
+        Produce a prompt template suitable for use with a particular model, to
+        produce output in a desired format.
+        """
+
+        if format in self.prompt_templates:
+            return self.prompt_templates[format]
+        else:
+            return self.prompt_templates["text"]  # Default to plain format
+
+    @property
+    def is_chat_provider(self):
+        return isinstance(self, BaseChatModel)
+
 
 class AI21Provider(BaseProvider, AI21):
     id = "ai21"
@@ -183,8 +224,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"]
@@ -524,16 +585,49 @@ class BedrockProvider(BaseProvider, Bedrock):
     id = "bedrock"
     name = "Amazon Bedrock"
     models = [
-        "amazon.titan-tg1-large",
+        "amazon.titan-text-express-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 = [
         "anthropic.claude-v1",
-        "anthropic.claude-instant-v1",
         "anthropic.claude-v2",
-        "ai21.j2-jumbo-instruct",
-        "ai21.j2-grande-instruct",
+        "anthropic.claude-instant-v1",
     ]
     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.308",
     "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,8 +66,11 @@ 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"]
+bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics:OpenAIEmbeddingsProvider"

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/base.py

@@ -9,6 +9,7 @@
 from jupyter_ai.config_manager import ConfigManager, Logger
 from jupyter_ai.models import AgentChatMessage, HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
+from langchain.chat_models.base import BaseChatModel
 
 if TYPE_CHECKING:
     from jupyter_ai.handlers import RootChatHandler

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

@@ -2,15 +2,15 @@
 
 from jupyter_ai.models import ChatMessage, ClearMessage, HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
-from langchain import ConversationChain
+from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferWindowMemory
 from langchain.prompts import (
     ChatPromptTemplate,
     HumanMessagePromptTemplate,
     MessagesPlaceholder,
+    PromptTemplate,
     SystemMessagePromptTemplate,
 )
-from langchain.schema import AIMessage
 
 from .base import BaseChatHandler
 
@@ -25,6 +25,11 @@
 The following is a friendly conversation between you and a human.
 """.strip()
 
+DEFAULT_TEMPLATE = """Current conversation:
+{history}
+Human: {input}
+AI:"""
+
 
 class DefaultChatHandler(BaseChatHandler):
     def __init__(self, chat_history: List[ChatMessage], *args, **kwargs):
@@ -36,16 +41,29 @@ 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(
+
+        if llm.is_chat_provider:
+            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}"),
+                ]
+            )
+            self.memory = ConversationBufferWindowMemory(return_messages=True, k=2)
+        else:
+            prompt_template = PromptTemplate(
+                input_variables=["history", "input"],
+                template=SYSTEM_PROMPT.format(
                     provider_name=llm.name, local_model_id=llm.model_id
-                ),
-                MessagesPlaceholder(variable_name="history"),
-                HumanMessagePromptTemplate.from_template("{input}"),
-                AIMessage(content=""),
-            ]
-        )
+                )
+                + "\n\n"
+                + DEFAULT_TEMPLATE,
+            )
+            self.memory = ConversationBufferWindowMemory(k=2)
+
         self.llm = llm
         self.llm_chain = ConversationChain(
             llm=llm, prompt=prompt_template, verbose=True, memory=self.memory