Use openai connector

Author: BenConstable9

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/open_ai.py

@@ -33,18 +33,28 @@ def get_authentication_properties(cls) -> dict:
 
         return token_provider, api_key
 
-    async def run_completion_request(self, messages: list[dict], temperature=0):
+    async def run_completion_request(
+        self, messages: list[dict], temperature=0, max_tokens=2000, model="4o-mini"
+    ) -> str:
+        if model == "4o-mini":
+            model_deployment = os.environ["OpenAI__MiniCompletionDeployment"]
+        elif model == "4o":
+            model_deployment = os.environ["OpenAI__CompletionDeployment"]
+        else:
+            raise ValueError(f"Model {model} not found")
+
         token_provider, api_key = self.get_authentication_properties()
         async with AsyncAzureOpenAI(
-            azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
+            azure_deployment=model_deployment,
             api_version=os.environ["OpenAI__ApiVersion"],
             azure_endpoint=os.environ["OpenAI__Endpoint"],
             azure_ad_token_provider=token_provider,
             api_key=api_key,
         ) as open_ai_client:
             response = await open_ai_client.chat.completions.create(
-                model=os.environ["OpenAI__MiniCompletionDeployment"],
+                model=model_deployment,
                 messages=messages,
                 temperature=temperature,
+                max_tokens=max_tokens,
             )
         return response.choices[0].message.content

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -8,14 +8,12 @@
 import logging
 from pydantic import BaseModel, Field, ConfigDict, computed_field
 from typing import Optional
-from text_2_sql_core.utils.environment import IdentityType, get_identity_type
-from openai import AsyncAzureOpenAI
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import random
 import re
 import networkx as nx
 from text_2_sql_core.utils.database import DatabaseEngine
 from tenacity import retry, stop_after_attempt, wait_exponential
+from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
 logging.basicConfig(level=logging.INFO)
 
@@ -279,6 +277,8 @@ def __init__(
         if output_directory is None:
             self.output_directory = "."
 
+        self.open_ai_connector = OpenAIConnector()
+
         load_dotenv(find_dotenv())
 
     @property
@@ -627,9 +627,6 @@ async def extract_columns_with_definitions(
 
         return columns
 
-    @retry(
-        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
-    )
     async def send_request_to_llm(self, system_prompt: str, input: str):
         """A method to use GPT to generate a definition for an entity.
 
@@ -640,55 +637,23 @@ async def send_request_to_llm(self, system_prompt: str, input: str):
         Returns:
             str: The generated definition."""
 
-        MAX_TOKENS = 2000
-
-        api_version = os.environ["OpenAI__ApiVersion"]
-        model = os.environ["OpenAI__CompletionDeployment"]
-
-        if get_identity_type() in [
-            IdentityType.SYSTEM_ASSIGNED,
-            IdentityType.USER_ASSIGNED,
-        ]:
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
-            )
-            api_key = None
-        else:
-            token_provider = None
-            api_key = os.environ["OpenAI__ApiKey"]
-
-        try:
-            async with AsyncAzureOpenAI(
-                api_key=api_key,
-                api_version=api_version,
-                azure_ad_token_provider=token_provider,
-                azure_endpoint=os.environ.get("OpenAI__Endpoint"),
-            ) as client:
-                response = await client.chat.completions.create(
-                    model=model,
-                    messages=[
-                        {
-                            "role": "system",
-                            "content": system_prompt,
-                        },
-                        {
-                            "role": "user",
-                            "content": [
-                                {
-                                    "type": "text",
-                                    "text": input,
-                                },
-                            ],
-                        },
-                    ],
-                    max_tokens=MAX_TOKENS,
-                )
+        messages = [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": input,
+                    },
+                ],
+            },
+        ]
 
-            return response.choices[0].message.content
-        except Exception as e:
-            logging.error(f"Unable to generate definition for {input}")
-            logging.error(f"Error generating definition: {e}")
-            return None
+        return await self.open_ai_connector.run_completion_request(messages)
 
     async def generate_entity_definition(self, entity: EntityItem):
         """A method to generate a definition for an entity.