How to use Azure OpenAI and Azure AI Studio with Weights & Biases Weave

!pip install weave openai

import weave

weave.init('azure-weave-cookbook')

from openai import AzureOpenAI

client = AzureOpenAI(
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        api_version="2024-02-01",
        azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    )

from openai import OpenAI

client = OpenAI(
        base_url=f"{os.getenv('AZURE_AI_STUDIO_API_ENDPOINT')}/v1",
        api_key=os.getenv('AZURE_AI_STUDIO_API_KEY')
    )

@weave.op()
def call_azure_chat(model_id: str, messages: list, max_tokens: int = 1000, temperature: float = 0.5):
    response = client.chat.completions.create(
        model=model_id,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature
    )
    return {"status": "success", "response": response.choices[0].message.content}

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Create a snack recipe for a dish called the Azure Weav-e-ohs"}
]
if "mistral" in model_id.lower():
    messages = format_messages_for_mistral(messages)
result = call_azure_chat(model_id, messages)
print(result)

@weave.op()
def format_prompt(prompt: str):
    "A formatting function for OpenAI models"
    system_prompt_formatted = "You are a helpful assistant."
    human_prompt = "{prompt}"
    human_prompt_formatted = human_prompt.format(prompt=prompt)
    messages = [{"role":"system", "content":system_prompt_formatted}, {"role":"user", "content":human_prompt_formatted}]
    return messages

@weave.op()
def run_chat(model_id: str, prompt: str):
    formatted_messages = format_prompt(prompt=prompt)
    if "mistral" in model_id.lower():
        formatted_messages = format_messages_for_mistral(formatted_messages)
    result = call_azure_chat(model_id, formatted_messages, max_tokens=1000)
    return result

from dataclasses import dataclass

@dataclass
class PromptTemplate:
    system_prompt: str
    human_prompt: str

    @weave.op()
    def format_prompt(self, email_content: str):
        system_prompt_formatted = self.system_prompt.format()
        human_prompt_formatted = self.human_prompt.format(email_content=email_content)
        messages = [{"role":"system", "content":system_prompt_formatted}, {"role":"user", "content":human_prompt_formatted}]
        return messages

from weave import Model

class AzureEmailAssistant(Model):
    model_id: str = model_id
    prompt_template: PromptTemplate
    max_tokens: int = 2048
    temperature: float = 0.0

    @weave.op()
    def format_doc(self, doc: str) -> list:
        messages = self.prompt_template.format_prompt(doc)
        return messages

    @weave.op()
    def respond(self, doc: str) -> dict:
        messages = self.format_doc(doc)
        if "mistral" in self.model_id.lower():
            messages = format_messages_for_mistral(messages)
        output = call_azure_chat(
            self.model_id,
            messages=messages,
            max_tokens=self.max_tokens,
            temperature=self.temperature)
        return output

    @weave.op()
    async def predict(self, email_content: str) -> str:
        return self.respond(email_content)["response"]

from weave import Dataset

dataset = Dataset(name=eval_dataset_name, rows=[
    {'id': '1', 'email_content': 'Subject: Inquiry about Order Delay\n\nHello,\n\nI placed an order last week for the new UltraGlow Skin Serum, but I have not received a shipping update yet. My order number is 12345. Could you please update me on the status of my shipment?\n\nThank you,\nJane Doe'},
    {'id': '2', 'email_content': 'Subject: Damaged Item Received\n\nHello,\n\nI received my order yesterday, but one of the items, a glass vase, was broken. My order number is 67890. How can I get a replacement or a refund?\n\nBest regards,\nJohn Smith'},
    # Additional rows
])
weave.publish(dataset)

dataset_uri = f"weave:///{wandb_entity}/{weave_project}/object/{eval_dataset_name}:latest"
dataset = weave.ref(dataset_uri).get()

@weave.op()
def check_conciseness(model_output: str) -> dict:
    result = len(model_output.split()) < 300
    return {'conciseness': result}

evaluation = weave.Evaluation(
    dataset=dataset, scorers=[check_conciseness],
)
await evaluation.evaluate(model)