Spaces:
Running
Running
File size: 1,878 Bytes
c151c44 |
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 52 53 54 55 56 57 58 59 |
from openai import AsyncOpenAI, OpenAI
import asyncio
import json
from typing import AsyncIterator
from typing import Union, List, Dict
from dotenv import load_dotenv
import os
load_dotenv()
# Initialize the async client
client = AsyncOpenAI(
base_url=os.getenv("BASE_URL"),
api_key=os.getenv("API_KEY"),
)
def process_input(data: Union[str, List[Dict[str, str]]]) -> Union[str, List[Dict[str, str]]]:
"""
Processes input to either uppercase a string or modify the 'content' field
of a list of dictionaries.
"""
if isinstance(data, str):
return data.strip() # Ensures prompt is cleaned up (optional)
elif isinstance(data, list):
# Ensure each item in the list is a dictionary with a 'content' key
return [
{**item, "content": item["content"].strip()} # Trims whitespace in 'content'
for item in data if isinstance(item, dict) and "content" in item
]
else:
raise TypeError("Input must be a string or a list of dictionaries with a 'content' field")
async def get_completions(
prompt: Union[str, List[Dict[str, str]]],
instructions: str
) -> str:
processed_prompt = process_input(prompt) # Ensures the input format is correct
if isinstance(processed_prompt, str):
messages = [
{"role": "system", "content": instructions},
{"role": "user", "content": processed_prompt}
]
elif isinstance(processed_prompt, list):
messages = [{"role": "system", "content": instructions}] + processed_prompt
else:
raise TypeError("Unexpected processed input type.")
response = await client.chat.completions.create(
model=os.getenv("MODEL"),
messages=messages,
response_format={"type": "json_object"}
)
output: str = response.choices[0].message.content
return output
|