src/models.py

from pydantic import BaseModel, Field

class ResponseState(BaseModel):
    isNeedContext: bool = Field(default=False, description="Whether the model needs additional context from data from or pdf")
    response: str = Field(default="", description="The response from the model")
    contextType: str = Field(default="", description="If isNeedContext true then The type of context needed, e.g., 'data' or 'pdf' or 'both'")
    retriverKey: str = Field(default="", description="If isNeedContext true and contextType is pdf or both then The key to use to retrieve the context from the pdf for vector db. The key is in plain text.")
    code: str = Field(default="", description="The data already loaded in a dataframe named df, you only need to write code to gain insights from the data. Insert all necessary info inside string variable named result")