Update app/core/clients.py

app/core/clients.py

@@ -1,30 +1,32 @@
-import os
 import logging
+from google import genai
+from huggingface_hub import InferenceClient
 from langchain_groq import ChatGroq
 from langchain_core.prompts import PromptTemplate
-from google import genai
+
 from .config import settings
 
 logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
-# --- LLM (LangChain) Client ---
 
+# --- LangChain (Groq LLM) ---
 def get_llm():
-    """Initializes and returns the ChatGroq LLM instance."""
     logger.info("Initializing ChatGroq LLM...")
     return ChatGroq(
         model="meta-llama/llama-4-scout-17b-16e-instruct",
         temperature=1,
         max_tokens=1024,
-        api_key=settings.LANGCHAIN_API_KEY
+        api_key=settings.LANGCHAIN_API_KEY,
     )
 
+
 def get_prompt_template():
     """Returns the prompt enhancement template."""
     return PromptTemplate.from_template('''
 You are a Prompt Enhancement AI Assistant. Your task is to take the
 user's raw image prompt and convert it into a detailed,
-professional prompt optimized for generating high-quality AI image.
+professional prompt optimized for generating high-quality AI images.
 Enhance the prompt by including relevant details such as:
 - Camera specifications (e.g., lens type, aperture, focal length)
 - Lighting setup (e.g., natural light, studio lighting, soft shadows)
@@ -37,21 +39,51 @@ Raw Prompt:
 Enhanced Prompt:
 ''')
 
-# --- Google GenAI Client (FIXED) ---
+# --- Google GenAI (Gemini) ---
 def get_genai_client():
-    """Initializes and returns the Google GenAI client (older SDK version)."""
+    """Initializes and returns the Google GenAI client."""
+    if not settings.GOOGLE_API_KEY:
+        raise RuntimeError("Missing GOOGLE_API_KEY in .env")
     try:
-        logger.info("Initializing Google GenAI Client...")
-        # Use the genai.Client from your original code
+        logger.info("Initializing Google GenAI client...")
         client = genai.Client(api_key=settings.GOOGLE_API_KEY)
-        logger.info("Google GenAI Client initialized successfully.")
+        logger.info("Google GenAI client initialized successfully.")
         return client
     except Exception as e:
-        logger.error(f"Failed to initialize GenAI client: {e}")
+        logger.exception("Failed to initialize GenAI client.")
         raise
 
-# --- Initialize Clients ---
+def get_genai_client1():
+    """Initializes and returns the Google GenAI client."""
+    if not settings.GOOGLE_API_KEY:
+        raise RuntimeError("Missing GOOGLE_API_KEY in .env")
+    try:
+        logger.info("Initializing Google GenAI client...")
+        client = genai.Client(api_key=settings.GOOGLE_API_KEY1)
+        logger.info("Google GenAI client initialized successfully.")
+        return client
+    except Exception as e:
+        logger.exception("Failed to initialize GenAI client.")
+        raise
+
+# --- Hugging Face (FLUX) ---
+def get_flux_client():
+    """Initializes and returns the Hugging Face FLUX client."""
+    if not settings.HF_API_KEY:
+        raise RuntimeError("Missing HF_API_KEY in .env")
+    try:
+        logger.info("Initializing Hugging Face InferenceClient...")
+        client = InferenceClient(provider="hf-inference", api_key=settings.HF_API_KEY)
+        logger.info("Hugging Face InferenceClient initialized successfully.")
+        return client
+    except Exception as e:
+        logger.exception("Failed to initialize FLUX client.")
+        raise
+
+
+# --- Export initialized clients ---
 llm = get_llm()
 prompt_template = get_prompt_template()
-# Export the client object
-genai_client = get_genai_client()
+genai_client = get_genai_client()
+flux_client = get_flux_client()
+genai_client1 = get_genai_client1()