Update utility.py

utility.py

@@ -1,16 +1,21 @@
 import json
 import os
+import re
 from datetime import datetime
 from google.cloud import firestore
 import pandas as pd
 from langchain_google_genai import ChatGoogleGenerativeAI
 import google.generativeai as genai
 import logging
+from typing import List, Dict, Union, Optional
 
+# Configure logging
 logger = logging.getLogger(__name__)
 
+# Initialize Firestore
 db = firestore.Client.from_service_account_json("firestore-key.json")
 
+# Initialize Gemini LLM
 llm = ChatGoogleGenerativeAI(
     model="gemini-2.0-flash-thinking-exp",
     max_output_tokens=1024,
@@ -19,8 +24,11 @@ llm = ChatGoogleGenerativeAI(
     top_p=0.01,
 )
 
-def generateResponse(prompt):
-    system_prompt = """You MUST format responses EXACTLY like this:
+def generateResponse(prompt: str) -> str:
+    """
+    Generate structured transaction response from user input using Gemini
+    """
+    system_prompt = """You MUST format responses EXACTLY as follows:
 
 *Intent*: [create/read/update/delete]
 *Transaction Type*: [purchase/sale/inventory/etc]
@@ -29,7 +37,7 @@ def generateResponse(prompt):
 - Field2: Value2
 - Field3: Value3
 
-For multiple transactions repeat this pattern exactly."""
+For multiple transactions, repeat this pattern for each transaction."""
     
     genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
     model = genai.GenerativeModel(
@@ -46,27 +54,33 @@ For multiple transactions repeat this pattern exactly."""
         response = model.generate_content([system_prompt, prompt])
         return response.text
     except Exception as e:
-        logger.error(f"Generation error: {str(e)}")
-        return None
+        logger.error(f"Response generation failed: {str(e)}")
+        return "Sorry, I couldn't process that request."
 
-def parse_multiple_transactions(response_text):
+def parse_multiple_transactions(response_text: str) -> List[Dict]:
+    """
+    Parse Gemini response into structured transaction data
+    Returns list of transaction dictionaries
+    """
     transactions = []
     current_trans = {}
     
+    # First try JSON parsing if response is JSON
     try:
-        # First try JSON parsing
         parsed = json.loads(response_text)
         if isinstance(parsed, dict):
             return [add_timestamp(parsed)]
-        return [add_timestamp(t) for t in parsed]
+        elif isinstance(parsed, list):
+            return [add_timestamp(t) for t in parsed]
     except json.JSONDecodeError:
-        pass  # Proceed to text parsing
+        pass  # Proceed to text parsing if not JSON
 
+    # Text-based parsing
     lines = [line.strip() for line in response_text.split('\n') if line.strip()]
     
     for line in lines:
         if line.startswith('*Intent*:'):
-            if current_trans:
+            if current_trans:  # Save previous transaction if exists
                 transactions.append(add_timestamp(current_trans))
                 current_trans = {}
             current_trans['intent'] = line.split(':', 1)[1].strip().lower()
@@ -81,16 +95,22 @@ def parse_multiple_transactions(response_text):
                 value = key_value[1].strip()
                 current_trans['details'][key] = value
     
+    # Add the last transaction if exists
     if current_trans:
         transactions.append(add_timestamp(current_trans))
     
     return transactions if transactions else []
 
-def add_timestamp(trans):
-    trans['created_at'] = datetime.now().isoformat()
-    return trans
+def add_timestamp(transaction: Dict) -> Dict:
+    """Add timestamp to transaction"""
+    transaction['created_at'] = datetime.now().isoformat()
+    return transaction
 
-def create_inventory(user_phone, transaction_data):
+def create_inventory(user_phone: str, transaction_data: List[Dict]) -> bool:
+    """
+    Create or update inventory items
+    Returns True if successful, False otherwise
+    """
     for transaction in transaction_data:
         item_name = transaction['details'].get('item')
         if not item_name:
@@ -98,74 +118,190 @@ def create_inventory(user_phone, transaction_data):
             
         quantity = int(transaction['details'].get('quantity', 0))
         doc_ref = db.collection("users").document(user_phone).collection("inventory").document(item_name)
-        existing_item = doc_ref.get()
         
-        if existing_item.exists:
-            existing_data = existing_item.to_dict()
-            current_qty = int(existing_data['details'].get('quantity', 0))
-            existing_data['details']['quantity'] = current_qty + quantity
-            doc_ref.set(existing_data)
-        else:
-            transaction['details']['quantity'] = quantity
-            doc_ref.set(transaction)
+        try:
+            # Update existing or create new
+            doc_ref.set({
+                'intent': 'create',
+                'transaction_type': 'inventory',
+                'details': {
+                    'item': item_name,
+                    'quantity': firestore.Increment(quantity),
+                    **{k:v for k,v in transaction['details'].items() if k not in ['item', 'quantity']}
+                },
+                'last_updated': datetime.now().isoformat()
+            }, merge=True)
+        except Exception as e:
+            logger.error(f"Inventory update failed for {user_phone}: {str(e)}")
+            return False
+    
     return True
 
-def create_sale(user_phone, transaction_data):
+def create_sale(user_phone: str, transaction_data: List[Dict]) -> bool:
+    """
+    Process sales transactions with inventory validation
+    Returns True if successful, False otherwise
+    """
+    batch = db.batch()
+    inventory_ref = db.collection("users").document(user_phone).collection("inventory")
+    sales_ref = db.collection("users").document(user_phone).collection("sales")
+    
     for transaction in transaction_data:
         item_name = transaction['details'].get('item')
         if not item_name:
             continue
             
         quantity = int(transaction['details'].get('quantity', 0))
-        inventory = fetch_transaction(user_phone, item_name)
+        item_doc = inventory_ref.document(item_name).get()
         
-        if not inventory or int(inventory['details'].get('quantity', 0)) < quantity:
+        if not item_doc.exists:
+            logger.error(f"Item {item_name} not found in inventory")
+            return False
+            
+        current_stock = int(item_doc.to_dict().get('details', {}).get('quantity', 0))
+        if current_stock < quantity:
+            logger.error(f"Insufficient stock for {item_name}")
             return False
             
-        new_stock = int(inventory['details']['quantity']) - quantity
-        inventory['details']['quantity'] = new_stock
-        db.collection("users").document(user_phone).collection("inventory").document(item_name).set(inventory)
+        # Update inventory
+        batch.update(inventory_ref.document(item_name), {
+            'details.quantity': firestore.Increment(-quantity),
+            'last_updated': datetime.now().isoformat()
+        })
         
-        sale_ref = db.collection("users").document(user_phone).collection("sales").document()
-        sale_ref.set(transaction)
-    return True
+        # Record sale
+        sale_doc = sales_ref.document()
+        batch.set(sale_doc, {
+            **transaction,
+            'timestamp': datetime.now().isoformat()
+        })
+    
+    try:
+        batch.commit()
+        return True
+    except Exception as e:
+        logger.error(f"Sales transaction failed: {str(e)}")
+        return False
 
-def read_datalake(user_phone, user_question):
-    inventory_ref = db.collection("users").document(user_phone).collection("inventory")
-    sales_ref = db.collection("users").document(user_phone).collection("sales")
-    inventory_list = [doc.to_dict() for doc in inventory_ref.stream()]
-    sales_list = [doc.to_dict() for doc in sales_ref.stream()]
+def read_datalake(user_phone: str, query: str) -> Union[str, Dict]:
+    """
+    Query user's transaction data
+    Returns formatted response or file path for visualizations
+    """
+    from pandasai import SmartDatalake
+    
+    try:
+        # Get all inventory and sales data
+        inventory = [doc.to_dict() for doc in 
+                   db.collection("users").document(user_phone).collection("inventory").stream()]
+        sales = [doc.to_dict() for doc in 
+                db.collection("users").document(user_phone).collection("sales").stream()]
+        
+        if not inventory and not sales:
+            return "No transaction data found"
+        
+        # Create DataFrames
+        inventory_df = pd.DataFrame(inventory)
+        sales_df = pd.DataFrame(sales)
+        
+        # Initialize SmartDatalake
+        lake = SmartDatalake(
+            [inventory_df, sales_df],
+            config={
+                "llm": llm,
+                "enable_cache": False,
+                "save_logs": False
+            }
+        )
+        
+        # Process query
+        response = lake.chat(query)
+        
+        # Handle different response types
+        if isinstance(response, str):
+            return response
+        elif hasattr(response, 'to_dict'):
+            return response.to_dict()
+        return str(response)
+        
+    except Exception as e:
+        logger.error(f"Data query failed: {str(e)}")
+        return "Sorry, I couldn't retrieve that data."
+
+def update_transaction(user_phone: str, transaction_id: str, 
+                     updates: Dict, collection: str = "inventory") -> bool:
+    """
+    Update existing transaction
+    Returns True if successful, False otherwise
+    """
+    doc_ref = db.collection("users").document(user_phone).collection(collection).document(transaction_id)
     
-    if not inventory_list and not sales_list:
-        return "No data found"
+    if not doc_ref.get().exists:
+        return False
+        
+    update_data = {
+        f"details.{k}": v for k,v in updates.items()
+    }
+    update_data['last_updated'] = datetime.now().isoformat()
     
-    inventory_df = pd.DataFrame(inventory_list)
-    sales_df = pd.DataFrame(sales_list)
+    try:
+        doc_ref.update(update_data)
+        return True
+    except Exception as e:
+        logger.error(f"Update failed: {str(e)}")
+        return False
+
+def delete_transaction(user_phone: str, transaction_data: List[Dict]) -> bool:
+    """
+    Delete specified transactions
+    Returns True if successful, False otherwise
+    """
+    batch = db.batch()
+    collection_map = {
+        'purchase': 'inventory',
+        'sale': 'sales',
+        'inventory': 'inventory'
+    }
     
-    from pandasai import SmartDatalake
-    lake = SmartDatalake(
-        [inventory_df, sales_df],
-        config={"llm": llm}
-    )
+    for transaction in transaction_data:
+        item_name = transaction['details'].get('item')
+        if not item_name:
+            continue
+            
+        trans_type = transaction.get('transaction_type', 'inventory').lower()
+        collection = collection_map.get(trans_type, 'inventory')
+        
+        doc_ref = db.collection("users").document(user_phone).collection(collection).document(item_name)
+        batch.delete(doc_ref)
     
     try:
-        return str(lake.chat(user_question))
+        batch.commit()
+        return True
     except Exception as e:
-        logger.error(f"Datalake error: {str(e)}")
-        return "Error generating response"
+        logger.error(f"Deletion failed: {str(e)}")
+        return False
 
-def persist_temporary_transaction(transactions, mobile):
+def persist_temporary_transaction(transactions: List[Dict], mobile: str) -> bool:
+    """
+    Store transactions temporarily before confirmation
+    Returns True if successful, False otherwise
+    """
     try:
-        db.collection("users").document(mobile).collection("temp_transactions").document("pending").set({
+        doc_ref = db.collection("users").document(mobile).collection("temp_transactions").document("pending")
+        doc_ref.set({
             "transactions": transactions,
             "timestamp": datetime.now().isoformat()
         })
         return True
     except Exception as e:
-        logger.error(f"Temp save error: {str(e)}")
+        logger.error(f"Temp storage failed: {str(e)}")
         return False
 
-def format_transaction_response(transactions):
+def format_transaction_response(transactions: Union[List[Dict], Dict]) -> str:
+    """
+    Format transaction data for user display
+    Returns formatted string
+    """
     if not transactions:
         return "No transaction data"
     
@@ -173,20 +309,97 @@ def format_transaction_response(transactions):
         transactions = [transactions]
     
     output = []
-    for trans in transactions:
-        output.append(f"Intent: {trans.get('intent', 'unknown')}")
-        output.append(f"Type: {trans.get('transaction_type', 'unknown')}")
+    for idx, trans in enumerate(transactions, 1):
+        output.append(f"Transaction {idx}:" if len(transactions) > 1 else "Transaction Details:")
+        output.append(f"• Intent: {trans.get('intent', 'unknown').title()}")
+        output.append(f"• Type: {trans.get('transaction_type', 'unknown').title()}")
+        
         if 'details' in trans:
-            output.append("Details:")
+            output.append("• Details:")
             for k, v in trans['details'].items():
-                output.append(f"- {k}: {v}")
-        output.append("")
+                if k == 'currency':
+                    continue
+                if 'price' in k or 'amount' in k or 'cost' in k:
+                    currency = trans['details'].get('currency', '$')
+                    output.append(f"  - {k.title()}: {currency}{v}")
+                else:
+                    output.append(f"  - {k.title()}: {v}")
+        output.append("")  # Blank line between transactions
     
     return "\n".join(output)
 
-def fetch_transaction(user_phone, item_name=None):
+def fetch_transaction(user_phone: str, 
+                    item_name: Optional[str] = None,
+                    collection: str = "inventory") -> Union[Dict, List[Dict]]:
+    """
+    Retrieve transaction(s) from Firestore
+    Returns single dict for specific item or list of all transactions
+    """
+    col_ref = db.collection("users").document(user_phone).collection(collection)
+    
     if item_name:
-        doc = db.collection("users").document(user_phone).collection("inventory").document(item_name).get()
+        doc = col_ref.document(item_name).get()
         return doc.to_dict() if doc.exists else None
     else:
-        return [doc.to_dict() for doc in db.collection("users").document(user_phone).collection("inventory").stream()]
+        return [doc.to_dict() for doc in col_ref.stream()]
+
+        
+def process_intent(parsed_trans_data: List[Dict], mobile: str) -> str:
+    """
+    Route transactions to appropriate CRUD operation based on intent
+    Returns formatted response message for user
+    """
+    if not parsed_trans_data:
+        return "Error: No transaction data provided"
+    
+    try:
+        intent = parsed_trans_data[0].get('intent', '').lower()
+        trans_type = parsed_trans_data[0].get('transaction_type', '').lower()
+        transaction_summary = format_transaction_response(parsed_trans_data)
+        
+        # Route to appropriate operation
+        if intent == 'create':
+            if trans_type in ('purchase', 'inventory'):
+                success = create_inventory(mobile, parsed_trans_data)
+                return ("Inventory updated successfully!\n\n" + transaction_summary) if success else \
+                       ("Failed to update inventory!\n\n" + transaction_summary)
+            elif trans_type == 'sale':
+                success = create_sale(mobile, parsed_trans_data)
+                return ("Sale recorded successfully!\n\n" + transaction_summary) if success else \
+                       ("Failed to record sale! Check inventory.\n\n" + transaction_summary)
+            else:
+                return f"Unsupported transaction type: {trans_type}\n\n{transaction_summary}"
+                
+        elif intent == 'read':
+            if 'item' in parsed_trans_data[0].get('details', {}):
+                item = parsed_trans_data[0]['details']['item']
+                result = fetch_transaction(mobile, item)
+                if result:
+                    return f"Item details:\n\n{format_transaction_response(result)}"
+                return f"Item '{item}' not found"
+            else:
+                results = fetch_transaction(mobile)
+                return f"All inventory items:\n\n{format_transaction_response(results)}" if results else "Inventory is empty"
+                
+        elif intent == 'update':
+            item = parsed_trans_data[0]['details'].get('item')
+            if not item:
+                return "Error: No item specified for update"
+                
+            updates = {k:v for k,v in parsed_trans_data[0]['details'].items() if k != 'item'}
+            success = update_transaction(mobile, item, updates, trans_type if trans_type else 'inventory')
+            return (f"Updated {item} successfully!\n\n{transaction_summary}") if success else \
+                   (f"Failed to update {item}!\n\n{transaction_summary}")
+                   
+        elif intent == 'delete':
+            item = parsed_trans_data[0]['details'].get('item', 'item')
+            success = delete_transaction(mobile, parsed_trans_data)
+            return (f"Deleted {item} successfully!") if success else \
+                   (f"Failed to delete {item}!")
+                   
+        else:
+            return f"Unsupported intent: {intent}\n\n{transaction_summary}"
+            
+    except Exception as e:
+        logger.error(f"Intent processing failed: {str(e)}")
+        return "Sorry, something went wrong while processing your request."