693ed39de4
Since the only usage of prompt is within LLMs, it is reasonable to keep it within the LLM module. This way, it would be easier to discover module, and make the code base less complicated. Changes: * Move prompt components into llms * Bump version 0.3.1 * Make pip install dependencies in eager mode --------- Co-authored-by: ian <ian@cinnamon.is>
64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
import pytest
|
|
|
|
from kotaemon.base import Document
|
|
from kotaemon.llms import BasePromptComponent, PromptTemplate
|
|
from kotaemon.post_processing.extractor import RegexExtractor
|
|
|
|
|
|
def test_set_attributes():
|
|
template = PromptTemplate("str = {s}, int = {i}, doc = {doc}, comp = {comp}")
|
|
doc = Document(text="Helloo, Alice!")
|
|
comp = RegexExtractor(
|
|
pattern=r"\d+", output_map={"1": "One", "2": "Two", "3": "Three"}
|
|
)
|
|
comp.set_run(kwargs={"text": "This is a test. 1 2 3"}, temp=True)
|
|
|
|
prompt = BasePromptComponent(template=template, s="Alice", i=30, doc=doc, comp=comp)
|
|
assert prompt.s == "Alice"
|
|
assert prompt.i == 30
|
|
assert prompt.doc == doc
|
|
assert prompt.comp == comp
|
|
|
|
|
|
def test_check_redundant_kwargs():
|
|
template = PromptTemplate("Hello, {name}!")
|
|
prompt = BasePromptComponent(template, name="Alice")
|
|
with pytest.warns(UserWarning, match="Keys provided but not in template: age"):
|
|
prompt._BasePromptComponent__check_redundant_kwargs(name="Alice", age=30)
|
|
|
|
|
|
def test_check_unset_placeholders():
|
|
template = PromptTemplate("Hello, {name}! I'm {age} years old.")
|
|
prompt = BasePromptComponent(template, name="Alice")
|
|
with pytest.raises(ValueError):
|
|
prompt._BasePromptComponent__check_unset_placeholders()
|
|
|
|
|
|
def test_validate_value_type():
|
|
template = PromptTemplate("Hello, {name}!")
|
|
prompt = BasePromptComponent(template)
|
|
with pytest.raises(ValueError):
|
|
prompt._BasePromptComponent__validate_value_type(name={})
|
|
|
|
|
|
def test_run():
|
|
template = PromptTemplate("str = {s}, int = {i}, doc = {doc}, comp = {comp}")
|
|
doc = Document(text="Helloo, Alice!")
|
|
comp = RegexExtractor(
|
|
pattern=r"\d+", output_map={"1": "One", "2": "Two", "3": "Three"}
|
|
)
|
|
comp.set_run(kwargs={"text": "This is a test. 1 2 3"}, temp=True)
|
|
|
|
prompt = BasePromptComponent(template=template, s="Alice", i=30, doc=doc, comp=comp)
|
|
|
|
result = prompt()
|
|
|
|
assert result.text == "str = Alice, int = 30, doc = Helloo, Alice!, comp = ['One']"
|
|
|
|
|
|
def test_set_method():
|
|
template = PromptTemplate("Hello, {name}!")
|
|
prompt = BasePromptComponent(template)
|
|
prompt.set(name="Alice")
|
|
assert prompt.name == "Alice"
|