From f6b80e01a108259e384920032124a52194203f1a Mon Sep 17 00:00:00 2001 From: Deshraj Yadav Date: Wed, 22 Nov 2023 01:06:33 -0800 Subject: [PATCH] [Feature] Add support for custom streaming callback (#971) --- embedchain/apps/app.py | 2 +- embedchain/bots/base.py | 4 +- embedchain/bots/discord.py | 2 +- embedchain/bots/poe.py | 2 +- embedchain/bots/slack.py | 2 +- embedchain/bots/whatsapp.py | 2 +- embedchain/chunkers/base_chunker.py | 2 +- embedchain/chunkers/common_chunker.py | 2 +- embedchain/chunkers/discourse.py | 2 +- embedchain/chunkers/docs_site.py | 2 +- embedchain/chunkers/docx_file.py | 2 +- embedchain/chunkers/gmail.py | 2 +- embedchain/chunkers/json.py | 2 +- embedchain/chunkers/mdx.py | 2 +- embedchain/chunkers/mysql.py | 2 +- embedchain/chunkers/notion.py | 2 +- embedchain/chunkers/pdf_file.py | 2 +- embedchain/chunkers/postgres.py | 2 +- embedchain/chunkers/qna_pair.py | 2 +- embedchain/chunkers/sitemap.py | 2 +- embedchain/chunkers/slack.py | 2 +- embedchain/chunkers/substack.py | 2 +- embedchain/chunkers/text.py | 2 +- embedchain/chunkers/unstructured_file.py | 2 +- embedchain/chunkers/web_page.py | 2 +- embedchain/chunkers/xml.py | 2 +- embedchain/chunkers/youtube_video.py | 2 +- embedchain/config/add_config.py | 2 +- embedchain/config/apps/app_config.py | 2 +- embedchain/config/apps/base_app_config.py | 2 +- embedchain/config/base_config.py | 2 +- embedchain/config/embedder/base.py | 2 +- embedchain/config/llm/base.py | 8 +- embedchain/config/pipeline_config.py | 2 +- embedchain/config/vectordb/chroma.py | 2 +- embedchain/config/vectordb/elasticsearch.py | 2 +- embedchain/config/vectordb/opensearch.py | 2 +- embedchain/config/vectordb/pinecone.py | 2 +- embedchain/config/vectordb/qdrant.py | 2 +- embedchain/config/vectordb/weaviate.py | 2 +- embedchain/config/vectordb/zilliz.py | 2 +- embedchain/data_formatter/data_formatter.py | 2 +- embedchain/embedchain.py | 2 +- embedchain/embedder/base.py | 4 +- embedchain/helpers/__init__.py | 0 embedchain/helpers/callbacks.py | 73 +++++++++++++++++++ .../{helper => helpers}/json_serializable.py | 0 embedchain/llm/anthropic.py | 2 +- embedchain/llm/azure_openai.py | 2 +- embedchain/llm/base.py | 2 +- embedchain/llm/cohere.py | 2 +- embedchain/llm/gpt4all.py | 2 +- embedchain/llm/huggingface.py | 2 +- embedchain/llm/jina.py | 2 +- embedchain/llm/llama2.py | 2 +- embedchain/llm/openai.py | 5 +- embedchain/llm/vertex_ai.py | 2 +- embedchain/loaders/base_loader.py | 2 +- embedchain/loaders/docs_site_loader.py | 2 +- embedchain/loaders/docx_file.py | 2 +- embedchain/loaders/local_qna_pair.py | 2 +- embedchain/loaders/local_text.py | 2 +- embedchain/loaders/mdx.py | 2 +- embedchain/loaders/notion.py | 2 +- embedchain/loaders/pdf_file.py | 2 +- embedchain/loaders/sitemap.py | 2 +- embedchain/loaders/substack.py | 2 +- embedchain/loaders/unstructured_file.py | 2 +- embedchain/loaders/web_page.py | 2 +- embedchain/loaders/xml.py | 2 +- embedchain/loaders/youtube_video.py | 2 +- embedchain/memory/message.py | 2 +- embedchain/pipeline.py | 2 +- embedchain/vectordb/base.py | 2 +- embedchain/vectordb/chroma.py | 2 +- embedchain/vectordb/elasticsearch.py | 2 +- embedchain/vectordb/opensearch.py | 2 +- embedchain/vectordb/pinecone.py | 2 +- embedchain/vectordb/weaviate.py | 2 +- embedchain/vectordb/zilliz.py | 2 +- pyproject.toml | 2 +- .../helper_classes/test_json_serializable.py | 4 +- 82 files changed, 162 insertions(+), 84 deletions(-) create mode 100644 embedchain/helpers/__init__.py create mode 100644 embedchain/helpers/callbacks.py rename embedchain/{helper => helpers}/json_serializable.py (100%) diff --git a/embedchain/apps/app.py b/embedchain/apps/app.py index 03dd3027f..56d1353ea 100644 --- a/embedchain/apps/app.py +++ b/embedchain/apps/app.py @@ -10,7 +10,7 @@ from embedchain.embedchain import EmbedChain from embedchain.embedder.base import BaseEmbedder from embedchain.embedder.openai import OpenAIEmbedder from embedchain.factory import EmbedderFactory, LlmFactory, VectorDBFactory -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm from embedchain.llm.openai import OpenAILlm from embedchain.utils import validate_yaml_config diff --git a/embedchain/bots/base.py b/embedchain/bots/base.py index a804b71b0..384ca47e3 100644 --- a/embedchain/bots/base.py +++ b/embedchain/bots/base.py @@ -3,8 +3,8 @@ from typing import Any from embedchain import Pipeline as App from embedchain.config import AddConfig, BaseLlmConfig, PipelineConfig from embedchain.embedder.openai import OpenAIEmbedder -from embedchain.helper.json_serializable import (JSONSerializable, - register_deserializable) +from embedchain.helpers.json_serializable import (JSONSerializable, + register_deserializable) from embedchain.llm.openai import OpenAILlm from embedchain.vectordb.chroma import ChromaDB diff --git a/embedchain/bots/discord.py b/embedchain/bots/discord.py index 2d2d482e4..adbf3b7c7 100644 --- a/embedchain/bots/discord.py +++ b/embedchain/bots/discord.py @@ -2,7 +2,7 @@ import argparse import logging import os -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .base import BaseBot diff --git a/embedchain/bots/poe.py b/embedchain/bots/poe.py index ebd71824e..762d090b1 100644 --- a/embedchain/bots/poe.py +++ b/embedchain/bots/poe.py @@ -3,7 +3,7 @@ import logging import os from typing import List, Optional -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .base import BaseBot diff --git a/embedchain/bots/slack.py b/embedchain/bots/slack.py index 21834d8d4..43e39f261 100644 --- a/embedchain/bots/slack.py +++ b/embedchain/bots/slack.py @@ -5,7 +5,7 @@ import signal import sys from embedchain import App -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .base import BaseBot diff --git a/embedchain/bots/whatsapp.py b/embedchain/bots/whatsapp.py index c9193f6f6..5106d40dc 100644 --- a/embedchain/bots/whatsapp.py +++ b/embedchain/bots/whatsapp.py @@ -4,7 +4,7 @@ import logging import signal import sys -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .base import BaseBot diff --git a/embedchain/chunkers/base_chunker.py b/embedchain/chunkers/base_chunker.py index f7b5cb09f..7130d30bf 100644 --- a/embedchain/chunkers/base_chunker.py +++ b/embedchain/chunkers/base_chunker.py @@ -1,6 +1,6 @@ import hashlib -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable from embedchain.models.data_type import DataType diff --git a/embedchain/chunkers/common_chunker.py b/embedchain/chunkers/common_chunker.py index 7d607357f..1527e3396 100644 --- a/embedchain/chunkers/common_chunker.py +++ b/embedchain/chunkers/common_chunker.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/discourse.py b/embedchain/chunkers/discourse.py index f78c616e1..14898bf01 100644 --- a/embedchain/chunkers/discourse.py +++ b/embedchain/chunkers/discourse.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/docs_site.py b/embedchain/chunkers/docs_site.py index 1b7c00d9c..d51dc8ee2 100644 --- a/embedchain/chunkers/docs_site.py +++ b/embedchain/chunkers/docs_site.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/docx_file.py b/embedchain/chunkers/docx_file.py index 77fee55cb..1452349e8 100644 --- a/embedchain/chunkers/docx_file.py +++ b/embedchain/chunkers/docx_file.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/gmail.py b/embedchain/chunkers/gmail.py index 59e87e7ec..6b804f546 100644 --- a/embedchain/chunkers/gmail.py +++ b/embedchain/chunkers/gmail.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/json.py b/embedchain/chunkers/json.py index 4eeee7ebe..ebc525419 100644 --- a/embedchain/chunkers/json.py +++ b/embedchain/chunkers/json.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/mdx.py b/embedchain/chunkers/mdx.py index 05225fac7..1c277dda7 100644 --- a/embedchain/chunkers/mdx.py +++ b/embedchain/chunkers/mdx.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/mysql.py b/embedchain/chunkers/mysql.py index 7a2ec7fc0..2b1c11ace 100644 --- a/embedchain/chunkers/mysql.py +++ b/embedchain/chunkers/mysql.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/notion.py b/embedchain/chunkers/notion.py index e0e2569cf..190d59b57 100644 --- a/embedchain/chunkers/notion.py +++ b/embedchain/chunkers/notion.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/pdf_file.py b/embedchain/chunkers/pdf_file.py index eb2355a4e..56bae064e 100644 --- a/embedchain/chunkers/pdf_file.py +++ b/embedchain/chunkers/pdf_file.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/postgres.py b/embedchain/chunkers/postgres.py index 168b6fcd5..7c6859bd0 100644 --- a/embedchain/chunkers/postgres.py +++ b/embedchain/chunkers/postgres.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/qna_pair.py b/embedchain/chunkers/qna_pair.py index 30f2b52f9..c0d8277b1 100644 --- a/embedchain/chunkers/qna_pair.py +++ b/embedchain/chunkers/qna_pair.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/sitemap.py b/embedchain/chunkers/sitemap.py index 640500113..64e773742 100644 --- a/embedchain/chunkers/sitemap.py +++ b/embedchain/chunkers/sitemap.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/slack.py b/embedchain/chunkers/slack.py index 93453b3b3..595682beb 100644 --- a/embedchain/chunkers/slack.py +++ b/embedchain/chunkers/slack.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/substack.py b/embedchain/chunkers/substack.py index 85f04fa62..92cacd6cb 100644 --- a/embedchain/chunkers/substack.py +++ b/embedchain/chunkers/substack.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/text.py b/embedchain/chunkers/text.py index 0c36c863b..f33d60c46 100644 --- a/embedchain/chunkers/text.py +++ b/embedchain/chunkers/text.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/unstructured_file.py b/embedchain/chunkers/unstructured_file.py index ab0322b02..d55f23ef0 100644 --- a/embedchain/chunkers/unstructured_file.py +++ b/embedchain/chunkers/unstructured_file.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/web_page.py b/embedchain/chunkers/web_page.py index b3da9ea3f..253b2b413 100644 --- a/embedchain/chunkers/web_page.py +++ b/embedchain/chunkers/web_page.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/xml.py b/embedchain/chunkers/xml.py index cae519ab9..c1bab0a77 100644 --- a/embedchain/chunkers/xml.py +++ b/embedchain/chunkers/xml.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/chunkers/youtube_video.py b/embedchain/chunkers/youtube_video.py index d2ca025c7..bde0a8f78 100644 --- a/embedchain/chunkers/youtube_video.py +++ b/embedchain/chunkers/youtube_video.py @@ -4,7 +4,7 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config.add_config import ChunkerConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/add_config.py b/embedchain/config/add_config.py index 16e99d687..669551184 100644 --- a/embedchain/config/add_config.py +++ b/embedchain/config/add_config.py @@ -3,7 +3,7 @@ from importlib import import_module from typing import Callable, Optional from embedchain.config.base_config import BaseConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/apps/app_config.py b/embedchain/config/apps/app_config.py index f0efaff9f..4ab90b022 100644 --- a/embedchain/config/apps/app_config.py +++ b/embedchain/config/apps/app_config.py @@ -1,6 +1,6 @@ from typing import Optional -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .base_app_config import BaseAppConfig diff --git a/embedchain/config/apps/base_app_config.py b/embedchain/config/apps/base_app_config.py index d12b1c630..f3a864700 100644 --- a/embedchain/config/apps/base_app_config.py +++ b/embedchain/config/apps/base_app_config.py @@ -2,7 +2,7 @@ import logging from typing import Optional from embedchain.config.base_config import BaseConfig -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/config/base_config.py b/embedchain/config/base_config.py index b02801f47..ff672f19b 100644 --- a/embedchain/config/base_config.py +++ b/embedchain/config/base_config.py @@ -1,6 +1,6 @@ from typing import Any, Dict -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable class BaseConfig(JSONSerializable): diff --git a/embedchain/config/embedder/base.py b/embedchain/config/embedder/base.py index 8e8501e74..9227e0791 100644 --- a/embedchain/config/embedder/base.py +++ b/embedchain/config/embedder/base.py @@ -1,6 +1,6 @@ from typing import Optional -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/llm/base.py b/embedchain/config/llm/base.py index a98c1b6d5..6dbfdb90f 100644 --- a/embedchain/config/llm/base.py +++ b/embedchain/config/llm/base.py @@ -1,9 +1,9 @@ import re from string import Template -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from embedchain.config.base_config import BaseConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable DEFAULT_PROMPT = """ Use the following pieces of context to answer the query at the end. @@ -68,6 +68,7 @@ class BaseLlmConfig(BaseConfig): system_prompt: Optional[str] = None, where: Dict[str, Any] = None, query_type: Optional[str] = None, + callbacks: Optional[List] = None, ): """ Initializes a configuration class instance for the LLM. @@ -98,6 +99,8 @@ class BaseLlmConfig(BaseConfig): :type system_prompt: Optional[str], optional :param where: A dictionary of key-value pairs to filter the database results., defaults to None :type where: Dict[str, Any], optional + :param callbacks: Langchain callback functions to use, defaults to None + :type callbacks: Optional[List], optional :raises ValueError: If the template is not valid as template should contain $context and $query (and optionally $history) :raises ValueError: Stream is not boolean @@ -113,6 +116,7 @@ class BaseLlmConfig(BaseConfig): self.deployment_name = deployment_name self.system_prompt = system_prompt self.query_type = query_type + self.callbacks = callbacks if type(template) is str: template = Template(template) diff --git a/embedchain/config/pipeline_config.py b/embedchain/config/pipeline_config.py index e46456a1b..8bfd4b4ea 100644 --- a/embedchain/config/pipeline_config.py +++ b/embedchain/config/pipeline_config.py @@ -1,6 +1,6 @@ from typing import Optional -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from .apps.base_app_config import BaseAppConfig diff --git a/embedchain/config/vectordb/chroma.py b/embedchain/config/vectordb/chroma.py index 38bf09213..d25de1c34 100644 --- a/embedchain/config/vectordb/chroma.py +++ b/embedchain/config/vectordb/chroma.py @@ -1,7 +1,7 @@ from typing import Optional from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/elasticsearch.py b/embedchain/config/vectordb/elasticsearch.py index 75498bd24..77d54a16c 100644 --- a/embedchain/config/vectordb/elasticsearch.py +++ b/embedchain/config/vectordb/elasticsearch.py @@ -2,7 +2,7 @@ import os from typing import Dict, List, Optional, Union from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/opensearch.py b/embedchain/config/vectordb/opensearch.py index 85d517b2d..d8dc9a109 100644 --- a/embedchain/config/vectordb/opensearch.py +++ b/embedchain/config/vectordb/opensearch.py @@ -1,7 +1,7 @@ from typing import Dict, Optional, Tuple from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/pinecone.py b/embedchain/config/vectordb/pinecone.py index 7bd462ae4..e9165fdca 100644 --- a/embedchain/config/vectordb/pinecone.py +++ b/embedchain/config/vectordb/pinecone.py @@ -1,7 +1,7 @@ from typing import Dict, Optional from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/qdrant.py b/embedchain/config/vectordb/qdrant.py index 4468c7b24..9802212e8 100644 --- a/embedchain/config/vectordb/qdrant.py +++ b/embedchain/config/vectordb/qdrant.py @@ -1,7 +1,7 @@ from typing import Dict, Optional from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/weaviate.py b/embedchain/config/vectordb/weaviate.py index 4035877b5..2db24134b 100644 --- a/embedchain/config/vectordb/weaviate.py +++ b/embedchain/config/vectordb/weaviate.py @@ -1,7 +1,7 @@ from typing import Dict, Optional from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/config/vectordb/zilliz.py b/embedchain/config/vectordb/zilliz.py index fbd6ec025..ba91922c0 100644 --- a/embedchain/config/vectordb/zilliz.py +++ b/embedchain/config/vectordb/zilliz.py @@ -2,7 +2,7 @@ import os from typing import Optional from embedchain.config.vectordb.base import BaseVectorDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable @register_deserializable diff --git a/embedchain/data_formatter/data_formatter.py b/embedchain/data_formatter/data_formatter.py index 1c218d8b8..368e43259 100644 --- a/embedchain/data_formatter/data_formatter.py +++ b/embedchain/data_formatter/data_formatter.py @@ -4,7 +4,7 @@ from typing import Any, Dict from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config import AddConfig from embedchain.config.add_config import ChunkerConfig, LoaderConfig -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable from embedchain.loaders.base_loader import BaseLoader from embedchain.models.data_type import DataType diff --git a/embedchain/embedchain.py b/embedchain/embedchain.py index 88c098a6d..dfd1627e2 100644 --- a/embedchain/embedchain.py +++ b/embedchain/embedchain.py @@ -13,7 +13,7 @@ from embedchain.config.apps.base_app_config import BaseAppConfig from embedchain.constants import SQLITE_PATH from embedchain.data_formatter import DataFormatter from embedchain.embedder.base import BaseEmbedder -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable from embedchain.llm.base import BaseLlm from embedchain.loaders.base_loader import BaseLoader from embedchain.models.data_type import (DataType, DirectDataType, diff --git a/embedchain/embedder/base.py b/embedchain/embedder/base.py index 50ed475b9..14941b2f1 100644 --- a/embedchain/embedder/base.py +++ b/embedchain/embedder/base.py @@ -3,12 +3,12 @@ from typing import Any, Callable, Optional from embedchain.config.embedder.base import BaseEmbedderConfig try: - from chromadb.api.types import Embeddings, Embeddable, EmbeddingFunction + from chromadb.api.types import Embeddable, EmbeddingFunction, Embeddings except RuntimeError: from embedchain.utils import use_pysqlite3 use_pysqlite3() - from chromadb.api.types import Embeddings, Embeddable, EmbeddingFunction + from chromadb.api.types import Embeddable, EmbeddingFunction, Embeddings class EmbeddingFunc(EmbeddingFunction): diff --git a/embedchain/helpers/__init__.py b/embedchain/helpers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/embedchain/helpers/callbacks.py b/embedchain/helpers/callbacks.py new file mode 100644 index 000000000..3c7ab356a --- /dev/null +++ b/embedchain/helpers/callbacks.py @@ -0,0 +1,73 @@ +import queue +from typing import Any, Dict, List, Union + +from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler +from langchain.schema import LLMResult + +STOP_ITEM = "[END]" +""" +This is a special item that is used to signal the end of the stream. +""" + + +class StreamingStdOutCallbackHandlerYield(StreamingStdOutCallbackHandler): + """ + This is a callback handler that yields the tokens as they are generated. + For a usage example, see the :func:`generate` function below. + """ + + q: queue.Queue + """ + The queue to write the tokens to as they are generated. + """ + + def __init__(self, q: queue.Queue) -> None: + """ + Initialize the callback handler. + q: The queue to write the tokens to as they are generated. + """ + super().__init__() + self.q = q + + def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None: + """Run when LLM starts running.""" + with self.q.mutex: + self.q.queue.clear() + + def on_llm_new_token(self, token: str, **kwargs: Any) -> None: + """Run on new LLM token. Only available when streaming is enabled.""" + self.q.put(token) + + def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + """Run when LLM ends running.""" + self.q.put(STOP_ITEM) + + def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: + """Run when LLM errors.""" + self.q.put("%s: %s" % (type(error).__name__, str(error))) + self.q.put(STOP_ITEM) + + +def generate(rq: queue.Queue): + """ + This is a generator that yields the items in the queue until it reaches the stop item. + + Usage example: + ``` + def askQuestion(callback_fn: StreamingStdOutCallbackHandlerYield): + llm = OpenAI(streaming=True, callbacks=[callback_fn]) + return llm(prompt="Write a poem about a tree.") + + @app.route("/", methods=["GET"]) + def generate_output(): + q = Queue() + callback_fn = StreamingStdOutCallbackHandlerYield(q) + threading.Thread(target=askQuestion, args=(callback_fn,)).start() + return Response(generate(q), mimetype="text/event-stream") + ``` + """ + while True: + result: str = rq.get() + if result == STOP_ITEM or result is None: + break + yield result diff --git a/embedchain/helper/json_serializable.py b/embedchain/helpers/json_serializable.py similarity index 100% rename from embedchain/helper/json_serializable.py rename to embedchain/helpers/json_serializable.py diff --git a/embedchain/llm/anthropic.py b/embedchain/llm/anthropic.py index 492854f6f..aab38640d 100644 --- a/embedchain/llm/anthropic.py +++ b/embedchain/llm/anthropic.py @@ -3,7 +3,7 @@ import os from typing import Optional from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/azure_openai.py b/embedchain/llm/azure_openai.py index 34c4f0815..99588aa0b 100644 --- a/embedchain/llm/azure_openai.py +++ b/embedchain/llm/azure_openai.py @@ -2,7 +2,7 @@ import logging from typing import Optional from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/base.py b/embedchain/llm/base.py index 8bb38833e..dc2471e42 100644 --- a/embedchain/llm/base.py +++ b/embedchain/llm/base.py @@ -7,7 +7,7 @@ from embedchain.config import BaseLlmConfig from embedchain.config.llm.base import (DEFAULT_PROMPT, DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE, DOCS_SITE_PROMPT_TEMPLATE) -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable from embedchain.memory.base import ECChatMemory from embedchain.memory.message import ChatMessage diff --git a/embedchain/llm/cohere.py b/embedchain/llm/cohere.py index 0811c0672..e996db6b7 100644 --- a/embedchain/llm/cohere.py +++ b/embedchain/llm/cohere.py @@ -5,7 +5,7 @@ from typing import Optional from langchain.llms import Cohere from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/gpt4all.py b/embedchain/llm/gpt4all.py index 699e9d339..73950bd7c 100644 --- a/embedchain/llm/gpt4all.py +++ b/embedchain/llm/gpt4all.py @@ -4,7 +4,7 @@ from langchain.callbacks.stdout import StdOutCallbackHandler from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/huggingface.py b/embedchain/llm/huggingface.py index 97adc67b9..4da6d5176 100644 --- a/embedchain/llm/huggingface.py +++ b/embedchain/llm/huggingface.py @@ -5,7 +5,7 @@ from typing import Optional from langchain.llms import HuggingFaceHub from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/jina.py b/embedchain/llm/jina.py index 2af5a7980..b09c0fd3e 100644 --- a/embedchain/llm/jina.py +++ b/embedchain/llm/jina.py @@ -5,7 +5,7 @@ from langchain.chat_models import JinaChat from langchain.schema import HumanMessage, SystemMessage from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/llama2.py b/embedchain/llm/llama2.py index d4f4aa2f5..fa2587bb8 100644 --- a/embedchain/llm/llama2.py +++ b/embedchain/llm/llama2.py @@ -5,7 +5,7 @@ from typing import Optional from langchain.llms import Replicate from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/llm/openai.py b/embedchain/llm/openai.py index 9e69085c0..9efa019b3 100644 --- a/embedchain/llm/openai.py +++ b/embedchain/llm/openai.py @@ -4,7 +4,7 @@ from langchain.chat_models import ChatOpenAI from langchain.schema import HumanMessage, SystemMessage from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm @@ -34,7 +34,8 @@ class OpenAILlm(BaseLlm): from langchain.callbacks.streaming_stdout import \ StreamingStdOutCallbackHandler - chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()]) + callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()] + chat = ChatOpenAI(**kwargs, streaming=config.stream, callbacks=callbacks) else: chat = ChatOpenAI(**kwargs) return chat(messages).content diff --git a/embedchain/llm/vertex_ai.py b/embedchain/llm/vertex_ai.py index 224c65f16..c453c67b5 100644 --- a/embedchain/llm/vertex_ai.py +++ b/embedchain/llm/vertex_ai.py @@ -3,7 +3,7 @@ import logging from typing import Optional from embedchain.config import BaseLlmConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm diff --git a/embedchain/loaders/base_loader.py b/embedchain/loaders/base_loader.py index bba58b813..26da91a2a 100644 --- a/embedchain/loaders/base_loader.py +++ b/embedchain/loaders/base_loader.py @@ -1,4 +1,4 @@ -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable class BaseLoader(JSONSerializable): diff --git a/embedchain/loaders/docs_site_loader.py b/embedchain/loaders/docs_site_loader.py index 2b880722c..ea76442e5 100644 --- a/embedchain/loaders/docs_site_loader.py +++ b/embedchain/loaders/docs_site_loader.py @@ -12,7 +12,7 @@ except ImportError: ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader diff --git a/embedchain/loaders/docx_file.py b/embedchain/loaders/docx_file.py index 59846c051..d72e64b29 100644 --- a/embedchain/loaders/docx_file.py +++ b/embedchain/loaders/docx_file.py @@ -6,7 +6,7 @@ except ImportError: raise ImportError( 'Docx file requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader diff --git a/embedchain/loaders/local_qna_pair.py b/embedchain/loaders/local_qna_pair.py index ffaa6feaa..1158d4a99 100644 --- a/embedchain/loaders/local_qna_pair.py +++ b/embedchain/loaders/local_qna_pair.py @@ -1,6 +1,6 @@ import hashlib -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader diff --git a/embedchain/loaders/local_text.py b/embedchain/loaders/local_text.py index 118cbd3af..e03ee12b2 100644 --- a/embedchain/loaders/local_text.py +++ b/embedchain/loaders/local_text.py @@ -1,6 +1,6 @@ import hashlib -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader diff --git a/embedchain/loaders/mdx.py b/embedchain/loaders/mdx.py index 9d73b6cad..45b112f14 100644 --- a/embedchain/loaders/mdx.py +++ b/embedchain/loaders/mdx.py @@ -1,6 +1,6 @@ import hashlib -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader diff --git a/embedchain/loaders/notion.py b/embedchain/loaders/notion.py index 7ff84ed58..e0e981e80 100644 --- a/embedchain/loaders/notion.py +++ b/embedchain/loaders/notion.py @@ -10,7 +10,7 @@ except ImportError: ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/loaders/pdf_file.py b/embedchain/loaders/pdf_file.py index 6b03554ba..03495edbe 100644 --- a/embedchain/loaders/pdf_file.py +++ b/embedchain/loaders/pdf_file.py @@ -6,7 +6,7 @@ except ImportError: raise ImportError( 'PDF File requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/loaders/sitemap.py b/embedchain/loaders/sitemap.py index 8b449f212..06e7e2390 100644 --- a/embedchain/loaders/sitemap.py +++ b/embedchain/loaders/sitemap.py @@ -13,7 +13,7 @@ except ImportError: 'Sitemap requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.loaders.web_page import WebPageLoader from embedchain.utils import is_readable diff --git a/embedchain/loaders/substack.py b/embedchain/loaders/substack.py index 4dcc609ee..0d46b6d1a 100644 --- a/embedchain/loaders/substack.py +++ b/embedchain/loaders/substack.py @@ -4,7 +4,7 @@ import time import requests -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import is_readable diff --git a/embedchain/loaders/unstructured_file.py b/embedchain/loaders/unstructured_file.py index be8cd931f..294c596c2 100644 --- a/embedchain/loaders/unstructured_file.py +++ b/embedchain/loaders/unstructured_file.py @@ -1,6 +1,6 @@ import hashlib -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/loaders/web_page.py b/embedchain/loaders/web_page.py index 931031826..ecf03e9df 100644 --- a/embedchain/loaders/web_page.py +++ b/embedchain/loaders/web_page.py @@ -10,7 +10,7 @@ except ImportError: 'Webpage requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/loaders/xml.py b/embedchain/loaders/xml.py index d200ffb25..00fe47701 100644 --- a/embedchain/loaders/xml.py +++ b/embedchain/loaders/xml.py @@ -6,7 +6,7 @@ except ImportError: raise ImportError( 'XML file requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/loaders/youtube_video.py b/embedchain/loaders/youtube_video.py index c3262822b..2aa080247 100644 --- a/embedchain/loaders/youtube_video.py +++ b/embedchain/loaders/youtube_video.py @@ -6,7 +6,7 @@ except ImportError: raise ImportError( 'YouTube video requires extra dependencies. Install with `pip install --upgrade "embedchain[dataloaders]"`' ) from None -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils import clean_string diff --git a/embedchain/memory/message.py b/embedchain/memory/message.py index 383b7c0f3..99081d2d5 100644 --- a/embedchain/memory/message.py +++ b/embedchain/memory/message.py @@ -1,7 +1,7 @@ import logging from typing import Any, Dict, Optional -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable class BaseMessage(JSONSerializable): diff --git a/embedchain/pipeline.py b/embedchain/pipeline.py index 91faf83c4..00a8a92df 100644 --- a/embedchain/pipeline.py +++ b/embedchain/pipeline.py @@ -15,7 +15,7 @@ from embedchain.embedchain import EmbedChain from embedchain.embedder.base import BaseEmbedder from embedchain.embedder.openai import OpenAIEmbedder from embedchain.factory import EmbedderFactory, LlmFactory, VectorDBFactory -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm from embedchain.llm.openai import OpenAILlm from embedchain.telemetry.posthog import AnonymousTelemetry diff --git a/embedchain/vectordb/base.py b/embedchain/vectordb/base.py index 77d6efc7a..12f9693f9 100644 --- a/embedchain/vectordb/base.py +++ b/embedchain/vectordb/base.py @@ -1,6 +1,6 @@ from embedchain.config.vectordb.base import BaseVectorDbConfig from embedchain.embedder.base import BaseEmbedder -from embedchain.helper.json_serializable import JSONSerializable +from embedchain.helpers.json_serializable import JSONSerializable class BaseVectorDB(JSONSerializable): diff --git a/embedchain/vectordb/chroma.py b/embedchain/vectordb/chroma.py index f5f1cbbe1..f32cf525a 100644 --- a/embedchain/vectordb/chroma.py +++ b/embedchain/vectordb/chroma.py @@ -6,7 +6,7 @@ from langchain.docstore.document import Document from tqdm import tqdm from embedchain.config import ChromaDbConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB try: diff --git a/embedchain/vectordb/elasticsearch.py b/embedchain/vectordb/elasticsearch.py index b27370802..5ae6fd7c4 100644 --- a/embedchain/vectordb/elasticsearch.py +++ b/embedchain/vectordb/elasticsearch.py @@ -10,7 +10,7 @@ except ImportError: ) from None from embedchain.config import ElasticsearchDBConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/vectordb/opensearch.py b/embedchain/vectordb/opensearch.py index 1ffe881a8..a86f6292b 100644 --- a/embedchain/vectordb/opensearch.py +++ b/embedchain/vectordb/opensearch.py @@ -16,7 +16,7 @@ from langchain.embeddings.openai import OpenAIEmbeddings from langchain.vectorstores import OpenSearchVectorSearch from embedchain.config import OpenSearchDBConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/vectordb/pinecone.py b/embedchain/vectordb/pinecone.py index 86a817ac9..c3420c092 100644 --- a/embedchain/vectordb/pinecone.py +++ b/embedchain/vectordb/pinecone.py @@ -9,7 +9,7 @@ except ImportError: ) from None from embedchain.config.vectordb.pinecone import PineconeDBConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/vectordb/weaviate.py b/embedchain/vectordb/weaviate.py index fde91caf6..6ff329cbc 100644 --- a/embedchain/vectordb/weaviate.py +++ b/embedchain/vectordb/weaviate.py @@ -10,7 +10,7 @@ except ImportError: ) from None from embedchain.config.vectordb.weaviate import WeaviateDBConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/vectordb/zilliz.py b/embedchain/vectordb/zilliz.py index 7779f3449..0608c12fe 100644 --- a/embedchain/vectordb/zilliz.py +++ b/embedchain/vectordb/zilliz.py @@ -2,7 +2,7 @@ import logging from typing import Dict, List, Optional, Tuple, Union from embedchain.config import ZillizDBConfig -from embedchain.helper.json_serializable import register_deserializable +from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB try: diff --git a/pyproject.toml b/pyproject.toml index 167459bd8..dc6167b62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "embedchain" -version = "0.1.18" +version = "0.1.19" description = "Data platform for LLMs - Load, index, retrieve and sync any unstructured data" authors = [ "Taranjeet Singh ", diff --git a/tests/helper_classes/test_json_serializable.py b/tests/helper_classes/test_json_serializable.py index ba06a005f..3cbe27634 100644 --- a/tests/helper_classes/test_json_serializable.py +++ b/tests/helper_classes/test_json_serializable.py @@ -4,8 +4,8 @@ from string import Template from embedchain import App from embedchain.config import AppConfig, BaseLlmConfig -from embedchain.helper.json_serializable import (JSONSerializable, - register_deserializable) +from embedchain.helpers.json_serializable import (JSONSerializable, + register_deserializable) class TestJsonSerializable(unittest.TestCase):