graph.py

# -*- coding: utf-8 -*-
"""
Fintech Orchestrator Graph — HuggingFace Version

Adapted from orchestrator_v3.py for HuggingFace Spaces deployment.
Uses HF Inference API (Gemma 3) instead of local Qwen.
Uses mocked banking data instead of A2A remote agent.
"""

from __future__ import annotations

from typing import Optional
from pydantic import BaseModel, Field

from hf_model import generate_response


# ---------------------------------------------------------------------------
# Mock Banking Data (replaces A2A agent)
# ---------------------------------------------------------------------------

MOCK_BANKING_DATA = {
    "net_worth": {
        "total": 87500.00,
        "assets": 142000.00,
        "liabilities": 54500.00,
        "currency": "USD",
    },
    "assets": {
        "checking": 12500.00,
        "savings": 35000.00,
        "investments": 89500.00,
        "other": 5000.00,
    },
    "liabilities": {
        "credit_cards": 4500.00,
        "student_loans": 25000.00,
        "auto_loan": 15000.00,
        "other": 10000.00,
    },
    "portfolio": {
        "AAPL": 15200,
        "GOOGL": 12800,
        "MSFT": 18500,
        "AMZN": 9000,
        "bonds": 14000,
        "ETFs": 18000,
    },
    "transactions": [
        {"date": "2026-02-08", "description": "Salary Deposit", "amount": 5200.00},
        {"date": "2026-02-07", "description": "Grocery Store", "amount": -127.43},
        {"date": "2026-02-06", "description": "Electric Bill", "amount": -145.00},
        {"date": "2026-02-05", "description": "Restaurant", "amount": -68.50},
        {"date": "2026-02-04", "description": "Gas Station", "amount": -52.30},
    ],
}


def get_banking_data(query: str) -> str:
    """Mock banking query - returns relevant data based on query."""
    query_lower = query.lower()
    
    if "net worth" in query_lower or "toplam" in query_lower:
        data = MOCK_BANKING_DATA["net_worth"]
        return f"""Net Worth Summary:
- Total Net Worth: ${data['total']:,.2f}
- Total Assets: ${data['assets']:,.2f}
- Total Liabilities: ${data['liabilities']:,.2f}"""
    
    elif "portfolio" in query_lower or "stocks" in query_lower:
        portfolio = MOCK_BANKING_DATA["portfolio"]
        lines = [f"- {k}: ${v:,.2f}" for k, v in portfolio.items()]
        total = sum(portfolio.values())
        return f"Portfolio (Total: ${total:,.2f}):\n" + "\n".join(lines)
    
    elif "assets" in query_lower or "varlık" in query_lower:
        assets = MOCK_BANKING_DATA["assets"]
        lines = [f"- {k.title()}: ${v:,.2f}" for k, v in assets.items()]
        total = sum(assets.values())
        return f"Assets (Total: ${total:,.2f}):\n" + "\n".join(lines)
    
    elif "liabilities" in query_lower or "borç" in query_lower:
        liabilities = MOCK_BANKING_DATA["liabilities"]
        lines = [f"- {k.replace('_', ' ').title()}: ${v:,.2f}" for k, v in liabilities.items()]
        total = sum(liabilities.values())
        return f"Liabilities (Total: ${total:,.2f}):\n" + "\n".join(lines)
    
    elif "transaction" in query_lower or "işlem" in query_lower:
        transactions = MOCK_BANKING_DATA["transactions"]
        lines = [f"- {t['date']}: {t['description']} (${t['amount']:+,.2f})" for t in transactions]
        return "Recent Transactions:\n" + "\n".join(lines)
    
    else:
        return f"""Account Summary:
- Net Worth: ${MOCK_BANKING_DATA['net_worth']['total']:,.2f}
- Checking: ${MOCK_BANKING_DATA['assets']['checking']:,.2f}
- Savings: ${MOCK_BANKING_DATA['assets']['savings']:,.2f}
- Investments: ${MOCK_BANKING_DATA['assets']['investments']:,.2f}"""


# ---------------------------------------------------------------------------
# Router Decision Schema
# ---------------------------------------------------------------------------

class RouterDecision(BaseModel):
    """Router decision with multi-step planning"""
    needs_banking: bool = Field(
        default=False,
        description="True if real account data is needed"
    )
    needs_calculation: bool = Field(
        default=False,
        description="True if financial calculation is needed"
    )
    needs_graph: bool = Field(
        default=False,
        description="True if visualization/chart is needed"
    )
    task_description: str = Field(
        default="",
        description="Description of what needs to be done"
    )


def route_query(query: str) -> RouterDecision:
    """Simple keyword-based router (fast, no LLM call)."""
    query_lower = query.lower()
    
    needs_banking = any(word in query_lower for word in [
        'balance', 'net worth', 'portfolio', 'assets', 'liabilities',
        'account', 'transaction', 'hesap', 'bakiye', 'varlık', 'borç'
    ])
    
    needs_calculation = any(word in query_lower for word in [
        'calculate', 'compute', 'roi', 'interest', 'compound', 'projection',
        'hesapla', 'faiz', 'getiri'
    ])
    
    needs_graph = any(word in query_lower for word in [
        'chart', 'graph', 'visualize', 'plot', 'pie', 'bar', 'line',
        'grafik', 'görselleştir', 'çiz'
    ])
    
    return RouterDecision(
        needs_banking=needs_banking,
        needs_calculation=needs_calculation,
        needs_graph=needs_graph,
        task_description=query
    )


# ---------------------------------------------------------------------------
# Calculator
# ---------------------------------------------------------------------------

def calculate(expression: str, banking_data: str = "") -> str:
    """Perform financial calculation using LLM."""
    prompt = f"""You are a financial calculator. Perform the calculation requested.

Request: {expression}

{"Available account data:\n" + banking_data if banking_data else ""}

Provide:
1. The calculation formula used
2. Step-by-step calculation
3. Final result

For compound interest: A = P(1 + r)^t
For ROI: ((Final - Initial) / Initial) * 100
"""
    
    messages = [{"role": "user", "content": prompt}]
    return generate_response(messages, max_tokens=600)


# ---------------------------------------------------------------------------
# Orchestrator State
# ---------------------------------------------------------------------------

class OrchestratorState:
    """State container for orchestrator."""
    def __init__(self):
        self.banking_data: Optional[str] = None
        self.calculation_result: Optional[str] = None
        self.graph_data: Optional[dict] = None
        self.output: str = ""


# ---------------------------------------------------------------------------
# Main Orchestrator Function
# ---------------------------------------------------------------------------

def run_orchestrator(query: str) -> tuple[str, Optional[dict]]:
    """
    Main entry point for the orchestrator.
    
    Args:
        query: User's query string
        
    Returns:
        Tuple of (response_text, chart_data_dict_or_none)
    """
    # Route the query
    decision = route_query(query)
    state = OrchestratorState()
    response_parts = []
    chart_data = None
    
    # Step 1: Get banking data if needed
    if decision.needs_banking:
        state.banking_data = get_banking_data(query)
        response_parts.append(f"📊 **Account Data:**\n{state.banking_data}")
    
    # Step 2: Perform calculation if needed
    if decision.needs_calculation:
        state.calculation_result = calculate(query, state.banking_data or "")
        response_parts.append(f"\n🧮 **Calculation:**\n{state.calculation_result}")
    
    # Step 3: Prepare chart data if needed
    if decision.needs_graph:
        query_lower = query.lower()
        
        if 'portfolio' in query_lower or 'pie' in query_lower:
            chart_data = {
                "type": "pie",
                "title": "Portfolio Distribution",
                "data": MOCK_BANKING_DATA["portfolio"]
            }
        elif 'assets' in query_lower:
            chart_data = {
                "type": "bar",
                "title": "Assets Breakdown",
                "data": MOCK_BANKING_DATA["assets"]
            }
        elif 'liabilities' in query_lower:
            chart_data = {
                "type": "bar",
                "title": "Liabilities Breakdown",
                "data": MOCK_BANKING_DATA["liabilities"]
            }
        else:
            # Default: net worth projection
            initial = MOCK_BANKING_DATA["net_worth"]["total"]
            rate = 0.08
            chart_data = {
                "type": "line",
                "title": "Net Worth Projection (8% Annual Growth)",
                "data": {f"Year {i}": initial * (1 + rate) ** i for i in range(6)}
            }
        
        response_parts.append(f"\n📈 **Chart:** {chart_data['title']}")
    
    # Step 4: If no specific action, use LLM for general response
    if not any([decision.needs_banking, decision.needs_calculation, decision.needs_graph]):
        context = f"""You are a fintech assistant. Answer the user's question about finance.

Available account data:
- Net Worth: ${MOCK_BANKING_DATA['net_worth']['total']:,.2f}
- Assets: ${MOCK_BANKING_DATA['net_worth']['assets']:,.2f}  
- Liabilities: ${MOCK_BANKING_DATA['net_worth']['liabilities']:,.2f}

User question: {query}

Provide a helpful, concise response."""
        
        messages = [{"role": "user", "content": context}]
        llm_response = generate_response(messages, max_tokens=800)
        response_parts.append(llm_response)
    
    # Combine response
    state.output = "\n\n".join(response_parts)
    
    return state.output, chart_data