performance improvement: initialize llm agent only once after startup

classifier/classifier.py

@@ -1,69 +1,79 @@
 import logging
-from fastapi import FastAPI
-from pydantic import BaseModel
-from typing import List
-import os
-from string import Formatter
-
 import os
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional
 
 import outlines
 from outlines.models import openai
 from outlines.generate import choice
 
 # Configure logger
-tools = logging.getLogger("classifier")
-tools.setLevel(logging.DEBUG)
-ch = logging.StreamHandler()
-ch.setLevel(logging.DEBUG)
-formatter = logging.Formatter("%(levelname)s:%(name)s:%(message)s")
-ch.setFormatter(formatter)
-tools.addHandler(ch)
-
-# Configure logger
-logging.basicConfig(
-    format="%(asctime)s %(levelname)s:%(name)s: %(message)s",
-    level=logging.DEBUG,
-)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger("classifier")
 
 app = FastAPI()
 
-# Pydantic model for incoming requests; prompt_template added
-class Req(BaseModel):
-    message: str
+# Global variables for shared config and classifier
+clf = None
+config_set = False
+
+class Config(BaseModel):
     model_name: str
     base_url: str
     class_set: List[str]
-    prompt_template: str  # template with {message} placeholder
+    prompt_template: str
+
+class Req(BaseModel):
+    message: str
 
 class Resp(BaseModel):
     result: str
 
-# Helper for safe formatting of {message} only
-class SafeFormatDict(dict):
-    def __missing__(self, key):
-        return '{' + key + '}'
+@app.post("/config")
+def configure(req: Config):
+    """Receive and initialize classifier configuration."""
+    global clf, config_set
 
-@app.post("/classify", response_model=Resp)
-def classify(req: Req):
-    logger.debug(f"Received request args: {req.dict()}")
-
-    prompt = req.prompt_template.replace("{message}", req.message)
-    logger.debug(f"Rendered prompt: {prompt!r}")
+    if config_set:
+        logger.warning("Classifier already configured. Ignoring new config.")
+        return {"status": "already_configured"}
 
     api_key = os.getenv("TOGETHERAI_API_KEY")
     logger.debug(f"Using API_KEY: {'set' if api_key else 'missing'}")
-    llm = openai(req.model_name, api_key=api_key, base_url=req.base_url)
-    clf = choice(llm, req.class_set)
-    logger.debug(f"Choice classifier created with labels: {req.class_set}")
 
     try:
-        result = clf(prompt)
-        # If it's a coroutine, run it; otherwise use result
-        logger.debug(f"Classifier returned: {result}")
+        llm = openai(req.model_name, api_key=api_key, base_url=req.base_url)
+        clf = choice(llm, req.class_set)
+        clf.class_set = req.class_set
+        clf.prompt_template = req.prompt_template
+        config_set = True
+        logger.info("Classifier configured successfully.")
+        return {"status": "configured"}
+    except Exception as e:
+        logger.error(f"Failed to configure classifier: {e}")
+        raise HTTPException(status_code=500, detail="Classifier configuration failed")
+
+@app.post("/classify", response_model=Resp)
+def classify(req: Req):
+    global clf
+    if clf is None or not config_set:
+        raise HTTPException(status_code=503, detail="Classifier not configured yet")
+
+    # Render the prompt using the template
+    try:
+        prompt = clf.prompt_template.replace("{message}", req.message)
+        logger.debug(f"Rendered prompt: {prompt!r}")
+    except Exception as e:
+        logger.warning(f"Prompt rendering failed: {e}")
+        prompt = req.message
+
+    # Run classifier
+    try:
+        result =  clf(prompt)
+        logger.debug(f"Classification result: {result}")
     except Exception as e:
-        result = req.class_set[-1]
-        logger.error(f"Classification error: {e}. Falling back to: {result}")
+        logger.error(f"Classification error: {e}. Falling back to: {clf.class_set[-1]}")
+        result = clf.class_set[-1]
 
     return Resp(result=result)

custom_components/llm_intent_classifier_client.py

@@ -30,6 +30,7 @@ class LlmIntentClassifier(IntentClassifier):
         component_config: Optional[Dict[Text, Any]] = None,
     ) -> None:
         super().__init__(component_config or {})
+
         self.url: str = self.component_config.get("classifier_url")
         self.timeout: float = float(self.component_config.get("timeout"))
         self.model_name: Optional[Text] = self.component_config.get("model_name")
@@ -52,6 +53,26 @@ class LlmIntentClassifier(IntentClassifier):
                 f"Missing configuration for {', '.join(missing)} in LlmIntentClassifier"
             )
 
+        # Push config to classifier backend
+        self._configure_remote_classifier()
+
+    def _configure_remote_classifier(self) -> None:
+        """Send configuration to the classifier backend to initialize the model."""
+        payload = {
+            "model_name": self.model_name,
+            "base_url": self.base_url,
+            "class_set": self.class_set,
+            "prompt_template": self.prompt_template,
+        }
+        try:
+            config_url = self.url.replace("/classify", "/config")
+            logger.debug(f"Sending classifier config to: {config_url}")
+            response = requests.post(config_url, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            logger.info("Remote classifier initialized successfully.")
+        except Exception as e:
+            logger.warning(f"Failed to initialize remote classifier: {e}")
+
     def train(
         self,
         training_data: TrainingData,
@@ -67,13 +88,7 @@ class LlmIntentClassifier(IntentClassifier):
         confidence: float = 0.0
 
         if text:
-            payload: Dict[str, Any] = {
-                "message": text,
-                "model_name": self.model_name,
-                "base_url": self.base_url,
-                "class_set": self.class_set,
-                "prompt_template": self.prompt_template,
-            }
+            payload: Dict[str, Any] = {"message": text}
             try:
                 resp = requests.post(self.url, json=payload, timeout=self.timeout)
                 resp.raise_for_status()