code upated with model lodaers

sunnysavita10 · sunnysavita10 · commit 1cd0a999ccd6 · 2025-08-24T12:41:35.000+05:30
diff --git a/.github/workflows/task_definition.json b/.github/workflows/task_definition.json
@@ -28,12 +28,7 @@
       ],
       "secrets": [
         {
-          "name": "GROQ_API_KEY",
-          "valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"
-        },
-        
-        {
-          "name": "GOOGLE_API_KEY",
+          "name": "API_KEYS",
           "valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"
         }
         
diff --git a/.gitignore b/.gitignore
@@ -7,4 +7,5 @@ logs/
 *.pyc
 faiss_index/
 main_archive/
-data/
+data/
+archive/
diff --git a/test.py b/test.py
@@ -241,3 +241,58 @@
 # #     secret = get_secret_value_response['SecretString']
 
 # #     # Your code goes here.
+
+
+
+# {
+#   "family": "documentportaltd",
+#   "networkMode": "awsvpc",
+#   "executionRoleArn": "arn:aws:iam::459497895986:role/ecsTaskExecutionRole",
+#   "requiresCompatibilities": ["FARGATE"],
+#   "cpu": "1024",
+#   "memory": "8192",
+#   "containerDefinitions": [
+#     {
+#       "name": "document-portal-container",
+#       "image": "459497895986.dkr.ecr.ap-southeast-2.amazonaws.com/documentportalliveclass",
+#       "cpu": 1024,
+#       "essential": true,
+#       "portMappings": [
+#         {
+#           "containerPort": 8080,
+#           "hostPort": 8080,
+#           "protocol": "tcp",
+#           "name": "document-portal-container-8080-tcp",
+#           "appProtocol": "http"
+#         }
+#       ],
+#       "environment": [
+#         {
+#           "name": "ENV",
+#           "value": "production"
+#         }
+#       ],
+#       "secrets": [
+#         {
+#           "name": "GROQ_API_KEY",
+#           "valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"
+#         },
+        
+#         {
+#           "name": "GOOGLE_API_KEY",
+#           "valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"
+#         }
+        
+#       ],
+#       "logConfiguration": {
+#         "logDriver": "awslogs",
+#         "options": {
+#           "awslogs-group": "/ecs/documentportaltd",
+#           "awslogs-region": "ap-southeast-2",
+#           "awslogs-stream-prefix": "ecs",
+#           "awslogs-create-group": "true"
+#         }
+#       }
+#     }
+#   ]
+# }
diff --git a/utils/model_loader.py b/utils/model_loader.py
@@ -1,127 +1,142 @@
-
 import os
 import sys
+import json
 from dotenv import load_dotenv
 from utils.config_loader import load_config
-from .config_loader import load_config
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-#from langchain_openai import ChatOpenAI
 from logger import GLOBAL_LOGGER as log
 from exception.custom_exception import DocumentPortalException
 
+
+class ApiKeyManager:
+    REQUIRED_KEYS = ["GROQ_API_KEY", "GOOGLE_API_KEY"]
+
+    def __init__(self):
+        self.api_keys = {}
+        raw = os.getenv("API_KEYS")
+
+        if raw:
+            try:
+                parsed = json.loads(raw)
+                if not isinstance(parsed, dict):
+                    raise ValueError("API_KEYS is not a valid JSON object")
+                self.api_keys = parsed
+                log.info("Loaded API_KEYS from ECS secret")
+            except Exception as e:
+                log.warning("Failed to parse API_KEYS as JSON", error=str(e))
+
+        # Fallback to individual env vars
+        for key in self.REQUIRED_KEYS:
+            if not self.api_keys.get(key):
+                env_val = os.getenv(key)
+                if env_val:
+                    self.api_keys[key] = env_val
+                    log.info(f"Loaded {key} from individual env var")
+
+        # Final check
+        missing = [k for k in self.REQUIRED_KEYS if not self.api_keys.get(k)]
+        if missing:
+            log.error("Missing required API keys", missing_keys=missing)
+            raise DocumentPortalException("Missing API keys", sys)
+
+        log.info("API keys loaded", keys={k: v[:6] + "..." for k, v in self.api_keys.items()})
+
+
+    def get(self, key: str) -> str:
+        val = self.api_keys.get(key)
+        if not val:
+            raise KeyError(f"API key for {key} is missing")
+        return val
+
+
 class ModelLoader:
-    
     """
-    A utility class to load embedding models and LLM models.
+    Loads embedding models and LLMs based on config and environment.
     """
-    
+
     def __init__(self):
-        
         if os.getenv("ENV", "local").lower() != "production":
             load_dotenv()
-            log.info("Running in LOCAL mode: .env file loaded")
+            log.info("Running in LOCAL mode: .env loaded")
         else:
-            log.info("Running in PRODUCTION mode: .env not loaded")
-        self._validate_env()
-        self.config=load_config()
-        log.info("Configuration loaded successfully", config_keys=list(self.config.keys()))
-        
-    def _validate_env(self):
-        """
-        Validate necessary environment variables.
-        Ensure API keys exist.
-        """
-        required_vars=["GOOGLE_API_KEY","GROQ_API_KEY"]
-        self.api_keys={key:os.getenv(key) for key in required_vars}
-        missing = [k for k, v in self.api_keys.items() if not v]
-        if missing:
-            log.error("Missing environment variables", missing_vars=missing)
-            raise DocumentPortalException("Missing environment variables", sys)
-        log.info("Environment variables validated", available_keys=[k for k in self.api_keys if self.api_keys[k]])
-        log.info("Environment variables validated", available_keys={k: v[:30] + "..." if v else None for k, v in self.api_keys.items()})
+            log.info("Running in PRODUCTION mode")
+
+        self.api_key_mgr = ApiKeyManager()
+        self.config = load_config()
+        log.info("YAML config loaded", config_keys=list(self.config.keys()))
 
-        
     def load_embeddings(self):
         """
-        Load and return the embedding model.
+        Load and return embedding model from Google Generative AI.
         """
         try:
-            log.info("Loading embedding model...")
             model_name = self.config["embedding_model"]["model_name"]
-            return GoogleGenerativeAIEmbeddings(model=model_name)
+            log.info("Loading embedding model", model=model_name)
+            return GoogleGenerativeAIEmbeddings(model=model_name,
+                                                google_api_key=self.api_key_mgr.get("GOOGLE_API_KEY")) #type: ignore
         except Exception as e:
             log.error("Error loading embedding model", error=str(e))
             raise DocumentPortalException("Failed to load embedding model", sys)
-        
+
     def load_llm(self):
         """
-        Load and return the LLM model.
+        Load and return the configured LLM model.
         """
-        """Load LLM dynamically based on provider in config."""
-        
         llm_block = self.config["llm"]
+        provider_key = os.getenv("LLM_PROVIDER", "google")
 
-        log.info("Loading LLM...")
-
-        provider_key = os.getenv("LLM_PROVIDER", "google")  # Default google
         if provider_key not in llm_block:
-            log.error("LLM provider not found in config", provider_key=provider_key)
-            raise ValueError(f"Provider '{provider_key}' not found in config")
+            log.error("LLM provider not found in config", provider=provider_key)
+            raise ValueError(f"LLM provider '{provider_key}' not found in config")
 
         llm_config = llm_block[provider_key]
         provider = llm_config.get("provider")
         model_name = llm_config.get("model_name")
         temperature = llm_config.get("temperature", 0.2)
         max_tokens = llm_config.get("max_output_tokens", 2048)
-        
-        log.info("Loading LLM", provider=provider, model=model_name, temperature=temperature, max_tokens=max_tokens)
+
+        log.info("Loading LLM", provider=provider, model=model_name)
 
         if provider == "google":
-            llm=ChatGoogleGenerativeAI(
+            return ChatGoogleGenerativeAI(
                 model=model_name,
+                google_api_key=self.api_key_mgr.get("GOOGLE_API_KEY"),
                 temperature=temperature,
                 max_output_tokens=max_tokens
             )
-            return llm
 
         elif provider == "groq":
-            llm=ChatGroq(
+            return ChatGroq(
                 model=model_name,
-                api_key=self.api_keys["GROQ_API_KEY"], #type: ignore
+                api_key=self.api_key_mgr.get("GROQ_API_KEY"), #type: ignore
                 temperature=temperature,
             )
-            return llm
-            
+
         # elif provider == "openai":
         #     return ChatOpenAI(
         #         model=model_name,
-        #         api_key=self.api_keys["OPENAI_API_KEY"],
+        #         api_key=self.api_key_mgr.get("OPENAI_API_KEY"),
         #         temperature=temperature,
         #         max_tokens=max_tokens
         #     )
+
         else:
             log.error("Unsupported LLM provider", provider=provider)
             raise ValueError(f"Unsupported LLM provider: {provider}")
-        
-    
-    
+
+
 if __name__ == "__main__":
     loader = ModelLoader()
-    
-    # Test embedding model loading
+
+    # Test Embedding
     embeddings = loader.load_embeddings()
-    print(f"Embedding Model Loaded: {embeddings}")
-    
-    # Test the ModelLoader
-    result=embeddings.embed_query("Hello, how are you?")
-    print(f"Embedding Result: {result}")
-    
-    # Test LLM loading based on YAML config
+    print(f"✅ Embedding Model Loaded: {embeddings}")
+    result = embeddings.embed_query("Hello, how are you?")
+    print(f"✅ Embedding Result: {result}")
+
+    # Test LLM
     llm = loader.load_llm()
-    print(f"LLM Loaded: {llm}")
-    
-    # Test the ModelLoader
-    result=llm.invoke("Hello, how are you?")
-    print(f"LLM Result: {result.content}")
+    print(f"✅ LLM Loaded: {llm}")
+    result = llm.invoke("Hello, how are you?")
+    print(f"✅ LLM Result: {result.content}")

Original file line number	Diff line number	Diff line change
`@@ -28,12 +28,7 @@`
`28`	`28`	`],`
`29`	`29`	`"secrets": [`
`30`	`30`	`{`
`31`		`- "name": "GROQ_API_KEY",`
`32`		`- "valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"`
`33`		`- },`
`34`		`-`
`35`		`- {`
`36`		`- "name": "GOOGLE_API_KEY",`
	`31`	`+ "name": "API_KEYS",`
`37`	`32`	`"valueFrom": "arn:aws:secretsmanager:ap-southeast-2:459497895986:secret:api_keys-nZTtj8"`
`38`	`33`	`}`
`39`	`34`