[AUR-390] Add prompt template and prompt component (#24)
* Export pipeline to config * Export the input to config * Preliminary creating UI dynamically * Add test for config export * Try out prompt UI * Add example projects * Fix test errors * Standardize interface for retrieval * Finalize the UI demo * Update README.md * Update README * Refactor according to main * Fix typing issue * Add openai key to git-secret * Add prompt template and prompt component * Update test * update tests * revert docstring --------- Co-authored-by: trducng <trungduc1992@gmail.com> Co-authored-by: Nguyen Trung Duc (john) <john@cinnamon.is>
This commit is contained in:
67
tests/test_prompt.py
Normal file
67
tests/test_prompt.py
Normal file
@@ -0,0 +1,67 @@
|
||||
import pytest
|
||||
|
||||
from kotaemon.documents.base import Document
|
||||
from kotaemon.post_processing.extractor import RegexExtractor
|
||||
from kotaemon.prompt.base import BasePrompt
|
||||
from kotaemon.prompt.template import PromptTemplate
|
||||
|
||||
|
||||
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 = BasePrompt(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 = BasePrompt(template, name="Alice")
|
||||
with pytest.raises(ValueError):
|
||||
prompt._BasePrompt__check_redundant_kwargs(name="Alice", age=30)
|
||||
|
||||
|
||||
def test_check_unset_placeholders():
|
||||
template = PromptTemplate("Hello, {name}! I'm {age} years old.")
|
||||
prompt = BasePrompt(template, name="Alice")
|
||||
with pytest.raises(ValueError):
|
||||
prompt._BasePrompt__check_unset_placeholders()
|
||||
|
||||
|
||||
def test_validate_value_type():
|
||||
template = PromptTemplate("Hello, {name}!")
|
||||
prompt = BasePrompt(template)
|
||||
with pytest.raises(ValueError):
|
||||
prompt._BasePrompt__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 = BasePrompt(template=template, s="Alice", i=30, doc=doc, comp=comp)
|
||||
|
||||
result = prompt()
|
||||
|
||||
assert (
|
||||
result
|
||||
== "str = Alice, int = 30, doc = Helloo, Alice!, comp = ['One', 'Two', 'Three']"
|
||||
)
|
||||
|
||||
|
||||
def test_set_method():
|
||||
template = PromptTemplate("Hello, {name}!")
|
||||
prompt = BasePrompt(template)
|
||||
prompt.set(name="Alice")
|
||||
assert prompt.name == "Alice"
|
53
tests/test_template.py
Normal file
53
tests/test_template.py
Normal file
@@ -0,0 +1,53 @@
|
||||
import pytest
|
||||
|
||||
from kotaemon.prompt.template import PromptTemplate
|
||||
|
||||
|
||||
def test_prompt_template_creation():
|
||||
# Test case 1: Ensure the PromptTemplate object is created correctly
|
||||
template_string = "This is a template"
|
||||
template = PromptTemplate(template_string)
|
||||
assert template.template == template_string
|
||||
|
||||
template_string = "Hello, {name}! Today is {day}."
|
||||
template = PromptTemplate(template_string)
|
||||
assert template.template == template_string
|
||||
assert template.placeholders == {"name", "day"}
|
||||
|
||||
|
||||
def test_prompt_template_addition():
|
||||
# Test case 2: Ensure the __add__ method concatenates the templates correctly
|
||||
template1 = PromptTemplate("Hello, ")
|
||||
template2 = PromptTemplate("world!")
|
||||
result = template1 + template2
|
||||
assert result.template == "Hello, \nworld!"
|
||||
|
||||
template1 = PromptTemplate("Hello, {name}!")
|
||||
template2 = PromptTemplate("Today is {day}.")
|
||||
result = template1 + template2
|
||||
assert result.template == "Hello, {name}!\nToday is {day}."
|
||||
|
||||
|
||||
def test_prompt_template_extract_placeholders():
|
||||
# Test case 3: Ensure the extract_placeholders method extracts placeholders
|
||||
# correctly
|
||||
template_string = "Hello, {name}! Today is {day}."
|
||||
result = PromptTemplate.extract_placeholders(template_string)
|
||||
assert result == {"name", "day"}
|
||||
|
||||
|
||||
def test_prompt_template_populate():
|
||||
# Test case 4: Ensure the populate method populates the template correctly
|
||||
template_string = "Hello, {name}! Today is {day}."
|
||||
template = PromptTemplate(template_string)
|
||||
result = template.populate(name="John", day="Monday")
|
||||
assert result == "Hello, John! Today is Monday."
|
||||
|
||||
|
||||
def test_prompt_template_unknown_placeholder():
|
||||
# Test case 5: Ensure the populate method raises an exception for unknown
|
||||
# placeholders
|
||||
template_string = "Hello, {name}! Today is {day}."
|
||||
template = PromptTemplate(template_string)
|
||||
with pytest.raises(ValueError):
|
||||
template.populate(name="John", month="January")
|
@@ -56,7 +56,7 @@ class TestChromaVectorStore:
|
||||
db.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
|
||||
|
||||
_, sim, out_ids = db.query(embedding=[0.1, 0.2, 0.3], top_k=1)
|
||||
assert sim == [0.0]
|
||||
assert sim == [1.0]
|
||||
assert out_ids == ["a"]
|
||||
|
||||
_, _, out_ids = db.query(embedding=[0.42, 0.52, 0.53], top_k=1)
|
||||
|
Reference in New Issue
Block a user