116 lines
2.6 KiB
Python
116 lines
2.6 KiB
Python
from enum import Enum
|
|
from typing import List, Optional
|
|
from pydantic import BaseModel, Field, model_validator
|
|
|
|
|
|
class PointModel(BaseModel):
|
|
x: float
|
|
y: float
|
|
|
|
def to_list(self) -> List[float]:
|
|
return [self.x, self.y]
|
|
|
|
|
|
class PointWithRadius(PointModel):
|
|
radius: Optional[float] = None
|
|
|
|
|
|
class BarSeriesModel(BaseModel):
|
|
name: str
|
|
data: List[float] = Field(
|
|
description="Only numbers should be given out in data. Don't include text/string in data."
|
|
)
|
|
|
|
|
|
class ScatterSeriesModel(BaseModel):
|
|
name: str
|
|
points: List[PointModel]
|
|
|
|
|
|
class BubbleSeriesModel(BaseModel):
|
|
name: str
|
|
points: List[PointWithRadius]
|
|
|
|
|
|
class LineSeriesModel(BaseModel):
|
|
name: str
|
|
data: List[float] = Field(
|
|
description="Only numbers should be given out in data. Don't include text/string in data."
|
|
)
|
|
|
|
|
|
class PieChartSeriesModel(BaseModel):
|
|
data: List[float]
|
|
|
|
|
|
class BarGraphDataModel(BaseModel):
|
|
categories: List[str]
|
|
series: List[BarSeriesModel] = Field(
|
|
description="There should be no more than 3 series"
|
|
)
|
|
|
|
|
|
class ScatterChartDataModel(BaseModel):
|
|
series: List[ScatterSeriesModel]
|
|
|
|
|
|
class BubbleChartDataModel(BaseModel):
|
|
series: List[BubbleSeriesModel]
|
|
|
|
|
|
class LineChartDataModel(BaseModel):
|
|
categories: List[str]
|
|
series: List[LineSeriesModel] = Field(
|
|
description="There should be no more than 3 series"
|
|
)
|
|
|
|
|
|
class PieChartDataModel(BaseModel):
|
|
categories: List[str]
|
|
series: List[PieChartSeriesModel] = Field(
|
|
description="One series model with list of data",
|
|
min_length=1,
|
|
)
|
|
|
|
@model_validator(mode="after")
|
|
def limit_series(self):
|
|
self.series = self.series[:1]
|
|
return self
|
|
|
|
|
|
class GraphTypeEnum(Enum):
|
|
pie = "pie"
|
|
bar = "bar"
|
|
line = "line"
|
|
|
|
|
|
class LLMGraphModel(BaseModel):
|
|
name: str
|
|
type: GraphTypeEnum
|
|
unit: Optional[str] = Field(
|
|
description="Unit of the data in the graph. Example: %, kg, million USD, tonnes, etc."
|
|
)
|
|
data: PieChartDataModel | LineChartDataModel | BarGraphDataModel
|
|
|
|
|
|
class GraphModel(LLMGraphModel):
|
|
style: Optional[dict] = {}
|
|
|
|
@classmethod
|
|
def from_llm_graph_model(
|
|
cls, llm_graph_model: LLMGraphModel, style: Optional[dict] = {}
|
|
):
|
|
return cls(
|
|
name=llm_graph_model.name,
|
|
type=llm_graph_model.type,
|
|
unit=llm_graph_model.unit,
|
|
data=llm_graph_model.data,
|
|
style=style,
|
|
)
|
|
|
|
|
|
GRAPH_TYPE_MAPPING = {
|
|
GraphTypeEnum.pie: PieChartDataModel,
|
|
GraphTypeEnum.bar: BarGraphDataModel,
|
|
GraphTypeEnum.line: LineChartDataModel,
|
|
}
|