embedchain.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776
  1. import hashlib
  2. import json
  3. import logging
  4. from typing import Any, Optional, Union
  5. from dotenv import load_dotenv
  6. from langchain.docstore.document import Document
  7. from embedchain.cache import adapt, get_gptcache_session, gptcache_data_convert, gptcache_update_cache_callback
  8. from embedchain.chunkers.base_chunker import BaseChunker
  9. from embedchain.config import AddConfig, BaseLlmConfig, ChunkerConfig
  10. from embedchain.config.base_app_config import BaseAppConfig
  11. from embedchain.core.db.models import ChatHistory, DataSource
  12. from embedchain.data_formatter import DataFormatter
  13. from embedchain.embedder.base import BaseEmbedder
  14. from embedchain.helpers.json_serializable import JSONSerializable
  15. from embedchain.llm.base import BaseLlm
  16. from embedchain.loaders.base_loader import BaseLoader
  17. from embedchain.models.data_type import DataType, DirectDataType, IndirectDataType, SpecialDataType
  18. from embedchain.utils.misc import detect_datatype, is_valid_json_string
  19. from embedchain.vectordb.base import BaseVectorDB
  20. load_dotenv()
  21. logger = logging.getLogger(__name__)
  22. class EmbedChain(JSONSerializable):
  23. def __init__(
  24. self,
  25. config: BaseAppConfig,
  26. llm: BaseLlm,
  27. db: BaseVectorDB = None,
  28. embedder: BaseEmbedder = None,
  29. system_prompt: Optional[str] = None,
  30. ):
  31. """
  32. Initializes the EmbedChain instance, sets up a vector DB client and
  33. creates a collection.
  34. :param config: Configuration just for the app, not the db or llm or embedder.
  35. :type config: BaseAppConfig
  36. :param llm: Instance of the LLM you want to use.
  37. :type llm: BaseLlm
  38. :param db: Instance of the Database to use, defaults to None
  39. :type db: BaseVectorDB, optional
  40. :param embedder: instance of the embedder to use, defaults to None
  41. :type embedder: BaseEmbedder, optional
  42. :param system_prompt: System prompt to use in the llm query, defaults to None
  43. :type system_prompt: Optional[str], optional
  44. :raises ValueError: No database or embedder provided.
  45. """
  46. self.config = config
  47. self.cache_config = None
  48. self.memory_config = None
  49. self.mem0_client = None
  50. # Llm
  51. self.llm = llm
  52. # Database has support for config assignment for backwards compatibility
  53. if db is None and (not hasattr(self.config, "db") or self.config.db is None):
  54. raise ValueError("App requires Database.")
  55. self.db = db or self.config.db
  56. # Embedder
  57. if embedder is None:
  58. raise ValueError("App requires Embedder.")
  59. self.embedder = embedder
  60. # Initialize database
  61. self.db._set_embedder(self.embedder)
  62. self.db._initialize()
  63. # Set collection name from app config for backwards compatibility.
  64. if config.collection_name:
  65. self.db.set_collection_name(config.collection_name)
  66. # Add variables that are "shortcuts"
  67. if system_prompt:
  68. self.llm.config.system_prompt = system_prompt
  69. # Fetch the history from the database if exists
  70. self.llm.update_history(app_id=self.config.id)
  71. # Attributes that aren't subclass related.
  72. self.user_asks = []
  73. self.chunker: Optional[ChunkerConfig] = None
  74. @property
  75. def collect_metrics(self):
  76. return self.config.collect_metrics
  77. @collect_metrics.setter
  78. def collect_metrics(self, value):
  79. if not isinstance(value, bool):
  80. raise ValueError(f"Boolean value expected but got {type(value)}.")
  81. self.config.collect_metrics = value
  82. @property
  83. def online(self):
  84. return self.llm.config.online
  85. @online.setter
  86. def online(self, value):
  87. if not isinstance(value, bool):
  88. raise ValueError(f"Boolean value expected but got {type(value)}.")
  89. self.llm.config.online = value
  90. def add(
  91. self,
  92. source: Any,
  93. data_type: Optional[DataType] = None,
  94. metadata: Optional[dict[str, Any]] = None,
  95. config: Optional[AddConfig] = None,
  96. dry_run=False,
  97. loader: Optional[BaseLoader] = None,
  98. chunker: Optional[BaseChunker] = None,
  99. **kwargs: Optional[dict[str, Any]],
  100. ):
  101. """
  102. Adds the data from the given URL to the vector db.
  103. Loads the data, chunks it, create embedding for each chunk
  104. and then stores the embedding to vector database.
  105. :param source: The data to embed, can be a URL, local file or raw content, depending on the data type.
  106. :type source: Any
  107. :param data_type: Automatically detected, but can be forced with this argument. The type of the data to add,
  108. defaults to None
  109. :type data_type: Optional[DataType], optional
  110. :param metadata: Metadata associated with the data source., defaults to None
  111. :type metadata: Optional[dict[str, Any]], optional
  112. :param config: The `AddConfig` instance to use as configuration options., defaults to None
  113. :type config: Optional[AddConfig], optional
  114. :raises ValueError: Invalid data type
  115. :param dry_run: Optional. A dry run displays the chunks to ensure that the loader and chunker work as intended.
  116. defaults to False
  117. :type dry_run: bool
  118. :param loader: The loader to use to load the data, defaults to None
  119. :type loader: BaseLoader, optional
  120. :param chunker: The chunker to use to chunk the data, defaults to None
  121. :type chunker: BaseChunker, optional
  122. :param kwargs: To read more params for the query function
  123. :type kwargs: dict[str, Any]
  124. :return: source_hash, a md5-hash of the source, in hexadecimal representation.
  125. :rtype: str
  126. """
  127. if config is not None:
  128. pass
  129. elif self.chunker is not None:
  130. config = AddConfig(chunker=self.chunker)
  131. else:
  132. config = AddConfig()
  133. try:
  134. DataType(source)
  135. logger.warning(
  136. 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
  137. )
  138. logger.warning(
  139. "Embedchain is swapping the arguments for you. This functionality might be deprecated in the future, so please adjust your code." # noqa #E501
  140. )
  141. source, data_type = data_type, source
  142. except ValueError:
  143. pass
  144. if data_type:
  145. try:
  146. data_type = DataType(data_type)
  147. except ValueError:
  148. logger.info(
  149. f"Invalid data_type: '{data_type}', using `custom` instead.\n Check docs to pass the valid data type: `https://docs.embedchain.ai/data-sources/overview`" # noqa: E501
  150. )
  151. data_type = DataType.CUSTOM
  152. if not data_type:
  153. data_type = detect_datatype(source)
  154. # `source_hash` is the md5 hash of the source argument
  155. source_hash = hashlib.md5(str(source).encode("utf-8")).hexdigest()
  156. self.user_asks.append([source, data_type.value, metadata])
  157. data_formatter = DataFormatter(data_type, config, loader, chunker)
  158. documents, metadatas, _ids, new_chunks = self._load_and_embed(
  159. data_formatter.loader, data_formatter.chunker, source, metadata, source_hash, config, dry_run, **kwargs
  160. )
  161. if data_type in {DataType.DOCS_SITE}:
  162. self.is_docs_site_instance = True
  163. # Convert the source to a string if it is not already
  164. if not isinstance(source, str):
  165. source = str(source)
  166. # Insert the data into the 'ec_data_sources' table
  167. self.db_session.add(
  168. DataSource(
  169. hash=source_hash,
  170. app_id=self.config.id,
  171. type=data_type.value,
  172. value=source,
  173. metadata=json.dumps(metadata),
  174. )
  175. )
  176. try:
  177. self.db_session.commit()
  178. except Exception as e:
  179. logger.error(f"Error adding data source: {e}")
  180. self.db_session.rollback()
  181. if dry_run:
  182. data_chunks_info = {"chunks": documents, "metadata": metadatas, "count": len(documents), "type": data_type}
  183. logger.debug(f"Dry run info : {data_chunks_info}")
  184. return data_chunks_info
  185. # Send anonymous telemetry
  186. if self.config.collect_metrics:
  187. # it's quicker to check the variable twice than to count words when they won't be submitted.
  188. word_count = data_formatter.chunker.get_word_count(documents)
  189. # Send anonymous telemetry
  190. event_properties = {
  191. **self._telemetry_props,
  192. "data_type": data_type.value,
  193. "word_count": word_count,
  194. "chunks_count": new_chunks,
  195. }
  196. self.telemetry.capture(event_name="add", properties=event_properties)
  197. return source_hash
  198. def _get_existing_doc_id(self, chunker: BaseChunker, src: Any):
  199. """
  200. Get id of existing document for a given source, based on the data type
  201. """
  202. # Find existing embeddings for the source
  203. # Depending on the data type, existing embeddings are checked for.
  204. if chunker.data_type.value in [item.value for item in DirectDataType]:
  205. # DirectDataTypes can't be updated.
  206. # Think of a text:
  207. # Either it's the same, then it won't change, so it's not an update.
  208. # Or it's different, then it will be added as a new text.
  209. return None
  210. elif chunker.data_type.value in [item.value for item in IndirectDataType]:
  211. # These types have an indirect source reference
  212. # As long as the reference is the same, they can be updated.
  213. where = {"url": src}
  214. if chunker.data_type == DataType.JSON and is_valid_json_string(src):
  215. url = hashlib.sha256((src).encode("utf-8")).hexdigest()
  216. where = {"url": url}
  217. if self.config.id is not None:
  218. where.update({"app_id": self.config.id})
  219. existing_embeddings = self.db.get(
  220. where=where,
  221. limit=1,
  222. )
  223. if len(existing_embeddings.get("metadatas", [])) > 0:
  224. return existing_embeddings["metadatas"][0]["doc_id"]
  225. else:
  226. return None
  227. elif chunker.data_type.value in [item.value for item in SpecialDataType]:
  228. # These types don't contain indirect references.
  229. # Through custom logic, they can be attributed to a source and be updated.
  230. if chunker.data_type == DataType.QNA_PAIR:
  231. # QNA_PAIRs update the answer if the question already exists.
  232. where = {"question": src[0]}
  233. if self.config.id is not None:
  234. where.update({"app_id": self.config.id})
  235. existing_embeddings = self.db.get(
  236. where=where,
  237. limit=1,
  238. )
  239. if len(existing_embeddings.get("metadatas", [])) > 0:
  240. return existing_embeddings["metadatas"][0]["doc_id"]
  241. else:
  242. return None
  243. else:
  244. raise NotImplementedError(
  245. f"SpecialDataType {chunker.data_type} must have a custom logic to check for existing data"
  246. )
  247. else:
  248. raise TypeError(
  249. f"{chunker.data_type} is type {type(chunker.data_type)}. "
  250. "When it should be DirectDataType, IndirectDataType or SpecialDataType."
  251. )
  252. def _load_and_embed(
  253. self,
  254. loader: BaseLoader,
  255. chunker: BaseChunker,
  256. src: Any,
  257. metadata: Optional[dict[str, Any]] = None,
  258. source_hash: Optional[str] = None,
  259. add_config: Optional[AddConfig] = None,
  260. dry_run=False,
  261. **kwargs: Optional[dict[str, Any]],
  262. ):
  263. """
  264. Loads the data from the given URL, chunks it, and adds it to database.
  265. :param loader: The loader to use to load the data.
  266. :type loader: BaseLoader
  267. :param chunker: The chunker to use to chunk the data.
  268. :type chunker: BaseChunker
  269. :param src: The data to be handled by the loader. Can be a URL for
  270. remote sources or local content for local loaders.
  271. :type src: Any
  272. :param metadata: Metadata associated with the data source.
  273. :type metadata: dict[str, Any], optional
  274. :param source_hash: Hexadecimal hash of the source.
  275. :type source_hash: str, optional
  276. :param add_config: The `AddConfig` instance to use as configuration options.
  277. :type add_config: AddConfig, optional
  278. :param dry_run: A dry run returns chunks and doesn't update DB.
  279. :type dry_run: bool, defaults to False
  280. :return: (list) documents (embedded text), (list) metadata, (list) ids, (int) number of chunks
  281. """
  282. existing_doc_id = self._get_existing_doc_id(chunker=chunker, src=src)
  283. app_id = self.config.id if self.config is not None else None
  284. # Create chunks
  285. embeddings_data = chunker.create_chunks(loader, src, app_id=app_id, config=add_config.chunker)
  286. # spread chunking results
  287. documents = embeddings_data["documents"]
  288. metadatas = embeddings_data["metadatas"]
  289. ids = embeddings_data["ids"]
  290. new_doc_id = embeddings_data["doc_id"]
  291. if existing_doc_id and existing_doc_id == new_doc_id:
  292. logger.info("Doc content has not changed. Skipping creating chunks and embeddings")
  293. return [], [], [], 0
  294. # this means that doc content has changed.
  295. if existing_doc_id and existing_doc_id != new_doc_id:
  296. logger.info("Doc content has changed. Recomputing chunks and embeddings intelligently.")
  297. self.db.delete({"doc_id": existing_doc_id})
  298. # get existing ids, and discard doc if any common id exist.
  299. where = {"url": src}
  300. if chunker.data_type == DataType.JSON and is_valid_json_string(src):
  301. url = hashlib.sha256((src).encode("utf-8")).hexdigest()
  302. where = {"url": url}
  303. # if data type is qna_pair, we check for question
  304. if chunker.data_type == DataType.QNA_PAIR:
  305. where = {"question": src[0]}
  306. if self.config.id is not None:
  307. where["app_id"] = self.config.id
  308. db_result = self.db.get(ids=ids, where=where) # optional filter
  309. existing_ids = set(db_result["ids"])
  310. if len(existing_ids):
  311. data_dict = {id: (doc, meta) for id, doc, meta in zip(ids, documents, metadatas)}
  312. data_dict = {id: value for id, value in data_dict.items() if id not in existing_ids}
  313. if not data_dict:
  314. src_copy = src
  315. if len(src_copy) > 50:
  316. src_copy = src[:50] + "..."
  317. logger.info(f"All data from {src_copy} already exists in the database.")
  318. # Make sure to return a matching return type
  319. return [], [], [], 0
  320. ids = list(data_dict.keys())
  321. documents, metadatas = zip(*data_dict.values())
  322. # Loop though all metadatas and add extras.
  323. new_metadatas = []
  324. for m in metadatas:
  325. # Add app id in metadatas so that they can be queried on later
  326. if self.config.id:
  327. m["app_id"] = self.config.id
  328. # Add hashed source
  329. m["hash"] = source_hash
  330. # Note: Metadata is the function argument
  331. if metadata:
  332. # Spread whatever is in metadata into the new object.
  333. m.update(metadata)
  334. new_metadatas.append(m)
  335. metadatas = new_metadatas
  336. if dry_run:
  337. return list(documents), metadatas, ids, 0
  338. # Count before, to calculate a delta in the end.
  339. chunks_before_addition = self.db.count()
  340. # Filter out empty documents and ensure they meet the API requirements
  341. valid_documents = [doc for doc in documents if doc and isinstance(doc, str)]
  342. documents = valid_documents
  343. # Chunk documents into batches of 2048 and handle each batch
  344. # helps wigth large loads of embeddings that hit OpenAI limits
  345. document_batches = [documents[i : i + 2048] for i in range(0, len(documents), 2048)]
  346. metadata_batches = [metadatas[i : i + 2048] for i in range(0, len(metadatas), 2048)]
  347. id_batches = [ids[i : i + 2048] for i in range(0, len(ids), 2048)]
  348. for batch_docs, batch_meta, batch_ids in zip(document_batches, metadata_batches, id_batches):
  349. try:
  350. # Add only valid batches
  351. if batch_docs:
  352. self.db.add(documents=batch_docs, metadatas=batch_meta, ids=batch_ids, **kwargs)
  353. except Exception as e:
  354. logger.info(f"Failed to add batch due to a bad request: {e}")
  355. # Handle the error, e.g., by logging, retrying, or skipping
  356. pass
  357. count_new_chunks = self.db.count() - chunks_before_addition
  358. logger.info(f"Successfully saved {str(src)[:100]} ({chunker.data_type}). New chunks count: {count_new_chunks}")
  359. return list(documents), metadatas, ids, count_new_chunks
  360. @staticmethod
  361. def _format_result(results):
  362. return [
  363. (Document(page_content=result[0], metadata=result[1] or {}), result[2])
  364. for result in zip(
  365. results["documents"][0],
  366. results["metadatas"][0],
  367. results["distances"][0],
  368. )
  369. ]
  370. def _retrieve_from_database(
  371. self,
  372. input_query: str,
  373. config: Optional[BaseLlmConfig] = None,
  374. where=None,
  375. citations: bool = False,
  376. **kwargs: Optional[dict[str, Any]],
  377. ) -> Union[list[tuple[str, str, str]], list[str]]:
  378. """
  379. Queries the vector database based on the given input query.
  380. Gets relevant doc based on the query
  381. :param input_query: The query to use.
  382. :type input_query: str
  383. :param config: The query configuration, defaults to None
  384. :type config: Optional[BaseLlmConfig], optional
  385. :param where: A dictionary of key-value pairs to filter the database results, defaults to None
  386. :type where: _type_, optional
  387. :param citations: A boolean to indicate if db should fetch citation source
  388. :type citations: bool
  389. :return: List of contents of the document that matched your query
  390. :rtype: list[str]
  391. """
  392. query_config = config or self.llm.config
  393. if where is not None:
  394. where = where
  395. else:
  396. where = {}
  397. if query_config is not None and query_config.where is not None:
  398. where = query_config.where
  399. if self.config.id is not None:
  400. where.update({"app_id": self.config.id})
  401. contexts = self.db.query(
  402. input_query=input_query,
  403. n_results=query_config.number_documents,
  404. where=where,
  405. citations=citations,
  406. **kwargs,
  407. )
  408. return contexts
  409. def query(
  410. self,
  411. input_query: str,
  412. config: BaseLlmConfig = None,
  413. dry_run=False,
  414. where: Optional[dict] = None,
  415. citations: bool = False,
  416. **kwargs: dict[str, Any],
  417. ) -> Union[tuple[str, list[tuple[str, dict]]], str, dict[str, Any]]:
  418. """
  419. Queries the vector database based on the given input query.
  420. Gets relevant doc based on the query and then passes it to an
  421. LLM as context to get the answer.
  422. :param input_query: The query to use.
  423. :type input_query: str
  424. :param config: The `BaseLlmConfig` instance to use as configuration options. This is used for one method call.
  425. To persistently use a config, declare it during app init., defaults to None
  426. :type config: BaseLlmConfig, optional
  427. :param dry_run: A dry run does everything except send the resulting prompt to
  428. the LLM. The purpose is to test the prompt, not the response., defaults to False
  429. :type dry_run: bool, optional
  430. :param where: A dictionary of key-value pairs to filter the database results., defaults to None
  431. :type where: dict[str, str], optional
  432. :param citations: A boolean to indicate if db should fetch citation source
  433. :type citations: bool
  434. :param kwargs: To read more params for the query function. Ex. we use citations boolean
  435. param to return context along with the answer
  436. :type kwargs: dict[str, Any]
  437. :return: The answer to the query, with citations if the citation flag is True
  438. or the dry run result
  439. :rtype: str, if citations is False and token_usage is False, otherwise if citations is true then
  440. tuple[str, list[tuple[str,str,str]]] and if token_usage is true then
  441. tuple[str, list[tuple[str,str,str]], dict[str, Any]]
  442. """
  443. contexts = self._retrieve_from_database(
  444. input_query=input_query, config=config, where=where, citations=citations, **kwargs
  445. )
  446. if citations and len(contexts) > 0 and isinstance(contexts[0], tuple):
  447. contexts_data_for_llm_query = list(map(lambda x: x[0], contexts))
  448. else:
  449. contexts_data_for_llm_query = contexts
  450. if self.cache_config is not None:
  451. logger.info("Cache enabled. Checking cache...")
  452. answer = adapt(
  453. llm_handler=self.llm.query,
  454. cache_data_convert=gptcache_data_convert,
  455. update_cache_callback=gptcache_update_cache_callback,
  456. session=get_gptcache_session(session_id=self.config.id),
  457. input_query=input_query,
  458. contexts=contexts_data_for_llm_query,
  459. config=config,
  460. dry_run=dry_run,
  461. )
  462. else:
  463. if self.llm.config.token_usage:
  464. answer, token_info = self.llm.query(
  465. input_query=input_query, contexts=contexts_data_for_llm_query, config=config, dry_run=dry_run
  466. )
  467. else:
  468. answer = self.llm.query(
  469. input_query=input_query, contexts=contexts_data_for_llm_query, config=config, dry_run=dry_run
  470. )
  471. # Send anonymous telemetry
  472. self.telemetry.capture(event_name="query", properties=self._telemetry_props)
  473. if citations:
  474. if self.llm.config.token_usage:
  475. return {"answer": answer, "contexts": contexts, "usage": token_info}
  476. return answer, contexts
  477. if self.llm.config.token_usage:
  478. return {"answer": answer, "usage": token_info}
  479. logger.warning(
  480. "Starting from v0.1.125 the return type of query method will be changed to tuple containing `answer`."
  481. )
  482. return answer
  483. def chat(
  484. self,
  485. input_query: str,
  486. config: Optional[BaseLlmConfig] = None,
  487. dry_run=False,
  488. session_id: str = "default",
  489. where: Optional[dict[str, str]] = None,
  490. citations: bool = False,
  491. **kwargs: dict[str, Any],
  492. ) -> Union[tuple[str, list[tuple[str, dict]]], str, dict[str, Any]]:
  493. """
  494. Queries the vector database on the given input query.
  495. Gets relevant doc based on the query and then passes it to an
  496. LLM as context to get the answer.
  497. Maintains the whole conversation in memory.
  498. :param input_query: The query to use.
  499. :type input_query: str
  500. :param config: The `BaseLlmConfig` instance to use as configuration options. This is used for one method call.
  501. To persistently use a config, declare it during app init., defaults to None
  502. :type config: BaseLlmConfig, optional
  503. :param dry_run: A dry run does everything except send the resulting prompt to
  504. the LLM. The purpose is to test the prompt, not the response., defaults to False
  505. :type dry_run: bool, optional
  506. :param session_id: The session id to use for chat history, defaults to 'default'.
  507. :type session_id: str, optional
  508. :param where: A dictionary of key-value pairs to filter the database results., defaults to None
  509. :type where: dict[str, str], optional
  510. :param citations: A boolean to indicate if db should fetch citation source
  511. :type citations: bool
  512. :param kwargs: To read more params for the query function. Ex. we use citations boolean
  513. param to return context along with the answer
  514. :type kwargs: dict[str, Any]
  515. :return: The answer to the query, with citations if the citation flag is True
  516. or the dry run result
  517. :rtype: str, if citations is False and token_usage is False, otherwise if citations is true then
  518. tuple[str, list[tuple[str,str,str]]] and if token_usage is true then
  519. tuple[str, list[tuple[str,str,str]], dict[str, Any]]
  520. """
  521. contexts = self._retrieve_from_database(
  522. input_query=input_query, config=config, where=where, citations=citations, **kwargs
  523. )
  524. if citations and len(contexts) > 0 and isinstance(contexts[0], tuple):
  525. contexts_data_for_llm_query = list(map(lambda x: x[0], contexts))
  526. else:
  527. contexts_data_for_llm_query = contexts
  528. memories = None
  529. if self.mem0_client:
  530. memories = self.mem0_client.search(
  531. query=input_query, agent_id=self.config.id, session_id=session_id, limit=self.memory_config.top_k
  532. )
  533. # Update the history beforehand so that we can handle multiple chat sessions in the same python session
  534. self.llm.update_history(app_id=self.config.id, session_id=session_id)
  535. if self.cache_config is not None:
  536. logger.debug("Cache enabled. Checking cache...")
  537. cache_id = f"{session_id}--{self.config.id}"
  538. answer = adapt(
  539. llm_handler=self.llm.chat,
  540. cache_data_convert=gptcache_data_convert,
  541. update_cache_callback=gptcache_update_cache_callback,
  542. session=get_gptcache_session(session_id=cache_id),
  543. input_query=input_query,
  544. contexts=contexts_data_for_llm_query,
  545. config=config,
  546. dry_run=dry_run,
  547. )
  548. else:
  549. logger.debug("Cache disabled. Running chat without cache.")
  550. if self.llm.config.token_usage:
  551. answer, token_info = self.llm.query(
  552. input_query=input_query,
  553. contexts=contexts_data_for_llm_query,
  554. config=config,
  555. dry_run=dry_run,
  556. memories=memories,
  557. )
  558. else:
  559. answer = self.llm.query(
  560. input_query=input_query,
  561. contexts=contexts_data_for_llm_query,
  562. config=config,
  563. dry_run=dry_run,
  564. memories=memories,
  565. )
  566. # Add to Mem0 memory if enabled
  567. # TODO: Might need to prepend with some text like:
  568. # "Remember user preferences from following user query: {input_query}"
  569. if self.mem0_client:
  570. self.mem0_client.add(data=input_query, agent_id=self.config.id, session_id=session_id)
  571. # add conversation in memory
  572. self.llm.add_history(self.config.id, input_query, answer, session_id=session_id)
  573. # Send anonymous telemetry
  574. self.telemetry.capture(event_name="chat", properties=self._telemetry_props)
  575. if citations:
  576. if self.llm.config.token_usage:
  577. return {"answer": answer, "contexts": contexts, "usage": token_info}
  578. return answer, contexts
  579. if self.llm.config.token_usage:
  580. return {"answer": answer, "usage": token_info}
  581. logger.warning(
  582. "Starting from v0.1.125 the return type of query method will be changed to tuple containing `answer`."
  583. )
  584. return answer
  585. def search(self, query, num_documents=3, where=None, raw_filter=None, namespace=None):
  586. """
  587. Search for similar documents related to the query in the vector database.
  588. Args:
  589. query (str): The query to use.
  590. num_documents (int, optional): Number of similar documents to fetch. Defaults to 3.
  591. where (dict[str, any], optional): Filter criteria for the search.
  592. raw_filter (dict[str, any], optional): Advanced raw filter criteria for the search.
  593. namespace (str, optional): The namespace to search in. Defaults to None.
  594. Raises:
  595. ValueError: If both `raw_filter` and `where` are used simultaneously.
  596. Returns:
  597. list[dict]: A list of dictionaries, each containing the 'context' and 'metadata' of a document.
  598. """
  599. # Send anonymous telemetry
  600. self.telemetry.capture(event_name="search", properties=self._telemetry_props)
  601. if raw_filter and where:
  602. raise ValueError("You can't use both `raw_filter` and `where` together.")
  603. filter_type = "raw_filter" if raw_filter else "where"
  604. filter_criteria = raw_filter if raw_filter else where
  605. params = {
  606. "input_query": query,
  607. "n_results": num_documents,
  608. "citations": True,
  609. "app_id": self.config.id,
  610. "namespace": namespace,
  611. filter_type: filter_criteria,
  612. }
  613. return [{"context": c[0], "metadata": c[1]} for c in self.db.query(**params)]
  614. def set_collection_name(self, name: str):
  615. """
  616. Set the name of the collection. A collection is an isolated space for vectors.
  617. Using `app.db.set_collection_name` method is preferred to this.
  618. :param name: Name of the collection.
  619. :type name: str
  620. """
  621. self.db.set_collection_name(name)
  622. # Create the collection if it does not exist
  623. self.db._get_or_create_collection(name)
  624. # TODO: Check whether it is necessary to assign to the `self.collection` attribute,
  625. # since the main purpose is the creation.
  626. def reset(self):
  627. """
  628. Resets the database. Deletes all embeddings irreversibly.
  629. `App` does not have to be reinitialized after using this method.
  630. """
  631. try:
  632. self.db_session.query(DataSource).filter_by(app_id=self.config.id).delete()
  633. self.db_session.query(ChatHistory).filter_by(app_id=self.config.id).delete()
  634. self.db_session.commit()
  635. except Exception as e:
  636. logger.error(f"Error deleting data sources: {e}")
  637. self.db_session.rollback()
  638. return None
  639. self.db.reset()
  640. self.delete_all_chat_history(app_id=self.config.id)
  641. # Send anonymous telemetry
  642. self.telemetry.capture(event_name="reset", properties=self._telemetry_props)
  643. def get_history(
  644. self,
  645. num_rounds: int = 10,
  646. display_format: bool = True,
  647. session_id: Optional[str] = "default",
  648. fetch_all: bool = False,
  649. ):
  650. history = self.llm.memory.get(
  651. app_id=self.config.id,
  652. session_id=session_id,
  653. num_rounds=num_rounds,
  654. display_format=display_format,
  655. fetch_all=fetch_all,
  656. )
  657. return history
  658. def delete_session_chat_history(self, session_id: str = "default"):
  659. self.llm.memory.delete(app_id=self.config.id, session_id=session_id)
  660. self.llm.update_history(app_id=self.config.id)
  661. def delete_all_chat_history(self, app_id: str):
  662. self.llm.memory.delete(app_id=app_id)
  663. self.llm.update_history(app_id=app_id)
  664. def delete(self, source_id: str):
  665. """
  666. Deletes the data from the database.
  667. :param source_hash: The hash of the source.
  668. :type source_hash: str
  669. """
  670. try:
  671. self.db_session.query(DataSource).filter_by(hash=source_id, app_id=self.config.id).delete()
  672. self.db_session.commit()
  673. except Exception as e:
  674. logger.error(f"Error deleting data sources: {e}")
  675. self.db_session.rollback()
  676. return None
  677. self.db.delete(where={"hash": source_id})
  678. logger.info(f"Successfully deleted {source_id}")
  679. # Send anonymous telemetry
  680. if self.config.collect_metrics:
  681. self.telemetry.capture(event_name="delete", properties=self._telemetry_props)