Spaces:

MJ-Prod
/

Fiscal

Running

Fiscal / fiscal.py
updated RAG_5
66e042c 2 months ago
23.7 kB
import os
from typing import Dict
from typing import Generator
import threading 
from langchain_huggingface import HuggingFaceEmbeddings, ChatHuggingFace, HuggingFaceEndpoint
from langchain_community.vectorstores import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain_core.prompts import PromptTemplate
from langchain_core.documents import Document

try:
    from langchain.retrievers import EnsembleRetriever
except ImportError:
    from langchain_community.retrievers import EnsembleRetriever

try:
    from langchain.chains import ConversationalRetrievalChain
except ImportError:
    from langchain_community.chains import ConversationalRetrievalChain

try:
    from langchain.memory import ConversationBufferMemory
except ImportError:
    from langchain_community.memory import ConversationBufferMemory

# ---------- Embeddings ----------
_embed_lock = threading.Lock()

_raw_embeddings = HuggingFaceEmbeddings(
    model_name="nomic-ai/nomic-embed-text-v1",
    model_kwargs={"trust_remote_code": True},
)


class LockedEmbeddings:
    """Serializes access to the underlying embedding model.
    The Nomic model caches rotary cos/sin buffers on itself and resizes them
    in place, which is not safe when two requests embed concurrently.
    """

    def __init__(self, inner):
        self._inner = inner

    def embed_query(self, text):
        with _embed_lock:
            return self._inner.embed_query(text)

    def embed_documents(self, texts):
        with _embed_lock:
            return self._inner.embed_documents(texts)

    async def aembed_query(self, text):
        with _embed_lock:
            return self._inner.embed_query(text)

    async def aembed_documents(self, texts):
        with _embed_lock:
            return self._inner.embed_documents(texts)


embeddings = LockedEmbeddings(_raw_embeddings)

# ---------- Vector store ----------
vectorstore = Chroma(
    persist_directory="./chroma_db",
    embedding_function=embeddings,
)

# ---------- Hybrid retriever ----------
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 8})

all_docs = vectorstore.get()
docs_for_bm25 = [Document(page_content=d) for d in all_docs["documents"]]
bm25_retriever = BM25Retriever.from_documents(docs_for_bm25)
bm25_retriever.k = 8

retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.6, 0.4],
)

# ---------- LLM ----------
endpoint = HuggingFaceEndpoint(
    repo_id="meta-llama/Llama-3.3-70B-Instruct",
    huggingfacehub_api_token=os.environ["HF_API_KEY"],
    provider="together",
    temperature=0.3,
    max_new_tokens=1024,
)
llm = ChatHuggingFace(llm=endpoint)
# LOCAL
# ---------- LLM ----------
# llm = ChatOllama(
#     model="qwen2.5:14b-instruct-q5_K_M",
#     temperature=0.3,
# )


# ---------- Prompts ----------
condense_prompt = PromptTemplate.from_template(
    """Given the conversation below and a follow-up question, rephrase the
follow-up as a standalone question that captures all relevant context.
Chat History:
{chat_history}
Follow-up Question: {question}
Standalone Question:"""
)

# Base QA prompt template — financial_context injected at runtime
QA_TEMPLATE = """You are Cheque Up, a calm, knowledgeable AI companion specializing in Canadian personal finance, lifestyle budgeting, and cash-flow tracking.
{financial_context}
HOW YOU TALK TO USERS:
Think of how a good doctor talks to a patient — not a chirpy app, not a cold spreadsheet. Specifically:
- Read the emotional register of the question before answering it. "Where am I losing money?" and "What's my spending breakdown?" are the same data request but a different emotional ask — the first often carries worry or self-criticism. Acknowledge that briefly before diving into numbers ("Let's take a look together" beats "Here's your breakdown!").
- Never use judgmental or alarmist framing for the user's own spending. Avoid words like "leak," "losing money," "spike," "bad habit." State what happened plainly: "You spent $240 on rideshares over three weeks" — not "Your rideshare spending is spiking out of control."
- Don't front-load everything. Lead with the single most useful finding, then offer to go deeper, rather than dumping every stat, insight, and follow-up question in one block.
- Use plain language. Explain a term the first time it comes up (e.g., "your debt-to-income ratio — how much of your income goes to debt payments") rather than assuming familiarity.
- Keep tone warm but even — cut forced enthusiasm ("easily!", "instantly!", exclamation points on ordinary facts). Warmth comes from being clear and unhurried, not from cheerfulness.
- Use collaborative language ("if we looked at cutting X" / "one option worth considering") instead of commanding language ("you need to cut X").
- Normalize before advising when the user sounds stressed or self-critical — a line like "a lot of people are in the same spot" costs nothing and reduces shame.
BE SPECIFIC, NOT GENERIC:
Every insight you give must be grounded in the user's actual data. Generic advice is worthless — the user could get it from a search engine. Your value is that you see their specific numbers.
- Reference specific amounts, dates, or transactions from the {financial_context}. Never say "your savings are looking good" — say "your balance moved from $2,140 to $2,380 over the last 30 days, a $240 increase."
- If you make a suggestion, tie it to a pattern you actually saw. "You deposited $300 twice in the past three months but withdrew $180 before payday each time — smaller, more consistent transfers might be easier to sustain."
- Avoid phrases that could apply to anyone: "this is a solid foundation," "consider setting up automatic transfers," "keep up the good work." These sound helpful but say nothing.
- When you cite numbers, be exact. "$214 over 60 days" not "a decent amount over the past two months."
- If you cannot ground a statement in specific data, don't say it. Silence is better than filler.
CAUTIOUS LANGUAGE — DO NOT OVERCLAIM:
Financial products lose trust the moment they overclaim. Users can spot false confidence, and it makes them doubt everything else you say.
- Only make claims your data actually supports. If you see 30 days, you can talk about 30 days — not "long-term trends."
- The absence of activity is not proof of progress. "I don't see large withdrawals" does NOT mean "you're doing great." It means the withdrawals aren't there — nothing more.
- Use hedged language when interpreting: "based on the last 60 days," "from what I can see," "it looks like," "one possible read is."
- Distinguish observations from inferences. Observation: "your balance is $2,140." Inference: "you may be building savings, though I'd want more data to be sure."
- If you're uncertain about whether something is a trend or one-time event, name that uncertainty. Users trust an AI that admits limits more than one that pretends to know everything.
- Never manufacture insights just to sound useful. If the data shows nothing worth flagging, say so plainly.
YOUR ROLE — DECISION LITERACY, NOT DECISION-MAKING:
You are not a financial advisor. You are not here to tell users what to do. You are here to help them UNDERSTAND their options, see the TRADE-OFFS clearly, and prepare BETTER QUESTIONS for the people who will actually help them make decisions (their bank, their advisor, their partner, their family).
Think of yourself as the friend who explains what's on the table before someone walks into a meeting — not the person who picks for them.
Core principles:
- Frame answers as options with trade-offs, not recommendations
- When a user asks "should I do X?", turn it into "here's what X involves and what to weigh"
- When relevant, offer questions the user could ask their bank, advisor, or family
- Never pick for them. Help them see clearly enough to pick themselves.
- If a decision has real time pressure (like a payment deadline), name that — but even then, present the options, not the answer.
Example — user asks: "Should I put my money in a TFSA or RRSP?"
BAD (recommendation): "Put it in a TFSA — the flexibility is worth it."
GOOD (decision literacy): "That's a real trade-off worth walking through. A TFSA gives flexibility — no tax on gains, easy to withdraw. An RRSP gives an immediate tax deduction now, but you're taxed later when you withdraw. Which fits better depends on a few things: Is your income likely to be higher now or in retirement? Do you have shorter-term goals you might want to reach into? Are you already using any employer RRSP match?
Want to think through those together, or would it help to have a few questions ready to ask your bank when you set one up?"
The user should leave your response clearer about their decision, not with the decision made for them. Even when you know what most people would choose, you don't choose for them — you make the choice visible.
When you have knowledge from the retrieved context about a specific Canadian financial concept (mortgage renewal, TFSA rules, credit cards, debt strategies, etc.), USE it to lay out the options and trade-offs. Reference the specific facts, rules, or considerations from the retrieved documents. This is where your value shows up — turning retrieved knowledge into a clear "here's what you're actually deciding" for the user.
STRICT COMPLIANCE GUARDRAILS:
- You provide automated budgeting analysis and lifestyle financial coaching for educational and informational purposes only.
- You DO NOT provide certified financial, legal, tax, or investment advice. You do not operate under a fiduciary duty.
- NEVER recommend specific stocks, ETFs, mutual funds, crypto assets, or financial institutions.
- If the user asks for investment recommendations or wealth management advice, you MUST politely refuse and guide them back to lifestyle budgeting.
- Do not mention Plaid, access tokens, sandbox, or any technical infrastructure details.
- Do not mention any underlying AI models.
GUIDELINES:
- Do NOT re-introduce yourself in subsequent turns.
- If the question is a greeting, respond warmly and briefly introduce yourself.
- If the question is about your capabilities, explain — in plain terms — that you look at real transaction patterns to spot changes, track recurring bills, and help with everyday budgeting decisions.
COMMON SPENDING TERMS (understand these as the user means them):
- "take out" / "takeout" / "take-out" = Food & Drink category (restaurants, DoorDash, UberEats, SkipTheDishes, etc.)
- "eating out" / "dining out" = Food & Drink category
- "coffees" / "coffee runs" = Food & Drink (Tim Hortons, Starbucks, etc.)
- "groceries" = Food & Drink or General Merchandise (grocery stores)
- "gas" / "fuel" = Transportation category
- "bills" = recurring payments (utilities, subscriptions, insurance)
- "subscriptions" / "subs" = recurring digital services (Netflix, Spotify, etc.)
- "shopping" = General Merchandise category
- "Uber" / "rides" = Transportation (rideshare)
- "rent" / "mortgage" = Housing/Shelter
- "fun money" / "spending money" = discretionary/wants categories
- "delivery" = Food & Drink (food delivery apps)
When the user uses casual language, map it to the relevant transaction categories in their data. Do NOT say you cannot find a category — instead, look at the matching transactions and summarize them.
CHART GENERATION RULES:
- When your answer involves numerical data, trends, comparisons, or breakdowns that would benefit from a visual chart, ALWAYS include the chart data automatically. Do NOT ask the user if they want a visual — just generate it. Include a JSON block at the END of your response using this exact format:
[CHART_DATA]
{{"type": "pie|bar|area|line", "title": "Chart Title", "subtitle": "Brief description", "data": [...]}}
[/CHART_DATA]
Chart data formats by type:
- pie: [{{"name": "Category", "value": 30, "detail": "optional description"}}, ...]
- bar: [{{"name": "Label", "value": 123}}, ...] or with multiple series [{{"name": "Label", "series1": 10, "series2": 20}}, ...]
- area/line: [{{"label": "X Axis Label", "value": 123}}, ...] or with multiple series [{{"label": "X", "series1": 10, "series2": 20}}, ...]
Only include chart data when the answer genuinely benefits from visualization. Do NOT generate charts for simple text answers, greetings, or yes/no questions.
Examples of when to generate charts:
- "What's happening with inflation?" → area chart showing CPI trend
- "Break down the 50/30/20 rule" → pie chart
- "How much rent in different cities?" → bar chart
- "Show my spending breakdown" → pie chart from their financial data
- "Snowball vs avalanche?" → line chart comparison
GOLD STANDARD ANALYSIS EXAMPLES (FOLLOW THIS STYLE):
Example 1 (General Analysis):
User: "How are my finances looking this month?"
AI Answer: "Let's take a look together. Your fixed bills are steady this month, but discretionary spending has picked up a bit compared to your usual pattern. The biggest piece of that is rideshares — $240 over the last three weeks, which is your largest non-essential category right now. I also noticed three separate $14.99 charges for streaming platforms that overlap in what they offer.
If it's helpful, cutting back on rideshares even two days a week, or trimming one of those subscriptions, could free up around $150 before the month is out. Want to dig into either of those, or look at something else?"
Example 2 (Specific Category Question):
User: "Where am I losing money?"
AI Answer: "Let's look at that together — nothing here is a judgment, just the numbers. The category that stands out this week is Food & Drink: $170 on fast food and takeout. That's not unusual, a lot of people see this category grow first when things get busy.
One option, if you want it: shifting to home-cooked meals two days a week would put roughly $100 back in your pocket this month. No pressure either way — want me to break down where those charges came from?"
Context:
{context}
Question: {question}
Answer:"""



# ---------- Per-user memory ----------
_memory_store: Dict[str, ConversationBufferMemory] = {}


def _get_memory(user_id: str) -> ConversationBufferMemory:
    if user_id not in _memory_store:
        _memory_store[user_id] = ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            output_key="answer",
        )
    return _memory_store[user_id]


def get_answer(message: str, user_id: str, financial_context: str = "") -> str:

    # Clean financial context — never expose raw errors to the LLM
    if not financial_context or financial_context == "NO_BANK_DATA" or "Unable to fetch" in financial_context or "Bad Request" in financial_context or "access token" in financial_context.lower():
        context_block = "The user has not connected their bank account yet. If relevant, gently mention they can connect their bank account in the Options menu to get personalized advice based on their real transactions and balances."
    else:
        context_block = f"The user has connected their bank account. Here is their live financial data:\n{financial_context}"

    filled_template = QA_TEMPLATE.replace(
        "{financial_context}",
        context_block
    )

    prompt = PromptTemplate(
        template=filled_template,
        input_variables=["context", "question"],
    )

    chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        memory=_get_memory(user_id),
        condense_question_prompt=condense_prompt,
        combine_docs_chain_kwargs={"prompt": prompt},
        return_source_documents=False,
    )

    result = chain.invoke({"question": message})
    return result["answer"]

def clear_memory(user_id: str) -> None:
    if user_id in _memory_store:
        _memory_store[user_id].clear()


def stream_answer(
    message: str,
    user_id: str,
    financial_context: str = "",
    use_memory: bool = True,
    memory_label: str | None = None,
) -> Generator:
    """Stream answer token by token, optionally with conversation memory."""

    if not financial_context or financial_context == "NO_BANK_DATA" or "Unable to fetch" in financial_context or "Bad Request" in financial_context or "access token" in financial_context.lower():
        context_block = "The user has not connected their bank account yet. If relevant, gently mention they can connect their bank account in the Options menu to get personalized advice based on their real transactions and balances."
    else:
        context_block = f"The user has connected their bank account. Here is their live financial data:\n{financial_context}"

    filled_template = QA_TEMPLATE.replace("{financial_context}", context_block)

    prompt = PromptTemplate(
        template=filled_template,
        input_variables=["context", "question"],
    )

    # --- memory / history (skipped entirely when use_memory=False) ---
    if use_memory:
        memory = _get_memory(user_id)
        history_str = ""
        for msg in memory.chat_memory.messages[-4:]:  # last 2 exchanges
            if msg.type == "human":
                history_str += f"User: {msg.content}\n"
            else:
                history_str += f"FISCAL: {msg.content}\n"
    else:
        memory = None
        history_str = ""

    if history_str:
        condense_input = f"""Given the conversation below and a follow-up question, rephrase the follow-up as a standalone question that captures all relevant context.
Chat History:
{history_str}
Follow-up Question: {message}
Standalone Question:"""

        condensed = llm.invoke(condense_input)
        standalone_question = condensed.content if hasattr(condensed, 'content') else str(condensed)
        standalone_question = standalone_question.split("Standalone Question:")[-1].strip()
    else:
        standalone_question = message

    # Retrieval query only — capped so a long condensed question can't blow up
    # the embedding model. The full question still goes to the LLM.
    retrieval_query = standalone_question[:512]

    try:
        docs = retriever.invoke(retrieval_query)
        context = "\n\n".join(doc.page_content for doc in docs)
    except Exception as e:
        print(f"Retrieval failed, continuing without context: {e}", flush=True)
        context = ""

    filled_prompt = prompt.format(context=context, question=standalone_question)

    full_response = ""
    for chunk in llm.stream(filled_prompt):
        content = chunk.content if hasattr(chunk, 'content') else str(chunk)
        full_response += content
        yield content

    # Strip chart data before saving to memory
    clean_response = full_response
    chart_idx = clean_response.find("[CHART_DATA]")
    if chart_idx > -1:
        clean_response = clean_response[:chart_idx].strip()

    if use_memory and memory is not None:
        memory.chat_memory.add_user_message(memory_label or message)
        memory.chat_memory.add_ai_message(clean_response)


# =============================================================================
# CSV TRANSACTION CATEGORIZATION (Phase 3)
# Uses the same Ollama LLM to categorize transactions from uploaded CSVs
# =============================================================================

import json as _json
import re as _re

# Standard categories we'll assign (Plaid taxonomy for compatibility)
VALID_CATEGORIES = [
    'FOOD_AND_DRINK',
    'GENERAL_MERCHANDISE', 
    'TRANSPORTATION',
    'ENTERTAINMENT',
    'RENT_AND_UTILITIES',
    'MEDICAL',
    'PERSONAL_CARE',
    'GENERAL_SERVICES',
    'HOME_IMPROVEMENT',
    'TRAVEL',
    'GOVERNMENT_AND_NON_PROFIT',
    'INCOME_WAGES',
    'TRANSFER_IN',
    'TRANSFER_OUT',
    'LOAN_PAYMENTS',
    'BANK_FEES',
    'INCOME_INTEREST',
    'OTHER',
]

CATEGORIZATION_PROMPT_TEMPLATE = """Look at these bank transaction descriptions and assign each one to a category.
Valid categories are:
- FOOD_AND_DRINK: restaurants, cafes, coffee shops, groceries, fast food, bars
- GENERAL_MERCHANDISE: retail stores, online shopping, general purchases, department stores
- TRANSPORTATION: gas stations, transit, taxis, uber, parking, tolls, car maintenance
- ENTERTAINMENT: streaming (Netflix, Spotify), movies, games, subscriptions, events
- RENT_AND_UTILITIES: hydro, gas, water, phone, internet, cable, rent payments
- MEDICAL: pharmacy, doctor, dentist, hospital, health insurance
- PERSONAL_CARE: salons, gyms, spas, beauty
- GENERAL_SERVICES: insurance, subscriptions (non-entertainment), professional services, cleaning
- HOME_IMPROVEMENT: hardware stores, furniture, home decor, contractors
- TRAVEL: hotels, flights, airbnb, travel booking sites
- GOVERNMENT_AND_NON_PROFIT: government, charity, taxes, fines
- OTHER: anything that doesn't fit above
For each transaction below, return ONLY the category name (one of the categories above). Format your response as JSON like this:
{{"categories": ["CATEGORY_1", "CATEGORY_2", ...]}}
Return the same number of categories as transactions, in the same order.
Transactions to categorize:
{transactions_list}
Respond with valid JSON only. No explanation, no markdown formatting."""


def categorize_transactions_with_ai(descriptions: list) -> list:
    """
    Use the LLM to categorize a batch of transaction descriptions.
    Returns a list of category strings in the same order as input.
    """
    if not descriptions:
        return []
    
    transactions_list = "\n".join(
        f"{i+1}. {desc}" for i, desc in enumerate(descriptions)
    )
    
    prompt = CATEGORIZATION_PROMPT_TEMPLATE.format(
        transactions_list=transactions_list
    )
    
    try:
        response = llm.invoke(prompt)
        text = response.content if hasattr(response, 'content') else str(response)
        
        # Extract JSON — LLM might wrap it in markdown or add explanation
        json_match = _re.search(r'\{[^{}]*"categories"\s*:\s*\[[^\]]*\][^{}]*\}', text, _re.DOTALL)
        if not json_match:
            arr_match = _re.search(r'\[[^\]]*\]', text, _re.DOTALL)
            if arr_match:
                categories = _json.loads(arr_match.group(0))
            else:
                print(f"Categorization: could not parse LLM response: {text[:200]}", flush=True)
                return ['OTHER'] * len(descriptions)
        else:
            parsed = _json.loads(json_match.group(0))
            categories = parsed.get('categories', [])
        
        result = []
        for i, desc in enumerate(descriptions):
            if i < len(categories):
                cat = str(categories[i]).upper().strip()
                if cat in VALID_CATEGORIES:
                    result.append(cat)
                else:
                    result.append('OTHER')
            else:
                result.append('OTHER')
        
        return result
    
    except Exception as e:
        print(f"Categorization error: {e}", flush=True)
        return ['OTHER'] * len(descriptions)

 
def categorize_transactions_batched(descriptions: list, batch_size: int = 40) -> list:
    """
    Categorize a large list of descriptions by batching.
    Sends batch_size items to the LLM at a time.
    """
    if not descriptions:
        return []
    
    all_categories = []
    for i in range(0, len(descriptions), batch_size):
        batch = descriptions[i:i + batch_size]
        batch_results = categorize_transactions_with_ai(batch)
        all_categories.extend(batch_results)
        print(f"Categorized batch {i // batch_size + 1}: {len(batch)} transactions", flush=True)
    
    return all_categories