Refactor to defer imports for LLM initialization, improving worker startup times and reducing cold-start race conditions.

Author: m-messer

src/agent/llm_factory.py

@@ -1,18 +1,14 @@
 import os
 from typing import Optional
 
-from langchain_openai import AzureChatOpenAI
-from langchain_openai import AzureOpenAIEmbeddings
-from langchain_community.llms import Ollama
-from langchain_community.embeddings import OllamaEmbeddings
 from langchain_openai import ChatOpenAI
-from langchain_openai import OpenAIEmbeddings
-from langchain_google_genai import ChatGoogleGenerativeAI
 from dotenv import load_dotenv
 load_dotenv()
 
 class AzureLLMs:
     def __init__(self, temperature: int = 0):
+        from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
+
         self._azure_llm = AzureChatOpenAI(
                         openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"],
                         azure_deployment=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"],
@@ -31,6 +27,9 @@ def get_embedding(self):
 
 class OllamaLLMs:
     def __init__(self):
+        from langchain_community.llms import Ollama
+        from langchain_community.embeddings import OllamaEmbeddings
+
         self._ollama_llm = Ollama(
             model=os.environ['OLLAMA_MODEL'],
             base_url=os.environ['OLLAMA_BASE_URL'],
@@ -62,19 +61,13 @@ def __init__(self, temperature: int = 0):
             api_key=os.environ["OPENAI_API_KEY"],
         )
 
-        self._openai_embedding = OpenAIEmbeddings(
-            model='text-embedding-ada-002',
-            api_key=os.environ['OPENAI_API_KEY'],
-        )
-
     def get_llm(self):
         return self._openai_llm
 
-    def get_embedding(self):
-        return self._openai_embedding
-
 class GoogleAILLMs:
     def __init__(self, temperature: int = 0):
+        from langchain_google_genai import ChatGoogleGenerativeAI
+
         self._google_llm = ChatGoogleGenerativeAI(
             model=os.environ['GOOGLE_AI_MODEL'],
             temperature=temperature,

src/module.py

@@ -5,7 +5,6 @@
 from lf_toolkit.shared.mued_api_v0_1_0 import DataPolicySupport, HealthStatus, Role
 
 from src.agent.context import parse_json_to_prompt
-from src.agent.agent import invoke_base_agent
 
 
 def chat_module(request: ChatRequest) -> ChatResponse:
@@ -21,6 +20,10 @@ def chat_module(request: ChatRequest) -> ChatResponse:
     Edit src/agent/prompts.py to change the chatbot's behaviour.
     Edit src/agent/agent.py to change the agent logic (summarisation threshold, LLM provider, etc.).
     """
+    # Deferred: this pulls in langgraph/langchain/the LLM provider SDKs and builds the agent.
+    # Importing it lazily keeps worker boot (and the RPC socket bind) fast, avoiding a cold-start
+    # dial race against shimmy's worker-start-timeout.
+    from src.agent.agent import invoke_base_agent
 
     conversation_id = request.conversationId