419 lines
17 KiB
Python
419 lines
17 KiB
Python
import hashlib
|
|
import importlib.metadata
|
|
import json
|
|
import logging
|
|
import os
|
|
import threading
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
import requests
|
|
from dotenv import load_dotenv
|
|
from langchain.docstore.document import Document
|
|
from tenacity import retry, stop_after_attempt, wait_fixed
|
|
|
|
from embedchain.chunkers.base_chunker import BaseChunker
|
|
from embedchain.config import AddConfig, BaseLlmConfig
|
|
from embedchain.config.apps.BaseAppConfig import BaseAppConfig
|
|
from embedchain.data_formatter import DataFormatter
|
|
from embedchain.embedder.base_embedder import BaseEmbedder
|
|
from embedchain.helper_classes.json_serializable import JSONSerializable
|
|
from embedchain.llm.base_llm import BaseLlm
|
|
from embedchain.loaders.base_loader import BaseLoader
|
|
from embedchain.models.data_type import DataType
|
|
from embedchain.utils import detect_datatype
|
|
from embedchain.vectordb.base_vector_db import BaseVectorDB
|
|
|
|
load_dotenv()
|
|
|
|
ABS_PATH = os.getcwd()
|
|
HOME_DIR = str(Path.home())
|
|
CONFIG_DIR = os.path.join(HOME_DIR, ".embedchain")
|
|
CONFIG_FILE = os.path.join(CONFIG_DIR, "config.json")
|
|
|
|
|
|
class EmbedChain(JSONSerializable):
|
|
def __init__(
|
|
self,
|
|
config: BaseAppConfig,
|
|
llm: BaseLlm,
|
|
db: BaseVectorDB = None,
|
|
embedder: BaseEmbedder = None,
|
|
system_prompt: Optional[str] = None,
|
|
):
|
|
"""
|
|
Initializes the EmbedChain instance, sets up a vector DB client and
|
|
creates a collection.
|
|
|
|
:param config: BaseAppConfig instance to load as configuration.
|
|
:param system_prompt: Optional. System prompt string.
|
|
"""
|
|
|
|
self.config = config
|
|
|
|
# Add subclasses
|
|
## Llm
|
|
self.llm = llm
|
|
## Database
|
|
# Database has support for config assignment for backwards compatibility
|
|
if db is None and (not hasattr(self.config, "db") or self.config.db is None):
|
|
raise ValueError("App requires Database.")
|
|
self.db = db or self.config.db
|
|
## Embedder
|
|
if embedder is None:
|
|
raise ValueError("App requires Embedder.")
|
|
self.embedder = embedder
|
|
|
|
# Initialize database
|
|
self.db._set_embedder(self.embedder)
|
|
self.db._initialize()
|
|
# Set collection name from app config for backwards compatibility.
|
|
if config.collection_name:
|
|
self.db.set_collection_name(config.collection_name)
|
|
|
|
# Add variables that are "shortcuts"
|
|
if system_prompt:
|
|
self.llm.config.system_prompt = system_prompt
|
|
|
|
# Attributes that aren't subclass related.
|
|
self.user_asks = []
|
|
|
|
# Send anonymous telemetry
|
|
self.s_id = self.config.id if self.config.id else str(uuid.uuid4())
|
|
self.u_id = self._load_or_generate_user_id()
|
|
# NOTE: Uncomment the next two lines when running tests to see if any test fires a telemetry event.
|
|
# if (self.config.collect_metrics):
|
|
# raise ConnectionRefusedError("Collection of metrics should not be allowed.")
|
|
thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("init",))
|
|
thread_telemetry.start()
|
|
|
|
def _load_or_generate_user_id(self):
|
|
"""
|
|
Loads the user id from the config file if it exists, otherwise generates a new
|
|
one and saves it to the config file.
|
|
"""
|
|
if not os.path.exists(CONFIG_DIR):
|
|
os.makedirs(CONFIG_DIR)
|
|
|
|
if os.path.exists(CONFIG_FILE):
|
|
with open(CONFIG_FILE, "r") as f:
|
|
data = json.load(f)
|
|
if "user_id" in data:
|
|
return data["user_id"]
|
|
|
|
u_id = str(uuid.uuid4())
|
|
with open(CONFIG_FILE, "w") as f:
|
|
json.dump({"user_id": u_id}, f)
|
|
|
|
return u_id
|
|
|
|
def add(
|
|
self,
|
|
source,
|
|
data_type: Optional[DataType] = None,
|
|
metadata: Optional[Dict] = None,
|
|
config: Optional[AddConfig] = None,
|
|
):
|
|
"""
|
|
Adds the data from the given URL to the vector db.
|
|
Loads the data, chunks it, create embedding for each chunk
|
|
and then stores the embedding to vector database.
|
|
|
|
:param source: The data to embed, can be a URL, local file or raw content, depending on the data type.
|
|
:param data_type: Optional. Automatically detected, but can be forced with this argument.
|
|
The type of the data to add.
|
|
:param metadata: Optional. Metadata associated with the data source.
|
|
:param config: Optional. The `AddConfig` instance to use as configuration
|
|
options.
|
|
:return: source_id, a md5-hash of the source, in hexadecimal representation.
|
|
"""
|
|
if config is None:
|
|
config = AddConfig()
|
|
|
|
try:
|
|
DataType(source)
|
|
logging.warning(
|
|
f"""Starting from version v0.0.40, Embedchain can automatically detect the data type. So, in the `add` method, the argument order has changed. You no longer need to specify '{source}' for the `source` argument. So the code snippet will be `.add("{data_type}", "{source}")`""" # noqa #E501
|
|
)
|
|
logging.warning(
|
|
"Embedchain is swapping the arguments for you. This functionality might be deprecated in the future, so please adjust your code." # noqa #E501
|
|
)
|
|
source, data_type = data_type, source
|
|
except ValueError:
|
|
pass
|
|
|
|
if data_type:
|
|
try:
|
|
data_type = DataType(data_type)
|
|
except ValueError:
|
|
raise ValueError(
|
|
f"Invalid data_type: '{data_type}'.",
|
|
f"Please use one of the following: {[data_type.value for data_type in DataType]}",
|
|
) from None
|
|
if not data_type:
|
|
data_type = detect_datatype(source)
|
|
|
|
# `source_id` is the hash of the source argument
|
|
hash_object = hashlib.md5(str(source).encode("utf-8"))
|
|
source_id = hash_object.hexdigest()
|
|
|
|
data_formatter = DataFormatter(data_type, config)
|
|
self.user_asks.append([source, data_type.value, metadata])
|
|
documents, _metadatas, _ids, new_chunks = self.load_and_embed(
|
|
data_formatter.loader, data_formatter.chunker, source, metadata, source_id
|
|
)
|
|
if data_type in {DataType.DOCS_SITE}:
|
|
self.is_docs_site_instance = True
|
|
|
|
# Send anonymous telemetry
|
|
if self.config.collect_metrics:
|
|
# it's quicker to check the variable twice than to count words when they won't be submitted.
|
|
word_count = sum([len(document.split(" ")) for document in documents])
|
|
|
|
extra_metadata = {"data_type": data_type.value, "word_count": word_count, "chunks_count": new_chunks}
|
|
thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("add", extra_metadata))
|
|
thread_telemetry.start()
|
|
|
|
return source_id
|
|
|
|
def add_local(self, source, data_type=None, metadata=None, config: AddConfig = None):
|
|
"""
|
|
Warning:
|
|
This method is deprecated and will be removed in future versions. Use `add` instead.
|
|
|
|
Adds the data from the given URL to the vector db.
|
|
Loads the data, chunks it, create embedding for each chunk
|
|
and then stores the embedding to vector database.
|
|
|
|
:param source: The data to embed, can be a URL, local file or raw content, depending on the data type.
|
|
:param data_type: Optional. Automatically detected, but can be forced with this argument.
|
|
The type of the data to add.
|
|
:param metadata: Optional. Metadata associated with the data source.
|
|
:param config: Optional. The `AddConfig` instance to use as configuration
|
|
options.
|
|
:return: md5-hash of the source, in hexadecimal representation.
|
|
"""
|
|
logging.warning(
|
|
"The `add_local` method is deprecated and will be removed in future versions. Please use the `add` method for both local and remote files." # noqa: E501
|
|
)
|
|
return self.add(source=source, data_type=data_type, metadata=metadata, config=config)
|
|
|
|
def load_and_embed(self, loader: BaseLoader, chunker: BaseChunker, src, metadata=None, source_id=None):
|
|
"""
|
|
Loads the data from the given URL, chunks it, and adds it to database.
|
|
|
|
:param loader: The loader to use to load the data.
|
|
:param chunker: The chunker to use to chunk the data.
|
|
:param src: The data to be handled by the loader. Can be a URL for
|
|
remote sources or local content for local loaders.
|
|
:param metadata: Optional. Metadata associated with the data source.
|
|
:param source_id: Hexadecimal hash of the source.
|
|
:return: (List) documents (embedded text), (List) metadata, (list) ids, (int) number of chunks
|
|
"""
|
|
embeddings_data = chunker.create_chunks(loader, src)
|
|
|
|
# spread chunking results
|
|
documents = embeddings_data["documents"]
|
|
metadatas = embeddings_data["metadatas"]
|
|
ids = embeddings_data["ids"]
|
|
|
|
# get existing ids, and discard doc if any common id exist.
|
|
where = {"app_id": self.config.id} if self.config.id is not None else {}
|
|
# where={"url": src}
|
|
existing_ids = self.db.get(
|
|
ids=ids,
|
|
where=where, # optional filter
|
|
)
|
|
|
|
if len(existing_ids):
|
|
data_dict = {id: (doc, meta) for id, doc, meta in zip(ids, documents, metadatas)}
|
|
data_dict = {id: value for id, value in data_dict.items() if id not in existing_ids}
|
|
|
|
if not data_dict:
|
|
print(f"All data from {src} already exists in the database.")
|
|
# Make sure to return a matching return type
|
|
return [], [], [], 0
|
|
|
|
ids = list(data_dict.keys())
|
|
documents, metadatas = zip(*data_dict.values())
|
|
|
|
# Loop though all metadatas and add extras.
|
|
new_metadatas = []
|
|
for m in metadatas:
|
|
# Add app id in metadatas so that they can be queried on later
|
|
if self.config.id:
|
|
m["app_id"] = self.config.id
|
|
|
|
# Add hashed source
|
|
m["hash"] = source_id
|
|
|
|
# Note: Metadata is the function argument
|
|
if metadata:
|
|
# Spread whatever is in metadata into the new object.
|
|
m.update(metadata)
|
|
|
|
new_metadatas.append(m)
|
|
metadatas = new_metadatas
|
|
|
|
# Count before, to calculate a delta in the end.
|
|
chunks_before_addition = self.db.count()
|
|
|
|
self.db.add(documents=documents, metadatas=metadatas, ids=ids)
|
|
count_new_chunks = self.db.count() - chunks_before_addition
|
|
print((f"Successfully saved {src} ({chunker.data_type}). New chunks count: {count_new_chunks}"))
|
|
return list(documents), metadatas, ids, count_new_chunks
|
|
|
|
def _format_result(self, results):
|
|
return [
|
|
(Document(page_content=result[0], metadata=result[1] or {}), result[2])
|
|
for result in zip(
|
|
results["documents"][0],
|
|
results["metadatas"][0],
|
|
results["distances"][0],
|
|
)
|
|
]
|
|
|
|
def retrieve_from_database(self, input_query, config: Optional[BaseLlmConfig] = None, where=None):
|
|
"""
|
|
Queries the vector database based on the given input query.
|
|
Gets relevant doc based on the query
|
|
|
|
:param input_query: The query to use.
|
|
:param config: The query configuration.
|
|
:param where: Optional. A dictionary of key-value pairs to filter the database results.
|
|
:return: The content of the document that matched your query.
|
|
"""
|
|
query_config = config or self.llm.config
|
|
|
|
if where is not None:
|
|
where = where
|
|
elif query_config is not None and query_config.where is not None:
|
|
where = query_config.where
|
|
else:
|
|
where = {}
|
|
|
|
if self.config.id is not None:
|
|
where.update({"app_id": self.config.id})
|
|
|
|
contents = self.db.query(
|
|
input_query=input_query,
|
|
n_results=query_config.number_documents,
|
|
where=where,
|
|
)
|
|
|
|
return contents
|
|
|
|
def query(self, input_query, config: BaseLlmConfig = None, dry_run=False, where=None):
|
|
"""
|
|
Queries the vector database based on the given input query.
|
|
Gets relevant doc based on the query and then passes it to an
|
|
LLM as context to get the answer.
|
|
|
|
:param input_query: The query to use.
|
|
:param config: Optional. The `LlmConfig` instance to use as configuration options.
|
|
This is used for one method call. To persistently use a config, declare it during app init.
|
|
:param dry_run: Optional. A dry run does everything except send the resulting prompt to
|
|
the LLM. The purpose is to test the prompt, not the response.
|
|
You can use it to test your prompt, including the context provided
|
|
by the vector database's doc retrieval.
|
|
The only thing the dry run does not consider is the cut-off due to
|
|
the `max_tokens` parameter.
|
|
:param where: Optional. A dictionary of key-value pairs to filter the database results.
|
|
:return: The answer to the query.
|
|
"""
|
|
contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where)
|
|
answer = self.llm.query(input_query=input_query, contexts=contexts, config=config, dry_run=dry_run)
|
|
|
|
# Send anonymous telemetry
|
|
thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("query",))
|
|
thread_telemetry.start()
|
|
|
|
return answer
|
|
|
|
def chat(self, input_query, config: BaseLlmConfig = None, dry_run=False, where=None):
|
|
"""
|
|
Queries the vector database on the given input query.
|
|
Gets relevant doc based on the query and then passes it to an
|
|
LLM as context to get the answer.
|
|
|
|
Maintains the whole conversation in memory.
|
|
:param input_query: The query to use.
|
|
:param config: Optional. The `LlmConfig` instance to use as configuration options.
|
|
This is used for one method call. To persistently use a config, declare it during app init.
|
|
:param dry_run: Optional. A dry run does everything except send the resulting prompt to
|
|
the LLM. The purpose is to test the prompt, not the response.
|
|
You can use it to test your prompt, including the context provided
|
|
by the vector database's doc retrieval.
|
|
The only thing the dry run does not consider is the cut-off due to
|
|
the `max_tokens` parameter.
|
|
:param where: Optional. A dictionary of key-value pairs to filter the database results.
|
|
:return: The answer to the query.
|
|
"""
|
|
contexts = self.retrieve_from_database(input_query=input_query, config=config, where=where)
|
|
answer = self.llm.chat(input_query=input_query, contexts=contexts, config=config, dry_run=dry_run)
|
|
|
|
# Send anonymous telemetry
|
|
thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("chat",))
|
|
thread_telemetry.start()
|
|
|
|
return answer
|
|
|
|
def set_collection(self, collection_name):
|
|
"""
|
|
Set the collection to use.
|
|
|
|
:param collection_name: The name of the collection to use.
|
|
"""
|
|
self.db.set_collection_name(collection_name)
|
|
# Create the collection if it does not exist
|
|
self.db._get_or_create_collection(collection_name)
|
|
# TODO: Check whether it is necessary to assign to the `self.collection` attribute,
|
|
# since the main purpose is the creation.
|
|
|
|
def count(self) -> int:
|
|
"""
|
|
Count the number of embeddings.
|
|
|
|
DEPRECATED IN FAVOR OF `db.count()`
|
|
|
|
:return: The number of embeddings.
|
|
"""
|
|
logging.warning("DEPRECATION WARNING: Please use `db.count()` instead of `count()`.")
|
|
return self.db.count()
|
|
|
|
def reset(self):
|
|
"""
|
|
Resets the database. Deletes all embeddings irreversibly.
|
|
`App` does not have to be reinitialized after using this method.
|
|
|
|
DEPRECATED IN FAVOR OF `db.reset()`
|
|
"""
|
|
# Send anonymous telemetry
|
|
thread_telemetry = threading.Thread(target=self._send_telemetry_event, args=("reset",))
|
|
thread_telemetry.start()
|
|
|
|
logging.warning("DEPRECATION WARNING: Please use `db.reset()` instead of `reset()`.")
|
|
self.db.reset()
|
|
|
|
@retry(stop=stop_after_attempt(3), wait=wait_fixed(1))
|
|
def _send_telemetry_event(self, method: str, extra_metadata: Optional[dict] = None):
|
|
if not self.config.collect_metrics:
|
|
return
|
|
|
|
with threading.Lock():
|
|
url = "https://api.embedchain.ai/api/v1/telemetry/"
|
|
metadata = {
|
|
"s_id": self.s_id,
|
|
"version": importlib.metadata.version(__package__ or __name__),
|
|
"method": method,
|
|
"language": "py",
|
|
"u_id": self.u_id,
|
|
}
|
|
if extra_metadata:
|
|
metadata.update(extra_metadata)
|
|
|
|
response = requests.post(url, json={"metadata": metadata})
|
|
if response.status_code != 200:
|
|
logging.warning(f"Telemetry event failed with status code {response.status_code}")
|