From 3ee63f521ed3c93db73eeba50ff547c5244d63f2 Mon Sep 17 00:00:00 2001
From: Mark Liffiton <liffiton@gmail.com>
Date: Thu, 28 Nov 2024 14:22:24 -0600
Subject: [PATCH] Lazy load LLM client objects only when needed.

---
 src/gened/openai.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/gened/openai.py b/src/gened/openai.py
index d422a05..ba3e6bc 100644
--- a/src/gened/openai.py
+++ b/src/gened/openai.py
@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: AGPL-3.0-only
 
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import wraps
 from sqlite3 import Row
 from typing import ParamSpec, TypeVar
@@ -29,11 +29,19 @@ class NoTokensError(Exception):
     pass
 
 
-@dataclass(frozen=True)
+@dataclass
 class LLMConfig:
-    client: AsyncOpenAI
+    api_key: str
     model: str
     tokens_remaining: int | None = None  # None if current user is not using tokens
+    _client: AsyncOpenAI | None = field(default=None, init=False, repr=False)
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        """ Lazy load the LLM client object only when requested. """
+        if self._client is None:
+            self._client = AsyncOpenAI(api_key=self.api_key)
+        return self._client
 
 
 def _get_llm(*, use_system_key: bool, spend_token: bool) -> LLMConfig:
@@ -63,10 +71,10 @@ def make_system_client(tokens_remaining: int | None = None) -> LLMConfig:
         """ Factory function to initialize a default client (using the system key)
             only if/when needed.
         """
-        system_model = current_app.config["SYSTEM_MODEL"]
         system_key = current_app.config["OPENAI_API_KEY"]
+        system_model = current_app.config["SYSTEM_MODEL"]
         return LLMConfig(
-            client=AsyncOpenAI(api_key=system_key),
+            api_key=system_key,
             model=system_model,
             tokens_remaining=tokens_remaining,
         )
@@ -104,7 +112,7 @@ def make_system_client(tokens_remaining: int | None = None) -> LLMConfig:
 
         api_key = class_row['openai_key']
         return LLMConfig(
-            client=AsyncOpenAI(api_key=api_key),
+            api_key=api_key,
             model=class_row['model'],
         )