from typing import List

from fastapi.params import Param
from llama_index.core import Document

from app.utils.milvus_utils import milvus_service
from app.utils.next_id import next_id


def add_knowledge_route(app):
  @app.get("/knowledge/search")
  async def knowledge_search(text: str = Param()):
    print("text", text)
    query_result = await  milvus_service.async_search(text or "hello")
    return query_result

  @app.post("/knowledge/embed_text")
  async def knowledge_embed_text(text_list: List[str]):
    id_list = await next_id(len(text_list))
    document_list = [
      Document(
        text=text,
        doc_id=id_list[index],
        metadata={"kb_id": "abc"}
      )
      for index, text in enumerate(text_list)
    ]
    await milvus_service.async_create_index_from_documents(document_list)
    search_result = await milvus_service.async_search("hello")
    return {
      "result": "嵌入成功:",
      "origin_documents": [document.to_dict() for document in document_list],
      "search_documents": search_result,
    }

  @app.post("/knowledge/delete")
  async def knowledge_embed_text(body: dict):
    await milvus_service.async_delete(body.get('id'))
    return {
      "result": "删除成功",
    }