Update app.py

app.py

@@ -149,23 +149,32 @@ def verify_user(username: str, password: str):
 
 @app.post("/validate-user")
 async def validate_user(request: Request, username: str = Form(...), password: str = Form(...)):
-    status,role = verify_user(username, password)
-    if status == 'success':
-        logging.info(f"user role {role}is rerturned")
+    status, role ,company_id= verify_user(username, password)
+    if status == 'success' and role and company_id:
+        logging.info(f"user role {role} is returned")
+        
+        # Set cookies and redirect to the dashboard
         response = RedirectResponse(url="/dashboard", status_code=302)
         response.set_cookie(key="role", value=role)
         response.set_cookie(key="username", value=username)
+        response.set_cookie(key="company_id",value=company_id)
         return response
-       
     else:
-        return templates.TemplateResponse("index.html", {"request": request})
-
+        # If login fails, redirect back to the index page with an error message
+        return templates.TemplateResponse("index.html", {
+            "request": request,
+            "error": "Invalid username or password"
+        })
+        
 @app.post("/submit_company_profile")
 async def submit_company_profile(request: Request,
                                  company_name: str = Form(...),
                                  company_code: str = Form(...),
                                  domain: str = Form(...),
-                                 llm_tools: List[str] = Form(...)):
+                                 llm_tools: List[str] = Form(...),
+                                 username:str=Form(...),
+                                 password:str=Form(...),
+                                 role:str=Form(...)):
     logging.info("Received form submission for company profile")
     logging.info(f"Form data - company_name: {company_name}, company_code: {company_code}, domain: {domain}, llm_tools: {llm_tools}")
 
@@ -176,6 +185,15 @@ async def submit_company_profile(request: Request,
         values = (company_name, company_code, domain, ",".join(llm_tools))
         logging.info(f"Executing query: {query} with values: {values}")
         cursor.execute(query, values)
+                # Retrieve the inserted company_id
+        company_id = cursor.lastrowid
+        logging.info(f"Company profile for {company_name} inserted successfully with company_id: {company_id}")
+
+        # Insert user details with the retrieved company_id
+        user_query = "INSERT INTO user_detail (company_id, username, password) VALUES (%s, %s, %s, %s)"
+        user_values = (company_id, username, password, role)
+        logging.info(f"Executing user detail query: {user_query} with values: {user_values}")
+        cursor.execute(user_query, user_values)
         cnx.commit()
         logging.info(f"Query executed successfully, {cursor.rowcount} row(s) affected")
         cursor.close()
@@ -201,37 +219,115 @@ async def get_companies():
     except mysql.connector.Error as err:
         logging.error(f"Database error: {err}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
+        
 @app.get("/dashboard")
 async def dashboard(request: Request):
     try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        username = request.cookies.get("username")
+        company_id=request._cookies.get("company_id")
+
+        # Establish database connection
         cnx = get_db_connection()
         cursor = cnx.cursor()
 
-        cursor.execute("show tables")
+        # Fetch all table names
+        cursor.execute("SHOW TABLES")
         all_tables = cursor.fetchall()
+        
+        # Dictionary to hold the count of records for each table
         table_count_of_each_table = {}
-        # Fetch count of each table in the database
-        for table in all_tables:
-            query = f"SELECT COUNT(*) FROM {table[0]}"
-            cursor.execute(query)
-            table_count_of_each_table[table[0]] = cursor.fetchone()[0]
 
+        # Fetch count of records for each table
+        for table in all_tables:
+            table_name = table[0]
+            query = f"SELECT COUNT(*) FROM {table_name} WHERE company_id = %s"
+            cursor.execute(query, (company_id,))
+        
+            count = cursor.fetchone()[0]
+            table_count_of_each_table[table_name] = count
+        query1=f"select company_name from company_detail where company_id = %s"
+        cursor.execute(query1,(company_id,))
+        company_name_result = cursor.fetchone()
+        
+        # Check if company_name_result is not None
+        if company_name_result:
+            company_name = company_name_result[0]
+        else:
+            company_name = "Unknown"  # Default
+        # Close cursor and connection
         cursor.close()
-        cnx.close() 
+        cnx.close()
         
+        # Log the counts for debugging purposes
         logging.info(table_count_of_each_table)
-       # return templates.TemplateResponse("dashboard.html", {"request": request,"title":"Dashboard" , "table_count_of_each_table": table_count_of_each_table})
+        
+        # Render the template with the data, role, and username
         return templates.TemplateResponse("dashboard.html", {
             "request": request,
             "title": "Dashboard",
-            "table_count_of_each_table": table_count_of_each_table
+            "table_count_of_each_table": table_count_of_each_table,
+            "role": role,
+            "username": username,
+            "company_id":company_id,
+            "company_name":company_name
         })
+    except mysql.connector.Error as err:
+        # Log the error and raise an HTTPException
+        logging.error(f"Database error: {err}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+
+@app.get("/api/company_record_count/{company_id}")
+async def get_company_record_count(company_id: int):
+    try:
+        # Establish database connection
+        cnx = get_db_connection()
+        cursor = cnx.cursor()
+
+        # List of tables to count records in
+        tables = ["knowledge_base", "data_connectors", "api_connectors", "prompt_templates"]
+
+        # Dictionary to hold the count of records for each table
+        table_counts = {}
+
+        # Fetch count of records for the selected company in each table
+        for table in tables:
+            query = f"SELECT COUNT(*) FROM {table} WHERE company_id = %s"
+            cursor.execute(query, (company_id,))
+            count = cursor.fetchone()[0]
+            table_counts[table] = count
+
+        # Close cursor and connection
+        cursor.close()
+        cnx.close()
+
+        return {"table_counts": table_counts}
+
     except mysql.connector.Error as err:
         logging.error(f"Database error: {err}")
         raise HTTPException(status_code=500, detail="Internal Server Error")
+
+        
 @app.get("/company_profile")
 async def company_profile(request: Request):
-    return templates.TemplateResponse("company_profile.html", {"request": request,"title":"Company Profile"})
+    try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        company_id = request.cookies.get("company_id")
+
+        # Render the template with the role and company_id
+        return templates.TemplateResponse("company_profile.html", {
+            "request": request,
+            "role": role,
+            "company_id": company_id,
+            "title":"Company Profile"
+        })
+    except Exception as e:
+        # Handle exceptions
+        logging.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+    #return templates.TemplateResponse("company_profile.html", {"request": request,"title":"Company Profile"})
 @app.get("/api/company_id")
 async def get_company_id(company_name: str):
     print(f"Received company_name: {company_name}")  # Debug statement
@@ -407,7 +503,23 @@ async def delete_company(company_id: int):
 
 @app.get("/knowledgebase")
 async def knowledgebase(request: Request):
-    return templates.TemplateResponse("knowledgebase.html", {"request": request,"title":"KnowledgeBase"})
+    try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        company_id = request.cookies.get("company_id")
+
+        # Render the template with the role and company_id
+        return templates.TemplateResponse("knowledgebase.html", {
+            "request": request,
+            "role": role,
+            "company_id": company_id,
+            "title":"KnowledgeBase"
+        })
+    except Exception as e:
+        # Handle exceptions
+        logging.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+   
 
 #to insert into knowledgebase
 @app.post("/upload_document")
@@ -707,8 +819,23 @@ async def delete_company(company_id: int):
  
 @app.get("/data_connectors")
 async def data_connectors(request: Request):
-    return templates.TemplateResponse("data_connectors.html", {"request": request, "title": "Data Connectors"})
+    try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        company_id = request.cookies.get("company_id")
 
+        # Render the template with the role and company_id
+        return templates.TemplateResponse("data_connectors.html", {
+            "request": request,
+            "role": role,
+            "company_id": company_id,
+            "title": "Data Connectors"
+        })
+    except Exception as e:
+        # Handle exceptions
+        logging.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+        
 #to insert into data_connectors
 @app.post("/save_data_connectors")
 async def save_data_connectors( request: Request,
@@ -822,7 +949,23 @@ async def get_data_connectors(company_id: str = Query(...), company_name: str =
  
 @app.get("/API_connectors")
 async def API_connectors(request: Request):
-    return templates.TemplateResponse("API_connectors.html", {"request": request,"title":"API Connectors"})
+    try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        company_id = request.cookies.get("company_id")
+
+        # Render the template with the role and company_id
+        return templates.TemplateResponse("API_connectors.html", {
+            "request": request,
+            "role": role,
+            "company_id": company_id,
+            "title":"API Connectors"
+        })
+    except Exception as e:
+        # Handle exceptions
+        logging.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
+  
 #save api connectors 
 @app.post("/api/save_api_details")
 async def API_saveconnectors(request: Request,
@@ -1081,7 +1224,22 @@ def delete_api_from_db(company_id: int) -> bool:
 
 @app.get("/prompt_template")
 async def prompt_template(request: Request):
-    return templates.TemplateResponse("prompt_template.html", {"request": request,"title":"Prompt Templates"})
+    try:
+        # Retrieve cookies
+        role = request.cookies.get("role")
+        company_id = request.cookies.get("company_id")
+
+        # Render the template with the role and company_id
+        return templates.TemplateResponse("prompt_template.html", {
+            "request": request,
+            "role": role,
+            "company_id": company_id,
+            "title":"Prompt Templates"
+        })
+    except Exception as e:
+        # Handle exceptions
+        logging.error(f"Error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
 # to insert into prompt templates
 @app.post("/api/save_prompt_details")
 async def prompt_saveconnectors(request: Request,