src/models/workflow_graph.py

from llama_index.core.workflow import Event
from pydantic import BaseModel
from typing import List


class Edge(BaseModel):
    source: str
    target: str
    desc: str


class Node(BaseModel):
    node_id: str
    name: str


class Graph(BaseModel):
    nodes: List[Node]
    edges: List[Edge]


class GetResponseEvent(Event):
    project_details: str
    rawResponse: str


class ConstructGraphEvent(Event):
    workflowGraph: Graph


class GraphInputSchema(BaseModel):
    desc: str


class GraphOutputSchema(BaseModel):
    graph: Graph