CSPON-TSC/aiAgent_gd/lianqi_server/run_server_async.py

import shutil
import tempfile
import time
import sys
from zipfile import ZipFile

import fiona
from shapely.geometry import shape
from sse_starlette.sse import EventSourceResponse
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
parent_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
sys.path.append(parent_dir)

from qwen_agent.planning.plan_executor import PlanExecutor
from qwen_agent.planning.plan_continue_executor import PlanContinueExecutor
from qwen_agent.llm.llm_client import LLMClient, LLMAsyncClient
from agent_config import LLMDict_Qwen_72B_1211, LLMDict_GPT4_TURBO
from agent_messages import BaseRequest

prompt_lan = "CN"
llm_name = "qwen-plus"
llm_turbo_name = "gpt-4-turbo"
max_ref_token = 4000
# model_server = "http://10.10.0.10:7907/v1"
model_server = "http://lq.lianqiai.cn:7905/v1"
api_key = ""
server_host = "0.0.0.0"
server_port = 8511

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

rspHeaders = {
    "Cache-Control": "no-cache",
    "Connection": "keep-alive",
    "Content-Type": "text/event-stream",
    "Transfer-Encoding": "chunked",
}

if model_server.startswith("http"):
    source = "local"
elif model_server.startswith("dashscope"):
    source = "dashscope"


def rm_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)


@app.post("/")
def index():
    return "Welcome to Lianqi AI"


@app.post("/subscribe/spatial")
async def upload_file(question: str = Form(...), file: UploadFile = File(...)):
    # 确保上传的文件是ZIP类型
    if file.content_type != "application/zip":
        return {"detail": "请上传ZIP文件"}

    # 将文件写入临时文件
    temp_dir = tempfile.TemporaryDirectory(dir=os.path.dirname(__file__) + "/upload_temps")

    with open(f"{temp_dir.name}/{file.filename}", "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
        print(buffer)

    # 解压缩
    with ZipFile(f"{temp_dir.name}/{file.filename}", "r") as zip_ref:
        zip_ref.extractall(temp_dir.name)

    for root, dirs, files in os.walk(temp_dir.name):
        for file in files:
            if file.endswith('.shp'):
                shp_file_path = os.path.join(root, file)

    with fiona.Collection(shp_file_path) as shp:
        # 文件上传逻辑将在此处编写
        for feature in shp:
            geom = shape(feature['geometry'])

    return EventSourceResponse(
        call_with_stream(f"原问题为:{question},上传的图形wkt为: {geom.wkt}"),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.post("/subscribe/", response_model=str)
async def subscribe(request: BaseRequest):
    print(request)
    return EventSourceResponse(
        call_with_stream(request.data, []),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.get("/subscribe/{question}", response_model=str)
async def subscribe(question: str):
    return EventSourceResponse(
        call_with_stream(question),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.post("/subscribeByTurbo/", response_model=str)
async def subscribeByTurbo(question: BaseRequest):
    return EventSourceResponse(
        call_with_stream(question.data, False, LLMDict_GPT4_TURBO),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.get("/subscribeByTurbo/{question}", response_model=str)
async def subscribeByTurbo(question: str):
    return EventSourceResponse(
        call_with_stream(question, False, LLMDict_GPT4_TURBO),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.post("/clarification/", response_model=str)
async def clarification(request: BaseRequest):
    print("clarification: ", request)
    return EventSourceResponse(
        call_with_stream(request.data, True),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


@app.post("/clarificationByTurbo/", response_model=str)
async def clarificationByTurbo(request: BaseRequest):
    print("clarificationByTurbo: ", request)
    return EventSourceResponse(
        call_with_stream(request.data, True, LLMDict_GPT4_TURBO),
        media_type="text/event-stream",
        headers=rspHeaders,
    )


llm_client = LLMClient(model=llm_name, model_server=model_server)
llm_client_async = LLMAsyncClient(model=llm_name, model_server=model_server)


async def call_with_stream(
        question,
        history=[],
        isClarification=False,
        llm_dict=LLMDict_Qwen_72B_1211,
):
    if isClarification:
        executor = PlanContinueExecutor(
            llm_dict=llm_dict, llm=llm_client_async, stream=True
        )
    else:
        executor = PlanExecutor(llm_dict=llm_dict, llm=llm_client_async, stream=True)

    async for rsp in executor.run(question, history):
        if not rsp:
            continue
        else:
            time.sleep(0.1)
            yield rsp
            yield "[DONE]"
    yield "[FINISH]"


if __name__ == "__main__":
    uvicorn.run(app, host=server_host, port=server_port, workers=1)