utils/ai.py

# pylint: disable=W0707
# pylint: disable=W0719

import os
import json
import tiktoken
import openai
from openai import OpenAI
import requests

from constants.cli import OPENAI_MODELS
from constants.ai import SYSTEM_PROMPT, PROMPT, API_URL


def retrieve(query, k=10, filters=None):
    """Retrieves and returns dict.

    Args:
        query (str): User query to pass in
        openai_api_key (str): openai api key. If not passed in, uses environment variable
        k (int, optional): number of results passed back. Defaults to 10.
        filters (dict, optional): Filters to apply to the query. You can filter based off
            any piece of metadata by passing in a dict of the format {metadata_name: filter_value}
            ie {"library_id": "1234"}.

            See the README for more details:
            https://github.com/fleet-ai/context/tree/main#using-fleet-contexts-rich-metadata

    Returns:
        list: List of queried results
    """

    url = f"{API_URL}/query"
    params = {
        "query": query,
        "dataset": "python_libraries",
        "n_results": k,
        "filters": filters,
    }
    return requests.post(url, json=params, timeout=120).json()


def retrieve_context(query, openai_api_key, k=10, filters=None):
    """Gets the context from our libraries vector db for a given query.

    Args:
        query (str): User input query
        k (int, optional): number of retrieved results. Defaults to 10.
    """

    # First, we query the API
    responses = retrieve(query, k=k, filters=filters)

    # Then, we build the prompt_with_context string
    prompt_with_context = ""
    for response in responses:
        prompt_with_context += f"\n\n### Context {response['metadata']['url']} ###\n{response['metadata']['text']}"
    return {"role": "user", "content": prompt_with_context}


def construct_prompt(
    messages,
    context_message,
    model="gpt-4-1106-preview",
    cite_sources=True,
    context_window=3000,
):
    """
    Constructs a RAG (Retrieval-Augmented Generation) prompt by balancing the token count of messages and context_message.
    If the total token count exceeds the maximum limit, it adjusts the token count of each to maintain a 1:1 proportion.
    It then combines both lists and returns the result.

    Parameters:
    messages (List[dict]): List of messages to be included in the prompt.
    context_message (dict): Context message to be included in the prompt.
    model (str): The model to be used for encoding, default is "gpt-4-1106-preview".

    Returns:
    List[dict]: The constructed RAG prompt.
    """
    # Get the encoding; default to cl100k_base
    if model in OPENAI_MODELS:
        encoding = tiktoken.encoding_for_model(model)
    else:
        encoding = tiktoken.get_encoding("cl100k_base")

    # 1) calculate tokens
    reserved_space = 1000
    max_messages_count = int((context_window - reserved_space) / 2)
    max_context_count = int((context_window - reserved_space) / 2)

    # 2) construct prompt
    prompts = messages.copy()
    prompts.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
    if cite_sources:
        prompts.insert(-1, {"role": "user", "content": PROMPT})

    # 3) find how many tokens each list has
    messages_token_count = len(
        encoding.encode(
            "\n".join(
                [
                    f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>"
                    for message in prompts
                ]
            )
        )
    )
    context_token_count = len(
        encoding.encode(
            f"<|im_start|>{context_message['role']}\n{context_message['content']}<|im_end|>"
        )
    )

    # 4) Balance the token count for each
    if (messages_token_count + context_token_count) > (context_window - reserved_space):
        # context has more than limit, messages has less than limit
        if (messages_token_count < max_messages_count) and (
            context_token_count > max_context_count
        ):
            max_context_count += max_messages_count - messages_token_count
        # messages has more than limit, context has less than limit
        elif (messages_token_count > max_messages_count) and (
            context_token_count < max_context_count
        ):
            max_messages_count += max_context_count - context_token_count

    # 5) Cut each list to the max count

    # Cut down messages
    while messages_token_count > max_messages_count:
        removed_encoding = encoding.encode(
            f"<|im_start|>{prompts[1]['role']}\n{prompts[1]['content']}<|im_end|>"
        )
        messages_token_count -= len(removed_encoding)
        if messages_token_count < max_messages_count:
            prompts = (
                [prompts[0]]
                + [
                    {
                        "role": prompts[1]["role"],
                        "content": encoding.decode(
                            removed_encoding[
                                : min(
                                    int(max_messages_count -
                                        messages_token_count),
                                    len(removed_encoding),
                                )
                            ]
                        )
                        .replace("<|im_start|>", "")
                        .replace("<|im_end|>", ""),
                    }
                ]
                + prompts[2:]
            )
        else:
            prompts = [prompts[0]] + prompts[2:]

    # Cut down context
    if context_token_count > max_context_count:
        # Taking a proportion of the content chars length
        reduced_chars_length = int(
            len(context_message["content"]) *
            (max_context_count / context_token_count)
        )
        context_message["content"] = context_message["content"][:reduced_chars_length]

    # 6) Combine both lists
    prompts.insert(-1, context_message)

    return prompts


def get_remote_chat_response(messages, model="gpt-4-1106-preview"):
    """
    Returns a streamed OpenAI chat response.

    Parameters:
    messages (List[dict]): List of messages to be included in the prompt.
    model (str): The model to be used for encoding, default is "gpt-4-1106-preview".

    Returns:
    str: The streamed OpenAI chat response.
    """
    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    try:
        response = client.chat.completions.create(
            model=model, messages=messages, temperature=0.2, stream=True
        )

        for chunk in response:
            current_context = chunk.choices[0].delta.content
            yield current_context

    except openai.AuthenticationError as error:
        print("401 Authentication Error:", error)
        raise Exception(
            "Invalid OPENAI_API_KEY. Please re-run with a valid key.")

    except Exception as error:
        print("Streaming Error:", error)
        raise Exception("Internal Server Error")


def get_other_chat_response(messages, model="local-model"):
    """
    Returns a streamed chat response from a local server.

    Parameters:
    messages (List[dict]): List of messages to be included in the prompt.
    model (str): The model to be used for encoding, default is "gpt-4-1106-preview".

    Returns:
    str: The streamed chat response.
    """
    try:
        if model == "local-model":
            url = "http://localhost:1234/v1/chat/completions"
            headers = {"Content-Type": "application/json"}
            data = {
                "messages": messages,
                "temperature": 0.2,
                "max_tokens": -1,
                "stream": True,
            }
            response = requests.post(
                url, headers=headers, data=json.dumps(data), stream=True, timeout=120
            )

            if response.status_code == 200:
                for chunk in response.iter_content(chunk_size=None):
                    decoded_chunk = chunk.decode()
                    if (
                        "data:" in decoded_chunk
                        and decoded_chunk.split("data:")[1].strip()
                    ):  # Check if the chunk is not empty
                        try:
                            chunk_dict = json.loads(
                                decoded_chunk.split("data:")[1].strip()
                            )
                            yield chunk_dict["choices"][0]["delta"].get("content", "")
                        except json.JSONDecodeError:
                            pass
            else:
                print(f"Error: {response.status_code}, {response.text}")
                raise Exception("Internal Server Error")
        else:
            if not os.environ.get("OPENROUTER_API_KEY"):
                raise Exception(
                    f"For non-OpenAI models, like {model}, set your OPENROUTER_API_KEY."
                )

            response = requests.post(
                url="https://openrouter.ai/api/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {os.environ.get('OPENROUTER_API_KEY')}",
                    "HTTP-Referer": os.environ.get(
                        "OPENROUTER_APP_URL", "https://fleet.so/context"
                    ),
                    "X-Title": os.environ.get("OPENROUTER_APP_TITLE", "Fleet Context"),
                    "Content-Type": "application/json",
                },
                data=json.dumps(
                    {"model": model, "messages": messages, "stream": True}),
                stream=True,
                timeout=120,
            )
            if response.status_code == 200:
                for chunk in response.iter_lines():
                    decoded_chunk = chunk.decode("utf-8")
                    if (
                        "data:" in decoded_chunk
                        and decoded_chunk.split("data:")[1].strip()
                    ):  # Check if the chunk is not empty
                        try:
                            chunk_dict = json.loads(
                                decoded_chunk.split("data:")[1].strip()
                            )
                            yield chunk_dict["choices"][0]["delta"].get("content", "")
                        except json.JSONDecodeError:
                            pass
            else:
                print(f"Error: {response.status_code}, {response.text}")
                raise Exception("Internal Server Error")

    except requests.exceptions.RequestException as error:
        print("Request Error:", error)
        raise Exception(
            "Invalid request. Please check your request parameters.")