Merge pull request #1101 from IlumCI/OpenAIFunctionCaller-Removal
[BUG-LiteLLM][Removed OpenAIFunctionCaller completely, and replaced with LiteLLM]pull/966/merge
Kye Gomez
1 week ago
committed by
GitHub
commit
a762c79d6e
@ -1,123 +0,0 @@
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
SUPPORTED_MODELS = [
|
||||
"o3-mini-2025-1-31",
|
||||
"o1-2024-12-17",
|
||||
"gpt-4o-mini-2024-07-18",
|
||||
"gpt-4o-2024-08-06",
|
||||
]
|
||||
|
||||
|
||||
def check_api_key():
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"API key is not set. Please set the API key using the api_key parameter."
|
||||
)
|
||||
return api_key
|
||||
|
||||
|
||||
class OpenAIFunctionCaller:
|
||||
"""
|
||||
A class to interact with the OpenAI API for generating text based on a system prompt and a task.
|
||||
|
||||
Attributes:
|
||||
- system_prompt (str): The system prompt to guide the AI's response.
|
||||
- api_key (str): The API key for the OpenAI service.
|
||||
- temperature (float): The temperature parameter for the AI model, controlling randomness.
|
||||
- base_model (BaseModel): The Pydantic model to parse the response into.
|
||||
- max_tokens (int): The maximum number of tokens in the response.
|
||||
- client (OpenAI): The OpenAI client instance.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
system_prompt: str,
|
||||
base_model: BaseModel,
|
||||
api_key: str = os.getenv("OPENAI_API_KEY"),
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 5000,
|
||||
model_name: str = "gpt-4o-2024-08-06",
|
||||
):
|
||||
self.system_prompt = system_prompt
|
||||
self.api_key = api_key
|
||||
self.temperature = temperature
|
||||
self.base_model = base_model
|
||||
self.max_tokens = max_tokens
|
||||
self.model_name = model_name
|
||||
|
||||
self.client = OpenAI(api_key=self.api_key)
|
||||
|
||||
def run(self, task: str):
|
||||
"""
|
||||
Run the OpenAI model with the system prompt and task to generate a response.
|
||||
|
||||
Args:
|
||||
- task (str): The task to be completed.
|
||||
- *args: Additional positional arguments for the OpenAI API.
|
||||
- **kwargs: Additional keyword arguments for the OpenAI API.
|
||||
|
||||
Returns:
|
||||
- BaseModel: The parsed response based on the base_model.
|
||||
"""
|
||||
try:
|
||||
completion = self.client.beta.chat.completions.parse(
|
||||
model=self.model_name,
|
||||
messages=[
|
||||
{"role": "system", "content": self.system_prompt},
|
||||
{"role": "user", "content": task},
|
||||
],
|
||||
response_format=self.base_model,
|
||||
max_tokens=self.max_tokens,
|
||||
temperature=self.temperature,
|
||||
)
|
||||
|
||||
return completion.choices[0].message.parsed
|
||||
|
||||
except Exception as e:
|
||||
print(f"There was an error: {e}")
|
||||
|
||||
def check_model_support(self):
|
||||
# need to print the supported models
|
||||
for model in SUPPORTED_MODELS:
|
||||
print(model)
|
||||
|
||||
return SUPPORTED_MODELS
|
||||
|
||||
def batch_run(self, tasks: List[str]) -> List[BaseModel]:
|
||||
"""
|
||||
Batch run the OpenAI model with the system prompt and task to generate a response.
|
||||
"""
|
||||
return [self.run(task) for task in tasks]
|
||||
|
||||
def concurrent_run(self, tasks: List[str]) -> List[BaseModel]:
|
||||
"""
|
||||
Concurrent run the OpenAI model with the system prompt and task to generate a response.
|
||||
"""
|
||||
with ThreadPoolExecutor(max_workers=len(tasks)) as executor:
|
||||
return list(executor.map(self.run, tasks))
|
||||
|
||||
|
||||
# class TestModel(BaseModel):
|
||||
# name: str
|
||||
# age: int
|
||||
|
||||
# # Example usage
|
||||
# model = OpenAIFunctionCaller(
|
||||
# system_prompt="You are a helpful assistant that returns structured data about people.",
|
||||
# base_model=TestModel,
|
||||
# api_key=os.getenv("OPENAI_API_KEY"),
|
||||
# temperature=0.7,
|
||||
# max_tokens=1000
|
||||
# )
|
||||
|
||||
# # Test with a more appropriate prompt for the TestModel schema
|
||||
# response = model.run("Tell me about a person named John who is 25 years old")
|
||||
# print(response)
|
||||
Loading…
Reference in new issue