File size: 1,660 Bytes
d184dfc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
from dotenv import load_dotenv
from openai import AsyncOpenAI, OpenAI
import openai
from typing import List, Dict, Any
import os
import asyncio


class ChatOpenAI:
    def __init__(self, model_name: str = "gpt-4-turbo-preview"):
        load_dotenv()
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        self.async_client = AsyncOpenAI()
        self.client = OpenAI()

        if self.openai_api_key is None:
            raise ValueError(
                "OPENAI_API_KEY environment variable is not set. Please set it to your OpenAI API key."
            )
        openai.api_key = self.openai_api_key
        self.model_name = model_name

    async def agenerate(self, messages: List[Dict[str, str]]) -> str:
        chat_completion = await self.async_client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=0.7,
        )

        return chat_completion.choices[0].message.content

    async def astream(self, messages: List[Dict[str, str]]):
        stream = await self.async_client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=0.7,
            stream=True,
        )

        async for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                yield chunk.choices[0].delta.content

    def generate(self, messages: List[Dict[str, str]]) -> str:
        chat_completion = self.client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            temperature=0.7,
        )

        return chat_completion.choices[0].message.content