Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplify the BaseComponent inteface (#64)

Browse Source 
This change remove `BaseComponent`'s:

- run_raw
- run_batch_raw
- run_document
- run_batch_document
- is_document
- is_batch

Each component is expected to support multiple types of inputs and a single type of output. Since we want the component to work out-of-the-box with both standardized and customized use cases, supporting multiple types of inputs are expected. At the same time, to reduce the complexity of understanding how to use a component, we restrict a component to only have a single output type.

To accommodate these changes, we also refactor some components to remove their run_raw, run_batch_raw... methods, and to decide the common output interface for those components.

Tests are updated accordingly.

Commit changes:

* Add kwargs to vector store's query
* Simplify the BaseComponent
* Update tests
* Remove support for Python 3.8 and 3.9
* Bump version 0.3.0
* Fix github PR caching still use old environment after bumping version

---------

Co-authored-by: ian <ian@cinnamon.is>
This commit is contained in:
Nguyen Trung Duc (john)
2023-11-13 15:10:18 +07:00
committed by GitHub
parent 6095526dc7
commit d79b3744cb
25 changed files with 280 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

71
tests/test_composite.py
View File

@@ -1,4 +1,7 @@
from copy import deepcopy
import pytest
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.composite import (
GatedBranchingPipeline,
@@ -10,6 +13,29 @@ from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.post_processing.extractor import RegexExtractor
from kotaemon.prompt.base import BasePromptComponent
_openai_chat_completion_response = ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "This is a test 123",
"finish_reason": "length",
"logprobs": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
@pytest.fixture
def mock_llm():
@@ -19,7 +45,6 @@ def mock_llm():
openai_api_version="OPENAI_API_VERSION",
deployment_name="dummy-q2-gpt35",
temperature=0,
request_timeout=600,
)
@@ -61,11 +86,12 @@ def mock_gated_linear_pipeline_negative(mock_prompt, mock_llm, mock_post_process
def test_simple_linear_pipeline_run(mocker, mock_simple_linear_pipeline):
openai_mocker = mocker.patch.object(
AzureChatOpenAI, "run", return_value="This is a test 123"
openai_mocker = mocker.patch(
"openai.resources.chat.completions.Completions.create",
return_value=_openai_chat_completion_response,
)
result = mock_simple_linear_pipeline.run(value="abc")
result = mock_simple_linear_pipeline(value="abc")
assert result.text == "123"
assert openai_mocker.call_count == 1
@@ -74,11 +100,12 @@ def test_simple_linear_pipeline_run(mocker, mock_simple_linear_pipeline):
def test_gated_linear_pipeline_run_positive(
mocker, mock_gated_linear_pipeline_positive
):
openai_mocker = mocker.patch.object(
AzureChatOpenAI, "run", return_value="This is a test 123."
openai_mocker = mocker.patch(
"openai.resources.chat.completions.Completions.create",
return_value=_openai_chat_completion_response,
)
result = mock_gated_linear_pipeline_positive.run(
result = mock_gated_linear_pipeline_positive(
value="abc", condition_text="positive condition"
)
@@ -89,11 +116,12 @@ def test_gated_linear_pipeline_run_positive(
def test_gated_linear_pipeline_run_negative(
mocker, mock_gated_linear_pipeline_positive
):
openai_mocker = mocker.patch.object(
AzureChatOpenAI, "run", return_value="This is a test 123."
openai_mocker = mocker.patch(
"openai.resources.chat.completions.Completions.create",
return_value=_openai_chat_completion_response,
)
result = mock_gated_linear_pipeline_positive.run(
result = mock_gated_linear_pipeline_positive(
value="abc", condition_text="negative condition"
)
@@ -102,14 +130,14 @@ def test_gated_linear_pipeline_run_negative(
def test_simple_branching_pipeline_run(mocker, mock_simple_linear_pipeline):
openai_mocker = mocker.patch.object(
AzureChatOpenAI,
"run",
side_effect=[
"This is a test 123.",
"a quick brown fox",
"jumps over the lazy dog 456",
],
response0: ChatCompletion = _openai_chat_completion_response
response1: ChatCompletion = deepcopy(_openai_chat_completion_response)
response1.choices[0].message.content = "a quick brown fox"
response2: ChatCompletion = deepcopy(_openai_chat_completion_response)
response2.choices[0].message.content = "jumps over the lazy dog 456"
openai_mocker = mocker.patch(
"openai.resources.chat.completions.Completions.create",
side_effect=[response0, response1, response2],
)
pipeline = SimpleBranchingPipeline()
for _ in range(3):
@@ -126,8 +154,11 @@ def test_simple_branching_pipeline_run(mocker, mock_simple_linear_pipeline):
def test_simple_gated_branching_pipeline_run(
mocker, mock_gated_linear_pipeline_positive, mock_gated_linear_pipeline_negative
):
openai_mocker = mocker.patch.object(
AzureChatOpenAI, "run", return_value="a quick brown fox"
response0: ChatCompletion = deepcopy(_openai_chat_completion_response)
response0.choices[0].message.content = "a quick brown fox"
openai_mocker = mocker.patch(
"openai.resources.chat.completions.Completions.create",
return_value=response0,
)
pipeline = GatedBranchingPipeline()

19
tests/test_embedding_models.py
View File

@@ -26,7 +26,8 @@ def test_azureopenai_embeddings_raw(openai_embedding_call):
)
output = model("Hello world")
assert isinstance(output, list)
assert isinstance(output[0], float)
assert isinstance(output[0], list)
assert isinstance(output[0][0], float)
openai_embedding_call.assert_called()
@@ -53,8 +54,8 @@ def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
side_effect=lambda *args, **kwargs: None,
)
@patch(
"langchain.embeddings.huggingface.HuggingFaceBgeEmbeddings.embed_query",
side_effect=lambda *args, **kwargs: [1.0, 2.1, 3.2],
"langchain.embeddings.huggingface.HuggingFaceBgeEmbeddings.embed_documents",
side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
)
def test_huggingface_embddings(
langchain_huggingface_embedding_call, sentence_transformers_init
@@ -67,21 +68,23 @@ def test_huggingface_embddings(
output = model("Hello World")
assert isinstance(output, list)
assert isinstance(output[0], float)
assert isinstance(output[0], list)
assert isinstance(output[0][0], float)
sentence_transformers_init.assert_called()
langchain_huggingface_embedding_call.assert_called()
@patch(
"langchain.embeddings.cohere.CohereEmbeddings.embed_query",
side_effect=lambda *args, **kwargs: [1.0, 2.1, 3.2],
"langchain.embeddings.cohere.CohereEmbeddings.embed_documents",
side_effect=lambda *args, **kwargs: [[1.0, 2.1, 3.2]],
)
def test_cohere_embddings(langchain_cohere_embedding_call):
def test_cohere_embeddings(langchain_cohere_embedding_call):
model = CohereEmbdeddings(
model="embed-english-light-v2.0", cohere_api_key="my-api-key"
)
output = model("Hello World")
assert isinstance(output, list)
assert isinstance(output[0], float)
assert isinstance(output[0], list)
assert isinstance(output[0][0], float)
langchain_cohere_embedding_call.assert_called()

7
tests/test_indexing_retrieval.py
View File

@@ -60,7 +60,8 @@ def test_retrieving(mock_openai_embedding, tmp_path):
)
index_pipeline(text=Document(text="Hello world"))
output = retrieval_pipeline(text=["Hello world", "Hello world"])
output = retrieval_pipeline(text="Hello world")
output1 = retrieval_pipeline(text="Hello world")
assert len(output) == 2, "Expect 2 results"
assert output[0] == output[1], "Expect identical results"
assert len(output) == 1, "Expect 1 results"
assert output == output1, "Expect identical results"

12
tests/test_llms_chat_models.py
View File

@@ -54,12 +54,6 @@ def test_azureopenai_model(openai_completion):
), "Output for single text is not LLMInterface"
openai_completion.assert_called()
# test for list[str] input - batch mode
output = model(["hello world"])
assert isinstance(output, list), "Output for batch string is not a list"
assert isinstance(output[0], LLMInterface), "Output for text is not LLMInterface"
openai_completion.assert_called()
# test for list[message] input - stream mode
messages = [
SystemMessage(content="You are a philosohper"),
@@ -73,9 +67,3 @@ def test_azureopenai_model(openai_completion):
output, LLMInterface
), "Output for single text is not LLMInterface"
openai_completion.assert_called()
# test for list[list[message]] input - batch mode
output = model([messages])
assert isinstance(output, list), "Output for batch string is not a list"
assert isinstance(output[0], LLMInterface), "Output for text is not LLMInterface"
openai_completion.assert_called()

10
tests/test_llms_completion_models.py
View File

@@ -44,11 +44,6 @@ def test_azureopenai_model(openai_completion):
model.agent, AzureOpenAILC
), "Agent not wrapped in Langchain's AzureOpenAI"
output = model(["hello world"])
assert isinstance(output, list), "Output for batch is not a list"
assert isinstance(output[0], LLMInterface), "Output for text is not LLMInterface"
openai_completion.assert_called()
output = model("hello world")
assert isinstance(
output, LLMInterface
@@ -72,11 +67,6 @@ def test_openai_model(openai_completion):
model.agent, OpenAILC
), "Agent is not wrapped in Langchain's OpenAI"
output = model(["hello world"])
assert isinstance(output, list), "Output for batch is not a list"
assert isinstance(output[0], LLMInterface), "Output for text is not LLMInterface"
openai_completion.assert_called()
output = model("hello world")
assert isinstance(
output, LLMInterface

14
tests/test_post_processing.py
View File

@@ -13,23 +13,13 @@ def regex_extractor():
def test_run_document(regex_extractor):
document = Document(text="This is a test. 1 2 3")
extracted_document = regex_extractor(document)
extracted_document = regex_extractor(document)[0]
assert extracted_document.text == "One"
assert extracted_document.matches == ["One", "Two", "Three"]
def test_is_document(regex_extractor):
assert regex_extractor.is_document(Document(text="Test"))
assert not regex_extractor.is_document("Test")
def test_is_batch(regex_extractor):
assert regex_extractor.is_batch([Document(text="Test")])
assert not regex_extractor.is_batch(Document(text="Test"))
def test_run_raw(regex_extractor):
output = regex_extractor("This is a test. 123")
output = regex_extractor("This is a test. 123")[0]
assert output.text == "123"
assert output.matches == ["123"]

2
tests/test_prompt.py
View File

@@ -54,7 +54,7 @@ def test_run():
result = prompt()
assert result.text == "str = Alice, int = 30, doc = Helloo, Alice!, comp = One"
assert result.text == "str = Alice, int = 30, doc = Helloo, Alice!, comp = ['One']"
def test_set_method():