Onboarding

.gitignore

@@ -43,6 +43,8 @@ next-env.d.ts
 /api/test_data
 api/__pycache__/
 api/utils/__pycache__/
+api/agent/__pycache__/
+api/celery/__pycache__/
 
 # log
 supabase_queries.log

.vscode/extensions.json

@@ -0,0 +1,3 @@
+{
+  "recommendations": ["denoland.vscode-deno"]
+}

.vscode/settings.json

@@ -0,0 +1,24 @@
+{
+  "deno.enablePaths": [
+    "supabase/functions"
+  ],
+  "deno.lint": true,
+  "deno.unstable": [
+    "bare-node-builtins",
+    "byonm",
+    "sloppy-imports",
+    "unsafe-proto",
+    "webgpu",
+    "broadcast-channel",
+    "worker-options",
+    "cron",
+    "kv",
+    "ffi",
+    "fs",
+    "http",
+    "net"
+  ],
+  "[typescript]": {
+    "editor.defaultFormatter": "denoland.vscode-deno"
+  }
+}

api/agent/chat.py

@@ -53,81 +53,66 @@ def log_supabase_query(operation: str, details: dict):
         log_entry += f"{'='*50}\n"
         logging.info(log_entry)
 
+def get_onboarding_stage(user: dict, supabase: Client):
+    # Fetch account count
+    accounts = supabase.table("Account").select("count", count="exact").filter(column="owner_id", operator="eq", criteria=user.id).execute()
+    account_check = accounts.count > 0
+
+    # Fetch transaction count
+    transactions = supabase.table("Transaction").select("count", count="exact").filter(column="owner_id", operator="eq", criteria=user.id).execute()
+    transaction_check = transactions.count > 0
+
+    # Fetch bucket count
+    buckets = supabase.table("TransactionBucket").select("count", count="exact").filter(column="owner_id", operator="eq", criteria=user.id).execute()
+    bucket_check = buckets.count > 0
+
+    # Fetch rule count
+    rules = supabase.table("Rule").select("count", count="exact").filter(column="owner_id", operator="eq", criteria=user.id).execute()
+    rule_check = rules.count > 0
+
+    onboarding_stage = 1
+    if account_check:
+        onboarding_stage = 2
+    if account_check and transaction_check:
+        onboarding_stage = 3
+    if account_check and transaction_check and bucket_check:
+        onboarding_stage = 4
+    if account_check and transaction_check and bucket_check and rule_check:
+        onboarding_stage = 5
+
+    return onboarding_stage
+
 # Constants and configurations
 OPENAI_MODEL = "o3-mini"
 GEMMA_MODEL = "gemma-3-27b-it"
 GROQ_MODEL = "qwen-qwq-32b"
 
 TEMPERATURE = 0
 
-async def process_chat_message_stream(message: str, user: dict, supabase: Client, buckets: List[Dict], account_types: List[Dict], accounts: List[Dict], chat_history: List[Dict]):
+async def process_chat_message_stream(message: str, user: dict, supabase: Client, buckets: List[Dict], account_types: List[Dict], accounts: List[Dict], chat_history: List[Dict], onboarding: bool = False):
     """
     Streaming version of process_chat_message that uses tool calling to execute SQL queries.
     The function follows this flow:
     1. Stream back text from the OpenAI response
     2. If a tool call is made (SQL query), collect all parameters
     3. Execute the SQL query with the parameters
     4. Return the visualization data
+    5. If onboarding and stage changes, generate a new response acknowledging the completion
     """
-    collected_text = ""
-    visualization_data = None
-    collected_tool_calls = {}
-    current_tool_call_id = None
 
-    try:
-        start_time = time.time()
+    async def generate_and_stream_response(messages, current_chat_history=None, is_stage_transition=False):
+        """Helper function to generate and stream AI responses"""
+        nonlocal chat_history
 
-        # Initialize OpenAI client
-        client = openai.OpenAI(
-            api_key=openai_api_key,
-            timeout=60.0  # 60 second timeout for all requests
-        )
-
-        # client = genai.Client(api_key=gemini_api_key)
-
-        # Send an initial event to indicate processing has started
-        start_data = json.dumps({
-            "type": "start",
-        })
-        yield f"data: {start_data}\n\n"
-
-        # Build initial conversation messages
-        messages = [
-            {
-                "role": "system",
-                "content": [
-                    {"type": "text", "text": get_chat_system_prompt(datetime.now(), buckets, account_types, accounts)}
-                ]
-            }
-        ]
-
-        # Append stored chat history if any
-        for stored_message in chat_history[-10:]:  # Only use last 10 messages for context
-            content = [{"type": "text", "text": stored_message.get("text", "")}]
+        if current_chat_history is not None:
+            chat_history = current_chat_history
 
-            # If there's visualization data, add it as a JSON string
-            if stored_message.get("visualization"):
-                content.append({
-                    "type": "text",
-                    "text": json_dumps(stored_message["visualization"])
-                })
-
-            formatted_message = {
-                "role": stored_message.get("role"),
-                "content": content
-            }
-            messages.append(formatted_message)
-
-        # Append the current user message
-        messages.append({
-            "role": "user",
-            "content": [
-                {"type": "text", "text": message}
-            ]
-        })
+        collected_text = ""
+        visualization_data = None
+        collected_tool_calls = {}
+        current_tool_call_id = None
 
         # Call OpenAI with streaming enabled and tool calling
-
         stream = client.chat.completions.create(
             model=OPENAI_MODEL,
             messages=messages,
@@ -136,16 +121,6 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
             stream=True,
             timeout=90.0
         )
-
-        # stream = client.models.generate_content_stream(
-        #     model='gemini-2.0-flash-001',
-        #     config=types.GenerateContentConfig(
-        #         system_instruction=get_chat_system_prompt(datetime.now(), buckets, account_types, accounts)
-        #     ),
-        #     contents=messages,
-        # )
-
-        # print(response)
 
         # Process the streaming response
         for chunk in stream:
@@ -186,7 +161,6 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
 
         # Process collected tool calls
         if collected_tool_calls:
-
             # Process each tool call
             for tool_call_id, tool_call in collected_tool_calls.items():
                 # All tool calls are processed through the function registry
@@ -198,13 +172,12 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
                         # Parse arguments
                         arguments = json.loads(tool_call["function"]["arguments"])
 
-                        function_start = time.time()
-
                         # Execute the function
                         result = await handler(arguments, user, supabase)
 
                         # Process result based on type
                         if result.get("type") == "text":
+                            collected_text += result.get("content", "")
                             # Send text response
                             text_event = json.dumps({
                                 "type": "text",
@@ -245,7 +218,10 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
         # Check if a chat history entry already exists for this user
         try:
             # Fetch the existing full chat history from the database
-            response = supabase.table("ChatHistory").select("*").eq("owner_id", user.id).maybe_single().execute()
+            if onboarding:
+                response = supabase.table("Onboarding").select("*").eq("owner_id", user.id).maybe_single().execute()
+            else:
+                response = supabase.table("ChatHistory").select("*").eq("owner_id", user.id).maybe_single().execute()
             if response.data:
                 full_history = response.data.get("chat_history", [])
             else:
@@ -255,12 +231,123 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
             full_history.append(assistant_message)
 
             # Update the chat history with the merged full history
-            supabase.table("ChatHistory").update({
-                "chat_history": full_history
-            }).eq("owner_id", user.id).execute()
+            if onboarding:
+                supabase.table("Onboarding").update({
+                    "chat_history": full_history
+                }).eq("owner_id", user.id).execute()
+            else:
+                supabase.table("ChatHistory").update({
+                    "chat_history": full_history
+                }).eq("owner_id", user.id).execute()
         except Exception as e:
             print(f"Error saving assistant message to chat history: {str(e)}")
 
+    try:
+        start_time = time.time()
+
+        # Initialize OpenAI client
+        client = openai.OpenAI(
+            api_key=openai_api_key,
+            timeout=60.0  # 60 second timeout for all requests
+        )
+
+        # Send an initial event to indicate processing has started
+        start_data = json.dumps({
+            "type": "start",
+        })
+        yield f"data: {start_data}\n\n"
+
+        if onboarding:
+            initial_stage = get_onboarding_stage(user, supabase)
+        else:
+            initial_stage = -1    
+
+        # Build initial conversation messages
+        messages = [
+            {
+                "role": "system",
+                "content": [
+                    {"type": "text", "text": get_chat_system_prompt(datetime.now(), buckets, account_types, accounts, initial_stage)}
+                ]
+            }
+        ]
+
+        # Append stored chat history if any
+        for stored_message in chat_history[-10:]:  # Only use last 10 messages for context
+            content = [{"type": "text", "text": stored_message.get("text", "")}]
+
+            # If there's visualization data, add it as a JSON string
+            if stored_message.get("visualization"):
+                content.append({
+                    "type": "text",
+                    "text": json_dumps(stored_message["visualization"])
+                })
+
+            formatted_message = {
+                "role": stored_message.get("role"),
+                "content": content
+            }
+            messages.append(formatted_message)
+
+        # Append the current user message
+        messages.append({
+            "role": "user",
+            "content": [
+                {"type": "text", "text": message}
+            ]
+        })
+
+        # Generate and stream the initial response
+        async for chunk in generate_and_stream_response(messages):
+            yield chunk
+
+        # Check if onboarding stage has changed and generate a new response if it has
+        if onboarding:
+            new_stage = get_onboarding_stage(user, supabase)
+            print(f"\n\n\n\n\nNew stage: {new_stage}\n\n\n\n\n")
+            print(f"\n\n\n\n\nInitial stage: {initial_stage}\n\n\n\n\n")
+            if new_stage != initial_stage:                
+                # Update chat history in database
+                if onboarding:
+                    supabase.table("Onboarding").update({
+                        "chat_history": chat_history
+                    }).eq("owner_id", user.id).execute()
+
+                # Create a new system message with updated onboarding stage
+                system_message = {
+                    "role": "system",
+                    "content": [{"type": "text", "text": get_chat_system_prompt(datetime.now(), buckets, account_types, accounts, new_stage)}]
+                }
+
+                # Create a new messages array with the updated system message and recent chat history
+                new_messages = [system_message]
+
+                # Add the most recent messages from chat history (excluding the transition message)
+                for stored_message in chat_history[-10:-1]:  # Last 10 messages excluding the transition
+                    content = [{"type": "text", "text": stored_message.get("text", "")}]
+                    if stored_message.get("visualization"):
+                        content.append({
+                            "type": "text",
+                            "text": json_dumps(stored_message["visualization"])
+                        })
+
+                    formatted_message = {
+                        "role": stored_message.get("role"),
+                        "content": content
+                    }
+                    new_messages.append(formatted_message)
+
+                # Add a special system message to acknowledge stage completion
+                stage_completion_message = {
+                    "role": "system",
+                    "content": [{"type": "text", "text": f"The user has just completed stage {initial_stage} and is now at stage {new_stage}. Acknowledge this achievement and guide them on what to do next."}]
+                }
+                new_messages.append(stage_completion_message)
+
+                # Generate and stream the stage transition response
+                async for chunk in generate_and_stream_response(new_messages, chat_history, True):
+                    yield chunk
+
     except Exception as e:
         error_message = str(e)
         print(f"Error processing chat: {error_message}")
@@ -282,10 +369,14 @@ async def process_chat_message_stream(message: str, user: dict, supabase: Client
         # Try to save the error message to chat history
         try:
             chat_history.append(assistant_message)
-
-            supabase.table("ChatHistory").update({
-                "chat_history": chat_history
-            }).eq("owner_id", user.id).execute()
+            if onboarding:
+                supabase.table("Onboarding").update({
+                    "chat_history": chat_history
+                }).eq("owner_id", user.id).execute()
+            else:
+                supabase.table("ChatHistory").update({
+                    "chat_history": chat_history
+                }).eq("owner_id", user.id).execute()
         except:
             pass  # Silently ignore errors in the error handler
 

api/agent/chat_tools.py

@@ -478,7 +478,7 @@ async def create_account_handler(arguments: Dict[str, Any], user: Dict, supabase
 
         return {
             "type": "text",
-            "content": f"Account '{account_data['name']}' created successfully!",
+            "content": f"Account '{account_data['name']}' created successfully! ",
             "account": result["account"]
         }