Fix transaction classification handling and improve error reporting in check_txn_and_decide method

app.py

@@ -37,11 +37,10 @@ ist_tz = timezone("Asia/Kolkata")
 class TransactionParser(BaseModel):
     """This Pydantic class is used to parse the transaction message. The message is taken and the output is structured in a specific format based upon below definitions."""
 
-    amount: str = Field(description="The amount of the transaction strictly in decimal format. Do not insert currency symbol.", example="123.45")
+    amount: str = Field(description="The amount of the transaction strictly in decimal format. Do not insert currency symbol.", example="123.45") # type: ignore
     dr_or_cr: str = Field(description="Identify if the transaction was debit (spent) or credit (received). Strictly choose one of the values - Debit or Credit")
     receiver: str = Field(description="The recipient of the transaction. Identify the Merchant Name from the message text.")
-    category: str = Field(description="The category of the transaction. The category of the transaction is linked to the Merchant Name. Strictly choose from one the of values - Shopping,EMI,Education,Miscellaneous,Grocery,Utility,House Help,Travel,Transport,Food")
-    # transaction_date: str = Field(description="Use today's date strictly in yyyy-mm-dd format.")
+    category: str = Field(description="The category of the transaction. The category of the transaction is linked to the Merchant Name. Strictly choose from one the of values - Shopping,EMI,Education,Miscellaneous,Grocery,Utility,House Help,Travel,Transport,Food,Insurance")
     transaction_origin: str = Field(description="The origin of the transaction. Provide the card or account number as well.")
 
 class TransactionClassification(BaseModel):
@@ -101,29 +100,29 @@ class Agent:
 
         SCOPES = ('https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive')
 
-        service_account_info = json.loads(GOOGLESHEETS_CREDENTIALS)
+        service_account_info = json.loads(GOOGLESHEETS_CREDENTIALS) # type: ignore
         credentials = service_account.Credentials.from_service_account_info(service_account_info, scopes=SCOPES)
         client = pygsheets.authorize(custom_credentials=credentials)
         worksheet = client.open_by_url(SHEET_URL)
         wk = worksheet[0]
         # Get number of rows in the worksheet
-        df = wk.get_as_df(start='A1', end='G999')
+        df = wk.get_as_df(start='A1', end='G999') # type: ignore
         nrows = df.shape[0]
-        wk.update_value(f'A{nrows+2}', result.amount)
-        wk.update_value(f'B{nrows+2}', result.dr_or_cr)
-        wk.update_value(f'C{nrows+2}', result.receiver)
-        wk.update_value(f'D{nrows+2}', result.category)
-        wk.update_value(f'E{nrows+2}', datetime.now(ist_tz).strftime("%Y-%m-%d"))
-        wk.update_value(f'F{nrows+2}', result.transaction_origin)
-        wk.update_value(f'G{nrows+2}', state["messages"][0])
+        wk.update_value(f'A{nrows+2}', result.amount) # type: ignore
+        wk.update_value(f'B{nrows+2}', result.dr_or_cr) # type: ignore
+        wk.update_value(f'C{nrows+2}', result.receiver) # type: ignore
+        wk.update_value(f'D{nrows+2}', result.category) # type: ignore
+        wk.update_value(f'E{nrows+2}', datetime.now(ist_tz).strftime("%Y-%m-%d")) # type: ignore
+        wk.update_value(f'F{nrows+2}', result.transaction_origin) # type: ignore
+        wk.update_value(f'G{nrows+2}', state["messages"][0]) # type: ignore
         print(f"{datetime.now(ist_tz)}: Writing transaction message to Google Sheets completed.")
-        return {"messages": ["Transaction Completed"]}
+        return {"messages": ["Transaction Completed"]} # type: ignore
         
     def check_txn_and_decide(self, state: AgentState):
         try:
-            result = json.loads(state['messages'][-1].content)['classification']
+            result = state['messages'][-1].classification # type: ignore
         except json.JSONDecodeError:
-            result = state['messages'][-1].content.strip()
+            raise HTTPException(status_code=400, detail="Invalid response format from model")
 
         return result == "Transaction"