Rename embedchain to mem0 and open sourcing code for long term memory (#1474)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Taranjeet Singh
co-authored by
Deshraj Yadav
Deshraj Yadav
parent
83e8c97295
commit
f842a92e25
@@ -0,0 +1,15 @@
|
||||
# flake8: noqa: F401
|
||||
|
||||
from .add_config import AddConfig, ChunkerConfig
|
||||
from .app_config import AppConfig
|
||||
from .base_config import BaseConfig
|
||||
from .cache_config import CacheConfig
|
||||
from .embedder.base import BaseEmbedderConfig
|
||||
from .embedder.base import BaseEmbedderConfig as EmbedderConfig
|
||||
from .embedder.ollama import OllamaEmbedderConfig
|
||||
from .llm.base import BaseLlmConfig
|
||||
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
|
||||
@@ -0,0 +1,79 @@
|
||||
import builtins
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from importlib import import_module
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class ChunkerConfig(BaseConfig):
|
||||
"""
|
||||
Config for the chunker used in `add` method
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunk_size: Optional[int] = 2000,
|
||||
chunk_overlap: Optional[int] = 0,
|
||||
length_function: Optional[Callable[[str], int]] = None,
|
||||
min_chunk_size: Optional[int] = 0,
|
||||
):
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_overlap = chunk_overlap
|
||||
self.min_chunk_size = min_chunk_size
|
||||
if self.min_chunk_size >= self.chunk_size:
|
||||
raise ValueError(f"min_chunk_size {min_chunk_size} should be less than chunk_size {chunk_size}")
|
||||
if self.min_chunk_size < self.chunk_overlap:
|
||||
logging.warning(
|
||||
f"min_chunk_size {min_chunk_size} should be greater than chunk_overlap {chunk_overlap}, otherwise it is redundant." # noqa:E501
|
||||
)
|
||||
|
||||
if isinstance(length_function, str):
|
||||
self.length_function = self.load_func(length_function)
|
||||
else:
|
||||
self.length_function = length_function if length_function else len
|
||||
|
||||
@staticmethod
|
||||
def load_func(dotpath: str):
|
||||
if "." not in dotpath:
|
||||
return getattr(builtins, dotpath)
|
||||
else:
|
||||
module_, func = dotpath.rsplit(".", maxsplit=1)
|
||||
m = import_module(module_)
|
||||
return getattr(m, func)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class LoaderConfig(BaseConfig):
|
||||
"""
|
||||
Config for the loader used in `add` method
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class AddConfig(BaseConfig):
|
||||
"""
|
||||
Config for the `add` method.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunker: Optional[ChunkerConfig] = None,
|
||||
loader: Optional[LoaderConfig] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the `add` method.
|
||||
|
||||
:param chunker: Chunker config, defaults to None
|
||||
:type chunker: Optional[ChunkerConfig], optional
|
||||
:param loader: Loader config, defaults to None
|
||||
:type loader: Optional[LoaderConfig], optional
|
||||
"""
|
||||
self.loader = loader
|
||||
self.chunker = chunker
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
from .base_app_config import BaseAppConfig
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class AppConfig(BaseAppConfig):
|
||||
"""
|
||||
Config to initialize an embedchain custom `App` instance, with extra config options.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_level: str = "WARNING",
|
||||
id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
collect_metrics: Optional[bool] = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for an App. This is the simplest form of an embedchain app.
|
||||
Most of the configuration is done in the `App` class itself.
|
||||
|
||||
:param log_level: Debug level ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], defaults to "WARNING"
|
||||
:type log_level: str, optional
|
||||
:param id: ID of the app. Document metadata will have this id., defaults to None
|
||||
:type id: Optional[str], optional
|
||||
:param collect_metrics: Send anonymous telemetry to improve embedchain, defaults to True
|
||||
:type collect_metrics: Optional[bool], optional
|
||||
"""
|
||||
self.name = name
|
||||
super().__init__(log_level=log_level, id=id, collect_metrics=collect_metrics, **kwargs)
|
||||
@@ -0,0 +1,58 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
from embedchain.helpers.json_serializable import JSONSerializable
|
||||
from embedchain.vectordb.base import BaseVectorDB
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseAppConfig(BaseConfig, JSONSerializable):
|
||||
"""
|
||||
Parent config to initialize an instance of `App`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
log_level: str = "WARNING",
|
||||
db: Optional[BaseVectorDB] = None,
|
||||
id: Optional[str] = None,
|
||||
collect_metrics: bool = True,
|
||||
collection_name: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for an App.
|
||||
Most of the configuration is done in the `App` class itself.
|
||||
|
||||
:param log_level: Debug level ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], defaults to "WARNING"
|
||||
:type log_level: str, optional
|
||||
:param db: A database class. It is recommended to set this directly in the `App` class, not this config,
|
||||
defaults to None
|
||||
:type db: Optional[BaseVectorDB], optional
|
||||
:param id: ID of the app. Document metadata will have this id., defaults to None
|
||||
:type id: Optional[str], optional
|
||||
:param collect_metrics: Send anonymous telemetry to improve embedchain, defaults to True
|
||||
:type collect_metrics: Optional[bool], optional
|
||||
:param collection_name: Default collection name. It's recommended to use app.db.set_collection_name() instead,
|
||||
defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
"""
|
||||
self.id = id
|
||||
self.collect_metrics = True if (collect_metrics is True or collect_metrics is None) else False
|
||||
self.collection_name = collection_name
|
||||
|
||||
if db:
|
||||
self._db = db
|
||||
logger.warning(
|
||||
"DEPRECATION WARNING: Please supply the database as the second parameter during app init. "
|
||||
"Such as `app(config=config, db=db)`."
|
||||
)
|
||||
|
||||
if collection_name:
|
||||
logger.warning("DEPRECATION WARNING: Please supply the collection name to the database config.")
|
||||
return
|
||||
|
||||
def _setup_logging(self, log_level):
|
||||
logger.basicConfig(format="%(asctime)s [%(name)s] [%(levelname)s] %(message)s", level=log_level)
|
||||
self.logger = logger.getLogger(__name__)
|
||||
@@ -0,0 +1,21 @@
|
||||
from typing import Any
|
||||
|
||||
from embedchain.helpers.json_serializable import JSONSerializable
|
||||
|
||||
|
||||
class BaseConfig(JSONSerializable):
|
||||
"""
|
||||
Base config.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initializes a configuration class for a class."""
|
||||
pass
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
"""Return config object as a dict
|
||||
|
||||
:return: config object as dict
|
||||
:rtype: dict[str, Any]
|
||||
"""
|
||||
return vars(self)
|
||||
@@ -0,0 +1,96 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class CacheSimilarityEvalConfig(BaseConfig):
|
||||
"""
|
||||
This is the evaluator to compare two embeddings according to their distance computed in embedding retrieval stage.
|
||||
In the retrieval stage, `search_result` is the distance used for approximate nearest neighbor search and have been
|
||||
put into `cache_dict`. `max_distance` is used to bound this distance to make it between [0-`max_distance`].
|
||||
`positive` is used to indicate this distance is directly proportional to the similarity of two entities.
|
||||
If `positive` is set `False`, `max_distance` will be used to subtract this distance to get the final score.
|
||||
|
||||
:param max_distance: the bound of maximum distance.
|
||||
:type max_distance: float
|
||||
:param positive: if the larger distance indicates more similar of two entities, It is True. Otherwise, it is False.
|
||||
:type positive: bool
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
strategy: Optional[str] = "distance",
|
||||
max_distance: Optional[float] = 1.0,
|
||||
positive: Optional[bool] = False,
|
||||
):
|
||||
self.strategy = strategy
|
||||
self.max_distance = max_distance
|
||||
self.positive = positive
|
||||
|
||||
@staticmethod
|
||||
def from_config(config: Optional[dict[str, Any]]):
|
||||
if config is None:
|
||||
return CacheSimilarityEvalConfig()
|
||||
else:
|
||||
return CacheSimilarityEvalConfig(
|
||||
strategy=config.get("strategy", "distance"),
|
||||
max_distance=config.get("max_distance", 1.0),
|
||||
positive=config.get("positive", False),
|
||||
)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class CacheInitConfig(BaseConfig):
|
||||
"""
|
||||
This is a cache init config. Used to initialize a cache.
|
||||
|
||||
:param similarity_threshold: a threshold ranged from 0 to 1 to filter search results with similarity score higher \
|
||||
than the threshold. When it is 0, there is no hits. When it is 1, all search results will be returned as hits.
|
||||
:type similarity_threshold: float
|
||||
:param auto_flush: it will be automatically flushed every time xx pieces of data are added, default to 20
|
||||
:type auto_flush: int
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
similarity_threshold: Optional[float] = 0.8,
|
||||
auto_flush: Optional[int] = 20,
|
||||
):
|
||||
if similarity_threshold < 0 or similarity_threshold > 1:
|
||||
raise ValueError(f"similarity_threshold {similarity_threshold} should be between 0 and 1")
|
||||
|
||||
self.similarity_threshold = similarity_threshold
|
||||
self.auto_flush = auto_flush
|
||||
|
||||
@staticmethod
|
||||
def from_config(config: Optional[dict[str, Any]]):
|
||||
if config is None:
|
||||
return CacheInitConfig()
|
||||
else:
|
||||
return CacheInitConfig(
|
||||
similarity_threshold=config.get("similarity_threshold", 0.8),
|
||||
auto_flush=config.get("auto_flush", 20),
|
||||
)
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class CacheConfig(BaseConfig):
|
||||
def __init__(
|
||||
self,
|
||||
similarity_eval_config: Optional[CacheSimilarityEvalConfig] = CacheSimilarityEvalConfig(),
|
||||
init_config: Optional[CacheInitConfig] = CacheInitConfig(),
|
||||
):
|
||||
self.similarity_eval_config = similarity_eval_config
|
||||
self.init_config = init_config
|
||||
|
||||
@staticmethod
|
||||
def from_config(config: Optional[dict[str, Any]]):
|
||||
if config is None:
|
||||
return CacheConfig()
|
||||
else:
|
||||
return CacheConfig(
|
||||
similarity_eval_config=CacheSimilarityEvalConfig.from_config(config.get("similarity_evaluation", {})),
|
||||
init_config=CacheInitConfig.from_config(config.get("init_config", {})),
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class BaseEmbedderConfig:
|
||||
def __init__(
|
||||
self,
|
||||
model: Optional[str] = None,
|
||||
deployment_name: Optional[str] = None,
|
||||
vector_dimension: Optional[int] = None,
|
||||
endpoint: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
model_kwargs: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""
|
||||
Initialize a new instance of an embedder config class.
|
||||
|
||||
:param model: model name of the llm embedding model (not applicable to all providers), defaults to None
|
||||
:type model: Optional[str], optional
|
||||
:param deployment_name: deployment name for llm embedding model, defaults to None
|
||||
:type deployment_name: Optional[str], optional
|
||||
:param vector_dimension: vector dimension of the embedding model, defaults to None
|
||||
:type vector_dimension: Optional[int], optional
|
||||
:param endpoint: endpoint for the embedding model, defaults to None
|
||||
:type endpoint: Optional[str], optional
|
||||
:param api_key: hugginface api key, defaults to None
|
||||
:type api_key: Optional[str], optional
|
||||
:param api_base: huggingface api base, defaults to None
|
||||
:type api_base: Optional[str], optional
|
||||
:param model_kwargs: key-value arguments for the embedding model, defaults a dict inside init.
|
||||
:type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init.
|
||||
"""
|
||||
self.model = model
|
||||
self.deployment_name = deployment_name
|
||||
self.vector_dimension = vector_dimension
|
||||
self.endpoint = endpoint
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
self.model_kwargs = model_kwargs or {}
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.embedder.base import BaseEmbedderConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class GoogleAIEmbedderConfig(BaseEmbedderConfig):
|
||||
def __init__(
|
||||
self,
|
||||
model: Optional[str] = None,
|
||||
deployment_name: Optional[str] = None,
|
||||
vector_dimension: Optional[int] = None,
|
||||
task_type: Optional[str] = None,
|
||||
title: Optional[str] = None,
|
||||
):
|
||||
super().__init__(model, deployment_name, vector_dimension)
|
||||
self.task_type = task_type or "retrieval_document"
|
||||
self.title = title or "Embeddings for Embedchain"
|
||||
@@ -0,0 +1,16 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.embedder.base import BaseEmbedderConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class OllamaEmbedderConfig(BaseEmbedderConfig):
|
||||
def __init__(
|
||||
self,
|
||||
model: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
vector_dimension: Optional[int] = None,
|
||||
):
|
||||
super().__init__(model=model, vector_dimension=vector_dimension)
|
||||
self.base_url = base_url or "http://localhost:11434"
|
||||
@@ -0,0 +1,2 @@
|
||||
from .base import (AnswerRelevanceConfig, ContextRelevanceConfig, # noqa: F401
|
||||
GroundednessConfig)
|
||||
@@ -0,0 +1,92 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
|
||||
ANSWER_RELEVANCY_PROMPT = """
|
||||
Please provide $num_gen_questions questions from the provided answer.
|
||||
You must provide the complete question, if are not able to provide the complete question, return empty string ("").
|
||||
Please only provide one question per line without numbers or bullets to distinguish them.
|
||||
You must only provide the questions and no other text.
|
||||
|
||||
$answer
|
||||
""" # noqa:E501
|
||||
|
||||
|
||||
CONTEXT_RELEVANCY_PROMPT = """
|
||||
Please extract relevant sentences from the provided context that is required to answer the given question.
|
||||
If no relevant sentences are found, or if you believe the question cannot be answered from the given context, return the empty string ("").
|
||||
While extracting candidate sentences you're not allowed to make any changes to sentences from given context or make up any sentences.
|
||||
You must only provide sentences from the given context and nothing else.
|
||||
|
||||
Context: $context
|
||||
Question: $question
|
||||
""" # noqa:E501
|
||||
|
||||
GROUNDEDNESS_ANSWER_CLAIMS_PROMPT = """
|
||||
Please provide one or more statements from each sentence of the provided answer.
|
||||
You must provide the symantically equivalent statements for each sentence of the answer.
|
||||
You must provide the complete statement, if are not able to provide the complete statement, return empty string ("").
|
||||
Please only provide one statement per line WITHOUT numbers or bullets.
|
||||
If the question provided is not being answered in the provided answer, return empty string ("").
|
||||
You must only provide the statements and no other text.
|
||||
|
||||
$question
|
||||
$answer
|
||||
""" # noqa:E501
|
||||
|
||||
GROUNDEDNESS_CLAIMS_INFERENCE_PROMPT = """
|
||||
Given the context and the provided claim statements, please provide a verdict for each claim statement whether it can be completely infered from the given context or not.
|
||||
Use only "1" (yes), "0" (no) and "-1" (null) for "yes", "no" or "null" respectively.
|
||||
You must provide one verdict per line, ONLY WITH "1", "0" or "-1" as per your verdict to the given statement and nothing else.
|
||||
You must provide the verdicts in the same order as the claim statements.
|
||||
|
||||
Contexts:
|
||||
$context
|
||||
|
||||
Claim statements:
|
||||
$claim_statements
|
||||
""" # noqa:E501
|
||||
|
||||
|
||||
class GroundednessConfig(BaseConfig):
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "gpt-4",
|
||||
api_key: Optional[str] = None,
|
||||
answer_claims_prompt: str = GROUNDEDNESS_ANSWER_CLAIMS_PROMPT,
|
||||
claims_inference_prompt: str = GROUNDEDNESS_CLAIMS_INFERENCE_PROMPT,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.answer_claims_prompt = answer_claims_prompt
|
||||
self.claims_inference_prompt = claims_inference_prompt
|
||||
|
||||
|
||||
class AnswerRelevanceConfig(BaseConfig):
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "gpt-4",
|
||||
embedder: str = "text-embedding-ada-002",
|
||||
api_key: Optional[str] = None,
|
||||
num_gen_questions: int = 1,
|
||||
prompt: str = ANSWER_RELEVANCY_PROMPT,
|
||||
):
|
||||
self.model = model
|
||||
self.embedder = embedder
|
||||
self.api_key = api_key
|
||||
self.num_gen_questions = num_gen_questions
|
||||
self.prompt = prompt
|
||||
|
||||
|
||||
class ContextRelevanceConfig(BaseConfig):
|
||||
def __init__(
|
||||
self,
|
||||
model: str = "gpt-4",
|
||||
api_key: Optional[str] = None,
|
||||
language: str = "en",
|
||||
prompt: str = CONTEXT_RELEVANCY_PROMPT,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.language = language
|
||||
self.prompt = prompt
|
||||
@@ -0,0 +1,275 @@
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from string import Template
|
||||
from typing import Any, Mapping, Optional, Dict, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_PROMPT = """
|
||||
You are a Q&A expert system. Your responses must always be rooted in the context provided for each query. Here are some guidelines to follow:
|
||||
|
||||
1. Refrain from explicitly mentioning the context provided in your response.
|
||||
2. The context should silently guide your answers without being directly acknowledged.
|
||||
3. Do not use phrases such as 'According to the context provided', 'Based on the context, ...' etc.
|
||||
|
||||
Context information:
|
||||
----------------------
|
||||
$context
|
||||
----------------------
|
||||
|
||||
Query: $query
|
||||
Answer:
|
||||
""" # noqa:E501
|
||||
|
||||
DEFAULT_PROMPT_WITH_HISTORY = """
|
||||
You are a Q&A expert system. Your responses must always be rooted in the context provided for each query. You are also provided with the conversation history with the user. Make sure to use relevant context from conversation history as needed.
|
||||
|
||||
Here are some guidelines to follow:
|
||||
|
||||
1. Refrain from explicitly mentioning the context provided in your response.
|
||||
2. The context should silently guide your answers without being directly acknowledged.
|
||||
3. Do not use phrases such as 'According to the context provided', 'Based on the context, ...' etc.
|
||||
|
||||
Context information:
|
||||
----------------------
|
||||
$context
|
||||
----------------------
|
||||
|
||||
Conversation history:
|
||||
----------------------
|
||||
$history
|
||||
----------------------
|
||||
|
||||
Query: $query
|
||||
Answer:
|
||||
""" # noqa:E501
|
||||
|
||||
DEFAULT_PROMPT_WITH_MEM0_MEMORY = """
|
||||
You are a Q&A expert system. Your responses must always be rooted in the context provided for each query. You are also provided with the conversation history and memories with the user. Make sure to use relevant context from conversation history and memories as needed.
|
||||
|
||||
Here are some guidelines to follow:
|
||||
|
||||
1. Refrain from explicitly mentioning the context provided in your response.
|
||||
2. Take into consideration the conversation history and memories provided.
|
||||
3. The context should silently guide your answers without being directly acknowledged.
|
||||
4. Do not use phrases such as 'According to the context provided', 'Based on the context, ...' etc.
|
||||
|
||||
Context information:
|
||||
----------------------
|
||||
$context
|
||||
----------------------
|
||||
|
||||
Conversation history:
|
||||
----------------------
|
||||
$history
|
||||
----------------------
|
||||
|
||||
Memories/Preferences:
|
||||
----------------------
|
||||
$memories
|
||||
----------------------
|
||||
|
||||
Query: $query
|
||||
Answer:
|
||||
""" # noqa:E501
|
||||
|
||||
DOCS_SITE_DEFAULT_PROMPT = """
|
||||
You are an expert AI assistant for developer support product. Your responses must always be rooted in the context provided for each query. Wherever possible, give complete code snippet. Dont make up any code snippet on your own.
|
||||
|
||||
Here are some guidelines to follow:
|
||||
|
||||
1. Refrain from explicitly mentioning the context provided in your response.
|
||||
2. The context should silently guide your answers without being directly acknowledged.
|
||||
3. Do not use phrases such as 'According to the context provided', 'Based on the context, ...' etc.
|
||||
|
||||
Context information:
|
||||
----------------------
|
||||
$context
|
||||
----------------------
|
||||
|
||||
Query: $query
|
||||
Answer:
|
||||
""" # noqa:E501
|
||||
|
||||
DEFAULT_PROMPT_TEMPLATE = Template(DEFAULT_PROMPT)
|
||||
DEFAULT_PROMPT_WITH_HISTORY_TEMPLATE = Template(DEFAULT_PROMPT_WITH_HISTORY)
|
||||
DEFAULT_PROMPT_WITH_MEM0_MEMORY_TEMPLATE = Template(DEFAULT_PROMPT_WITH_MEM0_MEMORY)
|
||||
DOCS_SITE_PROMPT_TEMPLATE = Template(DOCS_SITE_DEFAULT_PROMPT)
|
||||
query_re = re.compile(r"\$\{*query\}*")
|
||||
context_re = re.compile(r"\$\{*context\}*")
|
||||
history_re = re.compile(r"\$\{*history\}*")
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class BaseLlmConfig(BaseConfig):
|
||||
"""
|
||||
Config for the `query` method.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
number_documents: int = 3,
|
||||
template: Optional[Template] = None,
|
||||
prompt: Optional[Template] = None,
|
||||
model: Optional[str] = None,
|
||||
temperature: float = 0,
|
||||
max_tokens: int = 1000,
|
||||
top_p: float = 1,
|
||||
stream: bool = False,
|
||||
online: bool = False,
|
||||
token_usage: bool = False,
|
||||
deployment_name: Optional[str] = None,
|
||||
system_prompt: Optional[str] = None,
|
||||
where: dict[str, Any] = None,
|
||||
query_type: Optional[str] = None,
|
||||
callbacks: Optional[list] = None,
|
||||
api_key: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
endpoint: Optional[str] = None,
|
||||
model_kwargs: Optional[dict[str, Any]] = None,
|
||||
http_client_proxies: Optional[Union[Dict, str]] = None,
|
||||
http_async_client_proxies: Optional[Union[Dict, str]] = None,
|
||||
local: Optional[bool] = False,
|
||||
default_headers: Optional[Mapping[str, str]] = None,
|
||||
api_version: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the LLM.
|
||||
|
||||
Takes the place of the former `QueryConfig` or `ChatConfig`.
|
||||
|
||||
:param number_documents: Number of documents to pull from the database as
|
||||
context, defaults to 1
|
||||
:type number_documents: int, optional
|
||||
:param template: The `Template` instance to use as a template for
|
||||
prompt, defaults to None (deprecated)
|
||||
:type template: Optional[Template], optional
|
||||
:param prompt: The `Template` instance to use as a template for
|
||||
prompt, defaults to None
|
||||
:type prompt: Optional[Template], optional
|
||||
:param model: Controls the OpenAI model used, defaults to None
|
||||
:type model: Optional[str], optional
|
||||
:param temperature: Controls the randomness of the model's output.
|
||||
Higher values (closer to 1) make output more random, lower values make it more deterministic, defaults to 0
|
||||
:type temperature: float, optional
|
||||
:param max_tokens: Controls how many tokens are generated, defaults to 1000
|
||||
:type max_tokens: int, optional
|
||||
:param top_p: Controls the diversity of words. Higher values (closer to 1) make word selection more diverse,
|
||||
defaults to 1
|
||||
:type top_p: float, optional
|
||||
:param stream: Control if response is streamed back to user, defaults to False
|
||||
:type stream: bool, optional
|
||||
:param online: Controls whether to use internet for answering query, defaults to False
|
||||
:type online: bool, optional
|
||||
:param token_usage: Controls whether to return token usage in response, defaults to False
|
||||
:type token_usage: bool, optional
|
||||
:param deployment_name: t.b.a., defaults to None
|
||||
:type deployment_name: Optional[str], optional
|
||||
:param system_prompt: System prompt string, defaults to None
|
||||
: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 api_key: The api key of the custom endpoint, defaults to None
|
||||
:type api_key: Optional[str], optional
|
||||
:param endpoint: The api url of the custom endpoint, defaults to None
|
||||
:type endpoint: Optional[str], optional
|
||||
:param model_kwargs: A dictionary of key-value pairs to pass to the model, defaults to None
|
||||
:type model_kwargs: Optional[Dict[str, Any]], optional
|
||||
:param callbacks: Langchain callback functions to use, defaults to None
|
||||
:type callbacks: Optional[list], optional
|
||||
:param query_type: The type of query to use, defaults to None
|
||||
:type query_type: Optional[str], optional
|
||||
:param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None
|
||||
:type http_client_proxies: Optional[Dict | str], optional
|
||||
:param http_async_client_proxies: The proxy server settings for async calls used to create
|
||||
self.http_async_client, defaults to None
|
||||
:type http_async_client_proxies: Optional[Dict | str], optional
|
||||
:param local: If True, the model will be run locally, defaults to False (for huggingface provider)
|
||||
:type local: Optional[bool], optional
|
||||
:param default_headers: Set additional HTTP headers to be sent with requests to OpenAI
|
||||
:type default_headers: Optional[Mapping[str, str]], 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
|
||||
"""
|
||||
if template is not None:
|
||||
logger.warning(
|
||||
"The `template` argument is deprecated and will be removed in a future version. "
|
||||
+ "Please use `prompt` instead."
|
||||
)
|
||||
if prompt is None:
|
||||
prompt = template
|
||||
|
||||
if prompt is None:
|
||||
prompt = DEFAULT_PROMPT_TEMPLATE
|
||||
|
||||
self.number_documents = number_documents
|
||||
self.temperature = temperature
|
||||
self.max_tokens = max_tokens
|
||||
self.model = model
|
||||
self.top_p = top_p
|
||||
self.online = online
|
||||
self.token_usage = token_usage
|
||||
self.deployment_name = deployment_name
|
||||
self.system_prompt = system_prompt
|
||||
self.query_type = query_type
|
||||
self.callbacks = callbacks
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.endpoint = endpoint
|
||||
self.model_kwargs = model_kwargs
|
||||
self.http_client = httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
|
||||
self.http_async_client = (
|
||||
httpx.AsyncClient(proxies=http_async_client_proxies) if http_async_client_proxies else None
|
||||
)
|
||||
self.local = local
|
||||
self.default_headers = default_headers
|
||||
self.online = online
|
||||
self.api_version = api_version
|
||||
|
||||
if token_usage:
|
||||
f = open("model_prices_and_context_window.json")
|
||||
self.model_pricing_map = json.load(f)
|
||||
|
||||
if isinstance(prompt, str):
|
||||
prompt = Template(prompt)
|
||||
|
||||
if self.validate_prompt(prompt):
|
||||
self.prompt = prompt
|
||||
else:
|
||||
raise ValueError("The 'prompt' should have 'query' and 'context' keys and potentially 'history' (if used).")
|
||||
|
||||
if not isinstance(stream, bool):
|
||||
raise ValueError("`stream` should be bool")
|
||||
self.stream = stream
|
||||
self.where = where
|
||||
|
||||
@staticmethod
|
||||
def validate_prompt(prompt: Template) -> Optional[re.Match[str]]:
|
||||
"""
|
||||
validate the prompt
|
||||
|
||||
:param prompt: the prompt to validate
|
||||
:type prompt: Template
|
||||
:return: valid (true) or invalid (false)
|
||||
:rtype: Optional[re.Match[str]]
|
||||
"""
|
||||
return re.search(query_re, prompt.template) and re.search(context_re, prompt.template)
|
||||
|
||||
@staticmethod
|
||||
def _validate_prompt_history(prompt: Template) -> Optional[re.Match[str]]:
|
||||
"""
|
||||
validate the prompt with history
|
||||
|
||||
:param prompt: the prompt to validate
|
||||
:type prompt: Template
|
||||
:return: valid (true) or invalid (false)
|
||||
:rtype: Optional[re.Match[str]]
|
||||
"""
|
||||
return re.search(history_re, prompt.template)
|
||||
@@ -0,0 +1,21 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class Mem0Config(BaseConfig):
|
||||
def __init__(self, api_key: str, top_k: Optional[int] = 10):
|
||||
self.api_key = api_key
|
||||
self.top_k = top_k
|
||||
|
||||
@staticmethod
|
||||
def from_config(config: Optional[dict[str, Any]]):
|
||||
if config is None:
|
||||
return Mem0Config()
|
||||
else:
|
||||
return Mem0Config(
|
||||
api_key=config.get("api_key", ""),
|
||||
init_config=config.get("top_k", 10),
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.base_config import BaseConfig
|
||||
|
||||
|
||||
class BaseVectorDbConfig(BaseConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: str = "db",
|
||||
host: Optional[str] = None,
|
||||
port: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the vector database.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
: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 host: Database connection remote host. Use this if you run Embedchain as a client, defaults to None
|
||||
:type host: Optional[str], optional
|
||||
:param host: Database connection remote port. Use this if you run Embedchain as a client, defaults to None
|
||||
:type port: Optional[str], optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
"""
|
||||
self.collection_name = collection_name or "embedchain_store"
|
||||
self.dir = dir
|
||||
self.host = host
|
||||
self.port = port
|
||||
# Assign additional keyword arguments
|
||||
if kwargs:
|
||||
for key, value in kwargs.items():
|
||||
setattr(self, key, value)
|
||||
@@ -0,0 +1,41 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class ChromaDbConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
host: Optional[str] = None,
|
||||
port: Optional[str] = None,
|
||||
batch_size: Optional[int] = 100,
|
||||
allow_reset=False,
|
||||
chroma_settings: Optional[dict] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for ChromaDB.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
:param dir: Path to the database directory, where the database is stored, defaults to None
|
||||
:type dir: Optional[str], optional
|
||||
:param host: Database connection remote host. Use this if you run Embedchain as a client, defaults to None
|
||||
:type host: Optional[str], optional
|
||||
:param port: Database connection remote port. Use this if you run Embedchain as a client, defaults to None
|
||||
:type port: Optional[str], optional
|
||||
:param batch_size: Number of items to insert in one batch, defaults to 100
|
||||
:type batch_size: Optional[int], optional
|
||||
:param allow_reset: Resets the database. defaults to False
|
||||
:type allow_reset: bool
|
||||
:param chroma_settings: Chroma settings dict, defaults to None
|
||||
:type chroma_settings: Optional[dict], optional
|
||||
"""
|
||||
|
||||
self.chroma_settings = chroma_settings
|
||||
self.allow_reset = allow_reset
|
||||
self.batch_size = batch_size
|
||||
super().__init__(collection_name=collection_name, dir=dir, host=host, port=port)
|
||||
@@ -0,0 +1,56 @@
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class ElasticsearchDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
es_url: Union[str, list[str]] = None,
|
||||
cloud_id: Optional[str] = None,
|
||||
batch_size: Optional[int] = 100,
|
||||
**ES_EXTRA_PARAMS: dict[str, any],
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for an Elasticsearch client.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
:param dir: Path to the database directory, where the database is stored, defaults to None
|
||||
:type dir: Optional[str], optional
|
||||
:param es_url: elasticsearch url or list of nodes url to be used for connection, defaults to None
|
||||
:type es_url: Union[str, list[str]], optional
|
||||
:param cloud_id: cloud id of the elasticsearch cluster, defaults to None
|
||||
:type cloud_id: Optional[str], optional
|
||||
:param batch_size: Number of items to insert in one batch, defaults to 100
|
||||
:type batch_size: Optional[int], optional
|
||||
:param ES_EXTRA_PARAMS: extra params dict that can be passed to elasticsearch.
|
||||
:type ES_EXTRA_PARAMS: dict[str, Any], optional
|
||||
"""
|
||||
if es_url and cloud_id:
|
||||
raise ValueError("Only one of `es_url` and `cloud_id` can be set.")
|
||||
# self, es_url: Union[str, list[str]] = None, **ES_EXTRA_PARAMS: dict[str, any]):
|
||||
self.ES_URL = es_url or os.environ.get("ELASTICSEARCH_URL")
|
||||
self.CLOUD_ID = cloud_id or os.environ.get("ELASTICSEARCH_CLOUD_ID")
|
||||
if not self.ES_URL and not self.CLOUD_ID:
|
||||
raise AttributeError(
|
||||
"Elasticsearch needs a URL or CLOUD_ID attribute, "
|
||||
"this can either be passed to `ElasticsearchDBConfig` or as `ELASTICSEARCH_URL` or `ELASTICSEARCH_CLOUD_ID` in `.env`" # noqa: E501
|
||||
)
|
||||
self.ES_EXTRA_PARAMS = ES_EXTRA_PARAMS
|
||||
# Load API key from .env if it's not explicitly passed.
|
||||
# Can only set one of 'api_key', 'basic_auth', and 'bearer_auth'
|
||||
if (
|
||||
not self.ES_EXTRA_PARAMS.get("api_key")
|
||||
and not self.ES_EXTRA_PARAMS.get("basic_auth")
|
||||
and not self.ES_EXTRA_PARAMS.get("bearer_auth")
|
||||
):
|
||||
self.ES_EXTRA_PARAMS["api_key"] = os.environ.get("ELASTICSEARCH_API_KEY")
|
||||
|
||||
self.batch_size = batch_size
|
||||
super().__init__(collection_name=collection_name, dir=dir)
|
||||
@@ -0,0 +1,33 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class LanceDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
host: Optional[str] = None,
|
||||
port: Optional[str] = None,
|
||||
allow_reset=True,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for LanceDB.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
:param dir: Path to the database directory, where the database is stored, defaults to None
|
||||
:type dir: Optional[str], optional
|
||||
:param host: Database connection remote host. Use this if you run Embedchain as a client, defaults to None
|
||||
:type host: Optional[str], optional
|
||||
:param port: Database connection remote port. Use this if you run Embedchain as a client, defaults to None
|
||||
:type port: Optional[str], optional
|
||||
:param allow_reset: Resets the database. defaults to False
|
||||
:type allow_reset: bool
|
||||
"""
|
||||
|
||||
self.allow_reset = allow_reset
|
||||
super().__init__(collection_name=collection_name, dir=dir, host=host, port=port)
|
||||
@@ -0,0 +1,41 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class OpenSearchDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
opensearch_url: str,
|
||||
http_auth: tuple[str, str],
|
||||
vector_dimension: int = 1536,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
batch_size: Optional[int] = 100,
|
||||
**extra_params: dict[str, any],
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for an OpenSearch client.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
:param opensearch_url: URL of the OpenSearch domain
|
||||
:type opensearch_url: str, Eg, "http://localhost:9200"
|
||||
:param http_auth: Tuple of username and password
|
||||
:type http_auth: tuple[str, str], Eg, ("username", "password")
|
||||
:param vector_dimension: Dimension of the vector, defaults to 1536 (openai embedding model)
|
||||
:type vector_dimension: int, optional
|
||||
:param dir: Path to the database directory, where the database is stored, defaults to None
|
||||
:type dir: Optional[str], optional
|
||||
:param batch_size: Number of items to insert in one batch, defaults to 100
|
||||
:type batch_size: Optional[int], optional
|
||||
"""
|
||||
self.opensearch_url = opensearch_url
|
||||
self.http_auth = http_auth
|
||||
self.vector_dimension = vector_dimension
|
||||
self.extra_params = extra_params
|
||||
self.batch_size = batch_size
|
||||
|
||||
super().__init__(collection_name=collection_name, dir=dir)
|
||||
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class PineconeDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
index_name: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
vector_dimension: int = 1536,
|
||||
metric: Optional[str] = "cosine",
|
||||
pod_config: Optional[dict[str, any]] = None,
|
||||
serverless_config: Optional[dict[str, any]] = None,
|
||||
hybrid_search: bool = False,
|
||||
bm25_encoder: any = None,
|
||||
batch_size: Optional[int] = 100,
|
||||
**extra_params: dict[str, any],
|
||||
):
|
||||
self.metric = metric
|
||||
self.api_key = api_key
|
||||
self.index_name = index_name
|
||||
self.vector_dimension = vector_dimension
|
||||
self.extra_params = extra_params
|
||||
self.hybrid_search = hybrid_search
|
||||
self.bm25_encoder = bm25_encoder
|
||||
self.batch_size = batch_size
|
||||
if pod_config is None and serverless_config is None:
|
||||
# If no config is provided, use the default pod spec config
|
||||
pod_environment = os.environ.get("PINECONE_ENV", "gcp-starter")
|
||||
self.pod_config = {"environment": pod_environment, "metadata_config": {"indexed": ["*"]}}
|
||||
else:
|
||||
self.pod_config = pod_config
|
||||
self.serverless_config = serverless_config
|
||||
|
||||
if self.pod_config and self.serverless_config:
|
||||
raise ValueError("Only one of pod_config or serverless_config can be provided.")
|
||||
|
||||
if self.hybrid_search and self.metric != "dotproduct":
|
||||
raise ValueError(
|
||||
"Hybrid search is only supported with dotproduct metric in Pinecone. See full docs here: https://docs.pinecone.io/docs/hybrid-search#limitations"
|
||||
) # noqa:E501
|
||||
|
||||
super().__init__(collection_name=self.index_name, dir=None)
|
||||
@@ -0,0 +1,48 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class QdrantDBConfig(BaseVectorDbConfig):
|
||||
"""
|
||||
Config to initialize a qdrant client.
|
||||
:param: url. qdrant url or list of nodes url to be used for connection
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
hnsw_config: Optional[dict[str, any]] = None,
|
||||
quantization_config: Optional[dict[str, any]] = None,
|
||||
on_disk: Optional[bool] = None,
|
||||
batch_size: Optional[int] = 10,
|
||||
**extra_params: dict[str, any],
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for a qdrant client.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
:type collection_name: Optional[str], optional
|
||||
:param dir: Path to the database directory, where the database is stored, defaults to None
|
||||
:type dir: Optional[str], optional
|
||||
:param hnsw_config: Params for HNSW index
|
||||
:type hnsw_config: Optional[dict[str, any]], defaults to None
|
||||
:param quantization_config: Params for quantization, if None - quantization will be disabled
|
||||
:type quantization_config: Optional[dict[str, any]], defaults to None
|
||||
:param on_disk: If true - point`s payload will not be stored in memory.
|
||||
It will be read from the disk every time it is requested.
|
||||
This setting saves RAM by (slightly) increasing the response time.
|
||||
Note: those payload values that are involved in filtering and are indexed - remain in RAM.
|
||||
:type on_disk: bool, optional, defaults to None
|
||||
:param batch_size: Number of items to insert in one batch, defaults to 10
|
||||
:type batch_size: Optional[int], optional
|
||||
"""
|
||||
self.hnsw_config = hnsw_config
|
||||
self.quantization_config = quantization_config
|
||||
self.on_disk = on_disk
|
||||
self.batch_size = batch_size
|
||||
self.extra_params = extra_params
|
||||
super().__init__(collection_name=collection_name, dir=dir)
|
||||
@@ -0,0 +1,18 @@
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class WeaviateDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
batch_size: Optional[int] = 100,
|
||||
**extra_params: dict[str, any],
|
||||
):
|
||||
self.batch_size = batch_size
|
||||
self.extra_params = extra_params
|
||||
super().__init__(collection_name=collection_name, dir=dir)
|
||||
@@ -0,0 +1,49 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from embedchain.config.vector_db.base import BaseVectorDbConfig
|
||||
from embedchain.helpers.json_serializable import register_deserializable
|
||||
|
||||
|
||||
@register_deserializable
|
||||
class ZillizDBConfig(BaseVectorDbConfig):
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: Optional[str] = None,
|
||||
dir: Optional[str] = None,
|
||||
uri: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
vector_dim: Optional[str] = None,
|
||||
metric_type: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initializes a configuration class instance for the vector database.
|
||||
|
||||
:param collection_name: Default name for the collection, defaults to None
|
||||
: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 token: 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)
|
||||
Reference in New Issue
Block a user