From d1f169dead9d3fec10c8b804bf92dc0a2fe24288 Mon Sep 17 00:00:00 2001
From: Robert Szefler <robert@robusta.dev>
Date: Wed, 12 Jun 2024 12:30:42 +0200
Subject: [PATCH] token handling fixes

---
 holmes/core/robusta_ai.py   |  3 +++
 holmes/core/supabase_dal.py | 11 +++++++----
 server.py                   |  6 +-----
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/holmes/core/robusta_ai.py b/holmes/core/robusta_ai.py
index a0cc86df..dd53315f 100644
--- a/holmes/core/robusta_ai.py
+++ b/holmes/core/robusta_ai.py
@@ -22,6 +22,9 @@ def __init__(self, base_url: str, session_manager: SessionManager, runbook_manag
 
     def call(self, system_prompt: str, user_prompt: str) -> LLMResult:
         auth_token = self.session_manager.get_current_token()
+        if auth_token is None:
+            auth_token = self.session_manager.create_token()
+
         payload = {
             "auth": {"account_id": auth_token.account_id, "token": auth_token.token},
             "system_message": system_prompt,
diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py
index 60e77fa9..95e6bc95 100644
--- a/holmes/core/supabase_dal.py
+++ b/holmes/core/supabase_dal.py
@@ -172,7 +172,7 @@ def create_auth_token(self, token_type: str) -> AuthToken:
                 {
                     "account_id": self.account_id,
                     "user_id": self.user_id,
-                    "token": uuid4(),
+                    "token": str(uuid4()),
                     "type": token_type,
                 }
             )
@@ -180,17 +180,20 @@ def create_auth_token(self, token_type: str) -> AuthToken:
         )
         return AuthToken(**result.data[0])
 
-    def get_freshest_auth_token(self, token_type: str) -> AuthToken:
+    def get_freshest_auth_token(self, token_type: str) -> Optional[AuthToken]:
         result = (
             self.client.table(TOKENS_TABLE)
             .select("*")
-            .filter("token_type", "eq", token_type)
+            .filter("type", "eq", token_type)
             .filter("deleted", "eq", False)
             .order("created_at", desc=True)
             .limit(1)
             .execute()
         )
-        return AuthToken(**result.data[0])
+        if not result.data:
+            return None
+        else:
+            return AuthToken(**result.data[0])
 
     def invalidate_auth_token(self, token: AuthToken) -> None:
         (
diff --git a/server.py b/server.py
index ba00e8e8..f254bb1f 100644
--- a/server.py
+++ b/server.py
@@ -18,11 +18,7 @@
 from fastapi import FastAPI
 from rich.console import Console
 
-from holmes.common.env_vars import (
-    ALLOWED_TOOLSETS,
-    HOLMES_HOST,
-    HOLMES_PORT,
-)
+from holmes.common.env_vars import ALLOWED_TOOLSETS, HOLMES_HOST, HOLMES_PORT
 from holmes.config import BaseLLMConfig, LLMProviderType
 from holmes.core.issue import Issue
 from holmes.core.provider import LLMProviderFactory