|
|
@ -308,7 +308,7 @@ def test_prompt_jinja2_extra_input_variables(
|
|
|
|
).input_variables == ["bar", "foo"]
|
|
|
|
).input_variables == ["bar", "foo"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_few_shot_chat_message_prompt_template() -> None:
|
|
|
|
async def test_few_shot_chat_message_prompt_template() -> None:
|
|
|
|
"""Tests for few shot chat message template."""
|
|
|
|
"""Tests for few shot chat message template."""
|
|
|
|
examples = [
|
|
|
|
examples = [
|
|
|
|
{"input": "2+2", "output": "4"},
|
|
|
|
{"input": "2+2", "output": "4"},
|
|
|
@ -333,8 +333,7 @@ def test_few_shot_chat_message_prompt_template() -> None:
|
|
|
|
+ HumanMessagePromptTemplate.from_template("{input}")
|
|
|
|
+ HumanMessagePromptTemplate.from_template("{input}")
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
messages = final_prompt.format_messages(input="100 + 1")
|
|
|
|
expected = [
|
|
|
|
assert messages == [
|
|
|
|
|
|
|
|
SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
|
|
|
|
SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
|
|
|
|
HumanMessage(content="2+2", additional_kwargs={}, example=False),
|
|
|
|
HumanMessage(content="2+2", additional_kwargs={}, example=False),
|
|
|
|
AIMessage(content="4", additional_kwargs={}, example=False),
|
|
|
|
AIMessage(content="4", additional_kwargs={}, example=False),
|
|
|
@ -343,6 +342,11 @@ def test_few_shot_chat_message_prompt_template() -> None:
|
|
|
|
HumanMessage(content="100 + 1", additional_kwargs={}, example=False),
|
|
|
|
HumanMessage(content="100 + 1", additional_kwargs={}, example=False),
|
|
|
|
]
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
messages = final_prompt.format_messages(input="100 + 1")
|
|
|
|
|
|
|
|
assert messages == expected
|
|
|
|
|
|
|
|
messages = await final_prompt.aformat_messages(input="100 + 1")
|
|
|
|
|
|
|
|
assert messages == expected
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class AsIsSelector(BaseExampleSelector):
|
|
|
|
class AsIsSelector(BaseExampleSelector):
|
|
|
|
"""An example selector for testing purposes.
|
|
|
|
"""An example selector for testing purposes.
|
|
|
@ -355,11 +359,9 @@ class AsIsSelector(BaseExampleSelector):
|
|
|
|
self.examples = examples
|
|
|
|
self.examples = examples
|
|
|
|
|
|
|
|
|
|
|
|
def add_example(self, example: Dict[str, str]) -> Any:
|
|
|
|
def add_example(self, example: Dict[str, str]) -> Any:
|
|
|
|
"""Adds an example to the selector."""
|
|
|
|
raise NotImplementedError
|
|
|
|
raise NotImplementedError()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
|
|
|
|
def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
|
|
|
|
"""Select which examples to use based on the inputs."""
|
|
|
|
|
|
|
|
return list(self.examples)
|
|
|
|
return list(self.examples)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -387,8 +389,63 @@ def test_few_shot_chat_message_prompt_template_with_selector() -> None:
|
|
|
|
+ few_shot_prompt
|
|
|
|
+ few_shot_prompt
|
|
|
|
+ HumanMessagePromptTemplate.from_template("{input}")
|
|
|
|
+ HumanMessagePromptTemplate.from_template("{input}")
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
expected = [
|
|
|
|
|
|
|
|
SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
|
|
|
|
|
|
|
|
HumanMessage(content="2+2", additional_kwargs={}, example=False),
|
|
|
|
|
|
|
|
AIMessage(content="4", additional_kwargs={}, example=False),
|
|
|
|
|
|
|
|
HumanMessage(content="2+3", additional_kwargs={}, example=False),
|
|
|
|
|
|
|
|
AIMessage(content="5", additional_kwargs={}, example=False),
|
|
|
|
|
|
|
|
HumanMessage(content="100 + 1", additional_kwargs={}, example=False),
|
|
|
|
|
|
|
|
]
|
|
|
|
messages = final_prompt.format_messages(input="100 + 1")
|
|
|
|
messages = final_prompt.format_messages(input="100 + 1")
|
|
|
|
assert messages == [
|
|
|
|
assert messages == expected
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class AsyncAsIsSelector(BaseExampleSelector):
|
|
|
|
|
|
|
|
"""An example selector for testing purposes.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
This selector returns the examples as-is.
|
|
|
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def __init__(self, examples: Sequence[Dict[str, str]]) -> None:
|
|
|
|
|
|
|
|
"""Initializes the selector."""
|
|
|
|
|
|
|
|
self.examples = examples
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def add_example(self, example: Dict[str, str]) -> Any:
|
|
|
|
|
|
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
|
|
|
|
|
|
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def aselect_examples(self, input_variables: Dict[str, str]) -> List[dict]:
|
|
|
|
|
|
|
|
return list(self.examples)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
async def test_few_shot_chat_message_prompt_template_with_selector_async() -> None:
|
|
|
|
|
|
|
|
"""Tests for few shot chat message template with an async example selector."""
|
|
|
|
|
|
|
|
examples = [
|
|
|
|
|
|
|
|
{"input": "2+2", "output": "4"},
|
|
|
|
|
|
|
|
{"input": "2+3", "output": "5"},
|
|
|
|
|
|
|
|
]
|
|
|
|
|
|
|
|
example_selector = AsyncAsIsSelector(examples)
|
|
|
|
|
|
|
|
example_prompt = ChatPromptTemplate.from_messages(
|
|
|
|
|
|
|
|
[
|
|
|
|
|
|
|
|
HumanMessagePromptTemplate.from_template("{input}"),
|
|
|
|
|
|
|
|
AIMessagePromptTemplate.from_template("{output}"),
|
|
|
|
|
|
|
|
]
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
few_shot_prompt = FewShotChatMessagePromptTemplate(
|
|
|
|
|
|
|
|
input_variables=["input"],
|
|
|
|
|
|
|
|
example_prompt=example_prompt,
|
|
|
|
|
|
|
|
example_selector=example_selector,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
final_prompt: ChatPromptTemplate = (
|
|
|
|
|
|
|
|
SystemMessagePromptTemplate.from_template("You are a helpful AI Assistant")
|
|
|
|
|
|
|
|
+ few_shot_prompt
|
|
|
|
|
|
|
|
+ HumanMessagePromptTemplate.from_template("{input}")
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
expected = [
|
|
|
|
SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
|
|
|
|
SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
|
|
|
|
HumanMessage(content="2+2", additional_kwargs={}, example=False),
|
|
|
|
HumanMessage(content="2+2", additional_kwargs={}, example=False),
|
|
|
|
AIMessage(content="4", additional_kwargs={}, example=False),
|
|
|
|
AIMessage(content="4", additional_kwargs={}, example=False),
|
|
|
@ -396,3 +453,5 @@ def test_few_shot_chat_message_prompt_template_with_selector() -> None:
|
|
|
|
AIMessage(content="5", additional_kwargs={}, example=False),
|
|
|
|
AIMessage(content="5", additional_kwargs={}, example=False),
|
|
|
|
HumanMessage(content="100 + 1", additional_kwargs={}, example=False),
|
|
|
|
HumanMessage(content="100 + 1", additional_kwargs={}, example=False),
|
|
|
|
]
|
|
|
|
]
|
|
|
|
|
|
|
|
messages = await final_prompt.aformat_messages(input="100 + 1")
|
|
|
|
|
|
|
|
assert messages == expected
|
|
|
|