Better assistant caching

app.py

@@ -1,7 +1,8 @@
 import os
 import urllib
-from time import sleep
 import urllib.request
+from time import sleep
+from typing import Dict, List, Generator
 
 import gradio as gr
 import openai
@@ -11,70 +12,32 @@ load_dotenv()
 
 class MockInterviewer:
 
-    def __init__(self):
-        self.client = openai.OpenAI(api_key=os.environ['OPENAI_API_KEY'])
-        self.job_role = ''
-        self.company = ''
-        self.assistant_id = ''
-
-    def create_files(self, company):
-        if company.lower() == 'amazon':
-            url = 'https://www.aboutamazon.com/about-us/leadership-principles'
-            filename = 'leadership_principles.html'
-        else:
-            return []
-        
-        filename, headers = urllib.request.urlretrieve(url, filename)
-        with open(filename, 'rb') as file:
-            assistant_file = self.client.files.create(file=file, purpose='assistants')
-            file_ids = [assistant_file.id]
-        os.remove(filename)
-        return file_ids
-
-    def init_assistant(self, job_role, company):
-        if not job_role and not company:
-            raise gr.Error('Job Role and Company are required fields.')
-        if not job_role:
-            raise gr.Error('Job Role is a required field.')
-        if not company:
-            raise gr.Error('Company is a required field.')
-
-        if job_role != self.job_role or company != self.company:
-            file_ids = self.create_files(company)
+    def __init__(self) -> None:
+        self._client = openai.OpenAI(api_key=os.environ['OPENAI_API_KEY'])
+        self._assistant_id_cache: Dict[str, str] = {}
 
-            assistant = self.client.beta.assistants.create(
-                name='Mock Interviewer',
-                instructions=f'You are an AI mock interviewer for {job_role} roles at {company}. Please make it obvious this is your purpose. If you have been provided a file, use it as an interview guide.',
-                model='gpt-4-0125-preview',
-                tools=[
-                    {
-                        'type': 'retrieval'  # This adds the knowledge base as a tool
-                    }
-                ],
-                file_ids=file_ids)
-            
-            self.assistant_id = assistant.id
-
-    def chat(self, usr_message, history, job_role, company):
+    def chat(self, usr_message: Dict, history: List[List], job_role: str, company: str) -> Generator:
         print('Started function')
-        thread = self.client.beta.threads.create()
-        user_input = usr_message['text']
+        self._validate_fields(job_role, company)
+
+        thread = self._client.beta.threads.create()
+        user_input = usr_message.get('text')
 
-        self.init_assistant(job_role, company)
+        assistant_id = self._init_assistant(job_role, company)
 
         # Add the user's message to the thread
-        self.client.beta.threads.messages.create(thread_id=thread.id,
+        self._client.beta.threads.messages.create(thread_id=thread.id,
                                                  role="user",
                                                  content=user_input)
         print('Client made')
         # Run the Assistant
-        run = self.client.beta.threads.runs.create(thread_id=thread.id,
-                                                   assistant_id=self.assistant_id)
+        run = self._client.beta.threads.runs.create(thread_id=thread.id,
+                                                   assistant_id=assistant_id)
         print('Run created')
 
         # Check if the Run requires action (function call)
         while True:
-            run_status = self.client.beta.threads.runs.retrieve(thread_id=thread.id,
+            run_status = self._client.beta.threads.runs.retrieve(thread_id=thread.id,
                                                                 run_id=run.id)
             print(f"Run status: {run_status.status}")
             if run_status.status == 'completed':
@@ -83,12 +46,60 @@ class MockInterviewer:
             sleep(1)  # Wait for a second before checking again
 
         # Retrieve and return the latest message from the assistant
-        messages = self.client.beta.threads.messages.list(thread_id=thread.id)
+        messages = self._client.beta.threads.messages.list(thread_id=thread.id)
         response = messages.data[0].content[0].text.value
 
         print(f"Assistant response: {response}")  # Debugging line
         #return json.dumps({"response": response})
-        yield response    
+        yield response
+
+    def _validate_fields(self, job_role: str, company: str) -> None:
+        if not job_role and not company:
+            raise gr.Error('Job Role and Company are required fields.')
+        if not job_role:
+            raise gr.Error('Job Role is a required field.')
+        if not company:
+            raise gr.Error('Company is a required field.')
+
+    def _create_files(self, company: str) -> List[str]:
+        if company.lower() == 'amazon':
+            url = 'https://www.aboutamazon.com/about-us/leadership-principles'
+            filename = 'leadership_principles.html'
+        else:
+            return []
+        
+        filename, headers = urllib.request.urlretrieve(url, filename)
+        with open(filename, 'rb') as file:
+            assistant_file = self._client.files.create(file=file, purpose='assistants')
+            file_ids = [assistant_file.id]
+        os.remove(filename)
+        return file_ids
+
+    def _init_assistant(self, job_role: str, company: str) -> str:
+        cache_key = self._create_cache_key(job_role, company)
+        if cache_key in self._assistant_id_cache:
+            print(f'Fetched from cache for key {cache_key}')
+            return self._assistant_id_cache.get(cache_key)
+        else:
+            print(f'Initializing new assistant for key {cache_key}')
+            file_ids = self._create_files(company)
+
+            assistant = self._client.beta.assistants.create(
+                name='Mock Interviewer',
+                instructions=f'You are an AI mock interviewer for {job_role} roles at {company}. Please make it obvious this is your purpose. If you have been provided a file, use it as an interview guide.',
+                model='gpt-4-0125-preview',
+                tools=[
+                    {
+                        'type': 'retrieval'  # This adds the knowledge base as a tool
+                    }
+                ],
+                file_ids=file_ids)
+            
+            self._assistant_id_cache[cache_key] = assistant.id
+            return assistant.id
+
+    def _create_cache_key(self, job_role: str, company: str) -> str:
+        return f'{job_role.lower()}+{company.lower()}'
 
 # Creating the Gradio interface
 with gr.Blocks() as demo: