Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Expose templates for customisation in providers #581

Merged
merged 3 commits into from
Jan 25, 2024
Merged

Expose templates for customisation in providers #581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
31c572d
Expose chat template for customisation in providers
krassowski Jan 15, 2024
7eb8518
Enable completion template customization, including suffix
krassowski Jan 21, 2024
b3e507c
Remove `inline` from "inline completion" as per review request
krassowski Jan 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 99 additions & 1 deletion packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,13 @@
from langchain.chat_models.base import BaseChatModel
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.llms.utils import enforce_stop_tokens
from langchain.prompts import PromptTemplate
from langchain.prompts import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
MessagesPlaceholder,
PromptTemplate,
SystemMessagePromptTemplate,
)
from langchain.pydantic_v1 import BaseModel, Extra, root_validator
from langchain.schema import LLMResult
from langchain.utils import get_from_dict_or_env
Expand Down Expand Up @@ -42,6 +48,49 @@
from pydantic.main import ModelMetaclass


CHAT_SYSTEM_PROMPT = """
You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
You are not a language model, but rather an application built on a foundation model from {provider_name} called {local_model_id}.
You are talkative and you provide lots of specific details from the foundation model's context.
You may use Markdown to format your response.
Code blocks must be formatted in Markdown.
Math should be rendered with inline TeX markup, surrounded by $.
If you do not know the answer to a question, answer truthfully by responding that you do not know.
The following is a friendly conversation between you and a human.
""".strip()

CHAT_DEFAULT_TEMPLATE = """Current conversation:
{history}
Human: {input}
AI:"""


COMPLETION_SYSTEM_PROMPT = """
You are an application built to provide helpful code completion suggestions.
You should only produce code. Keep comments to minimum, use the
programming language comment syntax. Produce clean code.
The code is written in JupyterLab, a data analysis and code development
environment which can execute code extended with additional syntax for
interactive features, such as magics.
""".strip()

# only add the suffix bit if present to save input tokens/computation time
COMPLETION_DEFAULT_TEMPLATE = """
The document is called `{{filename}}` and written in {{language}}.
{% if suffix %}
The code after the completion request is:

```
{{suffix}}
```
{% endif %}

Complete the following code:

```
{{prefix}}"""


class EnvAuthStrategy(BaseModel):
"""Require one auth token via an environment variable."""

Expand Down Expand Up @@ -265,6 +314,55 @@ def get_prompt_template(self, format) -> PromptTemplate:
else:
return self.prompt_templates["text"] # Default to plain format

def get_chat_prompt_template(self) -> PromptTemplate:
"""
Produce a prompt template optimised for chat conversation.
The template should take two variables: history and input.
"""
name = self.__class__.name
if self.is_chat_provider:
return ChatPromptTemplate.from_messages(
[
SystemMessagePromptTemplate.from_template(
CHAT_SYSTEM_PROMPT
).format(provider_name=name, local_model_id=self.model_id),
MessagesPlaceholder(variable_name="history"),
HumanMessagePromptTemplate.from_template("{input}"),
]
)
else:
return PromptTemplate(
input_variables=["history", "input"],
template=CHAT_SYSTEM_PROMPT.format(
provider_name=name, local_model_id=self.model_id
)
+ "\n\n"
+ CHAT_DEFAULT_TEMPLATE,
)

def get_completion_prompt_template(self) -> PromptTemplate:
"""
Produce a prompt template optimised for inline code or text completion.
The template should take variables: prefix, suffix, language, filename.
"""
if self.is_chat_provider:
return ChatPromptTemplate.from_messages(
[
SystemMessagePromptTemplate.from_template(COMPLETION_SYSTEM_PROMPT),
HumanMessagePromptTemplate.from_template(
COMPLETION_DEFAULT_TEMPLATE, template_format="jinja2"
),
]
)
else:
return PromptTemplate(
input_variables=["prefix", "suffix", "language", "filename"],
template=COMPLETION_SYSTEM_PROMPT
+ "\n\n"
+ COMPLETION_DEFAULT_TEMPLATE,
template_format="jinja2",
)

@property
def is_chat_provider(self):
return isinstance(self, BaseChatModel)
Expand Down
48 changes: 4 additions & 44 deletions packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,32 +4,9 @@
from jupyter_ai_magics.providers import BaseProvider
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import (
ChatPromptTemplate,
HumanMessagePromptTemplate,
MessagesPlaceholder,
PromptTemplate,
SystemMessagePromptTemplate,
)

from .base import BaseChatHandler, SlashCommandRoutingType

SYSTEM_PROMPT = """
You are Jupyternaut, a conversational assistant living in JupyterLab to help users.
You are not a language model, but rather an application built on a foundation model from {provider_name} called {local_model_id}.
You are talkative and you provide lots of specific details from the foundation model's context.
You may use Markdown to format your response.
Code blocks must be formatted in Markdown.
Math should be rendered with inline TeX markup, surrounded by $.
If you do not know the answer to a question, answer truthfully by responding that you do not know.
The following is a friendly conversation between you and a human.
""".strip()

DEFAULT_TEMPLATE = """Current conversation:
{history}
Human: {input}
AI:"""


class DefaultChatHandler(BaseChatHandler):
id = "default"
Expand All @@ -49,27 +26,10 @@ def create_llm_chain(
model_parameters = self.get_model_parameters(provider, provider_params)
llm = provider(**provider_params, **model_parameters)

if llm.is_chat_provider:
prompt_template = ChatPromptTemplate.from_messages(
[
SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT).format(
provider_name=provider.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=provider.name, local_model_id=llm.model_id
)
+ "\n\n"
+ DEFAULT_TEMPLATE,
)
self.memory = ConversationBufferWindowMemory(k=2)
prompt_template = llm.get_chat_prompt_template()
self.memory = ConversationBufferWindowMemory(
return_messages=llm.is_chat_provider, k=2
)

self.llm = llm
self.llm_chain = ConversationChain(
Expand Down
43 changes: 2 additions & 41 deletions packages/jupyter-ai/jupyter_ai/completions/handlers/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,32 +18,6 @@
)
from .base import BaseInlineCompletionHandler

SYSTEM_PROMPT = """
You are an application built to provide helpful code completion suggestions.
You should only produce code. Keep comments to minimum, use the
programming language comment syntax. Produce clean code.
The code is written in JupyterLab, a data analysis and code development
environment which can execute code extended with additional syntax for
interactive features, such as magics.
""".strip()

AFTER_TEMPLATE = """
The code after the completion request is:

```
{suffix}
```
""".strip()

DEFAULT_TEMPLATE = """
The document is called `{filename}` and written in {language}.
{after}

Complete the following code:

```
{prefix}"""


class DefaultInlineCompletionHandler(BaseInlineCompletionHandler):
llm_chain: Runnable
Expand All @@ -57,18 +31,7 @@ def create_llm_chain(
model_parameters = self.get_model_parameters(provider, provider_params)
llm = provider(**provider_params, **model_parameters)

if llm.is_chat_provider:
prompt_template = ChatPromptTemplate.from_messages(
[
SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT),
HumanMessagePromptTemplate.from_template(DEFAULT_TEMPLATE),
]
)
else:
prompt_template = PromptTemplate(
input_variables=["prefix", "suffix", "language", "filename"],
template=SYSTEM_PROMPT + "\n\n" + DEFAULT_TEMPLATE,
)
prompt_template = llm.get_completion_prompt_template()

self.llm = llm
self.llm_chain = prompt_template | llm | StrOutputParser()
Expand Down Expand Up @@ -151,13 +114,11 @@ def _token_from_request(self, request: InlineCompletionRequest, suggestion: int)

def _template_inputs_from_request(self, request: InlineCompletionRequest) -> Dict:
suffix = request.suffix.strip()
# only add the suffix template if the suffix is there to save input tokens/computation time
after = AFTER_TEMPLATE.format(suffix=suffix) if suffix else ""
filename = request.path.split("/")[-1] if request.path else "untitled"

return {
"prefix": request.prefix,
"after": after,
"suffix": suffix,
"language": request.language,
"filename": filename,
"stop": ["\n```"],
Expand Down
Loading