diff --git a/knowledgehub/embeddings/cohere.py b/knowledgehub/embeddings/cohere.py new file mode 100644 index 0000000..e161787 --- /dev/null +++ b/knowledgehub/embeddings/cohere.py @@ -0,0 +1,12 @@ +from langchain.embeddings import CohereEmbeddings as LCCohereEmbeddings + +from kotaemon.embeddings.base import LangchainEmbeddings + + +class CohereEmbdeddings(LangchainEmbeddings): + """Cohere embeddings. + + This class wraps around the Langchain CohereEmbeddings class. + """ + + _lc_class = LCCohereEmbeddings diff --git a/knowledgehub/embeddings/huggingface.py b/knowledgehub/embeddings/huggingface.py new file mode 100644 index 0000000..37ba5e9 --- /dev/null +++ b/knowledgehub/embeddings/huggingface.py @@ -0,0 +1,12 @@ +from langchain.embeddings import HuggingFaceBgeEmbeddings as LCHuggingFaceEmbeddings + +from kotaemon.embeddings.base import LangchainEmbeddings + + +class HuggingFaceEmbeddings(LangchainEmbeddings): + """HuggingFace embeddings + + This class wraps around the Langchain HuggingFaceEmbeddings class + """ + + _lc_class = LCHuggingFaceEmbeddings diff --git a/knowledgehub/loaders/base.py b/knowledgehub/loaders/base.py index 70fead7..64dbd44 100644 --- a/knowledgehub/loaders/base.py +++ b/knowledgehub/loaders/base.py @@ -8,7 +8,7 @@ from ..base import BaseComponent from ..documents.base import Document -class AutoReader(BaseComponent, BaseReader): +class AutoReader(BaseComponent): """General auto reader for a variety of files. (based on llama-hub)""" def __init__(self, reader_type: Union[str, Type[BaseReader]]) -> None: @@ -31,7 +31,7 @@ class AutoReader(BaseComponent, BaseReader): return self.load_data(file=file, **kwargs) -class LIBaseReader(BaseComponent, BaseReader): +class LIBaseReader(BaseComponent): _reader_class: Type[BaseReader] def __init__(self, *args, **kwargs): diff --git a/knowledgehub/pipelines/agents/rewoo/planner.py b/knowledgehub/pipelines/agents/rewoo/planner.py index af2ddf3..89c1218 100644 --- a/knowledgehub/pipelines/agents/rewoo/planner.py +++ b/knowledgehub/pipelines/agents/rewoo/planner.py @@ -75,8 +75,8 @@ class Planner(BaseComponent): try: response = self.model(prompt) output.info("Planner run successful.") - except ValueError: + except ValueError as e: output.error("Planner failed to retrieve response from LLM") - raise ValueError("Planner failed to retrieve response from LLM") + raise ValueError("Planner failed to retrieve response from LLM") from e return response diff --git a/setup.py b/setup.py index 1b51fda..bb20417 100644 --- a/setup.py +++ b/setup.py @@ -59,6 +59,9 @@ setuptools.setup( "pytest-mock", "unstructured[pdf]", "farm-haystack==1.19.0", + "sentence_transformers", + "cohere", + "pypdf", ], }, entry_points={"console_scripts": ["kh=kotaemon.cli:main"]}, diff --git a/tests/test_agent.py b/tests/test_agent.py index 558b47e..a8e2eea 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -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): diff --git a/tests/test_cot.py b/tests/test_cot.py index f9423af..6583732 100644 --- a/tests/test_cot.py +++ b/tests/test_cot.py @@ -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): diff --git a/tests/test_embedding_models.py b/tests/test_embedding_models.py index d5a7642..6162353 100644 --- a/tests/test_embedding_models.py +++ b/tests/test_embedding_models.py @@ -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() diff --git a/tests/test_indexing_retrieval.py b/tests/test_indexing_retrieval.py index c253c25..fd4dd74 100644 --- a/tests/test_indexing_retrieval.py +++ b/tests/test_indexing_retrieval.py @@ -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): diff --git a/tests/test_llms_chat_models.py b/tests/test_llms_chat_models.py index 392d54e..7ce4ec3 100644 --- a/tests/test_llms_chat_models.py +++ b/tests/test_llms_chat_models.py @@ -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 diff --git a/tests/test_llms_completion_models.py b/tests/test_llms_completion_models.py index 495e5b4..a10d1f5 100644 --- a/tests/test_llms_completion_models.py +++ b/tests/test_llms_completion_models.py @@ -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): diff --git a/tests/test_qa.py b/tests/test_qa.py index 5f20cff..b6ff788 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -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.") diff --git a/tests/test_tools.py b/tests/test_tools.py index e5d22ae..771cb71 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -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):