file_rag/client.py

import argparse
import requests
import json

BASE_URL = "http://localhost:5666"

def chat_with_server(client_id: str, prompt: str, history: str):
    url = f"{BASE_URL}/chat"
    print(history)
    print(type(history))
    history = '[{"role":"user", "content":"你好"}]'
    data = {
        "client_id": client_id,
        "prompt": prompt,
        "history": history
    }
    response = requests.post(url, data=data)
    return response.json()

def upload_file_to_server(client_id: str, file_path: str):
    url = f"{BASE_URL}/uploadfile/"
    files = {"file": open(file_path, "rb")}
    data = {"client_id": client_id}
    response = requests.post(url, files=files, data=data)
    return response.json()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["chat", "upload"], required=True, help="选择调用的接口")
    parser.add_argument("--client_id", required=True, help="客户端ID")
    parser.add_argument("--prompt", help="用户输入的问题，仅用于chat模式")
    parser.add_argument("--history", default='[]', help="历史消息JSON字符串，仅用于chat模式")
    parser.add_argument("--file", help="上传的文件或文件夹路径，仅用于upload模式")
    
    args = parser.parse_args()
    
    if args.mode == "chat":
        if not args.prompt:
            print("Error: --prompt 参数在 chat 模式下是必须的！")
            return
        response = chat_with_server(args.client_id, args.prompt, args.history)
        print("Chat Response:", response)
    elif args.mode == "upload":
        if not args.file:
            print("Error: --file 参数在 upload 模式下是必须的！")
            return
            
        import os
        if os.path.isfile(args.file):
            files = [args.file]
        elif os.path.isdir(args.file):
            files = [os.path.join(args.file, f) for f in os.listdir(args.file) 
                    if os.path.isfile(os.path.join(args.file, f))]
        else:
            print(f"Error: 路径 {args.file} 不存在或不可访问！")
            return
            
        for file_path in files:
            try:
                print(f"\n正在上传文件: {file_path}")
                response = upload_file_to_server(args.client_id, file_path)
                print(f"上传结果: {response}")
            except Exception as e:
                print(f"文件 {file_path} 上传失败: {str(e)}")

if __name__ == "__main__":
    main()