model.py

import os
import json
from dotenv import load_dotenv
from openai import AzureOpenAI

load_dotenv('.env')

AZURE_OPENAI_API_KEY = os.environ.get('api_key')
OPENAI_API_VERSION = os.environ.get('api_version')
AZURE_OPENAI_ENDPOINT = os.environ.get('endpoint')
API_DEPLOYMENT_ID = {
    "o3-mini": "o3-mini-2025-01-31",
    "gpt-4.1": "gpt-4.1-2025-04-14",
    "gpt4o": "gpt-4o-2024-11-20"
}

client = AzureOpenAI(
    api_version=API_VERSION,
    azure_endpoint=RESOURCE_ENDPOINT,
    api_key=API_KEY
)

class LLM:
    def __init__(self, model_name):
        assert model_name in API_DEPLOYMENT_ID, f"""
        Model name not found in API_DEPLOYMENT_ID. Please check the model name.
        Available models: {list(API_DEPLOYMENT_ID.keys())}
        """
        self.model = API_DEPLOYMENT_ID[model_name]

    def _construct_message_chat(self, prompt, system_prompt=None):
        message = []
        if system_prompt:
            message.append({"role": "system", "content": system_prompt})
        else:
            message.append({"role": "system", "content": "Extract the relevant information."})
        message.append({
            "role": "user",
            "content": prompt
        })
        return message
    
    def _construct_message_structured(self, prompt, system_prompt=None):
        message = []
        if system_prompt:
            message.append({"role": "system", "content": system_prompt})
        else:
            message.append({"role": "system", "content": "Extract the relevant information."})
        message.append({"role": "user", "content": json.dumps(prompt)})      
        return message      
    
    def chat(self, prompt, system_prompt):
        message = self._construct_message_chat(prompt, system_prompt)
        response = client.chat.completions.create(
            model=self.model,
            messages=message
        )
        return response.choices[0].message.content

    def structured_chat(self, prompt, response_format, system_prompt=None):
        message = self._construct_message_structured(prompt, system_prompt)
        response = client.beta.chat.completions.parse(
            model=self.model,
            messages=message,
            response_format=response_format
        )
        return response.choices[0].message.parsed