diff --git a/easycompletion/model.py b/easycompletion/model.py index dc623af..e778d82 100644 --- a/easycompletion/model.py +++ b/easycompletion/model.py @@ -377,6 +377,7 @@ def function_completion( api_key=None, debug=DEBUG, temperature=0.0, + ): """ Send text and a list of functions to the model and return optional text and a function call. @@ -574,5 +575,40 @@ def function_completion( "arguments": arguments, "usage": usage, "finish_reason": finish_reason, + "model_used": model, # Include the model used in the response "error": None, } + + +def status_info(api_key=None, model=None, debug=DEBUG): + """ + Get status information about the API key and model. + + Parameters: + api_key (str, optional): OpenAI API key. If not provided, it uses the one defined in constants.py. + model (str, optional): The model to use. Default is the TEXT_MODEL defined in constants.py. + + Returns: + dict: A dictionary containing status information. + + Example: + >>> status_info(api_key='your_openai_api_key', model='gpt-3.5-turbo') + """ + # Validate the API key + if not validate_api_key(api_key): + return {"error": "Invalid OpenAI API key"} + + openai.api_key = api_key + + if model is None: + model = TEXT_MODEL + + # Get the model information + model_info = openai.Model.retrieve(model) + model_status = model_info.status + + return { + "api_key": api_key, + "model": model, + "model_status": model_status, + } diff --git a/easycompletion/prompt.py b/easycompletion/prompt.py deleted file mode 100644 index 8254b9e..0000000 --- a/easycompletion/prompt.py +++ /dev/null @@ -1,226 +0,0 @@ -import re -import tiktoken - -from .constants import TEXT_MODEL, DEFAULT_CHUNK_LENGTH, DEBUG -from .logger import log - - -def trim_prompt( - text, - max_tokens=DEFAULT_CHUNK_LENGTH, - model=TEXT_MODEL, - preserve_top=True, - debug=DEBUG, -): - """ - Trim the given text to a maximum number of tokens. - - Args: - text: Input text which needs to be trimmed. - max_tokens: Maximum number of tokens allowed in the trimmed text. - Default value is taken from the constants. - model: The model to use for tokenization. - preserve_top: If True, the function will keep the first 'max_tokens' tokens, - if False, it will keep the last 'max_tokens' tokens. - - Returns: - Trimmed text that fits within the specified token limit. - - Example: - trim_prompt("This is a test.", 3, preserve_top=True) - Output: "This is" - """ - # Encoding the text into tokens. - encoding = tiktoken.encoding_for_model(model) - tokens = encoding.encode(text) - if len(tokens) <= max_tokens: - return text # If text is already within limit, return as is. - - log(f"Trimming prompt, token len is {str(len(tokens))}", type="warning", log=debug) - - # If 'preserve_top' is True, keep the first 'max_tokens' tokens. - # Otherwise, keep the last 'max_tokens' tokens. - return encoding.decode( - tokens[:max_tokens] if preserve_top else tokens[-max_tokens:] - ) - - -def chunk_prompt(prompt, chunk_length=DEFAULT_CHUNK_LENGTH, debug=DEBUG): - """ - Split the given prompt into chunks where each chunk has a maximum number of tokens. - - Args: - prompt: Input text that needs to be split. - chunk_length: Maximum number of tokens allowed per chunk. - Default value is taken from the constants. - - Returns: - A list of string chunks where each chunk is within the specified token limit. - - Example: - chunk_prompt("This is a test. I am writing a function.", 4) - Output: ['This is', 'a test.', 'I am', 'writing a', 'function.'] - """ - if count_tokens(prompt) <= chunk_length: - return [prompt] - - # Splitting the prompt into sentences using regular expressions. - sentences = re.split(r"(?<=[.!?])\s+", prompt) - current_chunk = "" - prompt_chunks = [] - - # For each sentence in the input text. - for sentence in sentences: - # If adding a new sentence doesn't exceed the token limit, add it to the current chunk. - if count_tokens(current_chunk + sentence + " ") <= chunk_length: - current_chunk += sentence + " " - else: - # If adding a new sentence exceeds the token limit, add the current chunk to the list. - # Then, start a new chunk with the current sentence. - prompt_chunks.append(current_chunk.strip()) - current_chunk = sentence + " " - - # If there's any sentence left after looping through all sentences, add it to the list. - if current_chunk: - prompt_chunks.append(current_chunk.strip()) - - log( - f"Chunked prompt into {str(len(prompt_chunks))} chunks", - type="warning", - log=debug, - ) - - return prompt_chunks - - -def count_tokens(prompt: str, model=TEXT_MODEL) -> int: - """ - Count the number of tokens in a string. - - Args: - prompt: The string to be tokenized. - model: The model to use for tokenization. - - Returns: - The number of tokens in the input string. - - Example: - count_tokens("This is a test.") - Output: 5 - """ - if not isinstance(prompt, str): - prompt = str(prompt) - - encoding = tiktoken.encoding_for_model(model) - length = len( - encoding.encode(prompt) - ) # Encoding the text into tokens and counting the number of tokens. - return length - - -def get_tokens(prompt: str, model=TEXT_MODEL) -> list: - """ - Returns a list of tokens in a string. - - Args: - prompt: The string to be tokenized. - model: The model to use for tokenization. - - Returns: - A list of tokens in the input string. - - Example: - get_tokens("This is a test.") - Output: [This, is, a, test, .] - """ - encoding = tiktoken.encoding_for_model(model) - return encoding.encode( - prompt - ) # Encoding the text into tokens and returning the list of tokens. - - -def compose_prompt(prompt_template, parameters, debug=DEBUG): - """ - Composes a prompt using a template and parameters. - Parameter keys are enclosed in double curly brackets and replaced with parameter values. - - Args: - prompt_template: A template string that contains placeholders for the parameters. - parameters: A dictionary containing key-value pairs to replace the placeholders. - - Returns: - A string where all placeholders have been replaced with actual values from the parameters. - - Example: - compose_prompt("Hello {{name}}!", {"name": "John"}) - Output: "Hello John!" - """ - prompt = prompt_template # Initial prompt template. - - # Replacing placeholders in the template with the actual values from the parameters. - for key, value in parameters.items(): - # check if "{{" + key + "}}" is in prompt - # if not, continue - if "{{" + key + "}}" not in prompt: - continue - try: - if isinstance(value, str): - prompt = prompt.replace("{{" + key + "}}", value) - elif isinstance(value, int): - prompt = prompt.replace("{{" + key + "}}", str(value)) - elif isinstance(value, dict): - for k, v in value.items(): - prompt = prompt.replace("{{" + key + "}}", k + "::" + v) - elif isinstance(value, list): - for item in value: - prompt = prompt.replace("{{" + key + "}}", item + "\n") - elif value is None: - prompt = prompt.replace("{{" + key + "}}", "None") - else: - raise Exception(f"ERROR PARSING:\n{key}\n{value}") - except: - raise Exception(f"ERROR PARSING:\n{key}\n{value}") - - log(f"Composed prompt:\n{prompt}", log=debug) - - return prompt - - -def compose_function(name, description, properties, required_properties, debug=DEBUG): - """ - Composes a function object for function calling. - - Parameters: - name (str): The name of the function. - description (str): Description of the function. - properties (dict): Dictionary of property objects. - required_properties (list): List of property names that are required. - - Returns: - A dictionary representing a function. - - Usage: - summarization_function = compose_function( - name="summarize_text", - description="Summarize the text. Include the topic, subtopics.", - properties={ - "summary": { - "type": "string", - "description": "Detailed summary of the text.", - }, - }, - required_properties=["summary"], - ) - """ - function = { - "name": name, - "description": description, - "parameters": { - "type": "object", - "properties": properties, - "required": required_properties, - }, - } - log(f"Function:\n{str(function)}", type="info", log=debug) - return function - diff --git a/easycompletion/tests/model.py b/easycompletion/tests/model.py index eb41cfc..b1f7bcd 100644 --- a/easycompletion/tests/model.py +++ b/easycompletion/tests/model.py @@ -1,3 +1,4 @@ +import openai from easycompletion.model import ( chat_completion, parse_arguments, @@ -5,6 +6,10 @@ text_completion, ) +# Check if the OpenAI API key is set +if not openai.api_key: + raise ValueError("OpenAI API key is missing. Set your API key using 'openai.api_key = '.") + def test_parse_arguments(): test_input = '{"key1": "value1", "key2": 2}' diff --git a/easycompletion/tests/prompt.py b/easycompletion/tests/prompt.py deleted file mode 100644 index 95a2380..0000000 --- a/easycompletion/tests/prompt.py +++ /dev/null @@ -1,69 +0,0 @@ -from easycompletion.model import parse_arguments -from easycompletion.prompt import ( - compose_prompt, - trim_prompt, - chunk_prompt, - count_tokens, - get_tokens, - compose_function, -) - - -def test_chunk_prompt(): - test_text = "Write a song about AI" - chunks = chunk_prompt(test_text, chunk_length=2) - assert len(chunks) == 2, "Test chunk_prompt failed" - - -def test_trim_prompt_and_get_tokens(): - test_text = "Write a song about AI" - trimmed = trim_prompt(test_text, max_tokens=2) - count = count_tokens(trimmed) - assert count == 2, "Test trim_prompt failed" - - tokens = get_tokens(test_text) - assert len(tokens) == 5, "Test get_tokens failed" - - -def test_parse_arguments(): - test_input = '{"key1": "value1", "key2": 2}' - expected_output = {"key1": "value1", "key2": 2} - assert parse_arguments(test_input) == expected_output, "Test parse_arguments failed" - - -def test_compose_prompt(): - test_prompt = "I am a {{object}}" - test_dict = {"object": "towel"} - prompt = compose_prompt(test_prompt, test_dict) - assert prompt == "I am a towel", "Test compose_prompt failed" - - -def test_compose_function(): - summarization_function = { - "name": "summarize_text", - "description": "Summarize the text. Include the topic, subtopics.", - "parameters": { - "type": "object", - "properties": { - "summary": { - "type": "string", - "description": "Detailed summary of the text.", - }, - }, - "required": ["summary"], - }, - } - composed_summarization_function = compose_function( - name="summarize_text", - description="Summarize the text. Include the topic, subtopics.", - properties={ - "summary": { - "type": "string", - "description": "Detailed summary of the text.", - }, - }, - required_properties=["summary"], - ) - assert ( - composed_summarization_function == summarization_function - ), "Test compose_function failed"