{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d029ad67",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_milvus import BM25BuiltInFunction, Milvus\n",
    "from typing import List\n",
    "URI = \"http://10.10.10.9:19530\"\n",
    "tongyiKey = \"sk-9464b2498c184982a9fe9d2c2e725ab5\"\n",
    "from langchain_community.embeddings import DashScopeEmbeddings\n",
    "embeddings = DashScopeEmbeddings(\n",
    "    model=\"text-embedding-v3\",\n",
    "    dashscope_api_key= tongyiKey, \n",
    ")\n",
    "memVectorstore = Milvus(\n",
    "    embedding_function=embeddings,\n",
    "    connection_args={\"uri\": URI, \"token\": \"root:Milvus\", \"db_name\": \"bbit_ai_lab\"},\n",
    "    collection_name=\"memory\",\n",
    "    index_params={\"index_type\": \"FLAT\", \"metric_type\": \"L2\"},\n",
    "    consistency_level=\"Strong\",\n",
    "    auto_id=True,\n",
    "\n",
    "    primary_field = \"id\",\n",
    "    text_field=\"text\",\n",
    "    vector_field=\"vector\",\n",
    "    partition_key_field = \"ai_id\",\n",
    "    enable_dynamic_field = True,\n",
    "    drop_old=False,  # set to True if seeking to drop the collection with that name if it exists\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a480053b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str:\n",
    "    print(\"ai_id是：\" , ai_ids)\n",
    "    \"\"\"\n",
    "    根据关键词和 ai_ids 列表，在知识库中检索相关内容，并返回整理后的文本字符串\n",
    "    \"\"\"\n",
    "    # 构建过滤表达式：只查 kn_ids 范围内的\n",
    "    if ai_ids:\n",
    "        ids_expr = \" or \".join([f'ai_id == \"{kid}\"' for kid in ai_ids])\n",
    "        expr = f\"({ids_expr})\"\n",
    "    else:\n",
    "        expr = \"\"  # 不限制 kn_id todo 实际上应该不反悔任何内容\n",
    "    \n",
    "    result = knVectorstore.similarity_search(\n",
    "        query=key_words,\n",
    "        k=5,  # 可调节返回条数\n",
    "        expr=expr\n",
    "    )\n",
    "    \n",
    "    # 整理成字符串\n",
    "    doc_texts = []\n",
    "    for idx, doc in enumerate(result, start=1):\n",
    "        text = doc.page_content.strip()\n",
    "        if text:\n",
    "            # 可以加个编号，便于LLM区分\n",
    "            doc_texts.append(f\"[记忆{idx}]: {text}\")\n",
    "    \n",
    "    # 拼成一个大字符串，用换行隔开\n",
    "    combined_text = \"\\n\\n\".join(doc_texts)\n",
    "    return combined_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "36759de5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ai_id是： ['3730f279-8b56-46ec-bde9-8a9e6c27f021']\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'knVectorstore' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mget_memory_by_key_words\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m共育室 部署 地方\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m3730f279-8b56-46ec-bde9-8a9e6c27f021\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[2], line 13\u001b[0m, in \u001b[0;36mget_memory_by_key_words\u001b[0;34m(key_words, ai_ids)\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     11\u001b[0m     expr \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m  \u001b[38;5;66;03m# 不限制 kn_id todo 实际上应该不反悔任何内容\u001b[39;00m\n\u001b[0;32m---> 13\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mknVectorstore\u001b[49m\u001b[38;5;241m.\u001b[39msimilarity_search(\n\u001b[1;32m     14\u001b[0m     query\u001b[38;5;241m=\u001b[39mkey_words,\n\u001b[1;32m     15\u001b[0m     k\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m5\u001b[39m,  \u001b[38;5;66;03m# 可调节返回条数\u001b[39;00m\n\u001b[1;32m     16\u001b[0m     expr\u001b[38;5;241m=\u001b[39mexpr\n\u001b[1;32m     17\u001b[0m )\n\u001b[1;32m     19\u001b[0m \u001b[38;5;66;03m# 整理成字符串\u001b[39;00m\n\u001b[1;32m     20\u001b[0m doc_texts \u001b[38;5;241m=\u001b[39m []\n",
      "\u001b[0;31mNameError\u001b[0m: name 'knVectorstore' is not defined"
     ]
    }
   ],
   "source": [
    "get_memory_by_key_words(\"共育室 部署 地方\",[\"3730f279-8b56-46ec-bde9-8a9e6c27f021\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lang",
   "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.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}