ronantakizawa/codereview-bench huggingface dataset:
Repository: meilisearch/meilisearch
File: crates/meilisearch/src/search/federated/network.rs (Rust)

Code before:
use meilisearch_types::index_uid::IndexUid;
use meilisearch_types::network::Network;

use crate::search::{Federation, FederationOptions, SearchQuery, SearchQueryWithIndex};

pub fn network_partition<'a>(
    federation: &mut Federation,
    query: &'a SearchQuery,
    federation_options: Option<FederationOptions>,
    index_uid: &'a IndexUid,
    network: Network,
) -> impl Iterator<Item = SearchQueryWithIndex> + 'a {
    let federation_options = federation_options.unwrap_or_default();
    let mut query = SearchQueryWithIndex::from_index_query_federation(
        index_uid.clone(),
        query.clone(),
        Some(federation_options),
    );

    // Move federation options from the query to the federation struct
    let SearchQueryWithIndex {
        index_uid,
        q: _,
        vector: _,
        media: _,
        hybrid: _,
        offset,
        limit,
        page,
        hits_per_page,
        attributes_to_retrieve: _,
        retrieve_vectors: _,
        attributes_to_crop: _,
        crop_length: _,
        attributes_to_highlight: _,
        show_ranking_score: _,
        show_ranking_score_details: _,
        show_performance_details,
        use_network: _,
        show_matches_position: _,
        filter: _,
        sort: _,
        distinct: _,
        facets,
        highlight_pre_tag: _,

Review comment (suggestion):
```suggestion
    // Move query parameters that make sense at the federation level
    // from the `SearchQueryWithIndex` to the `Federation`
```

Code after:
use meilisearch_types::index_uid::IndexUid;
use meilisearch_types::network::Network;

use crate::search::{Federation, FederationOptions, SearchQuery, SearchQueryWithIndex};

pub fn network_partition<'a>(
    federation: &mut Federation,
    query: &'a SearchQuery,
    federation_options: Option<FederationOptions>,
    index_uid: &'a IndexUid,
    network: Network,
) -> impl Iterator<Item = SearchQueryWithIndex> + 'a {
    let federation_options = federation_options.unwrap_or_default();
    let mut query = SearchQueryWithIndex::from_index_query_federation(
        index_uid.clone(),
        query.clone(),
        Some(federation_options),
    );

    // Move query parameters that make sense at the federation level
    // from the `SearchQueryWithIndex` to the `Federation`
    let SearchQueryWithIndex {
        index_uid,
        q: _,
        vector: _,
        media: _,
        hybrid: _,
        offset,
        limit,
        page,
        hits_per_page,
        attributes_to_retrieve: _,
        retrieve_vectors: _,
        attributes_to_crop: _,
        crop_length: _,
        attributes_to_highlight: _,
        show_ranking_score: _,
        show_ranking_score_details: _,
        show_performance_details,
        use_network: _,
        show_matches_position: _,
        filter: _,
        sort: _,
        distinct: _,
        facets,
Repository: mem0ai/mem0
File: embedchain/vectordb/chroma_db.py (Python)

Code before:
import logging
import os  # noqa: F401

import chromadb
from chromadb.utils import embedding_functions  # noqa: F401

from embedchain.vectordb.base_vector_db import BaseVectorDB


class ChromaDB(BaseVectorDB):
    """Vector database using ChromaDB."""

    def __init__(self, db_dir=None, ef=None, host=None, port=None):
        self.ef = ef

        if host and port:
            logging.info(f"Connecting to ChromaDB server: {host}:{port}")
            self.client_settings = chromadb.config.Settings(
                chroma_api_impl="rest",
                chroma_server_host=host,
                chroma_server_http_port=port,
            )
        else:
            if db_dir is None:
                db_dir = "db"
            self.client_settings = chromadb.config.Settings(
                chroma_db_impl="duckdb+parquet",
                persist_directory=db_dir,
                anonymized_telemetry=False,
            )

Review comment (suggestion):
```suggestion
from chromadb.utils import embedding_functions
```

Code after:
import logging
import os

import chromadb
from chromadb.utils import embedding_functions

from embedchain.vectordb.base_vector_db import BaseVectorDB


class ChromaDB(BaseVectorDB):
    """Vector database using ChromaDB."""

    def __init__(self, db_dir=None, ef=None, host=None, port=None):
        self.ef = ef

        if host and port:
            logging.info(f"Connecting to ChromaDB server: {host}:{port}")
            self.client_settings = chromadb.config.Settings(
                chroma_api_impl="rest",
                chroma_server_host=host,
                chroma_server_http_port=port,
            )
        else:
            if db_dir is None:
                db_dir = "db"
            self.client_settings = chromadb.config.Settings(
                chroma_db_impl="duckdb+parquet",
                persist_directory=db_dir,
                anonymized_telemetry=False,
            )
Repository: mem0ai/mem0
File: embedchain/llm/hugging_face_hub.py (Python)

Code before:
import importlib
import os
from typing import Optional

from langchain.llms import HuggingFaceHub

from embedchain.config import BaseLlmConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.llm.base import BaseLlm


@register_deserializable
class HuggingFaceHubLlm(BaseLlm):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        if "HUGGINGFACEHUB_API_TOKEN" not in os.environ:
            raise ValueError("Please set the HUGGINGFACEHUB_API_TOKEN environment variable.")

        try:
            importlib.import_module("huggingface_hub")
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "The required dependencies for HuggingFaceHub are not installed."
                'Please install with `pip install --upgrade "embedchain[huggingface_hub]"`'
            ) from None

        super().__init__(config=config)

    def get_llm_model_answer(self, prompt):
        if self.config.system_prompt:
            raise ValueError("HuggingFaceHubLlm does not support `system_prompt`")
        return HuggingFaceHubLlm._get_answer(prompt=prompt, config=self.config)

    @staticmethod
    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
        model_kwargs = {
            "temperature": config.temperature or 0.1,
            "max_new_tokens": config.max_tokens,
        }

        if config.top_p > 0.0 and config.top_p < 1.0:

Review comment (suggestion):
```suggestion
        if "HUGGINGFACEHUB_ACCESS_TOKEN" not in os.environ:
```

Code after:
import importlib
import os
from typing import Optional

from langchain.llms import HuggingFaceHub

from embedchain.config import BaseLlmConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.llm.base import BaseLlm


@register_deserializable
class HuggingFaceHubLlm(BaseLlm):
    def __init__(self, config: Optional[BaseLlmConfig] = None):
        if "HUGGINGFACEHUB_ACCESS_TOKEN" not in os.environ:
            raise ValueError("Please set the HUGGINGFACEHUB_ACCESS_TOKEN environment variable.")

        try:
            importlib.import_module("huggingface_hub")
        except ModuleNotFoundError:
            raise ModuleNotFoundError(
                "The required dependencies for HuggingFaceHub are not installed."
                'Please install with `pip install --upgrade "embedchain[huggingface_hub]"`'
            ) from None

        super().__init__(config=config)

    def get_llm_model_answer(self, prompt):
        if self.config.system_prompt:
            raise ValueError("HuggingFaceHubLlm does not support `system_prompt`")
        return HuggingFaceHubLlm._get_answer(prompt=prompt, config=self.config)

    @staticmethod
    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
        model_kwargs = {
            "temperature": config.temperature or 0.1,
            "max_new_tokens": config.max_tokens,
        }

        if config.top_p > 0.0 and config.top_p < 1.0:
Repository: mem0ai/mem0
File: embedchain/llm/hugging_face_hub.py (Python)

Code before:
            raise ModuleNotFoundError(
                "The required dependencies for HuggingFaceHub are not installed."
                'Please install with `pip install --upgrade "embedchain[huggingface_hub]"`'
            ) from None

        super().__init__(config=config)

    def get_llm_model_answer(self, prompt):
        if self.config.system_prompt:
            raise ValueError("HuggingFaceHubLlm does not support `system_prompt`")
        return HuggingFaceHubLlm._get_answer(prompt=prompt, config=self.config)

    @staticmethod
    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
        model_kwargs = {
            "temperature": config.temperature or 0.1,
            "max_new_tokens": config.max_tokens,
        }

        if config.top_p > 0.0 and config.top_p < 1.0:
            model_kwargs["top_p"] = config.top_p
        else:
            raise ValueError("`top_p` must be > 0.0 and < 1.0")

        llm = HuggingFaceHub(
            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
            repo_id=config.model or "google/flan-t5-xxl",
            model_kwargs=model_kwargs,
        )

        return llm(prompt)


Review comment (suggestion):
```suggestion
            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"],
```

Code after:
            raise ModuleNotFoundError(
                "The required dependencies for HuggingFaceHub are not installed."
                'Please install with `pip install --upgrade "embedchain[huggingface_hub]"`'
            ) from None

        super().__init__(config=config)

    def get_llm_model_answer(self, prompt):
        if self.config.system_prompt:
            raise ValueError("HuggingFaceHubLlm does not support `system_prompt`")
        return HuggingFaceHubLlm._get_answer(prompt=prompt, config=self.config)

    @staticmethod
    def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
        model_kwargs = {
            "temperature": config.temperature or 0.1,
            "max_new_tokens": config.max_tokens,
        }

        if config.top_p > 0.0 and config.top_p < 1.0:
            model_kwargs["top_p"] = config.top_p
        else:
            raise ValueError("`top_p` must be > 0.0 and < 1.0")

        llm = HuggingFaceHub(
            huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"],
            repo_id=config.model or "google/flan-t5-xxl",
            model_kwargs=model_kwargs,
        )

        return llm(prompt)
Repository: mem0ai/mem0
File: tests/llm/test_hugging_face_hub.py (Python)

Code before:
import importlib
import os
import unittest
from unittest.mock import patch

from embedchain.config import BaseLlmConfig
from embedchain.llm.hugging_face_hub import HuggingFaceHubLlm


class TestHuggingFaceHubLlm(unittest.TestCase):
    def setUp(self):
        os.environ["HUGGINGFACEHUB_API_TOKEN"] = "test_api_key"
        self.config = BaseLlmConfig(model="google/flan-t5-xxl", max_tokens=50, temperature=0.7, top_p=0.8)

    def test_init_raises_value_error_without_api_key(self):
        os.environ.pop("HUGGINGFACEHUB_API_TOKEN")
        with self.assertRaises(ValueError):
            HuggingFaceHubLlm()

    def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
        llm = HuggingFaceHubLlm(self.config)
        llm.config.system_prompt = "system_prompt"
        with self.assertRaises(ValueError):
            llm.get_llm_model_answer("prompt")

    def test_top_p_value_within_range(self):
        config = BaseLlmConfig(top_p=1.0)
        with self.assertRaises(ValueError):
            HuggingFaceHubLlm._get_answer("test_prompt", config)

    def test_importlib_is_imported(self):
        importlib_installed = True
        try:
            importlib.import_module("huggingface_hub")
        except ImportError:
            importlib_installed = False
        self.assertTrue(importlib_installed)

Review comment (suggestion):
```suggestion
        os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"] = "test_access_token"
```

Code after:
import importlib
import os
import unittest
from unittest.mock import patch, MagicMock

from embedchain.config import BaseLlmConfig
from embedchain.llm.hugging_face_hub import HuggingFaceHubLlm


class TestHuggingFaceHubLlm(unittest.TestCase):
    def setUp(self):
        os.environ["HUGGINGFACEHUB_ACCESS_TOKEN"] = "test_access_token"
        self.config = BaseLlmConfig(model="google/flan-t5-xxl", max_tokens=50, temperature=0.7, top_p=0.8)

    def test_init_raises_value_error_without_api_key(self):
        os.environ.pop("HUGGINGFACEHUB_ACCESS_TOKEN")
        with self.assertRaises(ValueError):
            HuggingFaceHubLlm()

    def test_get_llm_model_answer_raises_value_error_for_system_prompt(self):
        llm = HuggingFaceHubLlm(self.config)
        llm.config.system_prompt = "system_prompt"
        with self.assertRaises(ValueError):
            llm.get_llm_model_answer("prompt")

    def test_top_p_value_within_range(self):
        config = BaseLlmConfig(top_p=1.0)
        with self.assertRaises(ValueError):
            HuggingFaceHubLlm._get_answer("test_prompt", config)

    def test_dependency_is_imported(self):
        importlib_installed = True
        try:
            importlib.import_module("huggingface_hub")
        except ImportError:
            importlib_installed = False
        self.assertTrue(importlib_installed)
Repository: mem0ai/mem0
File: embedchain/vectordb/zilliz.py (Python)

Code before:
from typing import Dict, List, Optional

from embedchain.config import ZillizDBConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.vectordb.base import BaseVectorDB

try:
    from pymilvus import MilvusClient
    from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
except ImportError:
    raise ImportError(
        "Zilliz requires extra dependencies. Install with `pip install --upgrade embedchain[pymilvus]`"
    ) from None


@register_deserializable
class ZillizVectorDB(BaseVectorDB):
    """Base class for vector database."""

    def __init__(self, config: ZillizDBConfig = None):
        """Initialize the database. Save the config and client as an attribute.

        :param config: Database configuration class instance.
        :type config: ZillizDBConfig
        """

        if config is None:
            self.config = ZillizDBConfig()
        else:
            self.config = config

        try:
            self.client = MilvusClient(
                uri=self.config.uri,
                token=self.config.token,
            )


Review comment (suggestion):
```suggestion
        "Zilliz requires extra dependencies. Install with `pip install --upgrade embedchain[milvus]`"
```

Code after:
from typing import Dict, List, Optional

from embedchain.config import ZillizDBConfig
from embedchain.helper.json_serializable import register_deserializable
from embedchain.vectordb.base import BaseVectorDB

try:
    from pymilvus import MilvusClient
    from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
except ImportError:
    raise ImportError(
        "Zilliz requires extra dependencies. Install with `pip install --upgrade embedchain[milvus]`"
    ) from None


@register_deserializable
class ZillizVectorDB(BaseVectorDB):
    """Base class for vector database."""

    def __init__(self, config: ZillizDBConfig = None):
        """Initialize the database. Save the config and client as an attribute.

        :param config: Database configuration class instance.
        :type config: ZillizDBConfig
        """

        if config is None:
            self.config = ZillizDBConfig()
        else:
            self.config = config

        self.client = MilvusClient(
            uri=self.config.uri,
            token=self.config.token,
        )

        self.connection = connections.connect(
Repository: mem0ai/mem0
File: embedchain/config/vectordb/zilliz.py (Python)

Code before:
        :type collection_name: Optional[str], optional
        :param dir: Path to the database directory, where the database is stored, defaults to "db"
        :type dir: str, optional
        :param uri: Cluster endpoint obtained from the Zilliz Console, defaults to None
        :type uri: Optional[str], optional
        :param token: API Key, if a Serverless Cluster, username:password, if a Dedicated Cluster, defaults to None
        :type port: Optional[str], optional
        """
        self.uri = uri or os.environ.get("ZILLIZ_CLOUD_URI")
        if not self.uri:
            raise AttributeError(
                "Zilliz needs a URI attribute, "
                "this can either be passed to `ZILLIZ_CLOUD_URI` or as `ZILLIZ_CLOUD_URI` in `.env`"
            )

        self.token = token or os.environ.get("ZILLIZ_CLOUD_TOKEN")
        if not self.token:
            raise AttributeError(
                "Zilliz needs a token attribute, "
                "this can either be passed to `ZILLIZ_CLOUD_TOKEN` or as `ZILLIZ_CLOUD_TOKEN` in `.env`,"
                "if having a username and password, pass it in the form 'username:password' to `ZILLIZ_CLOUD_TOKEN`"
            )

        self.metric_type = metric_type
        if not self.metric_type:
            self.metric_type = "L2"

        self.vector_dim = vector_dim
        super().__init__(collection_name=collection_name, dir=dir)


Review comment (suggestion):
Can convert these three lines into one line by doing:

```python
self.metric_type = metric_type if metric_type else "L2"

Code after:
        :type collection_name: Optional[str], optional
        :param dir: Path to the database directory, where the database is stored, defaults to "db"
        :type dir: str, optional
        :param uri: Cluster endpoint obtained from the Zilliz Console, defaults to None
        :type uri: Optional[str], optional
        :param token: API Key, if a Serverless Cluster, username:password, if a Dedicated Cluster, defaults to None
        :type port: Optional[str], optional
        """
        self.uri = uri or os.environ.get("ZILLIZ_CLOUD_URI")
        if not self.uri:
            raise AttributeError(
                "Zilliz needs a URI attribute, "
                "this can either be passed to `ZILLIZ_CLOUD_URI` or as `ZILLIZ_CLOUD_URI` in `.env`"
            )

        self.token = token or os.environ.get("ZILLIZ_CLOUD_TOKEN")
        if not self.token:
            raise AttributeError(
                "Zilliz needs a token attribute, "
                "this can either be passed to `ZILLIZ_CLOUD_TOKEN` or as `ZILLIZ_CLOUD_TOKEN` in `.env`,"
                "if having a username and password, pass it in the form 'username:password' to `ZILLIZ_CLOUD_TOKEN`"
            )

        self.metric_type = metric_type if metric_type else "L2"

        self.vector_dim = vector_dim
        super().__init__(collection_name=collection_name, dir=dir)
Repository: mem0ai/mem0
File: embedchain/chunkers/json.py (Python)

Code before:
from typing import Optional

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


@register_deserializable
class JsonChunker(BaseChunker):
    """Chunker for json."""

    def __init__(self, config: Optional[ChunkerConfig] = None):
        if config is None:
            config = ChunkerConfig(chunk_size=1000, chunk_overlap=0, length_function=len)
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=config.chunk_size,
            chunk_overlap=config.chunk_overlap,
            length_function=config.length_function,
        )
        super().__init__(text_splitter)


Review comment (suggestion):
```suggestion
class JSONChunker(BaseChunker):
```

Code after:
from typing import Optional

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


@register_deserializable
class JSONChunker(BaseChunker):
    """Chunker for json."""

    def __init__(self, config: Optional[ChunkerConfig] = None):
        if config is None:
            config = ChunkerConfig(chunk_size=1000, chunk_overlap=0, length_function=len)
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=config.chunk_size,
            chunk_overlap=config.chunk_overlap,
            length_function=config.length_function,
        )
        super().__init__(text_splitter)
Repository: mem0ai/mem0
File: embedchain/data_formatter/data_formatter.py (Python)

Code before:
    def _get_chunker(self, data_type: DataType, config: ChunkerConfig) -> BaseChunker:
        """Returns the appropriate chunker for the given data type.

        :param data_type: The type of the data to chunk.
        :type data_type: DataType
        :param config: Config to initialize the chunker with.
        :type config: ChunkerConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The chunker for the given data type.
        :rtype: BaseChunker
        """
        chunker_classes = {
            DataType.YOUTUBE_VIDEO: YoutubeVideoChunker,
            DataType.PDF_FILE: PdfFileChunker,
            DataType.WEB_PAGE: WebPageChunker,
            DataType.QNA_PAIR: QnaPairChunker,
            DataType.TEXT: TextChunker,
            DataType.DOCX: DocxFileChunker,
            DataType.DOCS_SITE: DocsSiteChunker,
            DataType.SITEMAP: SitemapChunker,
            DataType.NOTION: NotionChunker,
            DataType.CSV: TableChunker,
            DataType.MDX: MdxChunker,
            DataType.IMAGES: ImagesChunker,
            DataType.XML: XmlChunker,
            DataType.JSON: JsonChunker,
        }
        if data_type in chunker_classes:
            chunker_class: type = chunker_classes[data_type]
            chunker: BaseChunker = chunker_class(config)
            chunker.set_data_type(data_type)
            return chunker
        else:
            raise ValueError(f"Unsupported data type: {data_type}")


Review comment (suggestion):
```suggestion
            DataType.JSON: JSONChunker,
```

Code after:
    def _get_chunker(self, data_type: DataType, config: ChunkerConfig) -> BaseChunker:
        """Returns the appropriate chunker for the given data type.

        :param data_type: The type of the data to chunk.
        :type data_type: DataType
        :param config: Config to initialize the chunker with.
        :type config: ChunkerConfig
        :raises ValueError: If an unsupported data type is provided.
        :return: The chunker for the given data type.
        :rtype: BaseChunker
        """
        chunker_classes = {
            DataType.YOUTUBE_VIDEO: YoutubeVideoChunker,
            DataType.PDF_FILE: PdfFileChunker,
            DataType.WEB_PAGE: WebPageChunker,
            DataType.QNA_PAIR: QnaPairChunker,
            DataType.TEXT: TextChunker,
            DataType.DOCX: DocxFileChunker,
            DataType.DOCS_SITE: DocsSiteChunker,
            DataType.SITEMAP: SitemapChunker,
            DataType.NOTION: NotionChunker,
            DataType.CSV: TableChunker,
            DataType.MDX: MdxChunker,
            DataType.IMAGES: ImagesChunker,
            DataType.XML: XmlChunker,
            DataType.JSON: JSONChunker,
        }
        if data_type in chunker_classes:
            chunker_class: type = chunker_classes[data_type]
            chunker: BaseChunker = chunker_class(config)
            chunker.set_data_type(data_type)
            return chunker
        else:
            raise ValueError(f"Unsupported data type: {data_type}")