microsoft--promptflow/examples/flows/standard/gen-docstring/azure_open_ai.py

import asyncio
import logging
import time
import uuid
from typing import List
from openai.version import VERSION as OPENAI_VERSION
import os
from abc import ABC, abstractmethod
import tiktoken
from dotenv import load_dotenv
from prompt import PromptLimitException


class AOAI(ABC):
    def __init__(self, **kwargs):
        if OPENAI_VERSION.startswith("0."):
            raise Exception(
                "Please upgrade your OpenAI package to version >= 1.0.0 or "
                "using the command: pip install --upgrade openai."
            )
        init_params = {}
        api_type = os.environ.get("API_TYPE")
        if os.getenv("OPENAI_API_VERSION") is not None:
            init_params["api_version"] = os.environ.get("OPENAI_API_VERSION")
        if os.getenv("OPENAI_ORG_ID") is not None:
            init_params["organization"] = os.environ.get("OPENAI_ORG_ID")
        if os.getenv("OPENAI_API_KEY") is None:
            raise ValueError("OPENAI_API_KEY is not set in environment variables")
        if os.getenv("OPENAI_API_BASE") is not None:
            if api_type == "azure":
                init_params["azure_endpoint"] = os.environ.get("OPENAI_API_BASE")
            else:
                init_params["base_url"] = os.environ.get("OPENAI_API_BASE")
        init_params["api_key"] = os.environ.get("OPENAI_API_KEY")
        # A few sanity checks
        if api_type == "azure":
            if init_params.get("azure_endpoint") is None:
                raise ValueError(
                    "OPENAI_API_BASE is not set in environment variables, this is required when api_type==azure"
                )
            if init_params.get("api_version") is None:
                raise ValueError(
                    "OPENAI_API_VERSION is not set in environment variables, this is required when api_type==azure"
                )
            if init_params["api_key"].startswith("sk-"):
                raise ValueError(
                    "OPENAI_API_KEY should not start with sk- when api_type==azure, "
                    "are you using openai key by mistake?"
                )
            from openai import AzureOpenAI as Client
            from openai import AsyncAzureOpenAI as AsyncClient
        else:
            from openai import OpenAI as Client
            from openai import AsyncClient as AsyncClient
        self.client = Client(**init_params)
        self.async_client = AsyncClient(**init_params)
        self.default_engine = None
        self.engine = kwargs.pop('model', None) or os.environ.get("MODEL")
        self.total_tokens = 4000
        self.max_tokens = kwargs.pop('max_tokens', None) or os.environ.get("MAX_TOKENS") or 1200
        if self.engine == "gpt-4-32k":
            self.total_tokens = 31000
        if self.engine == "gpt-4":
            self.total_tokens = 7000
        if self.engine == "gpt-3.5-turbo-16k":
            self.total_tokens = 15000
        if self.max_tokens > self.total_tokens:
            raise ValueError(f"max_tokens must be less than total_tokens, "
                             f"total_tokens is {self.total_tokens}, max_tokens is {self.max_tokens}")
        self.tokens_limit = self.total_tokens - self.max_tokens

    def count_tokens(self, text: str) -> int:
        try:
            encoding = tiktoken.encoding_for_model(self.engine)
        except KeyError:
            encoding = tiktoken.encoding_for_model(self.default_engine)
        return len(encoding.encode(text))

    def query(self, text, **kwargs):
        stream = kwargs.pop("stream", False)
        for i in range(3):
            try:
                if not stream:
                    return self.query_with_no_stream(text, **kwargs)
                else:
                    return "".join(self.query_with_stream(text, **kwargs))
            except Exception as e:
                logging.error(f"Query failed, message={e}, "
                              f"will retry request llm after {(i + 1) * (i + 1)} seconds.")
                time.sleep((i + 1) * (i + 1))
        raise Exception("Query failed, and retry 3 times, but still failed.")

    async def async_query(self, text, **kwargs):
        stream = kwargs.pop("stream", False)
        for i in range(3):
            try:
                if not stream:
                    res = await self.async_query_with_no_stream(text, **kwargs)
                    return res
                else:
                    res = await self.async_query_with_stream(text, **kwargs)
                    return "".join(res)
            except Exception as e:
                logging.error(f"llm response error, message={e}, "
                              f"will retry request llm after {(i + 1) * (i + 1)} seconds.")
                await asyncio.sleep((i + 1) * (i + 1))
        raise Exception("llm response error, and retry 3 times, but still failed.")

    @abstractmethod
    def query_with_no_stream(self, text, **kwargs):
        pass

    @abstractmethod
    def query_with_stream(self, text, **kwargs):
        pass

    @abstractmethod
    async def async_query_with_no_stream(self, text, **kwargs):
        pass

    @abstractmethod
    async def async_query_with_stream(self, text, **kwargs):
        pass


class ChatLLM(AOAI):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.default_engine = "gpt-3.5-turbo"
        self.engine = self.engine or self.default_engine
        self.system_prompt = "You are a Python engineer."
        self.conversation = dict()

    def query_with_no_stream(self, text, **kwargs):
        conversation_id = kwargs.pop('conversation', None)
        messages = self.create_prompt(text, conversation_id)
        self.validate_tokens(messages)
        temperature = kwargs.pop("temperature", 0.1)
        response = self.client.chat.completions.create(
            model=self.engine,
            messages=messages,
            temperature=temperature,
            max_tokens=self.max_tokens,
            stream=False,
            **kwargs,
        )
        response_role = response.choices[0].message.role
        full_response = response.choices[0].message.content
        self.add_to_conversation(text, "user", conversation_id=conversation_id)
        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
        return full_response

    def query_with_stream(self, text, **kwargs):
        conversation_id = kwargs.pop('conversation', None)
        messages = self.create_prompt(text, conversation_id)
        self.validate_tokens(messages)
        temperature = kwargs.pop("temperature", 0.1)
        response = self.client.chat.completions.create(
            model=self.engine,
            messages=messages,
            temperature=temperature,
            max_tokens=self.max_tokens,
            stream=True,
            **kwargs,
        )

        response_role = None
        full_response = ""
        for chunk in response:
            delta = chunk.choices[0].delta
            response_role = delta.role
            if delta.content:
                content = delta.content
                full_response += content
                yield content
        self.add_to_conversation(text, "user", conversation_id=conversation_id)
        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)

    async def async_query_with_no_stream(self, text, **kwargs):
        conversation_id = kwargs.pop('conversation', None)
        messages = self.create_prompt(text, conversation_id)
        self.validate_tokens(messages)
        temperature = kwargs.pop("temperature", 0.1)
        response = await self.async_client.chat.completions.create(
            model=self.engine,
            messages=messages,
            temperature=temperature,
            max_tokens=self.max_tokens,
            stream=False,
            **kwargs,
        )
        response_role = response.choices[0].message.role
        full_response = response.choices[0].message.content
        self.add_to_conversation(text, "user", conversation_id=conversation_id)
        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
        return full_response

    async def async_query_with_stream(self, text, **kwargs):
        conversation_id = kwargs.pop('conversation', None)
        messages = self.create_prompt(text, conversation_id)
        self.validate_tokens(messages)
        temperature = kwargs.pop("temperature", 0.1)
        response = await self.async_client.chat.completions.create(
            model=self.engine,
            messages=messages,
            temperature=temperature,
            max_tokens=self.max_tokens,
            stream=True,
            **kwargs,
        )

        response_role = None
        full_response = ""
        for chunk in response:
            delta = chunk.choices[0].delta
            response_role = delta.role
            if delta.content:
                content = delta.content
                full_response += content
                yield content
        self.add_to_conversation(text, "user", conversation_id=conversation_id)
        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)

    def get_unique_conversation_id(self):
        return str(uuid.uuid4()).replace('-', '')

    def add_to_conversation(self, message: str, role: str, conversation_id: str) -> None:
        """
        Add a message to the conversation
        """
        if type(conversation_id) is str:
            self.conversation[conversation_id].append({"role": role, "content": message})

    def del_conversation(self, conversation_id: str) -> None:
        if conversation_id in self.conversation:
            del self.conversation[conversation_id]

    def init_conversation(self, conversation_id: str, system_prompt) -> None:
        """
        Init a new conversation
        """
        if type(conversation_id) is str:
            self.conversation[conversation_id] = [{"role": "system", "content": system_prompt}]

    def get_tokens_count(self, messages: List[dict]) -> int:
        """
        Get token count
        """
        num_tokens = 0
        for message in messages:
            # every message follows <im_start>{role/name}\n{content}<im_end>\n
            num_tokens += 5
            for key, value in message.items():
                if value:
                    num_tokens += self.count_tokens(value)
                if key == "name":  # if there's a name, the role is omitted
                    num_tokens += 5  # role is always required and always 1 token
        num_tokens += 5  # every reply is primed with <im_start>assistant
        return num_tokens

    def validate_tokens(self, messages: List[dict]) -> None:
        total_tokens = self.get_tokens_count(messages)
        if total_tokens > self.tokens_limit:
            message = f"token count {total_tokens} exceeds limit {self.tokens_limit}"
            raise PromptLimitException(message)

    def create_prompt(self, text: str, conversation_id: str = None):
        unique_conversation_id = self.get_unique_conversation_id()
        conversation_id = conversation_id or unique_conversation_id
        if conversation_id not in self.conversation:
            self.init_conversation(conversation_id=conversation_id, system_prompt=self.system_prompt)

        _conversation = self.conversation[conversation_id] + [{"role": "user", "content": text}]

        while self.get_tokens_count(_conversation) > self.tokens_limit and len(_conversation) > 2:
            _conversation.pop(1)

        if unique_conversation_id == conversation_id:
            self.del_conversation(conversation_id=unique_conversation_id)

        return _conversation


if __name__ == "__main__":
    load_dotenv()
    llm = ChatLLM()
    print(llm.query(text='how are you?'))
    res = llm.query_with_stream(text='how are you?')
    for item in res:
        print(item)