api.py

import json
import logging
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from agent_graph import app


logger = logging.getLogger("api_stream_logger")
api_app = FastAPI(title="Bank Responsibility Mapping API")
api_app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], # Restrict this to the specific UI port in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ChatResponse(BaseModel):
    query: str
    last_response: str = None
    user_id: str = None
    user_name: str = None
    updated_law_ids: str = None

@api_app.post("/api/chat")
async def chat_endpoint(request: ChatResponse):
    parsed_law_ids = []
    if request.updated_law_ids and request.updated_law_ids.strip():
        parsed_law_ids = [str(lid).strip() for lid in request.updated_law_ids.split("-") if lid.strip()]

    initial_state = {
        "updated_law_ids": None,
        "updated_laws_target_users": None,
        "updated_law_changes": None,

        "target_id": request.user_id,
        "target_name": request.user_name,
        "target_updated_law_ids": parsed_law_ids,
        "target_query": request.query,
        "target_last_response": request.last_response,
        "target_classification": None,
        "target_law_changes": None,
        "target_knowledge": None,
        "target_note": None,

        "final_output": ""
    }

    async def event_generator():
        try:
            # .astream() tells LangGraph to stream execution states asynchronously
            # stream_mode="values" tracks variable changes across the final node pipeline
            async for event in app.astream(initial_state, stream_mode="values"):
                # Check if the final analysis text string is accumulating content
                current_output = event.get("final_output", "")
                print(event.get("target_classification", ""))

                if current_output:
                    # We pass the current string snapshot down the port pipe
                    # Wrapping it in a JSON structure ensures frontend parameters match
                    payload = {"output": current_output}
                    yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"

        except Exception as stream_err:
            logger.error(f"Streaming anomaly detected: {str(stream_err)}")
            yield f"data: {json.dumps({'output': 'Error mid-stream'})}\n\n"

    # 4. Return the live network text stream container with Server-Sent Events content type
    return StreamingResponse(event_generator(), media_type="text/event-stream")


api_app.mount("/", StaticFiles(directory="static", html=True), name="static")

# uvicorn api:api_app --host 10.0.0.6 --port 8000 --reload