大模型RAG实战｜基于ThinkDoc文档解析与融合检索能力，提升RAG效果

import requestsdef upload_file_to_kb(api_key, kb_id, file_path):    """    上传文件到知识库    Args:        api_key: API密钥        kb_id: 知识库ID        file_path: 本地文件路径    Returns:        dict: 上传结果    """    url = f"https://doc.bluedigit.ai/api/kb/{kb_id}/files"    headers = {        "Authorization": f"Bearer {api_key}"    }    with open(file_path, 'rb') as f:        files = {            'file': (file_path.split('/')[-1], f, 'application/octet-stream')        }        response = requests.post(url, headers=headers, files=files)    if response.status_code == 200:        return response.json()    else:        print(f"上传失败: {response.status_code} - {response.text}")        return None# 使用示例api_key = "td-your-api-key"kb_id = "your-knowledge-base-id"file_path = "/path/to/your/document.pdf"result = upload_file_to_kb(api_key, kb_id, file_path)if result:    print(f"文件上传成功，文件ID: {result['uploaded_files'][0]['id']}")

def parse_file(api_key, kb_id, file_id):    """    启动文件解析    Args:        api_key: API密钥        kb_id: 知识库ID        file_id: 文件ID    Returns:        dict: 解析任务信息    """    url = f"https://doc.bluedigit.ai/api/kb/{kb_id}/files/{file_id}/parse"    headers = {        "Authorization": f"Bearer {api_key}",        "Content-Type": "application/json"    }    response = requests.post(url, headers=headers)    if response.status_code == 200:        return response.json()    else:        print(f"解析启动失败: {response.status_code} - {response.text}")        return None# 使用示例parse_result = parse_file(api_key, kb_id, "file_123456")if parse_result:    job_id = parse_result['job_id']    print(f"解析任务已启动，任务ID: {job_id}")

def retrieve_content(api_key, query, kb_ids=None, file_ids=None, top_k=10, score_threshold=0.5):    """    检索知识库内容    Args:        api_key: API密钥        query: 检索查询        kb_ids: 知识库ID列表        file_ids: 文件ID列表        top_k: 返回结果数量        score_threshold: 相关性阈值    Returns:        dict: 检索结果    """    url = "https://doc.bluedigit.ai/api/retrieve"    headers = {        "Authorization": f"Bearer {api_key}",        "Content-Type": "application/json"    }    data = {        "query": query,        "retrieval_setting": {            "top_k": top_k,            "score_threshold": score_threshold,            "query_type": "vector",            "rerank": False,            "top_n": 5        }    }    if kb_ids:        data["kb_ids"] = kb_ids    if file_ids:        data["file_ids"] = file_ids    response = requests.post(url, headers=headers, json=data)    if response.status_code == 200:        return response.json()    else:        print(f"检索失败: {response.status_code} - {response.text}")        return None# 使用示例# 在知识库中检索result = retrieve_content(    api_key=api_key,    query="人工智能的发展趋势",    kb_ids=["kb_123456"],    top_k=5,    score_threshold=0.7)if result:    print(f"找到 {len(result['records'])} 条相关记录:")    for record in result['records']:        print(f"- 文件: {record['file_name']}")        print(f"  相关性: {record['score']}")        print(f"  内容: {record['text'][:100]}...")        print()