refactor(singleton): wire OmniParser, RAG KB, and engine via app.state

Author: Santosh69

app/api/routes/evaluation.py

@@ -1,11 +1,11 @@
 import logging
 from fastapi import APIRouter, HTTPException, UploadFile, File, Body, Request
-from typing import Optional, List
-import json
+from typing import List
 
-from app.services.heuristic_engine import HeuristicEvaluationEngine
-from app.core.config import settings, ALLOWED_IMAGE_TYPES
-from app.services.exceptions import InvalidInputError
+from app.core.config import ALLOWED_IMAGE_TYPES
+from app.services.exceptions import InvalidInputError, ModelInferenceError, RAGKnowledgeBaseError
+from app.services.omniparser_client import UIElement
+from app.core.constants import HeuristicId, NIELSEN_HEURISTICS
 
 logger = logging.getLogger(__name__)
 router = APIRouter()
@@ -55,12 +55,13 @@ async def evaluate_heuristics(
         detection_client = request.app.state.omniparser_client
         contents = await image.read()
 
-        detection_result = await detection_client.detect_elements(contents)
-
-        # Initialize evaluation engine and evaluate
-        evaluation_engine = HeuristicEvaluationEngine()
-        await evaluation_engine.initialize()
+        detection_result = await detection_client.detect_elements(
+            contents,
+            content_type=content_type
+        )
 
+        # Use singleton evaluation engine and evaluate
+        evaluation_engine = request.app.state.heuristic_engine
         evaluation_result = await evaluation_engine.evaluate_interface(detection_result)
 
         return {
@@ -82,6 +83,7 @@ async def evaluate_heuristics(
 
 @router.post("/evaluate-legacy/{heuristic_id}")
 async def evaluate_legacy_format(
+    request: Request,
     heuristic_id: str,
     elements: List[dict] = Body(...)
 ):
@@ -100,7 +102,6 @@ async def evaluate_legacy_format(
         500: Unexpected server error
     """
     try:
-        from app.core.constants import HeuristicId
 
         # Normalize heuristic_id to uppercase
         normalized_id = heuristic_id.upper()
@@ -113,10 +114,8 @@ async def evaluate_legacy_format(
                 }
             )
 
-        evaluation_engine = HeuristicEvaluationEngine()
-        await evaluation_engine.initialize()
+        evaluation_engine = request.app.state.heuristic_engine
 
-        from app.services.omniparser_client import UIElement
         try:
             ui_elements = [UIElement.from_dict(e) for e in elements]
         except Exception as e:
@@ -168,8 +167,6 @@ async def evaluate_legacy_format(
 
 @router.get("/heuristics")
 async def get_heuristics():
-    from app.core.constants import NIELSEN_HEURISTICS
-
     return {
         "success": True,
         "data": {
@@ -183,7 +180,7 @@ async def get_heuristics():
     }
 
 @router.get("/knowledge-base/stats")
-async def get_knowledge_base_stats():
+async def get_knowledge_base_stats(request: Request):
     """Get statistics about the RAG knowledge base.
     
     Returns:
@@ -194,10 +191,7 @@ async def get_knowledge_base_stats():
         500: Unexpected server error
     """
     try:
-        from app.services.rag_knowledge_base import RAGKnowledgeBase
-
-        kb = RAGKnowledgeBase()
-        await kb.initialize()
+        kb = request.app.state.rag_knowledge_base
         stats = await kb.get_stats()
 
         return {

app/api/routes/heuristic.py

@@ -1,10 +1,7 @@
 import logging
 from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Request
 from typing import Optional
-from PIL import Image
-import io
 
-from app.services.omniparser_client import OmniParserClient, UIElementDetectionResult
 
 logger = logging.getLogger(__name__)
 router = APIRouter()
@@ -37,11 +34,11 @@ async def detect_ui_elements(
 
 @router.post("/analyze")
 async def analyze_interface(
+    request: Request,
     image: UploadFile = File(...)
 ):
     try:
-        client = OmniParserClient()
-        await client.initialize()
+        client = request.app.state.omniparser_client
 
         contents = await image.read()
         result = await client.detect_elements(contents)

app/services/heuristic_engine.py

@@ -2,14 +2,13 @@
 import json
 from typing import List, Dict, Any, Optional
 from datetime import datetime
-import asyncio
 from openai import AsyncOpenAI
 
 from app.core.constants import NIELSEN_HEURISTICS, HeuristicId, SeverityLevel
 from app.core.config import settings
 from app.services.omniparser_client import UIElementDetectionResult, UIElement
 from app.services.rag_knowledge_base import RAGKnowledgeBase
-from app.services.exceptions import ModelInferenceError, InvalidInputError
+from app.services.exceptions import ModelInferenceError
 
 logger = logging.getLogger(__name__)
 
@@ -111,24 +110,28 @@ class HeuristicEvaluationEngine:
     Evaluation Method: LLM-based with prompt engineering
     
     Example Usage:
-        engine = HeuristicEvaluationEngine()
+        engine = request.app.state.heuristic_engine
         result = await engine.evaluate_interface(detection_result)
     """
 
-    def __init__(self):
+    def __init__(self, rag_kb: Optional[RAGKnowledgeBase] = None):
         self.logger = logging.getLogger(__name__)
         self.llm_client = None
-        self.rag_kb = None
+        self.rag_kb = rag_kb
         self.initialized = False
 
     async def initialize(self):
+        if self.initialized:
+            return
         self.logger.info("Initializing Heuristic Evaluation Engine...")
         self.llm_client = AsyncOpenAI(
             api_key=settings.OPENAI_API_KEY,
             base_url=settings.OPENAI_BASE_URL
         )
-        self.rag_kb = RAGKnowledgeBase()
-        await self.rag_kb.initialize()
+        if self.rag_kb is None:
+            raise RuntimeError("RAGKnowledgeBase dependency missing. Inject via startup.")
+        if not self.rag_kb.index_initialized:
+            await self.rag_kb.initialize()
         self.initialized = True
         self.logger.info("Heuristic Evaluation Engine initialized")
 

main.py

@@ -1,9 +1,6 @@
-from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
 import logging
-import os
-from pathlib import Path
 
 from app.core.config import settings
 from app.services.heuristic_engine import HeuristicEvaluationEngine
@@ -32,24 +29,28 @@
 
 @app.on_event("startup")
 async def startup_event():
-    # Initialize OmniParser Client (Singleton)
-    app.state.omniparser_client = OmniParserClient()
-    await app.state.omniparser_client.initialize()
-    
     setup_logging()
     logger = logging.getLogger(__name__)
     logger.info("AI Heuristic Evaluation API starting up...")
 
-    # Initialize singleton OmniParser client to avoid re-initializing model on every request
+    # Initialize OmniParser Client (Singleton)
     app.state.omniparser_client = OmniParserClient()
     await app.state.omniparser_client.initialize()
     logger.info("OmniParser client initialized (singleton)")
 
-    heuristic_engine = HeuristicEvaluationEngine()
-    await heuristic_engine.initialize()
+    # Initialize RAG Knowledge Base (Singleton)
+    app.state.rag_knowledge_base = RAGKnowledgeBase()
+    await app.state.rag_knowledge_base.initialize()
+    logger.info("RAG knowledge base initialized")
 
+    # Initilize Hueristic Evaluation Engine (Singleton)
+    app.state.heuristic_engine = HeuristicEvaluationEngine(
+        rag_kb=app.state.rag_knowledge_base
+    )
+    await app.state.heuristic_engine.initialize()
     logger.info("Heuristic evaluation engine initialized")
 
+
 @app.on_event("shutdown")
 async def shutdown_event():
     logger = logging.getLogger(__name__)