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

Add Huggingface embeddings and Cohere embeddings (#63)

Browse Source 
* Add huggingface embeddings and cohere embeddings
* Update openai interface and the mock for newer OpenAI SDK

---------

Co-authored-by: trducng <trungduc1992@gmail.com>
This commit is contained in:
ian_Cin
2023-11-10 09:38:30 +07:00
committed by GitHub
parent 9035e25666
commit 6095526dc7
13 changed files with 249 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

129
tests/test_agent.py
View File

@@ -1,6 +1,7 @@
from unittest.mock import patch
import pytest
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.agents.react import ReactAgent
@@ -14,24 +15,30 @@ from kotaemon.pipelines.tools import (
FINAL_RESPONSE_TEXT = "Hello Cinnamon AI!"
_openai_chat_completion_responses_rewoo = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
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": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"#Plan1: Search for Cinnamon AI company on Google\n"
@@ -44,23 +51,28 @@ _openai_chat_completion_responses_rewoo = [
]
_openai_chat_completion_responses_react = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
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": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"I don't have prior knowledge about Cinnamon AI company, "
@@ -82,23 +94,28 @@ _openai_chat_completion_responses_react = [
]
_openai_chat_completion_responses_react_langchain_tool = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
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": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"I don't have prior knowledge about Cinnamon AI company, "
@@ -123,7 +140,7 @@ _openai_chat_completion_responses_react_langchain_tool = [
@pytest.fixture
def llm():
return AzureChatOpenAI(
openai_api_base="https://dummy.openai.azure.com/",
azure_endpoint="https://dummy.openai.azure.com/",
openai_api_key="dummy",
openai_api_version="2023-03-15-preview",
deployment_name="dummy-q2",
@@ -132,7 +149,7 @@ def llm():
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_rewoo,
)
def test_rewoo_agent(openai_completion, llm, mock_google_search):
@@ -150,7 +167,7 @@ def test_rewoo_agent(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react,
)
def test_react_agent(openai_completion, llm, mock_google_search):
@@ -167,7 +184,7 @@ def test_react_agent(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react,
)
def test_react_agent_langchain(openai_completion, llm, mock_google_search):
@@ -191,7 +208,7 @@ def test_react_agent_langchain(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react_langchain_tool,
)
def test_react_agent_with_langchain_tools(openai_completion, llm):

47
tests/test_cot.py
View File

@@ -1,32 +1,39 @@
from unittest.mock import patch
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.cot import ManualSequentialChainOfThought, Thought
_openai_chat_completion_response = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
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": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in ["Bonjour", "こんにちは (Konnichiwa)"]
]
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_plus_operator(openai_completion):
@@ -58,7 +65,7 @@ def test_cot_plus_operator(openai_completion):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_manual(openai_completion):
@@ -88,7 +95,7 @@ def test_cot_manual(openai_completion):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_with_termination_callback(openai_completion):

45
tests/test_embedding_models.py
View File

@@ -2,6 +2,8 @@ import json
from pathlib import Path
from unittest.mock import patch
from kotaemon.embeddings.cohere import CohereEmbdeddings
from kotaemon.embeddings.huggingface import HuggingFaceEmbeddings
from kotaemon.embeddings.openai import AzureOpenAIEmbeddings
with open(Path(__file__).parent / "resources" / "embedding_openai_batch.json") as f:
@@ -12,7 +14,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@patch(
"openai.api_resources.embedding.Embedding.create",
"openai.resources.embeddings.Embeddings.create",
side_effect=lambda *args, **kwargs: openai_embedding,
)
def test_azureopenai_embeddings_raw(openai_embedding_call):
@@ -29,7 +31,7 @@ def test_azureopenai_embeddings_raw(openai_embedding_call):
@patch(
"openai.api_resources.embedding.Embedding.create",
"openai.resources.embeddings.Embeddings.create",
side_effect=lambda *args, **kwargs: openai_embedding_batch,
)
def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
@@ -44,3 +46,42 @@ def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
assert isinstance(output[0], list)
assert isinstance(output[0][0], float)
openai_embedding_call.assert_called()
@patch(
"sentence_transformers.SentenceTransformer",
side_effect=lambda *args, **kwargs: None,
)
@patch(
"langchain.embeddings.huggingface.HuggingFaceBgeEmbeddings.embed_query",
side_effect=lambda *args, **kwargs: [1.0, 2.1, 3.2],
)
def test_huggingface_embddings(
langchain_huggingface_embedding_call, sentence_transformers_init
):
model = HuggingFaceEmbeddings(
model_name="intfloat/multilingual-e5-large",
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": False},
)
output = model("Hello World")
assert isinstance(output, list)
assert isinstance(output[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],
)
def test_cohere_embddings(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)
langchain_cohere_embedding_call.assert_called()

4
tests/test_indexing_retrieval.py
View File

@@ -3,7 +3,7 @@ from pathlib import Path
from typing import cast
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from kotaemon.docstores import InMemoryDocumentStore
from kotaemon.documents.base import Document
@@ -18,7 +18,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
def test_indexing(mock_openai_embedding, tmp_path):

43
tests/test_llms_chat_models.py
View File

@@ -2,31 +2,37 @@ from unittest.mock import patch
from langchain.chat_models import AzureChatOpenAI as AzureChatOpenAILC
from langchain.schema.messages import AIMessage, HumanMessage, SystemMessage
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.base import LLMInterface
from kotaemon.llms.chats.openai import AzureChatOpenAI
_openai_chat_completion_response = {
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
_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": "Hello! How can I assist you today?",
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_chat_completion_response,
)
def test_azureopenai_model(openai_completion):
@@ -36,7 +42,6 @@ def test_azureopenai_model(openai_completion):
openai_api_version="2023-03-15-preview",
deployment_name="gpt35turbo",
temperature=0,
request_timeout=60,
)
assert isinstance(
model.agent, AzureChatOpenAILC

33
tests/test_llms_completion_models.py
View File

@@ -2,24 +2,33 @@ from unittest.mock import patch
from langchain.llms import AzureOpenAI as AzureOpenAILC
from langchain.llms import OpenAI as OpenAILC
from openai.types.completion import Completion
from kotaemon.llms.base import LLMInterface
from kotaemon.llms.completions.openai import AzureOpenAI, OpenAI
_openai_completion_response = {
"id": "cmpl-7qyNoIo6gRSCJR0hi8o3ZKBH4RkJ0",
"object": "sample text_completion",
"created": 1392751226,
"model": "gpt-35-turbo",
"choices": [
{"text": "completion", "index": 0, "finish_reason": "length", "logprobs": None}
],
"usage": {"completion_tokens": 20, "prompt_tokens": 2, "total_tokens": 22},
}
_openai_completion_response = Completion.parse_obj(
{
"id": "cmpl-7qyNoIo6gRSCJR0hi8o3ZKBH4RkJ0",
"object": "text_completion",
"created": 1392751226,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"text": "completion",
"index": 0,
"finish_reason": "length",
"logprobs": None,
}
],
"usage": {"completion_tokens": 20, "prompt_tokens": 2, "total_tokens": 22},
}
)
@patch(
"openai.api_resources.completion.Completion.create",
"openai.resources.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_azureopenai_model(openai_completion):
@@ -47,7 +56,7 @@ def test_azureopenai_model(openai_completion):
@patch(
"openai.api_resources.completion.Completion.create",
"openai.resources.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_openai_model(openai_completion):

47
tests/test_qa.py
View File

@@ -3,7 +3,8 @@ from pathlib import Path
from unittest.mock import patch
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.ingest import ReaderIndexingPipeline
@@ -12,32 +13,37 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
openai_embedding = json.load(f)
_openai_chat_completion_response = {
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
_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": "Hello! How can I assist you today?",
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_chat_completion_response,
)
def test_ingest_pipeline(patch, mock_openai_embedding, tmp_path):
@@ -61,7 +67,6 @@ def test_ingest_pipeline(patch, mock_openai_embedding, tmp_path):
openai_api_version="2023-03-15-preview",
deployment_name="gpt35turbo",
temperature=0,
request_timeout=60,
)
qa_pipeline = indexing_pipeline.to_qa_pipeline(llm=llm, openai_api_key="some-key")
response = qa_pipeline("Summarize this document.")

4
tests/test_tools.py
View File

@@ -2,7 +2,7 @@ import json
from pathlib import Path
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from kotaemon.docstores import InMemoryDocumentStore
from kotaemon.documents.base import Document
@@ -18,7 +18,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
def test_google_tool(mock_google_search):