From 12d211897003163276609068066446d9bf94f8c3 Mon Sep 17 00:00:00 2001
From: Bharatdeep <140376418+Bharatdeep@users.noreply.github.com>
Date: Wed, 6 Nov 2024 22:18:46 +0530
Subject: [PATCH] Using base url in openai_api.py

---
 data_gemma/openai_api.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/data_gemma/openai_api.py b/data_gemma/openai_api.py
index a81a12d..8ded65e 100644
--- a/data_gemma/openai_api.py
+++ b/data_gemma/openai_api.py
@@ -30,11 +30,13 @@ class OpenAI(base.LLM):
   def __init__(
       self,
       model: str,
+      base_url: str,
       api_key: str,
       verbose: bool = True,
       session: requests.Session | None = None,
   ):
     self.key = api_key
+    self.base_url = base_url
     if not session:
       session = requests.Session()
     self.session: requests.Session = session
@@ -86,7 +88,7 @@ def _call_api(self, req_data: str) -> Any:
         'Authorization': f'Bearer {self.key}',
     }
     r = self.session.post(
-        'https://api.openai.com/v1/chat/completions',
+        self.base_url,
         data=req_data,
         headers=header,
     )