add llm service for generating income statements

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/api/routers/income_statement.py

@@ -3,7 +3,7 @@ from fastapi import APIRouter, Depends, HTTPException, status
 from sqlalchemy.ext.asyncio import AsyncSession
 from app.model.transaction import Transaction as TransactionModel
 from app.model.income_statement import IncomeStatement as IncomeStatementModel
-from app.schema.index import IncomeStatementCreate, IncomeStatementResponse
+from app.schema.index import IncomeStatementCreateRequest, IncomeStatementResponse
 from app.engine.postgresdb import get_db_session
 from app.service.income_statement import call_llm_to_create_income_statement
 
@@ -18,7 +18,7 @@ income_statement_router = r = APIRouter(prefix="/api/v1/income_statement", tags=
         500: {"description": "Internal server error"},
     },
 )
-async def create_income_statement(payload: IncomeStatementCreate, db: AsyncSession = Depends(get_db_session)) -> None:
+async def create_income_statement(payload: IncomeStatementCreateRequest, db: AsyncSession = Depends(get_db_session)) -> None:
     try:
         await call_llm_to_create_income_statement(payload, db)
 
@@ -27,7 +27,7 @@ async def create_income_statement(payload: IncomeStatementCreate, db: AsyncSessi
 
 
 @r.get(
-    "/{user_id}",
+    "/user/{user_id}",
     response_model=List[IncomeStatementResponse],
     responses={
         200: {"description": "New user created"},
@@ -43,14 +43,13 @@ async def get_income_statements(
     Retrieve all income statements.
     """
     result = await IncomeStatementModel.get_by_user(db, user_id)
-    all_rows = result.all()
-    if len(all_rows) == 0:
+    if len(result) == 0:
         raise HTTPException(status_code=status.HTTP_204_NO_CONTENT, detail="No income statements found for this user")
-    return all_rows
+    return result
 
 
 @r.get(
-    "/{report_id}",
+    "/report/{report_id}",
     response_model=IncomeStatementResponse,
     responses={
         200: {"description": "Income statement found"},

app/api/routers/user.py

@@ -27,7 +27,7 @@ async def create_user(user: UserCreate, db: AsyncSession = Depends(get_db_sessio
         if db_user and not db_user.is_deleted:
             raise HTTPException(status_code=409, detail="User already exists")
 
-        await UserModel.create(db, **user.dict())
+        await UserModel.create(db, **user.model_dump())
         user = await UserModel.get(db, email=user.email)
         return user
     except Exception as e:
@@ -64,7 +64,7 @@ async def update_user(email: str, user_payload: UserUpdate, db: AsyncSession = D
         user = await UserModel.get(db, email=email)
         if not user:
             raise HTTPException(status_code=404, detail="User not found")
-        await UserModel.update(db, id=user.id, **user_payload.dict())
+        await UserModel.update(db, id=user.id, **user_payload.model_dump())
         user = await UserModel.get(db, email=email)
         return user
     except Exception as e:

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/income_statement.py

@@ -8,6 +8,7 @@ from sqlalchemy.dialects.postgresql import JSON
 
 from app.model.base import BaseModel
 from app.engine.postgresdb import Base
+from app.schema.index import IncomeStatement as IncomeStatementSchema
 
 
 class IncomeStatement(Base, BaseModel):
@@ -21,19 +22,22 @@ class IncomeStatement(Base, BaseModel):
     user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
     user = relationship("User", back_populates="income_statements")
 
+    def __str__(self) -> str:
+        return f"IncomeStatement(id={self.id}, user_id={self.user_id}, date_from={self.date_from}, date_to={self.date_to}, income={self.income}, expenses={self.expenses})"
+
     @classmethod
-    def create(cls: "type[IncomeStatement]", db: AsyncSession, **kwargs) -> "IncomeStatement":
-        query = sql.insert(cls).values(**kwargs)
-        income_statements = db.execute(query)
-        income_statement = income_statements.first()
-        db.commit()
+    async def create(cls: "type[IncomeStatement]", db: AsyncSession, **kwargs: IncomeStatementSchema) -> "IncomeStatement":
+        income_statement = cls(**kwargs)
+        db.add(income_statement)
+        await db.commit()
+        await db.refresh(income_statement)
         return income_statement
 
     @classmethod
     async def get_by_user(cls: "type[IncomeStatement]", db: AsyncSession, user_id: int) -> "List[IncomeStatement]":
         query = sql.select(cls).where(cls.user_id == user_id)
         income_statements = await db.scalars(query)
-        return income_statements
+        return income_statements.all()
 
     @classmethod
     async def get(cls: "type[IncomeStatement]", db: AsyncSession, id: int) -> "IncomeStatement":

app/model/transaction.py

@@ -21,6 +21,9 @@ class Transaction(Base, BaseModel):
     user_id = mapped_column(ForeignKey("users.id"))
     user = relationship("User", back_populates="transactions")
 
+    def __str__(self) -> str:
+        return f"{self.transaction_date}, {self.category}, {self.name_description}, {self.amount}, {self.type}"
+
     @classmethod
     async def create(cls: "type[Transaction]", db: AsyncSession, **kwargs) -> "Transaction":
         query = sql.insert(cls).values(**kwargs)
@@ -31,7 +34,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()
 
@@ -55,4 +59,4 @@ class Transaction(Base, BaseModel):
     ) -> "List[Transaction]":
         query = sql.select(cls).where(cls.user_id == user_id).where(cls.transaction_date.between(start_date, end_date))
         transactions = await db.scalars(query)
-        return transactions
+        return transactions.all()

app/schema/index.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from datetime import datetime
-from typing import List, Optional
+from typing import Dict, List
+from typing_extensions import TypedDict
 
 from app.schema.base import BaseModel, PydanticBaseModel
 
@@ -62,15 +63,23 @@ class FileUploadCreate(PydanticBaseModel):
     type: str
 
 
-class IncomeStatementCreate(PydanticBaseModel):
+class IncomeStatementCreateRequest(PydanticBaseModel):
     user_id: int
     date_from: datetime
     date_to: datetime
 
 
-class IncomeStatementResponse(PydanticBaseModel):
-    id: int
-    date_from: datetime
-    date_to: datetime
-    income: dict
-    expenses: dict
+class IncomeStatementDetail(TypedDict):
+    total: float
+    category_totals: List[Dict[str, str | float]]
+
+class IncomeStatementLLMResponse(PydanticBaseModel):
+    income: IncomeStatementDetail
+    expenses: IncomeStatementDetail
+
+
+class IncomeStatement(IncomeStatementCreateRequest, IncomeStatementLLMResponse):
+    pass
+
+class IncomeStatementResponse(IncomeStatement):
+    id: int

app/service/income_statement.py

@@ -1,15 +1,37 @@
-from app.schema.index import IncomeStatementCreate
+import logging
+from app.schema.index import IncomeStatement, IncomeStatementCreateRequest
 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:
+
+async def call_llm_to_create_income_statement(payload: IncomeStatementCreateRequest, 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 = {}
-    await IncomeStatementModel.create(db, **payload, income=income, expenses=expenses)
+    if not transactions:
+        print("No transactions found")
+        return
+
+    response = await call_llm(transactions)
+
+    income = response.dict()['income']
+    expenses = response.dict()['expenses']
+
+    try:
+        income_statement_create_payload = IncomeStatement(
+            user_id=payload.user_id,
+            date_from=payload.date_from,
+            date_to=payload.date_to,
+            income=income,
+            expenses=expenses,
+        )
+
+        income_statement = await IncomeStatementModel.create(db, **income_statement_create_payload.model_dump())
+        print(f"Income statement created: {income_statement}")
+    except Exception as e:
+        print(e)
+        raise e

app/service/llm.py

@@ -0,0 +1,55 @@
+from typing import List
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables.base import RunnableSequence
+from langchain_openai import OpenAI
+from langchain.globals import set_llm_cache
+
+from app.model.transaction import Transaction
+from app.schema.index import IncomeStatementLLMResponse
+from config.index import config as env
+
+from langchain_core.output_parsers import PydanticOutputParser
+
+set_llm_cache(None)
+
+def income_statement_prompt () -> ChatPromptTemplate:
+    context_str = """
+        You are an accountant skilled at organizing transactions from multiple different bank
+        accounts and credit card statements to prepare an income statement.
+
+        Input data is in the below csv format:
+        transaction_date, category, name_description, amount, type\n
+        {input_data_csv}
+
+        Your task is to prepare an income statement. The output should be in the following format: {format_instructions}
+
+        """
+    prompt = ChatPromptTemplate.from_template(context_str)
+    return prompt
+
+async def call_llm(inputData: List[Transaction]) -> str:
+    input_data_csv = '\n'.join(str(x) for x in inputData)
+
+    output_parser = PydanticOutputParser(pydantic_object=IncomeStatementLLMResponse)
+
+    prompt = income_statement_prompt().partial(format_instructions=output_parser.get_format_instructions())
+
+    llm = OpenAI(name='Income Statement Generation Bot',
+                        api_key=env.OPENAI_API_KEY,
+                        # cache=True,
+                        temperature=0.7,
+                        verbose=True)
+
+    try:
+        runnable_chain = RunnableSequence(prompt, llm, output_parser)
+    except Exception as e:
+        print(f"runnable_chain error: {str(e)}")
+        raise e
+
+    try:
+        output_chunks = runnable_chain.invoke({"input_data_csv": input_data_csv})
+        return output_chunks
+
+    except Exception as e:
+        print(f"runnable_chain.invoke error: {str(e)}")
+        raise e

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,11 @@ 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 +75,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_income_statement.py

@@ -0,0 +1,51 @@
+from fastapi.testclient import TestClient
+import pytest
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from app.model.transaction import Transaction
+from app.model.user import User
+from tests.utils import get_fake_transactions
+
+@pytest.mark.asyncio
+async def test_income_statement(client: TestClient, get_db_session_fixture: AsyncSession) -> None:
+    session_override = get_db_session_fixture
+
+    # 1. Create a user
+    user = await User.create(session_override, name="user", email="email", hashed_password="password")
+
+    # 2. Create a bunch of transactions
+    fake_transactions = get_fake_transactions(user.id)
+    await Transaction.bulk_create(session_override, fake_transactions)
+
+    # 3. Create an income statement
+    min_date = min(t.transaction_date for t in fake_transactions)
+    max_date = max(t.transaction_date for t in fake_transactions)
+
+    print(f"min_date: {min_date}, max_date: {max_date}")
+    response = client.post(
+        "/api/v1/income_statement",
+        json={
+            "user_id": 1,
+            "date_from": str(min_date),
+            "date_to": str(max_date),
+        },
+    )
+
+    assert response.status_code == 200
+
+    # 4. Verify that the income statement matches the transactions
+    response = client.get(f"/api/v1/income_statement/user/1")
+    print(response.json())
+    assert response.status_code == 200
+    assert response.json()[0].get("income")
+    assert response.json()[0].get("expenses")
+
+    report_id = response.json()[0].get("id")
+
+    # # 5. Verify that the income statement can be retrieved
+    if report_id is not None:
+        response = client.get(f"/api/v1/income_statement/report/{report_id}")
+        assert response.status_code == 200
+        assert response.json().get("income")
+        assert response.json().get("expenses")
+        assert response.json().get("id") == report_id

tests/test_transactions.py

@@ -0,0 +1,25 @@
+from fastapi.testclient import TestClient
+import pytest
+
+from app.model.transaction import Transaction
+from app.model.user import User
+from sqlalchemy.ext.asyncio import AsyncSession
+from tests.utils import get_fake_transactions
+
+
+@pytest.mark.asyncio
+async def test_transactions(client: TestClient, get_db_session_fixture: AsyncSession) -> None:
+
+    session_override = get_db_session_fixture
+
+    # 1. Create a user
+    user = await User.create(session_override, name="user", email="email", hashed_password="password")
+
+    # 2. Create a bunch of transactions
+    fake_transactions = get_fake_transactions(user.id)
+    await Transaction.bulk_create(session_override, fake_transactions)
+
+    # 3. Verify that the transactions are returned
+    response = client.get("/api/v1/transactions/1")
+    assert response.status_code == 200
+    assert len(response.json()) == 10

tests/utils.py

@@ -0,0 +1,88 @@
+from datetime import datetime
+from typing import List
+from app.schema.index import TransactionCreate, TransactionType
+
+
+def get_fake_transactions(user_id: int) -> List[TransactionCreate]:
+    return [
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 1),
+            category="category1",
+            name_description="name_description",
+            amount=1.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 2),
+            category="category2",
+            name_description="name_description",
+            amount=2.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 3),
+            category="category3",
+            name_description="name_description",
+            amount=3.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 4),
+            category="category1",
+            name_description="name_description",
+            amount=4.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 5),
+            category="category2",
+            name_description="name_description",
+            amount=5.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 6),
+            category="category3",
+            name_description="name_description",
+            amount=6.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 7),
+            category="category1",
+            name_description="name_description",
+            amount=7.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 8),
+            category="category2",
+            name_description="name_description",
+            amount=8.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 9),
+            category="category3",
+            name_description="name_description",
+            amount=9.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 10),
+            category="category1",
+            name_description="name_description",
+            amount=10.0,
+            type=TransactionType.EXPENSE,
+        ),
+    ]