Moved seafile token handling to realtime user chat

02_oms_sql_pipeline.py

@@ -7,11 +7,13 @@ from openpyxl import load_workbook
 from openpyxl.styles import Alignment, Font
 from openpyxl.drawing.image import Image
 from copy import copy
+import pytz
 import logging
 import os
 import requests
 import json
 import time
+import datetime
 import ast
 
 logging.basicConfig(level=logging.DEBUG)
@@ -42,7 +44,7 @@ class Pipeline:
         TMP_STORAGE_DIR: str
 
     def __init__(self):
-        self.name = "01 PolarPress Database RAG"
+        self.name = "Østfold Milsim Medlemsdatabase"
         self.engine = None
         self.nlsql_response = ""
         self.last_emit_time = 0
@@ -256,6 +258,158 @@ class Pipeline:
             logging.error(f"Unable to share {file}. Error: {e}")
             return False
 
+    # def check_cloud_token
+    #
+    # userid (str): Brukerens UUID
+    #
+    # Returnerer: bool
+    # Returnerer True hvis brukerid ble funnet på keystore OG token
+    # er lagret OG updated nøkkelen i token ikke er eldre enn 45 minutter
+    #
+    # Beskrivelse: 
+    # Denne funksjonen sjekker om brukeren har en Outlands tilgangstoken
+    # for excel regneark opplasting og fildeling
+
+    def check_cloud_token(self, userid: str) -> bool:
+        try:
+            headers = {
+                'Content-Type': 'application/json',
+                'Accept': 'application/json; indent=4',
+                'X-API-KEY': self.valves.KEYSTORE_API
+            }
+
+            response = requests.get(
+                f"{self.valves.KEYSTORE_URL}/api/post/{userid}/store/seatoken",
+                headers=headers
+            )
+
+            if response.status_code == 404:
+                return False
+
+            res = response.json()
+
+            if 'updated_at' in res:
+                updated_str = res['updated']
+                updated_datetime = datetime.datetime.strptime(updated_str, '%Y-%m-%d %H:%M:%S')
+                timezone = pytz.timezone('Europe/Oslo')
+                updated_datetime = timezone.localize(updated_datetime)
+                now = datetime.datetime.now(timezone)
+                time_difference = now - updated_datetime
+                return time_difference > datetime.timedelta(minutes=45)
+
+            return False
+        except requests.RequestException as e:
+            logging.error(f"Error checking keystore for user auth token: {e}")
+            return False
+
+    # def new_cloud_token
+    #
+    # userid (str): Brukerens UUID
+    # token (str): Seafile autentiseringstoken
+    #
+    # Returnerer True hvis vellykket
+    #
+    # Beskrivelse:
+    # Denne funksjonen oppretter/oppdaterer
+    # brukerens seafile autentiseringstoken
+    # til nøkkelserveren
+
+    def update_cloud_token(self, userid: str, value_data: dict) ->bool:
+        try:
+            headers = {
+                'Content-Type': 'application/json',
+                'Accept': 'application/json; indent=4',
+                'X-API-KEY': self.valves.KEYSTORE_API
+            }
+            json_data = {
+                'body': userid,
+                'key': 'sea_token',
+                'value': json.dumps(value_data)
+            }
+
+            response = requests.post(
+                f"{self.valves.KEYSTORE_URL}/api/post/store",
+                json=json_data,
+                headers=headers
+            )
+
+            if response.status_code == 200:
+                return True
+                
+            return False
+        except requests.RequestException as e:
+            logging.error(f"Error updating keystore user auth token: {e}")
+            return False
+
+    # def check_cloud_token
+    #
+    # userid (str): Brukerens UUID
+    # token (str): Seafile autentiseringstoken
+    #
+    # Returnerer True hvis seafile token er gyldig
+    #
+    # Beskrivelse:
+    # Denne funksjonen sjekker seafile nettsky om brukerens
+    # autentiseringstoken er gyldig og kan brukes ved å
+    # hente brukerprofilen ved å bruke token
+
+    def check_cloud_user_profile(self, token: str) ->bool:
+        try:
+            headers = {
+                'Authorization': f"Token {token}",
+                'Accept': 'application/json; indent=4'
+            }
+
+            response = requests.get(
+                f"https://{self.valves.CLOUD_HOST}/{self.valves.CLOUD_API_VERSION}/user/",
+                headers=headers
+            )
+
+            if response.status_code == 200:
+                return True
+
+            return False
+        except requests.RequestException as e:
+            logging.error(f"Error trying user token on seafile server: {e}")
+            return False
+
+    # def get_cloud_token
+    #
+    # userid (str): Brukerens UUID
+    #
+    # Returnerer string
+    #
+    # Beskrivelse:
+    # Denne funksjonen henter brukerens autentiseringstoken fra
+    # nøkkelserveren
+
+    def get_cloud_token(self, userid: str) ->str:
+        try:
+            headers = {
+                'Content-Type': 'application/json',
+                'Accept': 'application/json; indent=4',
+                'X-API-KEY': self.valves.KEYSTORE_API
+            }
+
+            response = requests.get(
+                f"{self.valves.KEYSTORE_URL}/api/post/{userid}/store/sea_token",
+                headers=headers
+            )
+
+            if response.status_code == 404:
+                return False
+
+            res = response.json()
+
+            if 'value' in res:
+                value = json.loads(res['value'])
+                return value['token']
+
+            return False
+        except requests.RequestException as e:
+            logging.error(f"Error checking keystore for user auth token: {e}")
+            return False
+
     def run_llm_query(self, message: str):
         try:
             response = requests.post(
@@ -404,8 +558,29 @@ class Pipeline:
         MySQLQuery:
         """
 
+        keyserver_data = {
+            'user_name': body['user']['name'],
+            'user_email': body['user']['email'],
+            'user_role': body['user']['role'],
+            'token': self.token
+        }
+
         # Vi sjekker om body['stream'] er sann som indikerer at pipelinen kjører brukerens forespørsel
         if body['stream']:
+            # Først sjekker vi om vi har en gyldig autentiserings token på nettskyen vi kan lagre
+            # Excel regnearket i
+            if self.check_cloud_token(body['user']['id']):
+                current_token = self.get_cloud_token(body['user']['id']) # Hent nåværende autentiserings token
+                if self.check_cloud_user_profile(current_token):
+                    self.token = current_token # Den nåværende tokenen er gyldig, bruk denne til å laste opp og dele filer
+                else:
+                    self.sea_login()
+                    keyserver_data['token'] = self.token
+                    self.update_cloud_token(body['user']['id'], keyserver_data)
+            else:
+                self.sea_login()
+                keyserver_data['token'] = self.token
+                self.update_cloud_token(body['user']['id'], keyserver_data)
 
             # Vi kjører første inferens med instruksjonene i llm_initial_instructions for å fortelle
             # KI modellen at vi ønsker en SQL spørring til svar basert på user_message (det brukeren
@@ -502,6 +677,10 @@ class Pipeline:
                 if share_file:
                     data += f"\n\n## Eksport:\nI tilfelle resultatet fra spørringen overgår begrensningene i antall tokens i svaret, har jeg eksportert dataene du spurte om i et eget excel ark hvor du kan se hele datasettet. Du kan laste ned Excel arket her:\n1. [Regneark delelink på outlands.no]({share_file})"
 
+                # Debug data
+                # data += f"\n\n\n```{body}```"
+                print(str(body))
+
                 # Hvis resultatet er vellykket returner andre inferens trinn svar i chatten
                 if formatted_result["success"]:
                     return data