add llm service for generating income statements (#14)

- feat: add llm service to generate income statement (6d92adc7e96771ebc55f1859275ee4206a668883)
- test: pytest for transactions (cf9085da983a5159a7de0d43b0b12c16fdc32822)

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,7 +31,8 @@ class Transaction(Base, BaseModel):
 
     @classmethod
     async def bulk_create(cls: "type[Transaction]", db: AsyncSession, transactions: List[TransactionCreate]) -> None:
-        query = sql.insert(cls).values(transactions)
+        values = [transaction.model_dump() for transaction in transactions]
+        query = sql.insert(cls).values(values)
         await db.execute(query)
         await db.commit()
 

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"

tests/conftest.py

@@ -37,7 +37,7 @@ def client(app):
         yield c
 
 
-@pytest.mark.asyncio(scope="session")
+@pytest.fixture(scope="session")
 def event_loop(request):
     loop = asyncio.get_event_loop_policy().new_event_loop()
     yield loop
@@ -52,9 +52,9 @@ async def connection_test(test_db, event_loop):
     pg_db = test_db.dbname
     pg_password = test_db.password
 
-    with DatabaseJanitor(pg_user, pg_host, pg_port, pg_db, test_db.version, pg_password):
-        connection_str = f"postgresql+asyncpg://{pg_user}:@{pg_host}:{pg_port}/{pg_db}"
-        sessionmanager.init(connection_str)
+    with DatabaseJanitor(user=pg_user, host=pg_host, port=pg_port, dbname=pg_db, version=test_db.version, password=pg_password):
+        connection_str = f"postgresql+psycopg://{pg_user}:@{pg_host}:{pg_port}/{pg_db}"
+        sessionmanager.init(connection_str, {"echo": True, "future": True})
         yield
         await sessionmanager.close()
 
@@ -73,3 +73,9 @@ async def session_override(app, connection_test):
             yield session
 
     app.dependency_overrides[get_db_session] = get_db_session_override
+
+
+@pytest.fixture(scope="function", autouse=True)
+async def get_db_session_fixture():
+    async with sessionmanager.session() as session:
+        yield session

tests/test_transactions.py

@@ -0,0 +1,109 @@
+from datetime import datetime
+from typing import List
+from fastapi import Depends
+from fastapi.testclient import TestClient
+import pytest
+
+from app.model.transaction import Transaction
+from app.schema.index import TransactionType, TransactionCreate
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from app.engine.postgresdb import get_db_session
+from app.model.user import User
+
+def get_fake_transactions(user_id: int) -> List[TransactionCreate]:
+    return [
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 1),
+            category="category",
+            name_description="name_description",
+            amount=1.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 2),
+            category="category",
+            name_description="name_description",
+            amount=2.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 3),
+            category="category",
+            name_description="name_description",
+            amount=3.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 4),
+            category="category",
+            name_description="name_description",
+            amount=4.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 5),
+            category="category",
+            name_description="name_description",
+            amount=5.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 6),
+            category="category",
+            name_description="name_description",
+            amount=6.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 7),
+            category="category",
+            name_description="name_description",
+            amount=7.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 8),
+            category="category",
+            name_description="name_description",
+            amount=8.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 9),
+            category="category",
+            name_description="name_description",
+            amount=9.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 10),
+            category="category",
+            name_description="name_description",
+            amount=10.0,
+            type=TransactionType.EXPENSE,
+        ),
+    ]
+
+@pytest.mark.asyncio
+async def test_transactions(client: TestClient, get_db_session_fixture: AsyncSession) -> None:
+
+    session_override = get_db_session_fixture
+    user = await User.create(session_override, name="user", email="email", hashed_password="password")
+
+    fake_transactions = get_fake_transactions(user.id)
+    await Transaction.bulk_create(session_override, fake_transactions)
+
+    response = client.get("/api/v1/transactions/1")
+    assert response.status_code == 200
+    assert len(response.json()) == 10