Add transaction_origin field and expand classification options in TransactionParser

app.py

@@ -44,6 +44,11 @@ class TransactionParser(BaseModel):
     # transaction_date: str = Field(description="Use today's date strictly in yyyy-mm-dd format.")
     transaction_origin: str = Field(description="The origin of the transaction. Provide the card or account number as well.")
 
+class TransactionClassification(BaseModel):
+    """This Pydantic class is used to classify the transaction message. The message is taken and the output is structured in a specific format based upon below definition."""
+
+    classification: str = Field(description="Classification of the transaction. Strictly choose one of the values - Transaction, OTP, Scheduled, Reminder")
+
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], operator.add]
 
@@ -71,7 +76,7 @@ class Agent:
         if self.system:
             messages = [SystemMessage(content=self.system)] + messages
 
-        message = self.model.invoke(messages)
+        message = self.model.with_structured_output(TransactionClassification).invoke(messages)
         print(f"{datetime.now(ist_tz)}: Classifying transaction type completed.")
         return {"messages": [message]}
     
@@ -134,16 +139,10 @@ def write_message(data: dict, header: str = Header()):
     if header != HF_TOKEN:
         raise HTTPException(status_code=400, detail="Invalid header")
 
-    prompt = """You are a smart assistant adept at classifying different messages. \
-    You will be penalized heavily for incorrect classification. \
-    Your task is to classify the message into one of the following categories: \
-    Transaction or OTP. \
-    Output the classification in a structured format like below. \
-    {"classification": "OTP"} \
-    """
+    prompt = """You are a smart assistant adept at working with bank transaction messages."""
 
     message = data['message']
-    
+
     try:
         model = ChatGoogleGenerativeAI(model=GOOGLE_MODEL, max_retries=3, callbacks = [OpikTracer()])
     except Exception as e: #fallback model