from server.knowledge_base.kb_doc_api import *


def get_docs(query, top_k=3, score_threshold=1, knowledge_base_name='agriculture'):
    docs = search_docs(query, knowledge_base_name, top_k, score_threshold)
    context = "\n".join([doc.page_content for doc in docs])
    return context


def get_complete_docs(query, top_k=3, score_threshold=1, knowledge_base_name='agriculture'):
    docs = search_docs(query, knowledge_base_name, top_k, score_threshold)
    context = [doc.page_content for doc in docs]
    return context


if __name__ == '__main__':
    query = "玉米"
    top_k = 3
    score_threshold = 1
    knowledge_base_name = "agriculture"
    r = get_complete_docs(query, top_k, score_threshold, knowledge_base_name)
    print(r)