Add income statement answer endpoint

app/api/routers/income_statement.py

@@ -7,6 +7,9 @@ from app.schema.index import IncomeStatementCreateRequest, IncomeStatementRespon
 from app.engine.postgresdb import get_db_session
 from app.service.income_statement import call_llm_to_create_income_statement
 
+from app.service.query_rag import answer_query, fetch_income_statement_documents
+
+
 income_statement_router = r = APIRouter(prefix="/api/v1/income_statement", tags=["income_statement"])
 
 
@@ -48,6 +51,30 @@ async def get_income_statements(
     return result
 
 
+@r.get(
+    "/answer/{user_id}/{query}",
+    responses={
+        200: {"description": "Query answered"},
+        500: {"description": "Internal server error"},
+    },
+)
+async def answer_transactions_query(user_id: int, query: str, db: AsyncSession = Depends(get_db_session)):
+    """
+    Answer queries based on income statements
+    """
+    try:
+        result = await IncomeStatementModel.get_by_user(db, user_id)
+        if len(result) == 0:
+            raise HTTPException(status_code=500, detail="No transactions found for this user")
+        document_splits = await fetch_income_statement_documents(result)
+        answer = await answer_query(document_splits, query, user_id)
+        return answer
+    
+    except Exception as e:
+         raise HTTPException(status_code=500, detail=f"/answer endpoint error: {str(e)}")
+
+
+
 @r.get(
     "/report/{report_id}",
     response_model=IncomeStatementResponse,

app/api/routers/transaction.py

@@ -53,7 +53,7 @@ async def get_transactions(user_id: int, db: AsyncSession = Depends(get_db_sessi
 
 
 @r.get(
-    "/transaction_answer/{user_id}/{query}",
+    "/answer/{user_id}/{query}",
     responses={
         200: {"description": "Query answered"},
         500: {"description": "Internal server error"},
@@ -61,7 +61,7 @@ async def get_transactions(user_id: int, db: AsyncSession = Depends(get_db_sessi
 )
 async def answer_transactions_query(user_id: int, query: str, db: AsyncSession = Depends(get_db_session)):
     """
-    Retrieve all transactions.
+    Answer queries based on transactions.
     """
     try:
         result = await TransactionModel.get_by_user(db, user_id)
@@ -73,4 +73,4 @@ async def answer_transactions_query(user_id: int, query: str, db: AsyncSession =
         return answer
     
     except Exception as e:
-         raise HTTPException(status_code=500, detail=f"answer_transactions_query error: {str(e)}")
+         raise HTTPException(status_code=500, detail=f"/answer endpoint error: {str(e)}")

app/service/query_rag.py

@@ -11,6 +11,7 @@ from fastapi import HTTPException
 
 from typing import List
 from app.model.transaction import Transaction
+from app.model.income_statement import IncomeStatement
 
 import os
 
@@ -29,6 +30,21 @@ async def fetch_transaction_documents(transactions: List[Transaction]) -> List[D
         raise HTTPException(status_code = 500, detail=f"fetch_transaction_documents error: {str(e)}")
     
 
+async def fetch_income_statement_documents(income_statements: List[IncomeStatement]) -> List[Document]:
+    try:
+        document = ''.join(str(row)+'\n' for row in income_statements)
+
+        page_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
+        pages = page_splitter.split_text(document)
+
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=10)
+        document_splits = text_splitter.create_documents(pages)
+
+        return document_splits
+    except Exception as e:
+        raise HTTPException(status_code = 500, detail=f"fetch_income_statement_documents error: {str(e)}")
+        
+
 async def answer_query(document_splits: List[Document], query: str, user_id: int) -> str:
     """Creates an embedding of the transactions table and then returns the answer for the given query.
     Args:

tests/test_transactions.py

@@ -10,100 +10,21 @@ 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 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")
 
     fake_transactions = get_fake_transactions(user.id)
+    # 2. Create a bunch of transactions
     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