Upgraded LangChain, fixed prompts for Bedrock

docs/source/users/index.md

@@ -118,6 +118,7 @@ Jupyter AI supports the following model providers:
 | 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`                        |

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,6 +16,7 @@
     AnthropicProvider,
     AzureChatOpenAIProvider,
     BaseProvider,
+    BedrockChatProvider,
     BedrockProvider,
     ChatAnthropicProvider,
     ChatOpenAINewProvider,

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

@@ -536,7 +536,7 @@ def run_ai_cell(self, args: CellArgs, prompt: str):
         ip = get_ipython()
         prompt = prompt.format_map(FormatDict(ip.user_ns))
 
-        if provider_id == "anthropic-chat":
+        if provider.is_chat_provider:
             result = provider.generate([[HumanMessage(content=prompt)]])
         else:
             # generate output from model via provider

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

@@ -8,8 +8,13 @@
 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, ChatAnthropic, ChatOpenAI
+from langchain.chat_models import (
+    AzureChatOpenAI,
+    BedrockChat,
+    ChatAnthropic,
+    ChatOpenAI,
+)
+from langchain.chat_models.base import BaseChatModel
 from langchain.llms import (
     AI21,
     Anthropic,
@@ -23,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
 
@@ -187,6 +194,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._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.
@@ -204,6 +223,10 @@ def get_prompt_template(self, format) -> PromptTemplate:
         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"
@@ -596,14 +619,34 @@ 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",
-        "ai21.j2-mid",
-        "ai21.j2-ultra",
+        "anthropic.claude-instant-v1",
     ]
     model_id_key = "model_id"
     pypi_package_deps = ["boto3"]
@@ -619,3 +662,6 @@ class BedrockProvider(BaseProvider, Bedrock):
 
     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",
@@ -67,8 +67,10 @@ 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

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.308",
     "tiktoken", # required for OpenAIEmbeddings
     "jupyter_ai_magics",
     "dask[distributed]",