From a972d2fb0743a67a73049741c619325b8bbd593d Mon Sep 17 00:00:00 2001 From: Dev Khant Date: Sat, 7 Sep 2024 22:39:28 +0530 Subject: [PATCH] Code Formatting (#1828) --- .pre-commit-config.yaml | 2 +- cookbooks/add_memory_using_qdrant_cloud.py | 2 ++ embedchain/embedchain/app.py | 8 +++++-- embedchain/embedchain/bots/base.py | 6 +++-- embedchain/embedchain/cache.py | 10 ++++---- embedchain/embedchain/cli.py | 22 +++++++++++------ embedchain/embedchain/config/__init__.py | 2 +- .../embedchain/config/evaluation/__init__.py | 7 ++++-- embedchain/embedchain/config/llm/base.py | 4 ++-- embedchain/embedchain/embedchain.py | 14 +++++++++-- embedchain/embedchain/embedder/clarifai.py | 6 ++--- embedchain/embedchain/embedder/gpt4all.py | 4 +++- embedchain/embedchain/embedder/openai.py | 1 - embedchain/embedchain/llm/aws_bedrock.py | 5 ++-- embedchain/embedchain/llm/base.py | 2 +- embedchain/embedchain/llm/gpt4all.py | 3 +-- embedchain/embedchain/llm/jina.py | 5 ++-- embedchain/embedchain/loaders/pdf_file.py | 1 + embedchain/embedchain/loaders/rss_feed.py | 4 +++- .../embedchain/loaders/unstructured_file.py | 3 +-- embedchain/embedchain/vectordb/zilliz.py | 11 +++++++-- embedchain/examples/chat-pdf/app.py | 3 +-- embedchain/examples/sadhguru-ai/app.py | 3 +-- embedchain/examples/unacademy-ai/app.py | 3 +-- .../embedder/test_azure_openai_embedder.py | 2 +- .../embedder/test_huggingface_embedder.py | 1 + .../helper_classes/test_json_serializable.py | 6 +++-- embedchain/tests/llm/conftest.py | 2 ++ embedchain/tests/llm/test_azure_openai.py | 2 +- embedchain/tests/llm/test_ollama.py | 2 +- embedchain/tests/models/test_data_type.py | 8 +++++-- mem0/__init__.py | 2 +- mem0/configs/base.py | 7 +++--- mem0/configs/embeddings/base.py | 5 ++-- mem0/configs/llms/base.py | 5 ++-- mem0/configs/vector_stores/chroma.py | 2 +- mem0/configs/vector_stores/pgvector.py | 2 +- mem0/configs/vector_stores/qdrant.py | 3 ++- mem0/embeddings/base.py | 2 +- mem0/graphs/configs.py | 3 +++ mem0/llms/anthropic.py | 5 +--- mem0/llms/aws_bedrock.py | 8 +++---- mem0/llms/azure_openai.py | 4 ++-- mem0/llms/base.py | 2 +- mem0/llms/groq.py | 6 ++--- mem0/llms/litellm.py | 4 +--- mem0/llms/ollama.py | 4 +--- mem0/llms/openai.py | 4 ++-- mem0/llms/openai_structured.py | 5 ++-- mem0/llms/together.py | 6 ++--- mem0/memory/graph_memory.py | 15 ++++++++---- mem0/memory/main.py | 24 ++++++++++--------- mem0/memory/telemetry.py | 3 +++ mem0/proxy/main.py | 7 +++--- mem0/utils/factory.py | 2 +- mem0/vector_stores/chroma.py | 4 +--- mem0/vector_stores/configs.py | 3 ++- mem0/vector_stores/pgvector.py | 5 ++-- mem0/vector_stores/qdrant.py | 2 +- tests/llms/test_azure_openai.py | 5 ++-- tests/llms/test_groq.py | 7 ++++-- tests/llms/test_litellm.py | 6 +++-- tests/llms/test_ollama.py | 7 ++++-- tests/llms/test_openai.py | 7 ++++-- tests/llms/test_together.py | 7 ++++-- tests/test_proxy.py | 9 +++---- 66 files changed, 208 insertions(+), 138 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 29127126..f3cac96f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,7 +3,7 @@ repos: hooks: - id: ruff name: Ruff - entry: ruff + entry: ruff check language: system types: [python] args: [--fix] diff --git a/cookbooks/add_memory_using_qdrant_cloud.py b/cookbooks/add_memory_using_qdrant_cloud.py index c2500ff8..d7142752 100644 --- a/cookbooks/add_memory_using_qdrant_cloud.py +++ b/cookbooks/add_memory_using_qdrant_cloud.py @@ -1,6 +1,8 @@ # This example shows how to use vector config to use QDRANT CLOUD import os + from dotenv import load_dotenv + from mem0 import Memory # Loading OpenAI API Key diff --git a/embedchain/embedchain/app.py b/embedchain/embedchain/app.py index ede9ec75..b4d05160 100644 --- a/embedchain/embedchain/app.py +++ b/embedchain/embedchain/app.py @@ -9,7 +9,6 @@ import requests import yaml from tqdm import tqdm -from mem0 import Memory from embedchain.cache import ( Config, ExactMatchEvaluation, @@ -26,7 +25,11 @@ from embedchain.embedchain import EmbedChain from embedchain.embedder.base import BaseEmbedder from embedchain.embedder.openai import OpenAIEmbedder from embedchain.evaluation.base import BaseMetric -from embedchain.evaluation.metrics import AnswerRelevance, ContextRelevance, Groundedness +from embedchain.evaluation.metrics import ( + AnswerRelevance, + ContextRelevance, + Groundedness, +) from embedchain.factory import EmbedderFactory, LlmFactory, VectorDBFactory from embedchain.helpers.json_serializable import register_deserializable from embedchain.llm.base import BaseLlm @@ -36,6 +39,7 @@ from embedchain.utils.evaluation import EvalData, EvalMetric from embedchain.utils.misc import validate_config from embedchain.vectordb.base import BaseVectorDB from embedchain.vectordb.chroma import ChromaDB +from mem0 import Memory logger = logging.getLogger(__name__) diff --git a/embedchain/embedchain/bots/base.py b/embedchain/embedchain/bots/base.py index aa9cbdf9..4a817cc4 100644 --- a/embedchain/embedchain/bots/base.py +++ b/embedchain/embedchain/bots/base.py @@ -3,8 +3,10 @@ from typing import Any from embedchain import App from embedchain.config import AddConfig, AppConfig, BaseLlmConfig from embedchain.embedder.openai import OpenAIEmbedder -from embedchain.helpers.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/embedchain/cache.py b/embedchain/embedchain/cache.py index e3334fd8..765141c3 100644 --- a/embedchain/embedchain/cache.py +++ b/embedchain/embedchain/cache.py @@ -9,10 +9,12 @@ from gptcache.manager import get_data_manager from gptcache.manager.scalar_data.base import Answer from gptcache.manager.scalar_data.base import DataType as CacheDataType from gptcache.session import Session -from gptcache.similarity_evaluation.distance import \ - SearchDistanceEvaluation # noqa: F401 -from gptcache.similarity_evaluation.exact_match import \ - ExactMatchEvaluation # noqa: F401 +from gptcache.similarity_evaluation.distance import ( # noqa: F401 + SearchDistanceEvaluation, +) +from gptcache.similarity_evaluation.exact_match import ( # noqa: F401 + ExactMatchEvaluation, +) logger = logging.getLogger(__name__) diff --git a/embedchain/embedchain/cli.py b/embedchain/embedchain/cli.py index 3ced1afd..e4f0401d 100644 --- a/embedchain/embedchain/cli.py +++ b/embedchain/embedchain/cli.py @@ -14,13 +14,21 @@ import requests from rich.console import Console from embedchain.telemetry.posthog import AnonymousTelemetry -from embedchain.utils.cli import (deploy_fly, deploy_gradio_app, - deploy_hf_spaces, deploy_modal, - deploy_render, deploy_streamlit, - get_pkg_path_from_name, setup_fly_io_app, - setup_gradio_app, setup_hf_app, - setup_modal_com_app, setup_render_com_app, - setup_streamlit_io_app) +from embedchain.utils.cli import ( + deploy_fly, + deploy_gradio_app, + deploy_hf_spaces, + deploy_modal, + deploy_render, + deploy_streamlit, + get_pkg_path_from_name, + setup_fly_io_app, + setup_gradio_app, + setup_hf_app, + setup_modal_com_app, + setup_render_com_app, + setup_streamlit_io_app, +) console = Console() api_process = None diff --git a/embedchain/embedchain/config/__init__.py b/embedchain/embedchain/config/__init__.py index f497b2e6..768408b7 100644 --- a/embedchain/embedchain/config/__init__.py +++ b/embedchain/embedchain/config/__init__.py @@ -8,8 +8,8 @@ from .embedder.base import BaseEmbedderConfig from .embedder.base import BaseEmbedderConfig as EmbedderConfig from .embedder.ollama import OllamaEmbedderConfig from .llm.base import BaseLlmConfig +from .mem0_config import Mem0Config from .vector_db.chroma import ChromaDbConfig from .vector_db.elasticsearch import ElasticsearchDBConfig from .vector_db.opensearch import OpenSearchDBConfig from .vector_db.zilliz import ZillizDBConfig -from .mem0_config import Mem0Config diff --git a/embedchain/embedchain/config/evaluation/__init__.py b/embedchain/embedchain/config/evaluation/__init__.py index ebbfcd0d..67e78dad 100644 --- a/embedchain/embedchain/config/evaluation/__init__.py +++ b/embedchain/embedchain/config/evaluation/__init__.py @@ -1,2 +1,5 @@ -from .base import (AnswerRelevanceConfig, ContextRelevanceConfig, # noqa: F401 - GroundednessConfig) +from .base import ( # noqa: F401 + AnswerRelevanceConfig, + ContextRelevanceConfig, + GroundednessConfig, +) diff --git a/embedchain/embedchain/config/llm/base.py b/embedchain/embedchain/config/llm/base.py index 77e892b7..693d09c5 100644 --- a/embedchain/embedchain/config/llm/base.py +++ b/embedchain/embedchain/config/llm/base.py @@ -1,9 +1,9 @@ import json import logging import re -from string import Template -from typing import Any, Mapping, Optional, Dict, Union from pathlib import Path +from string import Template +from typing import Any, Dict, Mapping, Optional, Union import httpx diff --git a/embedchain/embedchain/embedchain.py b/embedchain/embedchain/embedchain.py index 661b3c8a..f7271d50 100644 --- a/embedchain/embedchain/embedchain.py +++ b/embedchain/embedchain/embedchain.py @@ -6,7 +6,12 @@ from typing import Any, Optional, Union from dotenv import load_dotenv from langchain.docstore.document import Document -from embedchain.cache import adapt, get_gptcache_session, gptcache_data_convert, gptcache_update_cache_callback +from embedchain.cache import ( + adapt, + get_gptcache_session, + gptcache_data_convert, + gptcache_update_cache_callback, +) from embedchain.chunkers.base_chunker import BaseChunker from embedchain.config import AddConfig, BaseLlmConfig, ChunkerConfig from embedchain.config.base_app_config import BaseAppConfig @@ -16,7 +21,12 @@ from embedchain.embedder.base import BaseEmbedder 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, IndirectDataType, SpecialDataType +from embedchain.models.data_type import ( + DataType, + DirectDataType, + IndirectDataType, + SpecialDataType, +) from embedchain.utils.misc import detect_datatype, is_valid_json_string from embedchain.vectordb.base import BaseVectorDB diff --git a/embedchain/embedchain/embedder/clarifai.py b/embedchain/embedchain/embedder/clarifai.py index 6f0b7c22..8f0bb2fe 100644 --- a/embedchain/embedchain/embedder/clarifai.py +++ b/embedchain/embedchain/embedder/clarifai.py @@ -1,18 +1,18 @@ import os from typing import Optional, Union +from chromadb import EmbeddingFunction, Embeddings + from embedchain.config import BaseEmbedderConfig from embedchain.embedder.base import BaseEmbedder -from chromadb import EmbeddingFunction, Embeddings - class ClarifaiEmbeddingFunction(EmbeddingFunction): def __init__(self, config: BaseEmbedderConfig) -> None: super().__init__() try: - from clarifai.client.model import Model from clarifai.client.input import Inputs + from clarifai.client.model import Model except ModuleNotFoundError: raise ModuleNotFoundError( "The required dependencies for ClarifaiEmbeddingFunction are not installed." diff --git a/embedchain/embedchain/embedder/gpt4all.py b/embedchain/embedchain/embedder/gpt4all.py index 95d7fd62..83123f49 100644 --- a/embedchain/embedchain/embedder/gpt4all.py +++ b/embedchain/embedchain/embedder/gpt4all.py @@ -9,7 +9,9 @@ class GPT4AllEmbedder(BaseEmbedder): def __init__(self, config: Optional[BaseEmbedderConfig] = None): super().__init__(config=config) - from langchain_community.embeddings import GPT4AllEmbeddings as LangchainGPT4AllEmbeddings + from langchain_community.embeddings import ( + GPT4AllEmbeddings as LangchainGPT4AllEmbeddings, + ) model_name = self.config.model or "all-MiniLM-L6-v2-f16.gguf" gpt4all_kwargs = {'allow_download': 'True'} diff --git a/embedchain/embedchain/embedder/openai.py b/embedchain/embedchain/embedder/openai.py index fc2c7d63..d9f4c742 100644 --- a/embedchain/embedchain/embedder/openai.py +++ b/embedchain/embedchain/embedder/openai.py @@ -3,7 +3,6 @@ from typing import Optional from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction - from embedchain.config import BaseEmbedderConfig from embedchain.embedder.base import BaseEmbedder from embedchain.models import VectorDimensions diff --git a/embedchain/embedchain/llm/aws_bedrock.py b/embedchain/embedchain/llm/aws_bedrock.py index 89b62032..7f916268 100644 --- a/embedchain/embedchain/llm/aws_bedrock.py +++ b/embedchain/embedchain/llm/aws_bedrock.py @@ -45,8 +45,9 @@ class AWSBedrockLlm(BaseLlm): } if config.stream: - from langchain.callbacks.streaming_stdout import \ - StreamingStdOutCallbackHandler + from langchain.callbacks.streaming_stdout import ( + StreamingStdOutCallbackHandler, + ) kwargs["streaming"] = True kwargs["callbacks"] = [StreamingStdOutCallbackHandler()] diff --git a/embedchain/embedchain/llm/base.py b/embedchain/embedchain/llm/base.py index bcbed4c8..ace4bb79 100644 --- a/embedchain/embedchain/llm/base.py +++ b/embedchain/embedchain/llm/base.py @@ -5,7 +5,6 @@ from typing import Any, Optional from langchain.schema import BaseMessage as LCBaseMessage -from embedchain.constants import SQLITE_PATH from embedchain.config import BaseLlmConfig from embedchain.config.llm.base import ( DEFAULT_PROMPT, @@ -13,6 +12,7 @@ from embedchain.config.llm.base import ( DEFAULT_PROMPT_WITH_MEM0_MEMORY_TEMPLATE, DOCS_SITE_PROMPT_TEMPLATE, ) +from embedchain.constants import SQLITE_PATH from embedchain.core.db.database import init_db, setup_engine from embedchain.helpers.json_serializable import JSONSerializable from embedchain.memory.base import ChatHistory diff --git a/embedchain/embedchain/llm/gpt4all.py b/embedchain/embedchain/llm/gpt4all.py index 833f3e26..76062b08 100644 --- a/embedchain/embedchain/llm/gpt4all.py +++ b/embedchain/embedchain/llm/gpt4all.py @@ -26,8 +26,7 @@ class GPT4ALLLlm(BaseLlm): @staticmethod def _get_instance(model): try: - from langchain_community.llms.gpt4all import \ - GPT4All as LangchainGPT4All + from langchain_community.llms.gpt4all import GPT4All as LangchainGPT4All except ModuleNotFoundError: raise ModuleNotFoundError( "The GPT4All python package is not installed. Please install it with `pip install --upgrade embedchain[opensource]`" # noqa E501 diff --git a/embedchain/embedchain/llm/jina.py b/embedchain/embedchain/llm/jina.py index 33796bc5..ac3a0e76 100644 --- a/embedchain/embedchain/llm/jina.py +++ b/embedchain/embedchain/llm/jina.py @@ -35,8 +35,9 @@ class JinaLlm(BaseLlm): if config.top_p: kwargs["model_kwargs"]["top_p"] = config.top_p if config.stream: - from langchain.callbacks.streaming_stdout import \ - StreamingStdOutCallbackHandler + from langchain.callbacks.streaming_stdout import ( + StreamingStdOutCallbackHandler, + ) chat = JinaChat(**kwargs, streaming=config.stream, callbacks=[StreamingStdOutCallbackHandler()]) else: diff --git a/embedchain/embedchain/loaders/pdf_file.py b/embedchain/embedchain/loaders/pdf_file.py index 78f93b5d..a7f6d554 100644 --- a/embedchain/embedchain/loaders/pdf_file.py +++ b/embedchain/embedchain/loaders/pdf_file.py @@ -1,6 +1,7 @@ import hashlib from langchain_community.document_loaders import PyPDFLoader + from embedchain.helpers.json_serializable import register_deserializable from embedchain.loaders.base_loader import BaseLoader from embedchain.utils.misc import clean_string diff --git a/embedchain/embedchain/loaders/rss_feed.py b/embedchain/embedchain/loaders/rss_feed.py index a4b3376f..bc17c68b 100644 --- a/embedchain/embedchain/loaders/rss_feed.py +++ b/embedchain/embedchain/loaders/rss_feed.py @@ -28,7 +28,9 @@ class RSSFeedLoader(BaseLoader): @staticmethod def get_rss_content(url: str): try: - from langchain_community.document_loaders import RSSFeedLoader as LangchainRSSFeedLoader + from langchain_community.document_loaders import ( + RSSFeedLoader as LangchainRSSFeedLoader, + ) except ImportError: raise ImportError( """RSSFeedLoader file requires extra dependencies. diff --git a/embedchain/embedchain/loaders/unstructured_file.py b/embedchain/embedchain/loaders/unstructured_file.py index 62c20ae4..856ac888 100644 --- a/embedchain/embedchain/loaders/unstructured_file.py +++ b/embedchain/embedchain/loaders/unstructured_file.py @@ -11,8 +11,7 @@ class UnstructuredLoader(BaseLoader): """Load data from an Unstructured file.""" try: import unstructured # noqa: F401 - from langchain_community.document_loaders import \ - UnstructuredFileLoader + from langchain_community.document_loaders import UnstructuredFileLoader except ImportError: raise ImportError( 'Unstructured file requires extra dependencies. Install with `pip install "unstructured[local-inference, all-docs]"`' # noqa: E501 diff --git a/embedchain/embedchain/vectordb/zilliz.py b/embedchain/embedchain/vectordb/zilliz.py index 0a30de9d..ca554473 100644 --- a/embedchain/embedchain/vectordb/zilliz.py +++ b/embedchain/embedchain/vectordb/zilliz.py @@ -6,8 +6,15 @@ from embedchain.helpers.json_serializable import register_deserializable from embedchain.vectordb.base import BaseVectorDB try: - from pymilvus import (Collection, CollectionSchema, DataType, FieldSchema, - MilvusClient, connections, utility) + from pymilvus import ( + Collection, + CollectionSchema, + DataType, + FieldSchema, + MilvusClient, + connections, + utility, + ) except ImportError: raise ImportError( "Zilliz requires extra dependencies. Install with `pip install --upgrade embedchain[milvus]`" diff --git a/embedchain/examples/chat-pdf/app.py b/embedchain/examples/chat-pdf/app.py index 0b2fefae..fb29b73e 100644 --- a/embedchain/examples/chat-pdf/app.py +++ b/embedchain/examples/chat-pdf/app.py @@ -8,8 +8,7 @@ import streamlit as st from embedchain import App from embedchain.config import BaseLlmConfig -from embedchain.helpers.callbacks import (StreamingStdOutCallbackHandlerYield, - generate) +from embedchain.helpers.callbacks import StreamingStdOutCallbackHandlerYield, generate def embedchain_bot(db_path, api_key): diff --git a/embedchain/examples/sadhguru-ai/app.py b/embedchain/examples/sadhguru-ai/app.py index 67d4c62d..5c123d60 100644 --- a/embedchain/examples/sadhguru-ai/app.py +++ b/embedchain/examples/sadhguru-ai/app.py @@ -8,8 +8,7 @@ import streamlit as st from embedchain import App from embedchain.config import BaseLlmConfig -from embedchain.helpers.callbacks import (StreamingStdOutCallbackHandlerYield, - generate) +from embedchain.helpers.callbacks import StreamingStdOutCallbackHandlerYield, generate @st.cache_resource diff --git a/embedchain/examples/unacademy-ai/app.py b/embedchain/examples/unacademy-ai/app.py index 5e552f1d..e31dba91 100644 --- a/embedchain/examples/unacademy-ai/app.py +++ b/embedchain/examples/unacademy-ai/app.py @@ -4,8 +4,7 @@ import streamlit as st from embedchain import App from embedchain.config import BaseLlmConfig -from embedchain.helpers.callbacks import (StreamingStdOutCallbackHandlerYield, - generate) +from embedchain.helpers.callbacks import StreamingStdOutCallbackHandlerYield, generate @st.cache_resource diff --git a/embedchain/tests/embedder/test_azure_openai_embedder.py b/embedchain/tests/embedder/test_azure_openai_embedder.py index 77357198..2667d01f 100644 --- a/embedchain/tests/embedder/test_azure_openai_embedder.py +++ b/embedchain/tests/embedder/test_azure_openai_embedder.py @@ -1,4 +1,4 @@ -from unittest.mock import patch, Mock +from unittest.mock import Mock, patch import httpx diff --git a/embedchain/tests/embedder/test_huggingface_embedder.py b/embedchain/tests/embedder/test_huggingface_embedder.py index 4760b926..ed97ccc9 100644 --- a/embedchain/tests/embedder/test_huggingface_embedder.py +++ b/embedchain/tests/embedder/test_huggingface_embedder.py @@ -1,5 +1,6 @@ from unittest.mock import patch + from embedchain.config import BaseEmbedderConfig from embedchain.embedder.huggingface import HuggingFaceEmbedder diff --git a/embedchain/tests/helper_classes/test_json_serializable.py b/embedchain/tests/helper_classes/test_json_serializable.py index adcfe5f7..e68e0507 100644 --- a/embedchain/tests/helper_classes/test_json_serializable.py +++ b/embedchain/tests/helper_classes/test_json_serializable.py @@ -4,8 +4,10 @@ from string import Template from embedchain import App from embedchain.config import AppConfig, BaseLlmConfig -from embedchain.helpers.json_serializable import (JSONSerializable, - register_deserializable) +from embedchain.helpers.json_serializable import ( + JSONSerializable, + register_deserializable, +) class TestJsonSerializable(unittest.TestCase): diff --git a/embedchain/tests/llm/conftest.py b/embedchain/tests/llm/conftest.py index edfafb95..6e3da3d5 100644 --- a/embedchain/tests/llm/conftest.py +++ b/embedchain/tests/llm/conftest.py @@ -1,7 +1,9 @@ from unittest import mock + import pytest + @pytest.fixture(autouse=True) def mock_alembic_command_upgrade(): with mock.patch("alembic.command.upgrade"): diff --git a/embedchain/tests/llm/test_azure_openai.py b/embedchain/tests/llm/test_azure_openai.py index 6aa69e6e..f5ac3874 100644 --- a/embedchain/tests/llm/test_azure_openai.py +++ b/embedchain/tests/llm/test_azure_openai.py @@ -1,4 +1,4 @@ -from unittest.mock import Mock, MagicMock, patch +from unittest.mock import MagicMock, Mock, patch import httpx import pytest diff --git a/embedchain/tests/llm/test_ollama.py b/embedchain/tests/llm/test_ollama.py index e45cb6bb..b0d93263 100644 --- a/embedchain/tests/llm/test_ollama.py +++ b/embedchain/tests/llm/test_ollama.py @@ -1,8 +1,8 @@ import pytest +from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler from embedchain.config import BaseLlmConfig from embedchain.llm.ollama import OllamaLlm -from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler @pytest.fixture diff --git a/embedchain/tests/models/test_data_type.py b/embedchain/tests/models/test_data_type.py index bf3d6e1e..60d66282 100644 --- a/embedchain/tests/models/test_data_type.py +++ b/embedchain/tests/models/test_data_type.py @@ -1,5 +1,9 @@ -from embedchain.models.data_type import (DataType, DirectDataType, - IndirectDataType, SpecialDataType) +from embedchain.models.data_type import ( + DataType, + DirectDataType, + IndirectDataType, + SpecialDataType, +) def test_subclass_types_in_data_type(): diff --git a/mem0/__init__.py b/mem0/__init__.py index 0433062e..66632ea0 100644 --- a/mem0/__init__.py +++ b/mem0/__init__.py @@ -2,5 +2,5 @@ import importlib.metadata __version__ = importlib.metadata.version("mem0ai") -from mem0.memory.main import Memory # noqa from mem0.client.main import MemoryClient # noqa +from mem0.memory.main import Memory # noqa diff --git a/mem0/configs/base.py b/mem0/configs/base.py index ee2a333d..42a0a2b2 100644 --- a/mem0/configs/base.py +++ b/mem0/configs/base.py @@ -1,12 +1,13 @@ import os from typing import Any, Dict, Optional + from pydantic import BaseModel, Field -from mem0.memory.setup import mem0_dir -from mem0.vector_stores.configs import VectorStoreConfig -from mem0.llms.configs import LlmConfig from mem0.embeddings.configs import EmbedderConfig from mem0.graphs.configs import GraphStoreConfig +from mem0.llms.configs import LlmConfig +from mem0.memory.setup import mem0_dir +from mem0.vector_stores.configs import VectorStoreConfig class MemoryItem(BaseModel): diff --git a/mem0/configs/embeddings/base.py b/mem0/configs/embeddings/base.py index 7014ea25..f4659dce 100644 --- a/mem0/configs/embeddings/base.py +++ b/mem0/configs/embeddings/base.py @@ -1,9 +1,10 @@ from abc import ABC -from mem0.configs.base import AzureConfig -from typing import Optional, Union, Dict +from typing import Dict, Optional, Union import httpx +from mem0.configs.base import AzureConfig + class BaseEmbedderConfig(ABC): """ diff --git a/mem0/configs/llms/base.py b/mem0/configs/llms/base.py index 0cf7cbb1..94c8e278 100644 --- a/mem0/configs/llms/base.py +++ b/mem0/configs/llms/base.py @@ -1,9 +1,10 @@ from abc import ABC -from mem0.configs.base import AzureConfig -from typing import Optional, Union, Dict +from typing import Dict, Optional, Union import httpx +from mem0.configs.base import AzureConfig + class BaseLlmConfig(ABC): """ diff --git a/mem0/configs/vector_stores/chroma.py b/mem0/configs/vector_stores/chroma.py index 23e182f1..ba58b3ac 100644 --- a/mem0/configs/vector_stores/chroma.py +++ b/mem0/configs/vector_stores/chroma.py @@ -1,6 +1,6 @@ import subprocess import sys -from typing import Optional, ClassVar, Dict, Any +from typing import Any, ClassVar, Dict, Optional from pydantic import BaseModel, Field, model_validator diff --git a/mem0/configs/vector_stores/pgvector.py b/mem0/configs/vector_stores/pgvector.py index 1ed59c75..df8dabf4 100644 --- a/mem0/configs/vector_stores/pgvector.py +++ b/mem0/configs/vector_stores/pgvector.py @@ -1,4 +1,4 @@ -from typing import Optional, Dict, Any +from typing import Any, Dict, Optional from pydantic import BaseModel, Field, model_validator diff --git a/mem0/configs/vector_stores/qdrant.py b/mem0/configs/vector_stores/qdrant.py index bb76cac3..10951db8 100644 --- a/mem0/configs/vector_stores/qdrant.py +++ b/mem0/configs/vector_stores/qdrant.py @@ -1,5 +1,6 @@ +from typing import Any, ClassVar, Dict, Optional + from pydantic import BaseModel, Field, model_validator -from typing import Optional, ClassVar, Dict, Any class QdrantConfig(BaseModel): diff --git a/mem0/embeddings/base.py b/mem0/embeddings/base.py index ef6800f3..c5d12b6e 100644 --- a/mem0/embeddings/base.py +++ b/mem0/embeddings/base.py @@ -1,5 +1,5 @@ -from typing import Optional from abc import ABC, abstractmethod +from typing import Optional from mem0.configs.embeddings.base import BaseEmbedderConfig diff --git a/mem0/graphs/configs.py b/mem0/graphs/configs.py index a93d27d9..033637c3 100644 --- a/mem0/graphs/configs.py +++ b/mem0/graphs/configs.py @@ -1,7 +1,10 @@ from typing import Optional + from pydantic import BaseModel, Field, field_validator, model_validator + from mem0.llms.configs import LlmConfig + class Neo4jConfig(BaseModel): url: Optional[str] = Field(None, description="Host address for the graph database") username: Optional[str] = Field(None, description="Username for the graph database") diff --git a/mem0/llms/anthropic.py b/mem0/llms/anthropic.py index 5e32d81c..fb390348 100644 --- a/mem0/llms/anthropic.py +++ b/mem0/llms/anthropic.py @@ -1,7 +1,4 @@ -import subprocess -import sys import os -import json from typing import Dict, List, Optional try: @@ -9,8 +6,8 @@ try: except ImportError: raise ImportError("The 'anthropic' library is required. Please install it using 'pip install anthropic'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class AnthropicLLM(LLMBase): diff --git a/mem0/llms/aws_bedrock.py b/mem0/llms/aws_bedrock.py index b5699a9e..5e7969c1 100644 --- a/mem0/llms/aws_bedrock.py +++ b/mem0/llms/aws_bedrock.py @@ -1,16 +1,14 @@ -import subprocess -import sys -import os import json -from typing import Dict, List, Optional, Any +import os +from typing import Any, Dict, List, Optional try: import boto3 except ImportError: raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class AWSBedrockLLM(LLMBase): diff --git a/mem0/llms/azure_openai.py b/mem0/llms/azure_openai.py index d18569a8..f093284b 100644 --- a/mem0/llms/azure_openai.py +++ b/mem0/llms/azure_openai.py @@ -1,11 +1,11 @@ -import os import json +import os from typing import Dict, List, Optional from openai import AzureOpenAI -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class AzureOpenAILLM(LLMBase): diff --git a/mem0/llms/base.py b/mem0/llms/base.py index 69e880b4..a1ed3c79 100644 --- a/mem0/llms/base.py +++ b/mem0/llms/base.py @@ -1,5 +1,5 @@ -from typing import Optional from abc import ABC, abstractmethod +from typing import Optional from mem0.configs.llms.base import BaseLlmConfig diff --git a/mem0/llms/groq.py b/mem0/llms/groq.py index c67436c0..e970a1ee 100644 --- a/mem0/llms/groq.py +++ b/mem0/llms/groq.py @@ -1,7 +1,5 @@ -import subprocess -import sys -import os import json +import os from typing import Dict, List, Optional try: @@ -9,8 +7,8 @@ try: except ImportError: raise ImportError("The 'groq' library is required. Please install it using 'pip install groq'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class GroqLLM(LLMBase): diff --git a/mem0/llms/litellm.py b/mem0/llms/litellm.py index 1a6fb7c0..cfdcba34 100644 --- a/mem0/llms/litellm.py +++ b/mem0/llms/litellm.py @@ -1,5 +1,3 @@ -import subprocess -import sys import json from typing import Dict, List, Optional @@ -8,8 +6,8 @@ try: except ImportError: raise ImportError("The 'litellm' library is required. Please install it using 'pip install litellm'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class LiteLLM(LLMBase): diff --git a/mem0/llms/ollama.py b/mem0/llms/ollama.py index e7acbdd7..6750f0dd 100644 --- a/mem0/llms/ollama.py +++ b/mem0/llms/ollama.py @@ -1,5 +1,3 @@ -import subprocess -import sys from typing import Dict, List, Optional try: @@ -7,8 +5,8 @@ try: except ImportError: raise ImportError("The 'ollama' library is required. Please install it using 'pip install ollama'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class OllamaLLM(LLMBase): diff --git a/mem0/llms/openai.py b/mem0/llms/openai.py index 68a36345..b988c33b 100644 --- a/mem0/llms/openai.py +++ b/mem0/llms/openai.py @@ -1,11 +1,11 @@ -import os import json +import os from typing import Dict, List, Optional from openai import OpenAI -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class OpenAILLM(LLMBase): diff --git a/mem0/llms/openai_structured.py b/mem0/llms/openai_structured.py index d2ea9c4e..ab56b85f 100644 --- a/mem0/llms/openai_structured.py +++ b/mem0/llms/openai_structured.py @@ -1,10 +1,11 @@ -import os, json +import json +import os from typing import Dict, List, Optional from openai import OpenAI -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class OpenAIStructuredLLM(LLMBase): diff --git a/mem0/llms/together.py b/mem0/llms/together.py index 816ed7a9..51ebac66 100644 --- a/mem0/llms/together.py +++ b/mem0/llms/together.py @@ -1,7 +1,5 @@ -import subprocess -import sys -import os import json +import os from typing import Dict, List, Optional try: @@ -9,8 +7,8 @@ try: except ImportError: raise ImportError("The 'together' library is required. Please install it using 'pip install together'.") -from mem0.llms.base import LLMBase from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.base import LLMBase class TogetherLLM(LLMBase): diff --git a/mem0/memory/graph_memory.py b/mem0/memory/graph_memory.py index 71352c94..9ec52e6b 100644 --- a/mem0/memory/graph_memory.py +++ b/mem0/memory/graph_memory.py @@ -1,10 +1,17 @@ -import json import logging + from langchain_community.graphs import Neo4jGraph from rank_bm25 import BM25Okapi -from mem0.utils.factory import LlmFactory, EmbedderFactory -from mem0.graphs.utils import get_update_memory_messages, EXTRACT_ENTITIES_PROMPT -from mem0.graphs.tools import UPDATE_MEMORY_TOOL_GRAPH, ADD_MEMORY_TOOL_GRAPH, NOOP_TOOL, ADD_MESSAGE_TOOL, SEARCH_TOOL + +from mem0.graphs.tools import ( + ADD_MEMORY_TOOL_GRAPH, + ADD_MESSAGE_TOOL, + NOOP_TOOL, + SEARCH_TOOL, + UPDATE_MEMORY_TOOL_GRAPH, +) +from mem0.graphs.utils import EXTRACT_ENTITIES_PROMPT, get_update_memory_messages +from mem0.utils.factory import EmbedderFactory, LlmFactory logger = logging.getLogger(__name__) diff --git a/mem0/memory/main.py b/mem0/memory/main.py index d5744b68..f158ec82 100644 --- a/mem0/memory/main.py +++ b/mem0/memory/main.py @@ -1,22 +1,24 @@ -import logging +import concurrent import hashlib -import uuid -import pytz import json +import logging +import threading +import uuid +import warnings from datetime import datetime from typing import Any, Dict -import warnings + +import pytz from pydantic import ValidationError + +from mem0.configs.base import MemoryConfig, MemoryItem +from mem0.configs.prompts import get_update_memory_messages from mem0.memory.base import MemoryBase from mem0.memory.setup import setup_config from mem0.memory.storage import SQLiteManager from mem0.memory.telemetry import capture_event from mem0.memory.utils import get_fact_retrieval_messages, parse_messages -from mem0.configs.prompts import get_update_memory_messages -from mem0.utils.factory import LlmFactory, EmbedderFactory, VectorStoreFactory -from mem0.configs.base import MemoryItem, MemoryConfig -import threading -import concurrent +from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory # Setup user config setup_config() @@ -154,7 +156,7 @@ class Memory(MemoryBase): logging.info(resp) try: if resp["event"] == "ADD": - memory_id = self._create_memory(data=resp["text"], metadata=metadata) + self._create_memory(data=resp["text"], metadata=metadata) elif resp["event"] == "UPDATE": self._update_memory(memory_id=resp["id"], data=resp["text"], metadata=metadata) elif resp["event"] == "DELETE": @@ -175,7 +177,7 @@ class Memory(MemoryBase): else: self.graph.user_id = "USER" data = "\n".join([msg["content"] for msg in messages if "content" in msg and msg["role"] != "system"]) - added_entities = self.graph.add(data, filters) + self.graph.add(data, filters) def get(self, memory_id): """ diff --git a/mem0/memory/telemetry.py b/mem0/memory/telemetry.py index 70336f78..ec819e93 100644 --- a/mem0/memory/telemetry.py +++ b/mem0/memory/telemetry.py @@ -1,3 +1,4 @@ +import logging import platform import sys @@ -5,6 +6,8 @@ from posthog import Posthog from mem0.memory.setup import get_user_id, setup_config +logging.getLogger('posthog').setLevel(logging.CRITICAL + 1) +logging.getLogger('urllib3').setLevel(logging.CRITICAL + 1) class AnonymousTelemetry: def __init__(self, project_api_key, host): diff --git a/mem0/proxy/main.py b/mem0/proxy/main.py index daf24ae1..bb614f4f 100644 --- a/mem0/proxy/main.py +++ b/mem0/proxy/main.py @@ -1,9 +1,10 @@ import logging import subprocess import sys -import httpx -from typing import Optional, List, Union import threading +from typing import List, Optional, Union + +import httpx try: import litellm @@ -20,9 +21,9 @@ except ImportError: raise ImportError("The required 'litellm' library is not installed.") sys.exit(1) -from mem0.memory.telemetry import capture_client_event from mem0 import Memory, MemoryClient from mem0.configs.prompts import MEMORY_ANSWER_PROMPT +from mem0.memory.telemetry import capture_client_event logger = logging.getLogger(__name__) diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index c672cf1d..aca433b0 100644 --- a/mem0/utils/factory.py +++ b/mem0/utils/factory.py @@ -1,7 +1,7 @@ import importlib -from mem0.configs.llms.base import BaseLlmConfig from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.configs.llms.base import BaseLlmConfig def load_class(class_type): diff --git a/mem0/vector_stores/chroma.py b/mem0/vector_stores/chroma.py index 21fc85c6..0dc97a3f 100644 --- a/mem0/vector_stores/chroma.py +++ b/mem0/vector_stores/chroma.py @@ -1,7 +1,5 @@ -import subprocess -import sys import logging -from typing import Optional, List, Dict +from typing import Dict, List, Optional from pydantic import BaseModel diff --git a/mem0/vector_stores/configs.py b/mem0/vector_stores/configs.py index 3b2eb293..2f052da5 100644 --- a/mem0/vector_stores/configs.py +++ b/mem0/vector_stores/configs.py @@ -1,4 +1,5 @@ -from typing import Optional, Dict +from typing import Dict, Optional + from pydantic import BaseModel, Field, model_validator diff --git a/mem0/vector_stores/pgvector.py b/mem0/vector_stores/pgvector.py index dae1e2d1..f9ec3f97 100644 --- a/mem0/vector_stores/pgvector.py +++ b/mem0/vector_stores/pgvector.py @@ -1,8 +1,7 @@ -import subprocess -import sys import json import logging -from typing import Optional, List +from typing import List, Optional + from pydantic import BaseModel try: diff --git a/mem0/vector_stores/qdrant.py b/mem0/vector_stores/qdrant.py index 0c7000b8..3ecb93f9 100644 --- a/mem0/vector_stores/qdrant.py +++ b/mem0/vector_stores/qdrant.py @@ -1,6 +1,6 @@ +import logging import os import shutil -import logging from qdrant_client import QdrantClient from qdrant_client.models import ( diff --git a/tests/llms/test_azure_openai.py b/tests/llms/test_azure_openai.py index effa19a1..63eb91b0 100644 --- a/tests/llms/test_azure_openai.py +++ b/tests/llms/test_azure_openai.py @@ -1,10 +1,11 @@ +from unittest.mock import Mock, patch + import httpx import pytest -from unittest.mock import Mock, patch -from mem0.llms.azure_openai import AzureOpenAILLM from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.azure_openai import AzureOpenAILLM MODEL = "gpt-4o" # or your custom deployment name TEMPERATURE = 0.7 diff --git a/tests/llms/test_groq.py b/tests/llms/test_groq.py index 52c3b222..e7d1f51c 100644 --- a/tests/llms/test_groq.py +++ b/tests/llms/test_groq.py @@ -1,7 +1,10 @@ -import pytest from unittest.mock import Mock, patch -from mem0.llms.groq import GroqLLM + +import pytest + from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.groq import GroqLLM + @pytest.fixture def mock_groq_client(): diff --git a/tests/llms/test_litellm.py b/tests/llms/test_litellm.py index b611d5c2..f4b265aa 100644 --- a/tests/llms/test_litellm.py +++ b/tests/llms/test_litellm.py @@ -1,8 +1,10 @@ -import pytest from unittest.mock import Mock, patch -from mem0.llms import litellm +import pytest + from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms import litellm + @pytest.fixture def mock_litellm(): diff --git a/tests/llms/test_ollama.py b/tests/llms/test_ollama.py index a781971c..d99fd2bc 100644 --- a/tests/llms/test_ollama.py +++ b/tests/llms/test_ollama.py @@ -1,9 +1,12 @@ -import pytest from unittest.mock import Mock, patch -from mem0.llms.ollama import OllamaLLM + +import pytest + from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.ollama import OllamaLLM from mem0.llms.utils.tools import ADD_MEMORY_TOOL + @pytest.fixture def mock_ollama_client(): with patch('mem0.llms.ollama.Client') as mock_ollama: diff --git a/tests/llms/test_openai.py b/tests/llms/test_openai.py index 81005079..204487c4 100644 --- a/tests/llms/test_openai.py +++ b/tests/llms/test_openai.py @@ -1,7 +1,10 @@ -import pytest from unittest.mock import Mock, patch -from mem0.llms.openai import OpenAILLM + +import pytest + from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.openai import OpenAILLM + @pytest.fixture def mock_openai_client(): diff --git a/tests/llms/test_together.py b/tests/llms/test_together.py index 1353416f..f317d106 100644 --- a/tests/llms/test_together.py +++ b/tests/llms/test_together.py @@ -1,7 +1,10 @@ -import pytest from unittest.mock import Mock, patch -from mem0.llms.together import TogetherLLM + +import pytest + from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.together import TogetherLLM + @pytest.fixture def mock_together_client(): diff --git a/tests/test_proxy.py b/tests/test_proxy.py index c8fe566f..c7776b84 100644 --- a/tests/test_proxy.py +++ b/tests/test_proxy.py @@ -1,10 +1,11 @@ -import pytest from unittest.mock import Mock, patch -from mem0.configs.prompts import MEMORY_ANSWER_PROMPT +import pytest + from mem0 import Memory, MemoryClient -from mem0.proxy.main import Mem0 -from mem0.proxy.main import Chat, Completions +from mem0.configs.prompts import MEMORY_ANSWER_PROMPT +from mem0.proxy.main import Chat, Completions, Mem0 + @pytest.fixture def mock_memory_client():