chore(pr/14): fix income statement endpoints, pytests

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/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)
@@ -56,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,19 +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: dict
-    expenses: dict
+    income: IncomeStatementDetail
+    expenses: IncomeStatementDetail
+
+
+class IncomeStatement(IncomeStatementCreateRequest, IncomeStatementLLMResponse):
+    pass
+
+class IncomeStatementResponse(IncomeStatement):
+    id: int

app/service/income_statement.py

@@ -1,4 +1,5 @@
-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
@@ -6,13 +7,31 @@ 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
     )
 
-    response = call_llm(transactions)
-    income = response.income
-    expenses = response.expenses
+    if not transactions:
+        print("No transactions found")
+        return
 
-    await IncomeStatementModel.create(db, **payload, income=income, expenses=expenses)
+    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

@@ -1,53 +1,55 @@
-import logging
-from llama_index.core.settings import Settings
-from llama_index.core import PromptTemplate
+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
 
-def income_statement_prompt (inputData: Transaction) -> PromptTemplate:
-    context_str = f"""
+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 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,
-            }
-        }
-        ```
+        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 = PromptTemplate(context_str, template_var_mappings={"inputData": inputData})
-    logging.info(f"Prompt: {prompt}")
+    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())
 
-async def call_llm(prompt: PromptTemplate) -> str:
-    llm = Settings.llm.copy()
-    prompt = income_statement_prompt()
-    llm.check_prompts(prompt)
+    llm = OpenAI(name='Income Statement Generation Bot',
+                        api_key=env.OPENAI_API_KEY,
+                        # cache=True,
+                        temperature=0.7,
+                        verbose=True)
 
-    llm.system_prompt = prompt
+    try:
+        runnable_chain = RunnableSequence(prompt, llm, output_parser)
+    except Exception as e:
+        print(f"runnable_chain error: {str(e)}")
+        raise e
 
-    output = await llm.astructured_predict(output_cls=IncomeStatementLLMResponse, prompt=prompt)
+    try:
+        output_chunks = runnable_chain.invoke({"input_data_csv": input_data_csv})
+        return output_chunks
 
-    logging.info(f"Output: {output}")
-    return output
+    except Exception as e:
+        print(f"runnable_chain.invoke error: {str(e)}")
+        raise e

tests/conftest.py

@@ -54,7 +54,9 @@ async def connection_test(test_db, event_loop):
 
     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})
+        sessionmanager.init(connection_str,
+                            # {"echo": True, "future": True}
+                            )
         yield
         await sessionmanager.close()
 

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

@@ -1,109 +1,25 @@
-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
+from sqlalchemy.ext.asyncio import AsyncSession
+from tests.utils import get_fake_transactions
 
-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
+
+    # 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,
+        ),
+    ]