Query Rag

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

@@ -6,6 +6,10 @@ from app.model.transaction import Transaction as TransactionModel
 from app.schema.index import TransactionResponse
 from app.engine.postgresdb import get_db_session
 
+from app.service.query_rag import answer_query, fetch_transaction_documents
+
+import pandas as pd
+
 transaction_router = r = APIRouter(prefix="/api/v1/transactions", tags=["transactions"])
 
 
@@ -46,3 +50,27 @@ async def get_transactions(user_id: int, db: AsyncSession = Depends(get_db_sessi
                 }
                 results.append(result)
             return results
+
+
+@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 transactions.
+    """
+    try:
+        result = await TransactionModel.get_by_user(db, user_id)
+        all_rows = result.all()
+        if len(all_rows) == 0:
+            raise HTTPException(status_code=500, detail="No transactions found for this user")
+        document_splits = await fetch_transaction_documents(all_rows)
+        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)}")

app/api/routers/user.py

@@ -6,6 +6,10 @@ from app.engine.postgresdb import get_db_session
 from app.schema.index import UserCreate, User as UserSchema, UserResponse, UserUpdate
 from app.model.user import User as UserModel
 
+from app.model.transaction import Transaction as TransactionModel
+
+from tests.utils import get_fake_transactions
+
 
 user_router = r = APIRouter(prefix="/api/v1/users", tags=["users"])
 logger = logging.getLogger(__name__)
@@ -24,12 +28,21 @@ logger = logging.getLogger(__name__)
 async def create_user(user: UserCreate, db: AsyncSession = Depends(get_db_session)):
     try:
         db_user = await UserModel.get(db, email=user.email)
+        print(f"db_user: {db_user}")
         if db_user and not db_user.is_deleted:
             raise HTTPException(status_code=409, detail="User already exists")
 
         await UserModel.create(db, **user.model_dump())
         user = await UserModel.get(db, email=user.email)
+        print(f"user: {user}")
+
+        # Supply fake transactions data when creating user to have data to query
+        fake_transactions = get_fake_transactions(user.id)
+        await TransactionModel.bulk_create(db, fake_transactions)
+        print(f"fake_transactions: {fake_transactions}")
+
         return user
+    
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 

app/model/transaction.py

@@ -23,6 +23,15 @@ class Transaction(Base, BaseModel):
 
     def __str__(self) -> str:
         return f"{self.transaction_date}, {self.category}, {self.name_description}, {self.amount}, {self.type}"
+    
+    def to_dict(self):
+        return {
+            'transaction_date': self.transaction_date,
+            'category': self.category,
+            'name_description': self.name_description,
+            'amount': self.amount,
+            'type': self.type
+        }
 
     @classmethod
     async def create(cls: "type[Transaction]", db: AsyncSession, **kwargs) -> "Transaction":

app/schema/index.py

@@ -53,6 +53,8 @@ class TransactionCreate(TransactionResponse):
 class Transaction(TransactionResponse):
     user: User
 
+class TransactionRagResponse(PydanticBaseModel):
+    answer: str
 
 class FileUploadCreate(PydanticBaseModel):
     source: str

app/service/query_rag.py

@@ -0,0 +1,158 @@
+from pinecone import Pinecone, ServerlessSpec
+from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_pinecone import PineconeVectorStore
+from langchain.vectorstores import Chroma
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain_core.documents import Document
+
+from fastapi import HTTPException
+
+from typing import List
+from app.model.transaction import Transaction
+from app.model.income_statement import IncomeStatement
+
+import os
+
+async def fetch_transaction_documents(transactions: List[Transaction]) -> List[Document]:
+    try:
+        document = ''.join(str(row)+'\n' for row in transactions)
+
+        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_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:
+        df (pd.DataFrame): DataFrame containing the transactions that a user has entered
+        query (str): The query the user will ask against said embedding
+
+    Returns:
+        str: Response to query
+    """
+    try:
+        openai_api_key = os.environ['OPENAI_API_KEY']
+
+        embeddings = OpenAIEmbeddings(
+            model="text-embedding-3-small",
+            openai_api_key=openai_api_key
+        )
+
+        print(f"answer_transactions_query: \n document_splits: {type(document_splits)}\n")
+
+        llm = ChatOpenAI(
+            openai_api_key=openai_api_key,
+            model_name="gpt-3.5-turbo",
+            temperature=0.0
+        )
+
+        prompt_template = """
+            Use the following pieces of context to answer the question at the end.
+            If you don't know the answer, please think rationally answer from your own knowledge base
+
+            {context}
+
+            Question: {question}
+        """
+        prompt = PromptTemplate(
+            template=prompt_template, input_variables=["context", "question"]
+        )
+
+        chain_type_kwargs = {"prompt": prompt}
+
+        vectorestore = Chroma.from_documents(documents=document_splits,embedding=embeddings)
+
+        qa = RetrievalQA.from_chain_type(
+            llm=llm,
+            retriever=vectorestore.as_retriever(),
+            chain_type_kwargs=chain_type_kwargs
+        )
+
+        print(f"answer_transactions_query: \n qa: {qa}\n")
+
+        answer = qa({"query": query})
+
+        print(f"answer_transactions_query: \n answer: {answer['result']}\n")
+        
+        result_split = answer['result'].split("\n")
+        result = ''.join(item for item in result_split)
+
+        return result
+    
+    except Exception as e:
+        raise HTTPException(status_code = 500, detail=f"answer_query error: {str(e)}")
+    
+
+        # Note: Attempted to use Pinecone however it failed to learn from documents so settled for Chroma
+
+        # batch_limit = 20
+
+        # pinecone_api_key = os.environ['PINECONE_API_KEY']
+
+        # namespace = "transactionsvector"
+        # index_name = "transactionsrag"
+
+        # pc = Pinecone(api_key=pinecone_api_key)
+        #  if index_name in pc.list_indexes().names():
+        #     pc.delete_index(index_name)
+
+        # pc.create_index(
+        #     name=index_name,
+        #     dimension=1536,
+        #     metric="cosine",
+        #     spec=ServerlessSpec(
+        #         cloud="aws", region="us-east-1"
+        #     )
+        # )
+
+        # index = pc.Index(index_name)
+        # transactions_search = PineconeVectorStore.from_documents(
+        #     documents=document_splits,
+        #     index_name=index_name,
+        #     embedding=embeddings,
+        #     namespace=namespace
+        # )
+
+        # print(f"answer_transactions_query: \n transactions_search: {transactions_search}\n")
+
+        # for ids in index.list(namespace=namespace):
+        #     query_item = index.query(
+        #         id=ids[0], 
+        #         namespace=namespace, 
+        #         top_k=1,
+        #         include_values=True,
+        #         include_metadata=True
+        #     )
+        #     print(f"answer_transactions_query: \n query_item: {query_item}\n")
+
+        # qa = RetrievalQA.from_chain_type(
+        #     llm=llm,
+        #     chain_type="stuff",
+        #     retriever=transactions_search.as_retriever(),
+        #     chain_type_kwargs=chain_type_kwargs
+        # )
+

tests/test_transactions.py

@@ -1,9 +1,15 @@
+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.model.user import User
+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
 
 
@@ -11,12 +17,11 @@ from tests.utils import get_fake_transactions
 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)
+    # 2. Create a bunch of transactions
     await Transaction.bulk_create(session_override, fake_transactions)
 
     # 3. Verify that the transactions are returned

tests/utils.py

@@ -8,80 +8,240 @@ def get_fake_transactions(user_id: int) -> List[TransactionCreate]:
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 1),
-            category="category1",
-            name_description="name_description",
-            amount=1.0,
+            category="shopping",
+            name_description="Amazon",
+            amount=11.0,
             type=TransactionType.EXPENSE,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 2),
-            category="category2",
-            name_description="name_description",
-            amount=2.0,
+            category="Streaming",
+            name_description="Netflix",
+            amount=4.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 3),
+            category="Other",
+            name_description="Consulting",
+            amount=3.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 4),
+            category="Other",
+            name_description="Selling Paintings",
+            amount=4.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 5),
+            category="Gym",
+            name_description="Gym membership",
+            amount=5.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 6),
+            category="Transportation",
+            name_description="Uber Taxi",
+            amount=6.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 7),
+            category="Other",
+            name_description="Freelancing",
+            amount=7.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 8),
+            category="Income",
+            name_description="Salary",
+            amount=8.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 9),
+            category="Shopping",
+            name_description="Amazon Lux",
+            amount=9.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 10),
+            category="Taxes",
+            name_description="CA Property Tax",
+            amount=10.0,
+            type=TransactionType.EXPENSE,
+        ),
+         TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 1),
+            category="shopping",
+            name_description="Amazon",
+            amount=11.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 2),
+            category="Streaming",
+            name_description="Netflix",
+            amount=4.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 3),
+            category="Other",
+            name_description="Consulting",
+            amount=3.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 4),
+            category="Other",
+            name_description="Selling Paintings",
+            amount=4.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 5),
+            category="Gym",
+            name_description="Gym membership",
+            amount=5.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 6),
+            category="Transportation",
+            name_description="Uber Taxi",
+            amount=6.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 7),
+            category="Other",
+            name_description="Freelancing",
+            amount=7.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 8),
+            category="Income",
+            name_description="Salary",
+            amount=8.0,
+            type=TransactionType.INCOME,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 9),
+            category="Shopping",
+            name_description="Amazon Lux",
+            amount=9.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 10),
+            category="Taxes",
+            name_description="CA Property Tax",
+            amount=10.0,
+            type=TransactionType.EXPENSE,
+        ),
+         TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 1),
+            category="shopping",
+            name_description="Amazon",
+            amount=11.0,
+            type=TransactionType.EXPENSE,
+        ),
+        TransactionCreate(
+            user_id=user_id,
+            transaction_date=datetime(2022, 1, 2),
+            category="Streaming",
+            name_description="Netflix",
+            amount=4.0,
             type=TransactionType.EXPENSE,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 3),
-            category="category3",
-            name_description="name_description",
+            category="Other",
+            name_description="Consulting",
             amount=3.0,
             type=TransactionType.INCOME,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 4),
-            category="category1",
-            name_description="name_description",
+            category="Other",
+            name_description="Selling Paintings",
             amount=4.0,
             type=TransactionType.INCOME,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 5),
-            category="category2",
-            name_description="name_description",
+            category="Gym",
+            name_description="Gym membership",
             amount=5.0,
             type=TransactionType.EXPENSE,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 6),
-            category="category3",
-            name_description="name_description",
+            category="Transportation",
+            name_description="Uber Taxi",
             amount=6.0,
             type=TransactionType.EXPENSE,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 7),
-            category="category1",
-            name_description="name_description",
+            category="Other",
+            name_description="Freelancing",
             amount=7.0,
             type=TransactionType.INCOME,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 8),
-            category="category2",
-            name_description="name_description",
+            category="Income",
+            name_description="Salary",
             amount=8.0,
             type=TransactionType.INCOME,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 9),
-            category="category3",
-            name_description="name_description",
+            category="Shopping",
+            name_description="Amazon Lux",
             amount=9.0,
             type=TransactionType.EXPENSE,
         ),
         TransactionCreate(
             user_id=user_id,
             transaction_date=datetime(2022, 1, 10),
-            category="category1",
-            name_description="name_description",
+            category="Taxes",
+            name_description="CA Property Tax",
             amount=10.0,
             type=TransactionType.EXPENSE,
         ),