68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
from dailyai.services.ai_services import AIService, TTSService, LLMService, ImageGenService
|
|
from typing import Generator
|
|
|
|
import requests
|
|
from PIL import Image
|
|
import io
|
|
from openai import OpenAI
|
|
|
|
import os
|
|
import json
|
|
|
|
class OpenAILLMService(LLMService):
|
|
def __init__(self, api_key=None, model=None):
|
|
super().__init__()
|
|
api_key = api_key or os.getenv("OPEN_AI_KEY")
|
|
self.model = model or os.getenv("OPEN_AI_MODEL")
|
|
self.client = OpenAI(api_key=api_key)
|
|
|
|
def get_response(self, messages, stream):
|
|
return self.client.chat.completions.create(
|
|
stream=stream,
|
|
messages=messages,
|
|
model=self.model
|
|
)
|
|
|
|
def run_llm_async(self, messages) -> Generator[str, None, None]:
|
|
messages_for_log = json.dumps(messages)
|
|
self.logger.debug(f"Generating chat via openai: {messages_for_log}")
|
|
|
|
response = self.get_response(messages, stream=True)
|
|
|
|
for chunk in response:
|
|
if len(chunk.choices) == 0:
|
|
continue
|
|
|
|
if chunk.choices[0].delta.content:
|
|
yield chunk.choices[0].delta.content
|
|
|
|
def run_llm(self, messages) -> str | None:
|
|
messages_for_log = json.dumps(messages)
|
|
self.logger.debug(f"Generating chat via azure: {messages_for_log}")
|
|
|
|
response = self.get_response(messages, stream=False)
|
|
if response and len(response.choices) > 0:
|
|
return response.choices[0].message.content
|
|
else:
|
|
return None
|
|
|
|
class OpenAIImageGenService(ImageGenService):
|
|
def __init__(self, api_key=None, model=None):
|
|
super().__init__()
|
|
api_key = api_key or os.getenv("OPEN_AI_KEY")
|
|
self.model = model or os.getenv("OPEN_AI_MODEL")
|
|
self.client = OpenAI(api_key=api_key)
|
|
|
|
def run_image_gen(self, sentence) -> tuple[str, Image.Image]:
|
|
image = self.client.images.generate(
|
|
prompt=sentence,
|
|
n=1,
|
|
size=f"1024x1024"
|
|
)
|
|
image_url = image.data[0].url
|
|
response = requests.get(image_url)
|
|
dalle_stream = io.BytesIO(response.content)
|
|
dalle_im = Image.open(dalle_stream)
|
|
|
|
return (image_url, dalle_im)
|