Upgraded LangChain, fixed prompts for Bedrock

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

@@ -15,6 +15,7 @@
     AnthropicProvider,
     AzureChatOpenAIProvider,
     BaseProvider,
+    BedrockChatProvider,
     BedrockProvider,
     ChatAnthropicProvider,
     ChatOpenAINewProvider,

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

@@ -417,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]
@@ -536,7 +539,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 self._is_chat_model(provider.id):
             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,14 @@
 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.llms import (
     AI21,
     Anthropic,
@@ -23,6 +29,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 +195,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.
@@ -596,14 +616,41 @@ 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",
-        "ai21.j2-mid",
-        "ai21.j2-ultra",
+        "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"]
@@ -619,3 +666,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.306",
     "typing_extensions>=4.5.0",
     "click~=8.0",
     "jsonpath-ng>=1.5.3,<2",
@@ -67,6 +67,7 @@ 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]",