support for openai and gemini models at API layer

Author: Khauneesh-AI

app/core/model_handlers.py

@@ -12,6 +12,10 @@
 from app.core.exceptions import APIError, InvalidModelError, ModelHandlerError, JSONParsingError
 from app.core.telemetry_integration import track_llm_operation
 from app.core.config import  _get_caii_token
+import os
+from dotenv import load_dotenv
+load_dotenv() 
+import google.generativeai as genai 
 
 
 
@@ -154,11 +158,21 @@ def _extract_json_from_text(self, text: str) -> List[Dict[str, Any]]:
 
 
     #@track_llm_operation("generate")
-    def generate_response(self, prompt: str, retry_with_reduced_tokens: bool = True, request_id = None) -> List[Dict[str, str]]:
+    def generate_response(
+        self,
+        prompt: str,
+        retry_with_reduced_tokens: bool = True,
+        request_id: Optional[str] = None,
+    ):
         if self.inference_type == "aws_bedrock":
             return self._handle_bedrock_request(prompt, retry_with_reduced_tokens)
-        elif self.inference_type == "CAII":
+        if self.inference_type == "CAII":
             return self._handle_caii_request(prompt)
+        if self.inference_type == "openai":
+            return self._handle_openai_request(prompt)
+        if self.inference_type == "gemini":
+            return self._handle_gemini_request(prompt)
+        raise ModelHandlerError(f"Unsupported inference_type={self.inference_type}", 400)
 
     def _handle_bedrock_request(self, prompt: str, retry_with_reduced_tokens: bool):
         """Handle Bedrock requests with retry logic"""
@@ -278,6 +292,50 @@ def _handle_bedrock_request(self, prompt: str, retry_with_reduced_tokens: bool):
             raise ModelHandlerError(f"Failed after {self.MAX_RETRIES} retries: {str(last_exception)}", status_code=500)
 
 
+    # ---------- OpenAI -------------------------------------------------------
+    def _handle_openai_request(self, prompt: str):
+        try:
+            client = OpenAI(
+                api_key=os.getenv("OPENAI_API_KEY"),
+                base_url=os.getenv("OPENAI_API_BASE", None) or None,
+            )
+            completion = client.chat.completions.create(
+                model=self.model_id,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=self.model_params.max_tokens,
+                temperature=self.model_params.temperature,
+                top_p=self.model_params.top_p,
+                stream=False,
+            )
+            text = completion.choices[0].message.content
+            return self._extract_json_from_text(text) if not self.custom_p else text
+        except Exception as e:
+            raise ModelHandlerError(f"OpenAI request failed: {e}", 500)
+
+    # ---------- Gemini -------------------------------------------------------
+    def _handle_gemini_request(self, prompt: str):
+        if genai is None:
+            raise ModelHandlerError(
+                "google-generativeai library not installed — `pip install google-generativeai`",
+                500,
+            )
+        try:
+            genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+            model = genai.GenerativeModel(self.model_id)  # e.g. 'gemini-1.5-pro-latest'
+            resp = model.generate_content(
+                prompt,
+                generation_config={
+                    "max_output_tokens": self.model_params.max_tokens,
+                    "temperature": self.model_params.temperature,
+                    "top_p": self.model_params.top_p,
+                },
+            )
+            text = resp.text
+            return self._extract_json_from_text(text) if not self.custom_p else text
+        except Exception as e:
+            raise ModelHandlerError(f"Gemini request failed: {e}", 500)
+
+
     def _handle_caii_request(self, prompt: str):
         """Original CAII implementation"""
         try:

app/main.py

@@ -731,7 +731,9 @@ async def get_model_id():
 
     models = {
         "aws_bedrock": bedrock_list,
-        "CAII": []
+        "CAII": [],
+        "OpenAI" : [],
+        "Google Gemini" : []
     }
 
     return {"models": models}

pyproject.toml

@@ -36,11 +36,14 @@ dependencies = [
     "psutil==5.9.8",
     "pandas>=2.2.3",
     # ── new packages for data‑analysis layer ─────────────────
-    "numpy>=1.24.0",            # explicit, for dcor/scipy (already a transitive dep of pandas)
-    "scipy>=1.12.0",            # chi‑square, ANOVA, etc.
-    "dcor>=0.6",                # distance‑correlation metric
-    "openpyxl>=3.1.2",          # read .xlsx files
-    "pyxlsb>=1.0.9",            # read .xlsb files
+    "numpy>=1.24.0", # explicit, for dcor/scipy (already a transitive dep of pandas)
+    "scipy>=1.12.0", # chi‑square, ANOVA, etc.
+    "dcor>=0.6", # distance‑correlation metric
+    "openpyxl>=3.1.2", # read .xlsx files
+    "pyxlsb>=1.0.9", # read .xlsb files
+    "google>=3.0.0",
+    "google-generativeai>=0.8.5",
+    "google-genai>=1.2.0",
 ]