Added documentation and demo3 work together

LexAI_iOS/LexAI_iOS/Managers/FirebaseManager.swift

@@ -19,6 +19,8 @@ class FirebaseManager: ObservableObject {
     @Published var errorMessage: String?
     @Published var isLoading = false
 
+    /// Listens for Firebase auth changes and updates published state.
+    /// Removed in `deinit` to avoid stale callbacks.
     private var authStateListener: AuthStateDidChangeListenerHandle?
     private lazy var db = Firestore.firestore()
     private let isPreview: Bool
@@ -35,6 +37,8 @@ class FirebaseManager: ObservableObject {
                 self?.isAuthenticated = (user?.isAnonymous == false)
             }
         }
+        // Keeps a valid auth token available for callable functions
+        // before the user performs explicit account auth.
         if Auth.auth().currentUser == nil {
             Task { @MainActor in
                 do {
@@ -223,6 +227,10 @@ class FirebaseManager: ObservableObject {
 
     // MARK: - Auth
 
+    /// Creates a new Firebase Auth account with email/password.
+    /// On success it sets `user` and `isAuthenticated`; on failure it maps errors.
+    /// - Parameter email: Email for the new account.
+    /// - Parameter password: Password for the new account.
     func signUp(email: String, password: String) async {
         isLoading = true
         errorMessage = nil
@@ -237,6 +245,10 @@ class FirebaseManager: ObservableObject {
     }
 
     @MainActor
+    /// Signs in an existing Firebase user with email/password.
+    /// Updates `user` and `isAuthenticated`; maps failures to friendly messages.
+    /// - Parameter email: Account email.
+    /// - Parameter password: Account password.
     func signIn(email: String, password: String) async {
         isLoading = true
         errorMessage = nil
@@ -251,6 +263,8 @@ class FirebaseManager: ObservableObject {
     }
 
     @MainActor
+    /// Signs out the current user and clears auth state.
+    /// Setting `user = nil` and `isAuthenticated = false` routes back to `AuthView`.
     func signOut() {
         do {
             try Auth.auth().signOut()
@@ -261,6 +275,10 @@ class FirebaseManager: ObservableObject {
         }
     }
 
+    /// Converts Firebase Auth errors into user-friendly UI text.
+    /// Maps common `AuthErrorCode` values and falls back to localized descriptions.
+    /// - Parameter error: Firebase auth error to convert.
+    /// - Returns: Message safe to show in UI.
     private func mapFirebaseError(_ error: Error) -> String {
         let nsError = error as NSError
         guard let errorCode = AuthErrorCode(rawValue: nsError.code) else {

LexAI_iOS/LexAI_iOS/Views/ChatView.swift

@@ -22,6 +22,8 @@ struct ChatView: View {
     @State private var inputText: String = ""
     @State private var showScanDocuments = false
     @State private var showFilePicker = false
+    // Tracks when a model response is in progress.
+    // Used for ThinkingBubbleView visibility and send-button disabling.
     @State private var isAwaitingReply = false
     @EnvironmentObject var firebaseManager: FirebaseManager
     @Environment(\.scenePhase) private var scenePhase
@@ -101,6 +103,8 @@ struct ChatView: View {
                     }
 
                     if isAwaitingReply {
+                        // Shows typing animation while waiting for LexAI.
+                        // Fades/scales out when the response is appended.
                         ThinkingBubbleView()
                             .transition(.opacity.combined(with: .scale(scale: 0.8)))
                     }

LexAI_iOS/LexAI_iOS/Views/Components/ThinkingBubbleView.swift

@@ -6,7 +6,10 @@
 //
 import SwiftUI
 
+/// Shows a 3-dot typing indicator while LexAI is generating a reply.
+/// Uses staggered `.easeInOut` animations on three circles that repeat indefinitely.
 struct ThinkingBubbleView: View {
+    /// Toggles the pulsing dots and is set to `true` in `onAppear`.
     @State private var animating = false
 
     var body: some View {

lexai-functions/main.py

@@ -1,11 +1,7 @@
-"""Firebase callable entrypoint for LexAI chat.
+"""Firebase Cloud Function backend for LexAI.
 
-Pipeline: optional translation of user + history into English, embed the English prompt,
-retrieve Michigan legislation chunks from Pinecone, call the English-only RunPod legal
-model, then translate the assistant reply back to the client's UI language when needed.
-
-Pinecone is initialized lazily so deploy-time import/discovery does not require
-``PINECONE_API_KEY`` until a request actually runs.
+Implements the RAG pipeline: retrieve legal context from Pinecone and
+generate answers through the RunPod-hosted model.
 """
 
 import os
@@ -73,6 +69,13 @@ def _get_pinecone() -> tuple[Pinecone, Any]:
 
 
 def embed_query(query_text: str) -> list:
+    """Convert user text into an embedding vector.
+
+    Args:
+        query_text: Raw query text.
+    Returns:
+        List of float values for Pinecone search.
+    """
     pc, _ = _get_pinecone()
     embeddings_response = pc.inference.embed(
         model=EMBED_MODEL,
@@ -125,6 +128,14 @@ def _extract_metadata(match: Any) -> Any:
 
 
 def query_pinecone(query_embedding: list, top_k: int = 5) -> list:
+    """Fetch top legislation chunks from Pinecone.
+
+    Args:
+        query_embedding: Embedding vector for lookup.
+        top_k: Number of matches to retrieve.
+    Returns:
+        List of chunk text strings from top metadata matches.
+    """
     _, index = _get_pinecone()
     # Prefer chunk vectors (embed_and_store sets chunk_idx). If index has no chunk_idx, fall back unfiltered.
     for use_chunk_filter in (True, False):
@@ -155,7 +166,13 @@ def query_pinecone(query_embedding: list, top_k: int = 5) -> list:
 
 
 def call_runpod(messages: list) -> str:
-    """POST to RunPod serverless ``/run``, then poll ``/status`` until COMPLETED or timeout."""
+    """Call RunPod chat completion and return generated text.
+
+    Args:
+        messages: OpenAI-style role/content messages.
+    Returns:
+        Generated text or an error string after polling for completion.
+    """
     endpoint_id = (os.environ.get("RUNPOD_ENDPOINT_ID") or "").strip()
     runpod_key = (os.environ.get("RUNPOD_API_KEY") or "").strip()
     if not endpoint_id:
@@ -173,6 +190,8 @@ def call_runpod(messages: list) -> str:
             "input": {
                 "openai_route": "/v1/chat/completions",
                 "openai_input": {
+                    # Fine-tuned LLaMA 3 8B via QLoRA (4-bit NF4 + LoRA adapters)
+                    # on Michigan legal Q&A + IDK examples, then merged for RunPod.
                     "model": "hbalkhafaji/llama3-8b-legal-merged",
                     "messages": messages,
                     "max_tokens": 1024,
@@ -222,9 +241,12 @@ def call_runpod(messages: list) -> str:
     secrets=[OPENAI_API_KEY, PINECONE_API_KEY, RUNPOD_API_KEY],
 )
 def chat(req: https_fn.CallableRequest) -> dict:
-    """HTTPS callable: ``req.data`` may include ``prompt``, ``chat_history``, ``language`` (default ``en``).
+    """Handle Firebase callable chat requests.
 
-    Returns ``{"response": str}`` on success or ``{"error": str}`` on validation/runtime failure.
+    Args:
+        req: Request with `prompt` and optional `chat_history`/`language` in `req.data`.
+    Returns:
+        `{"response": text}` on success or `{"error": message}` on failure.
     """
     try:
         prompt = req.data.get("prompt", "")
@@ -244,6 +266,8 @@ def chat(req: https_fn.CallableRequest) -> dict:
 
         context = "\n\n---\n\n".join(relevant_chunks)
 
+        # This prompt line reinforces fine-tuned IDK/unlearning behavior so the
+        # model declines out-of-scope questions when excerpts are insufficient.
         system_message = {
             "role": "system",
             "content": (