Create ai.py

ai.py

@@ -0,0 +1,291 @@
+# 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.")