{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b1f2ff89",
   "metadata": {},
   "outputs": [],
   "source": [
    "## 准备部分\n",
    "# 指定Embed对象\n",
    "embed_object = \"labors-law\"\n",
    "# 指定用户问题\n",
    "query_text = \"对于劳动合同的解除有什么规定?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b9d9fa0-1c73-4cef-abc4-397458215159",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入必要的包\n",
    "import time\n",
    "import os\n",
    "import sys\n",
    "import json\n",
    "from chromadb import Client\n",
    "\n",
    "# 获取当前脚本的目录\n",
    "current_dir = os.getcwd()\n",
    "\n",
    "# 读取设定文件\n",
    "config_file_path = os.path.join(current_dir, \"..\", \"setup\", \"config_embed.json\")\n",
    "\n",
    "# 读取配置项目\n",
    "try:\n",
    "    with open(config_file_path, \"r\", encoding=\"utf-8\") as f:\n",
    "        dict_config = json.load(f)\n",
    "    # 读取通用项目\n",
    "    API_URL = dict_config[\"API_URL\"]\n",
    "    cohere_access_token = dict_config[\"cohere_access_token\"]\n",
    "    custom_proxies = dict_config[\"custom_proxies\"]\n",
    "    list_chroma_dir = dict_config[\"list_chroma_dir\"]\n",
    "    list_embed_file_path = dict_config[\"list_embed_file_path\"]\n",
    "    # 读取Embed对象项目\n",
    "    model_name = dict_config[\"docs\"][embed_object][\"model_name\"]\n",
    "    embed_file_name = dict_config[\"docs\"][embed_object][\"file_name\"]\n",
    "    split_chunk_size = dict_config[\"docs\"][embed_object][\"split_chunk_size\"]\n",
    "    split_overlap = dict_config[\"docs\"][embed_object][\"split_overlap\"]\n",
    "    model_batch_size = dict_config[\"docs\"][embed_object][\"model_batch_size\"]\n",
    "    collection_ids_prefix = dict_config[\"docs\"][embed_object][\"collection_ids_prefix\"]\n",
    "    embed_collection_name = dict_config[\"docs\"][embed_object][\"collection_name\"]\n",
    "except Exception as e:\n",
    "    print(\"配置文件读取失败\", e)\n",
    "    sys.exit(\"Stop processing\")\n",
    "\n",
    "# 回退到上级目录并构建相对路径\n",
    "chroma_dir = os.path.join(current_dir, *list_chroma_dir)\n",
    "embed_file_path = os.path.join(current_dir, *list_embed_file_path, embed_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "630b5215-b3ad-43e6-8f48-83b78abf9c10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb\n",
    "\n",
    "# 初始化 Chroma 客户端，指定之前的数据存储目录\n",
    "client = chromadb.PersistentClient(path=chroma_dir)\n",
    "\n",
    "# 获取 \"document_embeddings\" 集合\n",
    "collection = client.get_collection(embed_collection_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8fc38844",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 装载本地的Hugging Face模型\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer(\"jinaai/jina-embeddings-v3\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "667253d4-1a55-406a-b3d4-51f4b9957aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 查询文本\n",
    "# query_embedding = get_embeddings([query_text])  # 使用之前的 get_embedding 函数生成嵌入\n",
    "query_embedding = model.encode([query_text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f3df429f-b4d6-4e66-803b-ae5e9ed95922",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document: ['一的，劳动者可以随时通知用人单位解除劳动合同：（一）在试用期内的；（二）用人单位以暴力、威胁或者非法限制人身自由的手段强迫劳动的；（三）用人单位未按照劳动合同约定支付劳动报酬或者提供劳动条件的。第三十三条\\u3000企业职工一方与企业可以就劳动报酬、工作时间、休息休假、劳动安全卫生、保险福利等事项，签订集体合同。集体合同草案应当提交职工代表大会或者全体职工讨论通过。集体合同由工会代表职工与企业签订；没有建立', '止条件出现，劳动合同即行终止。第二十四条\\u3000经劳动合同当事人协商一致，劳动合同可以解除。第二十五条\\u3000劳动者有下列情形之一的，用人单位可以解除劳动合同：（一）在试用期间被证明不符合录用条件的；（二）严重违反劳动纪律或者用人单位规章制度的；（三）严重失职，营私舞弊，对用人单位利益造成重大损害的；（四）被依法追究刑事责任的。第二十六条\\u3000有下列情形之一的，用人单位可以解除劳动合同，但是应当提前三十日以书面形', '在孕期、产期、哺乳期内的；（四）法律、行政法规规定的其他情形。第三十条\\u3000用人单位解除劳动合同，工会认为不适当的，有权提出意见。如果用人单位违反法律、法规或者劳动合同，工会有权要求重新处理；劳动者申请仲裁或者提起诉讼的，工会应当依法给予支持和帮助。第三十一条\\u3000劳动者解除劳动合同，应当提前三十日以书面形式通知用人单位。第三十二条\\u3000有下列情形之一的，劳动者可以随时通知用人单位解除劳动合同：（一）在试用期', '之一的，用人单位可以解除劳动合同，但是应当提前三十日以书面形式通知劳动者本人：（一）劳动者患病或者非因工负伤，医疗期满后，不能从事原工作也不能从事由用人单位另行安排的工作的；（二）劳动者不能胜任工作，经过培训或者调整工作岗位，仍不能胜任工作的；（三）劳动合同订立时所依据的客观情况发生重大变化，致使原劳动合同无法履行，经当事人协商不能就变更劳动合同达成协议的。第二十七条\\u3000用人单位濒临破产进行法定整顿', '动保护和劳动条件；（四）劳动报酬；（五）劳动纪律；（六）劳动合同终止的条件；（七）违反劳动合同的责任。劳动合同除前款规定的必备条款外，当事人可以协商约定其他内容。第二十条\\u3000劳动合同的期限分为有固定期限、无固定期限和以完成一定的工作为期限。11/22/24,']\n",
      "Metadata: [{'chunk': 21}, {'chunk': 15}, {'chunk': 20}, {'chunk': 16}, {'chunk': 12}]\n"
     ]
    }
   ],
   "source": [
    "# 查询集合，返回前5个匹配的文档\n",
    "results = collection.query(\n",
    "    # query_embeddings=query_embedding[\"embeddings\"],\n",
    "    query_embeddings=query_embedding,\n",
    "    n_results=5\n",
    ")\n",
    "\n",
    "# 显示查询结果\n",
    "for doc, metadata in zip(results['documents'], results['metadatas']):\n",
    "    print(\"Document:\", doc)\n",
    "    print(\"Metadata:\", metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "01ed14d1-71b5-42b3-82d6-47a053894a74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "    [\n",
      "        \"一的，劳动者可以随时通知用人单位解除劳动合同：（一）在试用期内的；（二）用人单位以暴力、威胁或者非法限制人身自由的手段强迫劳动的；（三）用人单位未按照劳动合同约定支付劳动报酬或者提供劳动条件的。第三十三条　企业职工一方与企业可以就劳动报酬、工作时间、休息休假、劳动安全卫生、保险福利等事项，签订集体合同。集体合同草案应当提交职工代表大会或者全体职工讨论通过。集体合同由工会代表职工与企业签订；没有建立\",\n",
      "        \"止条件出现，劳动合同即行终止。第二十四条　经劳动合同当事人协商一致，劳动合同可以解除。第二十五条　劳动者有下列情形之一的，用人单位可以解除劳动合同：（一）在试用期间被证明不符合录用条件的；（二）严重违反劳动纪律或者用人单位规章制度的；（三）严重失职，营私舞弊，对用人单位利益造成重大损害的；（四）被依法追究刑事责任的。第二十六条　有下列情形之一的，用人单位可以解除劳动合同，但是应当提前三十日以书面形\",\n",
      "        \"在孕期、产期、哺乳期内的；（四）法律、行政法规规定的其他情形。第三十条　用人单位解除劳动合同，工会认为不适当的，有权提出意见。如果用人单位违反法律、法规或者劳动合同，工会有权要求重新处理；劳动者申请仲裁或者提起诉讼的，工会应当依法给予支持和帮助。第三十一条　劳动者解除劳动合同，应当提前三十日以书面形式通知用人单位。第三十二条　有下列情形之一的，劳动者可以随时通知用人单位解除劳动合同：（一）在试用期\",\n",
      "        \"之一的，用人单位可以解除劳动合同，但是应当提前三十日以书面形式通知劳动者本人：（一）劳动者患病或者非因工负伤，医疗期满后，不能从事原工作也不能从事由用人单位另行安排的工作的；（二）劳动者不能胜任工作，经过培训或者调整工作岗位，仍不能胜任工作的；（三）劳动合同订立时所依据的客观情况发生重大变化，致使原劳动合同无法履行，经当事人协商不能就变更劳动合同达成协议的。第二十七条　用人单位濒临破产进行法定整顿\",\n",
      "        \"动保护和劳动条件；（四）劳动报酬；（五）劳动纪律；（六）劳动合同终止的条件；（七）违反劳动合同的责任。劳动合同除前款规定的必备条款外，当事人可以协商约定其他内容。第二十条　劳动合同的期限分为有固定期限、无固定期限和以完成一定的工作为期限。11/22/24,\"\n",
      "    ]\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "print(json.dumps(results[\"documents\"], indent=4, ensure_ascii=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3aa24e59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28863"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 释放内存\n",
    "del model\n",
    "import gc\n",
    "gc.collect()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}