feat: add llm service to generate income statement

app/api/routers/file_upload.py

@@ -1,11 +1,14 @@
 from typing import Annotated
-from fastapi import APIRouter, UploadFile
+from fastapi import APIRouter, UploadFile, Depends
 from app.categorization.file_processing import process_file, save_results
 from app.schema.index import FileUploadCreate
 import asyncio
 import os
 import csv
 
+from app.engine.postgresdb import get_db_session
+from sqlalchemy.ext.asyncio import AsyncSession
+
 file_upload_router = r = APIRouter(prefix="/api/v1/file_upload", tags=["file_upload"])
 
 @r.post(
@@ -16,7 +19,7 @@ file_upload_router = r = APIRouter(prefix="/api/v1/file_upload", tags=["file_upl
         500: {"description": "Internal server error"},
     },
 )
-async def create_file(input_file: UploadFile):
+async def create_file(input_file: UploadFile, db: AsyncSession = Depends(get_db_session)):
     try:
          # Create directory to store all uploaded .csv files
         file_upload_directory_path = "data/tx_data/input"
@@ -31,10 +34,10 @@ async def create_file(input_file: UploadFile):
         # With the newly created file and it's path, process and save it for embedding
         processed_file = process_file(os.path.realpath(input_file.filename))
         result = await asyncio.gather(processed_file)
-        save_results(result)
+        await save_results(db, result)
 
     except Exception:
         return {"message": "There was an error uploading this file. Ensure you have a .csv file with the following columns:"
-        "\n source, date, type, category, description, amount"}
+        "\n transaction_date, type, category, name_description, amount"}
 
     return {"message": f"Successfully uploaded {input_file.filename}"}

app/categorization/file_processing.py

@@ -14,6 +14,8 @@ from app.categorization.config import RESULT_OUTPUT_FILE, CATEGORY_REFERENCE_OUT
 from app.model.transaction import Transaction
 from app.schema.index import TransactionCreate
 
+from sqlalchemy.ext.asyncio import AsyncSession
+
 
 # Read file and process it (e.g. categorize transactions)
 async def process_file(file_path: str) -> Dict[str, Union[str, pd.DataFrame]]:
@@ -71,7 +73,7 @@ def standardize_csv_file(file_path: str) -> pd.DataFrame:
     return tx_list
 
 
-async def save_results(results: List) -> None:
+async def save_results(db: AsyncSession, results: List) -> None:
     """
     Merge all interim results in the input folder and write the merged results to the output file.
 
@@ -104,7 +106,7 @@ async def save_results(results: List) -> None:
     # Save to database
     # FIXME: get user_id from session
     txn_list_to_save = [TransactionCreate(**row.to_dict(), user_id=1) for _, row in tx_list.iterrows()]
-    await Transaction.bulk_create(txn_list_to_save)
+    await Transaction.bulk_create(db, txn_list_to_save)
 
     new_ref_data = tx_list[["name/description", "category"]]
     if os.path.exists(CATEGORY_REFERENCE_OUTPUT_FILE):

app/model/transaction.py

@@ -31,8 +31,12 @@ class Transaction(Base, BaseModel):
 
     @classmethod
     async def bulk_create(cls: "type[Transaction]", db: AsyncSession, transactions: List[TransactionCreate]) -> None:
-        query = sql.insert(cls).values(transactions)
-        await db.execute(query)
+        transactions_list = [cls(transaction_date=transaction.transaction_date,
+                                 category=transaction.category,
+                                 type=transaction.type,
+                                 amount= transaction.amount,
+                                 name_description=transaction.name_description) for transaction in transactions]
+        db.add(transactions_list)
         await db.commit()
 
     @classmethod

app/schema/index.py

@@ -74,3 +74,7 @@ class IncomeStatementResponse(PydanticBaseModel):
     date_to: datetime
     income: dict
     expenses: dict
+
+class IncomeStatementLLMResponse(PydanticBaseModel):
+    income: dict
+    expenses: dict

app/service/income_statement.py

@@ -3,13 +3,16 @@ from app.model.transaction import Transaction as TransactionModel
 from app.model.income_statement import IncomeStatement as IncomeStatementModel
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from app.service.llm import call_llm
+
 
 async def call_llm_to_create_income_statement(payload: IncomeStatementCreate, db: AsyncSession) -> None:
     transactions = await TransactionModel.get_by_user_between_dates(
         db, payload.user_id, payload.date_from, payload.date_to
     )
 
-    # TODO: Call LLM to generate income and expenses
-    income = {}
-    expenses = {}
+    response = call_llm(transactions)
+    income = response.income
+    expenses = response.expenses
+
     await IncomeStatementModel.create(db, **payload, income=income, expenses=expenses)

app/service/llm.py

@@ -0,0 +1,53 @@
+import logging
+from llama_index.core.settings import Settings
+from llama_index.core import PromptTemplate
+
+from app.model.transaction import Transaction
+from app.schema.index import IncomeStatementLLMResponse
+
+def income_statement_prompt (inputData: Transaction) -> PromptTemplate:
+    context_str = f"""
+        You are an accountant skilled at organizing transactions from multiple different bank
+        accounts and credit card statements to prepare an income statement.
+
+        Input data has the following format: transaction_date, type, category, name_description, amount.
+        The <IN> tag is prepended to the input data as follows:
+        ```<IN>
+        {inputData}
+        ```
+
+        Your task is to prepare an income statement. The income statement's output should be in a json format.
+        An example of the expected output is as follows with the <OUT> tag. Note that not all categories of the transactions
+        have been listed below. Use below <OUT> tag as a reference in preparing the income statement.
+        ```<OUT>
+        {
+            "REVENUE": {
+                "Gross Sales": 73351.11,
+                "Other Income": 0,
+                "Balance Dec 2022": 3987.39,
+            },
+            "EXPENSES": {
+                "Advertising": 0,
+                "Commissions": 0,
+                "Insurance": 0,
+                "Memberships": 0,
+                "Utilities": 0,
+            }
+        }
+        ```
+
+        """
+    prompt = PromptTemplate(context_str, template_var_mappings={"inputData": inputData})
+    logging.info(f"Prompt: {prompt}")
+
+async def call_llm(prompt: PromptTemplate) -> str:
+    llm = Settings.llm.copy()
+    prompt = income_statement_prompt()
+    llm.check_prompts(prompt)
+
+    llm.system_prompt = prompt
+
+    output = await llm.astructured_predict(output_cls=IncomeStatementLLMResponse, prompt=prompt)
+
+    logging.info(f"Output: {output}")
+    return output

pyproject.toml

@@ -52,6 +52,9 @@ version = "0.2.2"
 [tool.black]
 line-length = 119
 
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+
 [build-system]
 requires = [ "poetry-core" ]
 build-backend = "poetry.core.masonry.api"