Update app.py

app.py

@@ -1,117 +1,119 @@
-import streamlit as st
-import requests
-
-# ----------------- Streamlit Setup ----------------- #
-st.set_page_config(page_title="Lease Agreement Extractor (OCR + AI)", layout="centered")
-st.title("📄 Lease Agreement Extractor (OCR + AI)")
-
-# ----------------- Assistant & API Config ----------------- #
-OPENAI_ASSISTANT_ID = "asst_xBnNfiyWmVa4iF3CgXwJnmBt"  # Replace with your assistant ID
-OPENAI_API_KEY = "YOUR_OPENAI_API_KEY"  # Replace with your OpenAI key
-
-HEADERS = {
-    "Authorization": f"Bearer {OPENAI_API_KEY}",
-    "Content-Type": "application/json"
-}
-
-# ----------------- Upload Interface ----------------- #
-uploaded_file = st.file_uploader("Upload a lease agreement (PDF, PNG, JPG)", type=["pdf", "png", "jpg", "jpeg"])
-
-def upload_file_to_openai(file_bytes, filename):
-    files = {
-        "file": (filename, file_bytes, "application/octet-stream")
-    }
-    response = requests.post("https://api.openai.com/v1/files", headers={"Authorization": f"Bearer {OPENAI_API_KEY}"}, files=files)
-    return response.json().get("id")
-
-def create_thread_and_run(file_id):
-    # Create thread
-    thread_res = requests.post("https://api.openai.com/v1/threads", headers=HEADERS, json={})
-    thread_data = thread_res.json()
-    thread_id = thread_data.get("id")
-
-    # Run assistant with structured lease extraction prompt
-    run_payload = {
-        "assistant_id": OPENAI_ASSISTANT_ID,
-        "instructions": """
-You are an AI assistant designed to extract structured lease data from any uploaded agreement, even if it is scanned. Your responsibilities are:
-
-- Accept PDF or image-based lease documents.
-- Use OCR (Optical Character Recognition) if the document is non-selectable or scanned.
-- Parse and extract key structured data fields from lease agreements.
-- Output results in JSON format following a fixed schema.
-
-OCR & Extraction Rules:
-1. If text cannot be extracted normally, perform OCR using image-to-text on each page.
-2. Preserve paragraph structure, detect headers, signature blocks, tables, and dates.
-
-Expected JSON Output:
-{
-  "document_title": "Lease Agreement Title or File Name",
-  "parties": {
-    "lessor": "Name of Lessor",
-    "lessee": "Name of Lessee"
-  },
-  "property_description": "Detailed description of leased property",
-  "term": {
-    "start_date": "YYYY-MM-DD",
-    "end_date": "YYYY-MM-DD",
-    "renewal_options": "Yes/No or Clause Text"
-  },
-  "financials": {
-    "monthly_rent": "Amount",
-    "deposit": "Amount",
-    "escalation_clause": "Text if present"
-  },
-  "obligations": {
-    "maintenance": "Lessor/Lessee",
-    "insurance": "Lessor/Lessee",
-    "subletting": "Allowed/Not allowed"
-  },
-  "signatures": [
-    {
-      "party": "Party Name",
-      "signed_by": "Full Name",
-      "date": "YYYY-MM-DD",
-      "position": "Title if listed"
-    }
-  ],
-  "raw_text_per_page": {
-    "page_1": "Full OCR text of page 1",
-    "page_2": "Full OCR text of page 2"
-  }
-}
-
-📌 Behaviors:
-- Automatically detect if OCR is needed.
-- Normalize dates to YYYY-MM-DD.
-- Use "Not Found" or null where info is missing.
-- Return JSON only.
-- If OCR text is low-confidence, include: "low_confidence": true
-""",
-        "file_ids": [file_id]
-    }
-
-    run_res = requests.post(f"https://api.openai.com/v1/threads/{thread_id}/runs", headers=HEADERS, json=run_payload)
-    run_data = run_res.json()
-    run_id = run_data.get("id")
-
-    return thread_id, run_id
-
-# ----------------- Main UI Flow ----------------- #
-if uploaded_file is not None:
-    st.success("📄 File uploaded successfully.")
-
-    if st.button("🚀 Run Lease Extraction"):
-        with st.spinner("Uploading and invoking assistant..."):
-            file_id = upload_file_to_openai(uploaded_file.getvalue(), uploaded_file.name)
-            if file_id:
-                thread_id, run_id = create_thread_and_run(file_id)
-                st.success("✅ Assistant run started!")
-                st.code(f"Thread ID: {thread_id}", language="text")
-                st.code(f"Run ID: {run_id}", language="text")
-                st.markdown("📌 You can retrieve results via OpenAI API using these IDs.")
-            else:
-                st.error("❌ File upload to OpenAI failed.")
-else:
-    st.info("Please upload a lease agreement file to begin.")
+import streamlit as st 
+from openai import OpenAI
+import time
+import os
+import uuid
+import firebase_admin
+from firebase_admin import credentials, firestore
+
+# 🔐 Firebase setup
+if not firebase_admin._apps:
+    cred = credentials.Certificate("firebase-service-account.json")
+    firebase_admin.initialize_app(cred)
+
+db = firestore.client()
+
+# 🔐 OpenAI setup
+openai_key = os.getenv("openai_key")
+assistant_id = os.getenv("ASSISTANT_ID")
+client = OpenAI(api_key=openai_key)
+
+# 🌐 Streamlit Config
+st.set_page_config(page_title="LeaseTracker AI Assistant", layout="wide")
+
+# 🎯 Session + User ID
+if "user_id" not in st.session_state:
+    st.session_state["user_id"] = str(uuid.uuid4())
+user_id = st.session_state["user_id"]
+
+# 🖼️ LeaseTracker Branding + Styling
+st.markdown("""
+    <style>
+    .block-container {padding-top: 1rem; padding-bottom: 0rem;}
+    header {visibility: hidden;}
+    .stChatMessage { max-width: 85%; border-radius: 12px; padding: 8px; margin-bottom: 10px; }
+    .stChatMessage[data-testid="stChatMessage-user"] { background: #f0f0f0; color: #000000; }
+    .stChatMessage[data-testid="stChatMessage-assistant"] { background: #e3f2fd; color: #000000; }
+    .lt-logo { vertical-align: middle; }
+    </style>
+""", unsafe_allow_html=True)
+
+st.markdown("""
+    <div style='text-align: center; margin-top: 20px; margin-bottom: -10px;'>
+        <span style='display: inline-flex; align-items: center; gap: 8px;'>
+            <img src='https://leasetracker.co.za/wp-content/uploads/2020/10/Lease-Tracker-Logo.png' width='100' class='lt-logo'/>
+            <span style='font-size: 12px; color: gray;'>Powered by LeaseTracker</span>
+        </span>
+    </div>
+""", unsafe_allow_html=True)
+
+# 🔁 Get or create a thread ID
+def get_or_create_thread_id():
+    doc_ref = db.collection("users").document(user_id)
+    doc = doc_ref.get()
+    if doc.exists:
+        return doc.to_dict()["thread_id"]
+    else:
+        thread = client.beta.threads.create()
+        doc_ref.set({"thread_id": thread.id, "created_at": firestore.SERVER_TIMESTAMP})
+        return thread.id
+
+# 💾 Save a message
+def save_message(role, content):
+    db.collection("users").document(user_id).collection("messages").add({
+        "role": role,
+        "content": content,
+        "timestamp": firestore.SERVER_TIMESTAMP
+    })
+
+# 💬 Display chat history
+def display_chat_history():
+    messages = db.collection("users").document(user_id).collection("messages").order_by("timestamp").stream()
+    assistant_icon_html = "<img src='https://leasetracker.co.za/wp-content/uploads/2020/10/Lease-Tracker-Logo.png' width='20' style='vertical-align:middle;'/>"
+    for msg in list(messages)[::-1]:
+        data = msg.to_dict()
+        if data["role"] == "user":
+            st.markdown(f"<div class='stChatMessage' data-testid='stChatMessage-user'>👤 <strong>You:</strong> {data['content']}</div>", unsafe_allow_html=True)
+        else:
+            st.markdown(f"<div class='stChatMessage' data-testid='stChatMessage-assistant'>{assistant_icon_html} <strong>LeaseTracker Assistant:</strong> {data['content']}</div>", unsafe_allow_html=True)
+
+# 🚀 Main Chat UI
+input_col, clear_col = st.columns([9, 1])
+with input_col:
+    user_input = st.chat_input("Type your message here...")
+
+with clear_col:
+    if st.button("🗑️", key="clear-chat", help="Clear Chat"):
+        try:
+            user_doc_ref = db.collection("users").document(user_id)
+            for msg in user_doc_ref.collection("messages").stream():
+                msg.reference.delete()
+            user_doc_ref.delete()
+            st.session_state.clear()
+            st.rerun()
+        except Exception as e:
+            st.error(f"Failed to clear chat: {e}")
+
+thread_id = get_or_create_thread_id()
+display_chat_history()
+
+if user_input:
+    # Send user message to OpenAI thread
+    client.beta.threads.messages.create(thread_id=thread_id, role="user", content=user_input)
+    save_message("user", user_input)
+
+    with st.spinner("Thinking and typing... 💭"):
+        run = client.beta.threads.runs.create(thread_id=thread_id, assistant_id=assistant_id)
+        while True:
+            run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
+            if run_status.status == "completed":
+                break
+            time.sleep(1)
+
+    messages_response = client.beta.threads.messages.list(thread_id=thread_id)
+    latest_response = sorted(messages_response.data, key=lambda x: x.created_at)[-1]
+    assistant_message = latest_response.content[0].text.value
+
+    save_message("assistant", assistant_message)
+    time.sleep(0.5)
+    st.rerun()