main.py

from uuid import uuid4
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from langchain_core.messages import HumanMessage, ToolMessage
from graph import supervisor
import pprint
import dotenv
import json
from pydantic import BaseModel, Field


dotenv.load_dotenv()


class ChatRequest(BaseModel):
    user_input: str
    thread_id: str | None = Field(default_factory=lambda: uuid4().hex)


app = FastAPI()

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


def messages_to_json(messages):
    """Convert a list of message objects to JSON-serializable format"""
    result = []

    for message in messages:
        if isinstance(message, dict):
            # Already a dict, just add it
            result.append(message)
        else:
            # Message object, extract attributes
            message_dict = {
                "type": getattr(
                    message,
                    "type",
                    type(message).__name__.replace("Message", "").lower(),
                ),
                "content": getattr(message, "content", None),
                "id": getattr(message, "id", None),
            }

            if hasattr(message, "additional_kwargs"):
                message_dict["additional_kwargs"] = message.additional_kwargs

            if hasattr(message, "response_metadata"):
                message_dict["response_metadata"] = message.response_metadata

            if hasattr(message, "name"):
                message_dict["name"] = message.name

            if hasattr(message, "tool_calls"):
                message_dict["tool_calls"] = message.tool_calls

            if hasattr(message, "usage_metadata"):
                message_dict["usage_metadata"] = message.usage_metadata

            if hasattr(message, "tool_call_id"):
                message_dict["tool_call_id"] = message.tool_call_id

            result.append(message_dict)

    return result


async def event_generator(request: ChatRequest):
    config = {"configurable": {"thread_id": request.thread_id}}
    seen_message_ids = set()

    try:
        for chunk in supervisor.stream(
            {
                "messages": [HumanMessage(content=request.user_input)],
                "generated_image": None,
            },
            config=config,
            stream_mode="updates",
        ):
            node_name = list(chunk)[0]
            node_update = chunk[node_name]

            print(f"Node update keys: {list(node_update.keys())}")
            generated_image = node_update.get("generated_image")
            new_messages = node_update.get("messages", [])
            if new_messages is None:
                new_messages = []

            # Deduplicate messages
            unique_messages = []
            for msg in new_messages:
                msg_id = getattr(msg, "id", None)
                if msg_id:
                    if msg_id not in seen_message_ids:
                        seen_message_ids.add(msg_id)
                        unique_messages.append(msg)
                else:
                    # If message has no ID, we just pass it through
                    unique_messages.append(msg)

            if not unique_messages and not generated_image:
                continue

            print(25 * "+")
            print("node name: ", node_name)
            print(f"New messages ({len(unique_messages)}):", unique_messages)
            if generated_image:
                print("Generated image present:", bool(generated_image))
            print(25 * "-")

            event_data = {
                "node_name": node_name,
                "messages": messages_to_json(unique_messages),
                "generated_image": generated_image,
            }

            yield json.dumps(event_data) + "\n"

    except Exception as e:
        print(f"Error in event_generator: {str(e)}")
        error_event = {
            "node_name": "error",
            "messages": [
                {
                    "type": "ai",
                    "content": f"An error occurred during processing: {str(e)}",
                    "id": str(uuid4()),
                }
            ],
            "generated_image": None,
        }
        yield json.dumps(error_event) + "\n"


@app.post("/chat/stream")
async def chat_stream_endpoint(request: ChatRequest):
    return StreamingResponse(
        event_generator(request=request),
        media_type="text/event-stream",
    )


@app.get("/")
def read_root():
    return {"status": "RAG API is running, View at queai.vercel.app"}