main.py

from fastapi import HTTPException
from fastapi.responses import StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
import json, uuid, os 
from backend.tools.parse_adk_response import parse
from backend.tools.structures_io import get_data_for_viewer
from backend.defines.global_vars import http_client, app, security
from backend.defines.data_struct import SSE_Params, ProcessData_Params


CONFIG = {
    "base_url": "http://127.0.0.1:8000",
    "app_name": "MatCreator"
}


app.mount("/static", StaticFiles(directory="static"), name="static")
@app.get("/")
async def read_index():
    return FileResponse('static/index.html')

# 创建一个session
@app.get("/config")
async def create_session():
    random_str = str(uuid.uuid4())
    session_id = random_str
    user_id = "user" + random_str.split('-')[0]
    header = {
        "content-type": "application/json"
    }
    url = f"{CONFIG['base_url']}/apps/{CONFIG['app_name']}/users/{user_id}/sessions/{session_id}"
    try:
        async with http_client.session.post(url, headers=header) as response:
            if response.status == 200:
                data = await response.json()
                return {
                    "user_id": data["userId"],
                    "app_name": data["appName"],
                    "session_id": data["id"],
                    "base_url": CONFIG['base_url']
                }
            else:
                raise HTTPException(status_code=500, detail="Failed to create session")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# If you print chunk, it will be like
# chunk = b'data:{...}\n'
# chunk = b'\n'
# chunk = b'data:{...}\n'
# chunk = b'\n'
@app.post("/run_sse")
async def run_sse(request: SSE_Params):
    user_message = request.message
    user_config = request.config
    payload = {
        "appName": user_config["app_name"],
        "userId":  user_config["user_id"],
        "sessionId":user_config["session_id"],
        "newMessage": {
            "role": "user",
            "parts": [{"text": user_message}]
        }, 
        "streaming": False,
        "stateDelta": None
    } 
    headers = {
        "Accept": "text/event-stream",
        "Content-Type": "application/json",
    }
    url = f"{user_config['base_url']}/run_sse"

    async def generate():
        async with http_client.session.post(url,headers=headers, json=payload) as response:
            if response.status != 200:
                yield f"data: {json.dumps({'error': 'Backend request failed'})}\n\n"
                return
            
            buffer = ""
            async for chunk in response.content:
                if chunk:
                    buffer += chunk.decode('utf-8')

                    while '\n\n' in buffer:
                        event, buffer = buffer.split('\n\n', 1)
                        
                        if event.startswith('data: '):
                            try:
                                data_str = event[6:]  
                                if data_str.strip():
                                    data_obj = json.loads(data_str)
                                    
                                    if "content" in data_obj and "parts" in data_obj["content"]:
                                        for part in data_obj["content"]["parts"]:
                                            yield parse(part)
                                    else:
                                        yield f"data: {json.dumps(data_obj)}\n\n"
                            except json.JSONDecodeError:
                                continue
            yield "data: [DONE]\n\n" 
        
    return StreamingResponse(
            generate(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache",
                "Connection": "keep-alive"})


@app.post("/data-processing")
async def data_processing(request: ProcessData_Params):
    task = request.task 
    params = request.params
    if task == "view_structure":
        return get_data_for_viewer(params)
    else:
        raise NotImplementedError(f"Task {task} is not implemented.")

@app.get("/download/{token}")
async def download_file(token: str):
    print(token)
    file_path = security.decrypt_data(token)
    print(file_path)
    filename = os.path.basename(file_path)
    return FileResponse(file_path, filename=filename)

if __name__ == "__main__":
    import argparse, uvicorn

    parser = argparse.ArgumentParser(description="A custom ui of pfd-agent")
    parser.add_argument("--port", type=int, default=3000, help="The port of the ui client.")
    parser.add_argument("--host", type=str, default="localhost", help="The host of the ui client.")
    parser.add_argument("--base_url", type=str, default="http://127.0.0.1:8000", help="The url of the pfd agent.")
    parser.add_argument("--app_name", type=str, default="MatCreator", help="The name of the app.")
    args = parser.parse_args()
    CONFIG["app_name"] = args.app_name
    CONFIG["base_url"] = args.base_url

    uvicorn.run("main:app", host=args.host, port=args.port, reload=True)