Update utility.py

utility.py

@@ -1,20 +1,16 @@
-import re
-import os
 import json
+import os
 from datetime import datetime
-import openai
 from google.cloud import firestore
-from dotenv import load_dotenv
 import pandas as pd
-from pandasai.responses.response_parser import ResponseParser
 from langchain_google_genai import ChatGoogleGenerativeAI
 import google.generativeai as genai
+import logging
+
+logger = logging.getLogger(__name__)
 
-# Load environment variables and configure Firestore
-load_dotenv()
 db = firestore.Client.from_service_account_json("firestore-key.json")
 
-# Initialize the Gemini LLM client
 llm = ChatGoogleGenerativeAI(
     model="gemini-2.0-flash-thinking-exp",
     max_output_tokens=1024,
@@ -23,32 +19,21 @@ llm = ChatGoogleGenerativeAI(
     top_p=0.01,
 )
 
-class FlaskResponse(ResponseParser):
-    def __init__(self, context) -> None:
-        super().__init__(context)
-
-    def format_dataframe(self, result):
-        return result['value'].to_html()
-
-    def format_plot(self, result):
-        try:
-            img_path = result['value']
-        except ValueError:
-            img_path = str(result['value'])
-        return img_path
+def generateResponse(prompt):
+    system_prompt = """You MUST format responses EXACTLY like this:
 
-    def format_other(self, result):
-        return str(result['value'])
+*Intent*: [create/read/update/delete]
+*Transaction Type*: [purchase/sale/inventory/etc]
+*Details*:
+- Field1: Value1
+- Field2: Value2
+- Field3: Value3
 
-def generateResponse(prompt, model='gemini-2.0-flash-thinking-exp'):
-    relevant_info_template = """
-Intent: The CRUD operation (create, read, update, delete)
-Transaction Type: e.g. Purchase, Sale, Inventory
-Details: A list of key detail fields extracted.
-"""
+For multiple transactions repeat this pattern exactly."""
+    
     genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
-    model_instance = genai.GenerativeModel(
-        model_name=model,
+    model = genai.GenerativeModel(
+        'gemini-2.0-flash-thinking-exp',
         generation_config={
             "temperature": 0.1,
             "top_p": 0.01,
@@ -57,75 +42,53 @@ Details: A list of key detail fields extracted.
         }
     )
     
-    system_prompt = (
-        f"You are a transaction classification assistant. Format your output with: "
-        f"{relevant_info_template}"
-    )
-    
     try:
-        response = model_instance.generate_content([
-            {"role": "user", "parts": [system_prompt, prompt]}
-        ])
+        response = model.generate_content([system_prompt, prompt])
         return response.text
     except Exception as e:
-        logger.error(f'Error generating response: {str(e)}')
+        logger.error(f"Generation error: {str(e)}")
         return None
 
 def parse_multiple_transactions(response_text):
+    transactions = []
+    current_trans = {}
+    
     try:
+        # First try JSON parsing
         parsed = json.loads(response_text)
         if isinstance(parsed, dict):
-            parsed["created_at"] = datetime.now().isoformat()
-            return [parsed]
-        elif isinstance(parsed, list):
-            for item in parsed:
-                item["created_at"] = datetime.now().isoformat()
-            return parsed
-    except Exception:
-        data = {}
-        lines = response_text.splitlines()
-        for line in lines:
-            if "*Intent*:" in line:
-                data["intent"] = line.split(":", 1)[1].strip()
-            elif "*Transaction Type*:" in line:
-                data["transaction_type"] = line.split(":", 1)[1].strip()
-            elif "*Details*:" in line:
-                details_line = line.split(":", 1)[1].strip()
-                details = {}
-                parts = details_line.split(",-")
-                for part in parts:
-                    part = part.strip(" -")
-                    if ":" in part:
-                        key, value = part.split(":", 1)
-                        details[key.strip().lower()] = value.strip().strip(",")
-                data["details"] = details
-        if data:
-            data["created_at"] = datetime.now().isoformat()
-            return [data]
-    return []
+            return [add_timestamp(parsed)]
+        return [add_timestamp(t) for t in parsed]
+    except json.JSONDecodeError:
+        pass  # Proceed to text parsing
 
-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()]
-    inventory_df = pd.DataFrame(inventory_list)
-    sales_df = pd.DataFrame(sales_list)
-    from pandasai import SmartDatalake
-    lake = SmartDatalake(
-        [inventory_df, sales_df],
-        config={
-            "llm": llm,
-            "custom_whitelisted_dependencies": ["ast"],
-            "response_parser": FlaskResponse,
-            "enable_cache": False,
-            "save_logs": False
-        }
-    )
-    response = lake.chat(user_question)
-    if isinstance(response, dict):
-        return json.dumps(response)
-    return str(response)
+    lines = [line.strip() for line in response_text.split('\n') if line.strip()]
+    
+    for line in lines:
+        if line.startswith('*Intent*:'):
+            if current_trans:
+                transactions.append(add_timestamp(current_trans))
+                current_trans = {}
+            current_trans['intent'] = line.split(':', 1)[1].strip().lower()
+        elif line.startswith('*Transaction Type*:'):
+            current_trans['transaction_type'] = line.split(':', 1)[1].strip().lower()
+        elif line.startswith('*Details*:'):
+            current_trans['details'] = {}
+        elif line.startswith('-') and 'details' in current_trans:
+            key_value = line[1:].strip().split(':', 1)
+            if len(key_value) == 2:
+                key = key_value[0].strip().lower()
+                value = key_value[1].strip()
+                current_trans['details'][key] = value
+    
+    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 create_inventory(user_phone, transaction_data):
     for transaction in transaction_data:
@@ -167,132 +130,63 @@ def create_sale(user_phone, transaction_data):
         sale_ref.set(transaction)
     return True
 
-def update_transaction(user_phone, transaction_id, update_data, transaction_type="inventory"):
+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()]
+    
+    if not inventory_list and not sales_list:
+        return "No data found"
+    
+    inventory_df = pd.DataFrame(inventory_list)
+    sales_df = pd.DataFrame(sales_list)
+    
+    from pandasai import SmartDatalake
+    lake = SmartDatalake(
+        [inventory_df, sales_df],
+        config={"llm": llm}
+    )
+    
     try:
-        doc_ref = db.collection("users").document(user_phone).collection(transaction_type).document(transaction_id)
-        if not doc_ref.get().exists:
-            return False
-            
-        updates = {}
-        for key, value in update_data.items():
-            updates[f'details.{key}'] = value
-            
-        doc_ref.update(updates)
-        return True
+        return str(lake.chat(user_question))
     except Exception as e:
-        logger.error(f"Update failed: {str(e)}")
-        return False
-
-def fetch_transaction(user_phone, transaction_id=None, transaction_type="inventory"):
-    if transaction_id:
-        doc_ref = db.collection("users").document(user_phone).collection(transaction_type).document(transaction_id)
-        transaction = doc_ref.get()
-        if transaction.exists:
-            return transaction.to_dict()
-        return None
-    else:
-        collection_ref = db.collection("users").document(user_phone).collection(transaction_type)
-        return [doc.to_dict() for doc in collection_ref.stream()]
-
-def delete_transaction(user_phone, transaction_data):
-    transaction_id = transaction_data[0]['details'].get('item')
-    transaction_type = transaction_data[0].get('transaction_type', 'inventory').lower()
-    doc_ref = db.collection("users").document(user_phone).collection(transaction_type).document(transaction_id)
-    if doc_ref.get().exists:
-        doc_ref.delete()
-        return True
-    return False
+        logger.error(f"Datalake error: {str(e)}")
+        return "Error generating response"
 
 def persist_temporary_transaction(transactions, mobile):
-    temp_ref = db.collection("users").document(mobile).collection("temp_transactions").document('pending-user-action')
-    data = {
-        "transactions": transactions,
-        "status": "pending",
-        "created_at": datetime.now().isoformat()
-    }
-    temp_ref.set(data)
-    return True
+    try:
+        db.collection("users").document(mobile).collection("temp_transactions").document("pending").set({
+            "transactions": transactions,
+            "timestamp": datetime.now().isoformat()
+        })
+        return True
+    except Exception as e:
+        logger.error(f"Temp save error: {str(e)}")
+        return False
 
 def format_transaction_response(transactions):
     if not transactions:
-        return "No transaction data found."
+        return "No transaction data"
     
     if isinstance(transactions, dict):
         transactions = [transactions]
-    elif not isinstance(transactions, list):
-        return str(transactions)
     
-    result = []
-    for idx, trans in enumerate(transactions, 1):
-        if not isinstance(trans, dict):
-            continue
-            
-        intent = trans.get('intent', 'Unknown')
-        trans_type = trans.get('transaction_type', 'Unknown')
-        details = trans.get('details', {})
-        
-        if len(transactions) > 1:
-            result.append(f"Transaction {idx}:")
-        
-        result.append(f"Intent: {intent}")
-        result.append(f"Type: {trans_type}")
-        
-        if details:
-            result.append("Details:")
-            for key, value in details.items():
-                if key == 'currency':
-                    continue
-                if 'price' in key or 'cost' in key or 'amount' in key:
-                    currency = details.get('currency', '$')
-                    result.append(f"  - {key.replace('_', ' ').title()}: {currency}{value}")
-                else:
-                    result.append(f"  - {key.replace('_', ' ').title()}: {value}")
-        
-        result.append("")
+    output = []
+    for trans in transactions:
+        output.append(f"Intent: {trans.get('intent', 'unknown')}")
+        output.append(f"Type: {trans.get('transaction_type', 'unknown')}")
+        if 'details' in trans:
+            output.append("Details:")
+            for k, v in trans['details'].items():
+                output.append(f"- {k}: {v}")
+        output.append("")
     
-    return "\n".join(result)
+    return "\n".join(output)
 
-def process_intent(parsed_trans_data, mobile):
-    intent = parsed_trans_data[0]['intent'].lower()
-    trans_type = parsed_trans_data[0].get('transaction_type', '').lower()
-    transaction_summary = format_transaction_response(parsed_trans_data)
-
-    if intent == 'create':
-        if trans_type in ('purchase', 'inventory'):
-            if create_inventory(mobile, parsed_trans_data):
-                return f"Transaction recorded successfully!\n\n{transaction_summary}"
-            return f"Failed to record transaction!\n\n{transaction_summary}"
-        elif trans_type == 'sale':
-            if create_sale(mobile, parsed_trans_data):
-                return f"Sale recorded successfully!\n\n{transaction_summary}"
-            return f"Insufficient stock or invalid item!\n\n{transaction_summary}"
-        return f"Unsupported transaction type: {trans_type}\n\n{transaction_summary}"
-    
-    elif intent == 'update':
-        item = parsed_trans_data[0]['details'].get('item')
-        update_details = {k:v for k,v in parsed_trans_data[0]['details'].items() if k != 'item'}
-        
-        if update_transaction(mobile, item, update_details, trans_type):
-            return f"Updated {item} successfully!\n\n{transaction_summary}"
-        return f"Failed to update {item}!\n\n{transaction_summary}"
-    
-    elif intent == 'delete':
-        item = parsed_trans_data[0]['details'].get('item', 'item')
-        if delete_transaction(mobile, parsed_trans_data):
-            return f"Deleted {item} successfully!\n\n{transaction_summary}"
-        return f"Failed to delete {item}!\n\n{transaction_summary}"
-    
-    elif intent == 'read':
-        if 'item' in parsed_trans_data[0]['details']:
-            item = parsed_trans_data[0]['details']['item']
-            transaction = fetch_transaction(mobile, item)
-            if transaction:
-                return f"Details for {item}:\n\n{format_transaction_response(transaction)}"
-            return f"Item {item} not found!"
-        else:
-            transactions = fetch_transaction(mobile, transaction_type=trans_type if trans_type else "inventory")
-            if transactions:
-                return f"All {trans_type if trans_type else 'inventory'} items:\n\n{format_transaction_response(transactions)}"
-            return "No items found!"
-    
-    return f"Unsupported intent: {intent}\n\n{transaction_summary}"
+def fetch_transaction(user_phone, item_name=None):
+    if item_name:
+        doc = db.collection("users").document(user_phone).collection("inventory").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()]