langchain tests

api/tests/stampy_chat/test_callbacks.py

@@ -0,0 +1,32 @@
+from stampy_chat.callbacks import stream_callback
+
+
+def test_stream_callback_generates():
+    def caller_backer(callback):
+        for i in range(5):
+            callback(f'value no {i}')
+
+    assert list(stream_callback(caller_backer)) == [
+        f'value no {i}' for i in range(5)
+    ]
+
+
+def test_stream_callback_formatter():
+    def caller_backer(callback):
+        for i in range(5):
+            callback(f'value no {i}')
+
+    assert list(stream_callback(caller_backer, lambda val: 'formatted ' + val)) == [
+        f'formatted value no {i}' for i in range(5)
+    ]
+
+
+def test_stream_callback_on_error():
+    def caller_backer(callback):
+        for i in range(5):
+            callback(f'value no {i}')
+        raise ValueError('this is a pen')
+
+    assert list(stream_callback(caller_backer)) == [
+        f'value no {i}' for i in range(5)
+    ] + ['this is a pen']

api/tests/stampy_chat/test_chat.py

@@ -0,0 +1,142 @@
+from langchain.llms.fake import FakeListLLM
+from langchain.memory import ChatMessageHistory
+from langchain.prompts import ChatPromptTemplate
+from langchain.schema import ChatMessage, HumanMessage, SystemMessage
+
+from stampy_chat.callbacks import StampyCallbackHandler
+from stampy_chat.chat import (
+    LimitedConversationSummaryBufferMemory,
+    MessageBufferPromptTemplate,
+    PrefixedPrompt
+)
+
+
+def make_prompt_template(max_tokens, examples):
+    template = ChatPromptTemplate.from_template('{content}')
+    return MessageBufferPromptTemplate(
+        example_prompt=template,
+        get_num_tokens=lambda s: len(s),
+        max_tokens=max_tokens,
+        examples=examples,
+    )
+
+
+def test_MessageBufferPromptTemplate_format_messages():
+    template = make_prompt_template(100, [
+        {'content': 'bla bla bla'},
+        {'content': 'ble ble ble'},
+        {'content': 'and some more'},
+    ])
+    assert template.format_messages() == [
+        HumanMessage(content='bla bla bla'),
+        HumanMessage(content='ble ble ble'),
+        HumanMessage(content='and some more'),
+    ]
+
+
+def test_MessageBufferPromptTemplate_format_messages_truncated():
+    template = make_prompt_template(20, [
+        {'content': 'bla bla bla'},
+        {'content': 'ble ble ble'},
+        {'content': 'and some more'},
+    ])
+    assert template.format_messages() == [
+        HumanMessage(content='bla bla bla'),
+    ]
+
+
+def test_MessageBufferPromptTemplate_format_all_messages_truncated():
+    template = make_prompt_template(10, [
+        {'content': 'bla bla bla'},
+        {'content': 'ble ble ble'},
+        {'content': 'and some more'},
+    ])
+    assert template.format_messages() == []
+
+
+def test_PrefixedPrompt_format_messages():
+    prompt = PrefixedPrompt(messages_field='history', prompt='bla bla bla', input_variables=[])
+    history = [HumanMessage(content=f'human message {i}') for i in range(5)]
+    assert prompt.format_messages(history=history) == [
+        SystemMessage(content='bla bla bla'),
+        HumanMessage(content='human message 0'),
+        HumanMessage(content='human message 1'),
+        HumanMessage(content='human message 2'),
+        HumanMessage(content='human message 3'),
+        HumanMessage(content='human message 4'),
+    ]
+
+
+def test_PrefixedPrompt_format_messages_no_history():
+    prompt = PrefixedPrompt(messages_field='history', prompt='bla bla bla', input_variables=[])
+    assert prompt.format_messages(history=[]) == []
+
+
+def test_LimitedConversationSummaryBufferMemory_set_empty():
+    llm = FakeListLLM(responses=['this is a summary of what was before'])
+    memory = LimitedConversationSummaryBufferMemory(llm=llm)
+    memory.chat_memory = [{'content': 'bla bla bla', 'role': 'human'}]
+
+    memory.set_messages([])
+    assert memory.chat_memory == ChatMessageHistory(messages=[])
+
+
+def test_LimitedConversationSummaryBufferMemory_set():
+    llm = FakeListLLM(responses=['this is a summary of what was before'])
+    memory = LimitedConversationSummaryBufferMemory(llm=llm)
+
+    memory.set_messages([
+        {'content': 'a system message', 'role': 'system'},
+        {'content': 'bla bla bla', 'role': 'human'},
+    ])
+    assert memory.chat_memory == ChatMessageHistory(messages=[
+        ChatMessage(content='a system message', role='system'),
+        ChatMessage(content='bla bla bla', role='human'),
+    ])
+
+
+def test_LimitedConversationSummaryBufferMemory_set_more():
+    llm = FakeListLLM(responses=['this is a summary of what was before'])
+    memory = LimitedConversationSummaryBufferMemory(llm=llm, max_history=4)
+
+    memory.set_messages([
+        {'content': 'a system message', 'role': 'system'},
+        {'content': 'message 1 - should be summarized', 'role': 'human'},
+        {'content': 'message 2 - should be summarized', 'role': 'human'},
+        {'content': 'message 3 - should be kept', 'role': 'human'},
+        {'content': 'message 4 - should be kept', 'role': 'human'},
+        {'content': 'message 5 - should be kept', 'role': 'human'},
+    ])
+    assert memory.chat_memory == ChatMessageHistory(messages=[
+        ChatMessage(content='this is a summary of what was before', role='assistant'),
+        ChatMessage(content='message 3 - should be kept', role='human'),
+        ChatMessage(content='message 4 - should be kept', role='human'),
+        ChatMessage(content='message 5 - should be kept', role='human'),
+    ])
+
+
+def test_LimitedConversationSummaryBufferMemory_set_with_callbacks():
+    callback_calls = {}
+    class DummyCallback(StampyCallbackHandler):
+        def on_memory_set_start(self, history):
+            callback_calls['start'] = history
+
+        def on_memory_set_end(self, messages):
+            callback_calls['end'] = messages
+
+    llm = FakeListLLM(responses=['this is a summary of what was before'])
+    memory = LimitedConversationSummaryBufferMemory(llm=llm, callbacks=[DummyCallback()])
+    history = [
+        {'content': 'a system message', 'role': 'system'},
+        {'content': 'bla bla bla', 'role': 'human'},
+    ]
+
+    memory.set_messages(history)
+    assert memory.chat_memory == ChatMessageHistory(messages=[
+        ChatMessage(content='a system message', role='system'),
+        ChatMessage(content='bla bla bla', role='human'),
+    ])
+    assert callback_calls == {
+        'start': history,
+        'end': memory.chat_memory,
+    }

api/tests/stampy_chat/test_citations.py

@@ -0,0 +1,137 @@
+import pytest
+from datetime import datetime
+from unittest.mock import patch, Mock, call
+
+from langchain.schema.vectorstore import VectorStore
+
+from stampy_chat.citations import ReferencesSelector, format_block
+
+
+class DummyVectorStore(VectorStore):
+    """LangChain is very restrictive with validated fields, so this class mocks a VectorStore."""
+    def __init__(self, *args, similarity_search=None, similarity_search_return=None, **kwargs):
+        self.similarity_search_func = similarity_search
+        self.similarity_search_return_value = similarity_search_return
+
+    def add_texts(self, *args, **kwargs):
+        pass
+
+    def from_texts(self, *args, **kwargs):
+        pass
+
+    def similarity_search(self, *args, **kwargs):
+        if self.similarity_search_return_value:
+            return self.similarity_search_return_value
+        elif self.similarity_search_func:
+            return self.similarity_search_func(*args, **kwargs)
+        return []
+
+
+@pytest.fixture
+def selector():
+    examples = [
+        Mock(page_content=f'{i}', metadata={
+            'bla': f'bla {i}'
+        }) for i in range(5)
+    ]
+    return ReferencesSelector(vectorstore=DummyVectorStore(similarity_search_return=examples))
+
+
+@pytest.mark.parametrize('num, letter', (
+    (0, 'a'), (25, 'z'),
+    (26, '{'),  # this is a basic ASCII translator, so too many citations will result in fun
+))
+def test_ReferencesSelector_make_references(num, letter):
+    assert ReferencesSelector.make_reference(num) == letter
+
+
+def test_ReferencesSelector_select_examples(selector):
+    assert selector.select_examples(input_variables={}) == [
+        {'bla': 'bla 0', 'id': '0', 'reference': 'a'},
+        {'bla': 'bla 1', 'id': '1', 'reference': 'b'},
+        {'bla': 'bla 2', 'id': '2', 'reference': 'c'},
+        {'bla': 'bla 3', 'id': '3', 'reference': 'd'},
+        {'bla': 'bla 4', 'id': '4', 'reference': 'e'},
+    ]
+
+
+def test_ReferencesSelector_select_examples_callbacks(selector):
+    callback = Mock()
+    selector.callbacks = [callback]
+
+    expected_examples = [
+        {'bla': 'bla 0', 'id': '0', 'reference': 'a'},
+        {'bla': 'bla 1', 'id': '1', 'reference': 'b'},
+        {'bla': 'bla 2', 'id': '2', 'reference': 'c'},
+        {'bla': 'bla 3', 'id': '3', 'reference': 'd'},
+        {'bla': 'bla 4', 'id': '4', 'reference': 'e'},
+    ]
+    input_variables = {'var1': 'bla', 'var2': 'ble'}
+
+    assert selector.select_examples(input_variables=input_variables) == expected_examples
+    callback.on_context_fetch_start.assert_called_once_with(input_variables)
+    callback.on_context_fetch_end.assert_called_once_with(expected_examples)
+
+
+def test_ReferencesSelector_select_examples_removes_duplicates(selector):
+    selector.vectorstore.similarity_search_return_value = [
+        Mock(page_content=f'{i}', metadata={
+            'bla': f'bla {i}'
+        }) for i in range(5)
+    ] * 5
+
+    assert selector.select_examples(input_variables={}) == [
+        {'bla': 'bla 0', 'id': '0', 'reference': 'a'},
+        {'bla': 'bla 1', 'id': '1', 'reference': 'b'},
+        {'bla': 'bla 2', 'id': '2', 'reference': 'c'},
+        {'bla': 'bla 3', 'id': '3', 'reference': 'd'},
+        {'bla': 'bla 4', 'id': '4', 'reference': 'e'},
+    ]
+
+
+@pytest.mark.parametrize("overrides, expected", [
+    # Basic fields
+    ({}, {}),
+    ({'title': 'bla bla'}, {'title': 'bla bla'}),
+    ({'text': 'bla bla'}, {'text': 'bla bla'}),
+    ({'url': 'different.bla.bla'}, {'url': 'different.bla.bla'}),
+    ({'tags': 'a tag, and another one'}, {'tags': 'a tag, and another one'}),
+
+    # Id
+    ({'id': 'some id'}, {'id': 'some id'}),
+    ({'hash_id': 'some hash id'}, {'id': 'some hash id'}),
+    ({'id': 'some id', 'hash_id': 'some hash id'}, {'id': 'some hash id'}),
+
+    # Authors
+    ({'authors': 'mr blobby'}, {'authors': 'mr blobby'}),
+    ({'authors': ['mr blobby', 'john snow']}, {'authors': ['mr blobby', 'john snow']}),
+    ({'author': 'your momma'}, {'authors': ['your momma']}),
+    ({'author': 'your momma', 'authors': ['mr blobby', 'john snow']}, {'authors': ['mr blobby', 'john snow']}),
+
+    # Date field
+    ({'date_published': '2020-01-02'}, {'date': '2020-01-02'}),
+    ({'date': '2020-01-02'}, {'date': '2020-01-02'}),
+    ({'date': '1930-01-02', 'date_published': '2020-01-02'}, {'date': '2020-01-02'}),
+
+    # Date format
+    ({'date_published': datetime.fromisoformat('2020-01-02T01:02:03')}, {'date': '2020-01-02'}),
+    ({'date_published': datetime.fromisoformat('2020-01-02T01:02:03').date()}, {'date': '2020-01-02'}),
+    ({'date_published': datetime.fromisoformat('2020-01-02T01:02:03').timestamp()}, {'date': '2020-01-02'}),
+    ({'date_published': int(datetime.fromisoformat('2020-01-02T01:02:03').timestamp())}, {'date': '2020-01-02'}),
+])
+def test_format_block(overrides, expected):
+    defaults = {
+        'title': 'das kapititle',
+        'url': 'http://bla.bla',
+        'text': 'some text'
+    }
+    default_exceptions = {
+        'id': None,
+        'title': 'das kapititle',
+        'url': 'http://bla.bla',
+        'text': 'some text',
+        'authors': None,
+        'date': None,
+        'tags': None,
+    }
+    assert format_block(dict(defaults, **overrides)) == dict(default_exceptions, **expected)