agent/datastructures.py

from typing import List, Optional
from enum import Enum
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.output_parsers import PydanticOutputParser
from langchain_core.output_parsers import JsonOutputParser


## DEFINE INPUT FRAMEWORK
class InputSchema(BaseModel):
    """Expect the input from the frontend to be a JSON object with this structure"""
    question: str = Field(description="The enquiry that is passed from the user")

class ActionTypes(str, Enum):
    SuggestGoal = "SuggestGoal"
    SuggestActivity = "VisitPage"

# Define your desired data structure.
class FrontEndActions(BaseModel):
    """Structure to pass actions back to the frontend"""
    heading: str = Field(description="The heading text to display on the button")
    detail: str = Field(description="More detailed information, for instance explaining why you have chosen this action for the user")
    id: int = Field(description="The ID of the object that is referenced")
    type: ActionTypes = Field(description="This should be a string that identifies the type of action. It can be one of: SuggestGoal, SuggestActivity")

class ResponseSchema(BaseModel):
    """Always use this to format the final response to the user. This will be passed back to the frontend."""
    message: str = Field(description="final answer to respond to the user")
    thread_id: int = Field(description="The ID of the checkpointer memory thread that this response is associated with. This is used to keep track of the conversation.")
    tools: List[str] = Field(description="A list of the tools used to generate the response.") 
    actions: List[FrontEndActions] = Field(description="List of suggested actions that should be passed back to the frontend to display. The use will click these to enact them. ")

parser = JsonOutputParser(pydantic_object=ResponseSchema)