AI vector embeddings of Roblox informational documents for semantic search
import requests  # for fetching the docs
import tiktoken  # for counting tokens
import re  # for cutting metadata out of doc headers
from json import dumps as json_dumps  # For debug printing

import write
import config

# OpenAI's best embeddings as of Sept 2023
EMBEDDING_MODEL = "text-embedding-ada-002"
def count_tokens(text: str, model: str = EMBEDDING_MODEL) -> int:
    """Return the number of tokens in a string."""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

def fetch_tree_data():
    data_res = requests.get(
        "https://api.github.com/repos/MaximumADHD/Roblox-Client-Tracker/git/trees/roblox?recursive=true", headers=config.GH_REQ_HEADERS)
    data_res.raise_for_status()

    data = data_res.json()
    return data

SELFCLOSING_HTML_PATTERN = re.compile(r"<([^/]\w*)[^>]*/>", re.DOTALL)
HTML_PATTERN = re.compile(r"<([^/]\w*)[^>]*>.*?</\1>", re.DOTALL)
def prepare_document_for_ingest(document):
    # Replace escaped characters with their actual characters such as &mdash;
    document = document.replace("&mdash;", "—")
    document = document.replace("&ndash;", "–")
    document = document.replace("&nbsp;", " ")
    document = document.replace("&quot;", "\"")
    document = document.replace("&apos;", "'")
    document = document.replace("&lt;", "<")
    document = document.replace("&gt;", ">")
    document = document.replace("&amp;", "&")
    # Replace <li></li> with - and <ul></ul> with \n
    document = document.replace("<li>", "- ")
    document = document.replace("</li>", "\n")
    document = document.replace("<ul>", "\n")
    document = document.replace("</ul>", "\n")
    # Remove html elements like <img /> and <video></video>
    document = HTML_PATTERN.sub("", document)
    document = SELFCLOSING_HTML_PATTERN.sub("", document)

    return document


def is_api_identifier(x):
    return re.match(r"^@\w+/", x)


def replace_xml_tags(item):
    if isinstance(item, list):
        for i in range(len(item)):
            if isinstance(item[i], str):
                bolded = re.sub(r"<strong>(.*?)</strong>", r"**\1**", item[i])
                italicized = re.sub(r"<em>(.*?)</em>", r"*\1*", bolded)
                monospaced = re.sub(r"<code>(.*?)</code>", r"`\1`", italicized)
                item[i] = monospaced
            elif isinstance(item[i], list):
                replace_xml_tags(item[i])
            elif isinstance(item[i], dict):
                replace_xml_tags(item[i])

    elif isinstance(item, dict):
        for prop in item:
            if isinstance(item[prop], str):
                bolded = re.sub(r"<strong>(.*?)</strong>",
                                r"**\1**", item[prop])
                italicized = re.sub(r"<em>(.*?)</em>", r"*\1*", bolded)
                monospaced = re.sub(r"<code>(.*?)</code>", r"`\1`", italicized)
                item[prop] = monospaced
            elif isinstance(item[prop], list):
                replace_xml_tags(item[prop])
            elif isinstance(item[prop], dict):
                replace_xml_tags(item[prop])


def replace_api_identifiers(content, object):
    if isinstance(object, list):
        for key in range(len(object)):
            item = object[key]
            if isinstance(item, str) and is_api_identifier(item):
                ref_obj = content[item]
                if 'documentation' in ref_obj and len(ref_obj) == 1:
                    # If the ref object only has a documentation, replace the identifier with the documentation
                    object[key] = ref_obj['documentation']
                else:
                    object[key] = ref_obj
            else:
                replace_api_identifiers(content, item)

    elif isinstance(object, dict):
        for key in object:
            item = object[key]
            if isinstance(item, str) and is_api_identifier(item):
                ref_obj = content[item]
                if 'documentation' in ref_obj and len(ref_obj) == 1:
                    # If the ref object only has a documentation, replace the identifier with the documentation
                    object[key] = ref_obj['documentation']
                else:
                    object[key] = ref_obj
            else:
                replace_api_identifiers(
                    content, item)

def get_api_dump():
    res = requests.get(
        "https://raw.githubusercontent.com/MaximumADHD/Roblox-Client-Tracker/roblox/API-Dump.json"
    )
    res.raise_for_status()
    dump = res.json()
    return dump

def get_api_docstrings():
    res = requests.get(
        "https://raw.githubusercontent.com/MaximumADHD/Roblox-Client-Tracker/roblox/api-docs/mini/en-us.json"
    )
    res.raise_for_status()

    api_docstrings = {}
    used_identifiers = []

    content = res.json()

    # Replace xml tags with markdown
    replace_xml_tags(content)

    # Replace identifiers with their object
    replace_api_identifiers(content, content)

    # Build the API reference
    for key in content:
        if len(content[key]) == 0:
            continue
        if re.match(r"\.Connect$", key) or re.match(r"\.Wait$", key):
            continue
        api_docstrings[key] = content[key]

    return api_docstrings

def createEnumReference(enumObj, api_docstrings):
    enumKey = '@roblox/enum/' + enumObj["Name"]
    if not enumKey in api_docstrings:
        # print("Couldn't find any docsstring for Enum." + enumObj["Name"])
        return ""

    api_docstring = api_docstrings[enumKey]

    desc = api_docstring["documentation"] or ""

    code_sample = ""
    if 'code_sample' in api_docstring and api_docstring['code_sample'] != "":
        code_sample = f"\n```Lua\n{api_docstring['code_sample']}\n```"

    itemDescs = []
    if 'keys' in api_docstring and len(api_docstring['keys']) > 0:
        for itemName in api_docstring['keys']:
            itemDocstring = api_docstring['keys'][itemName]
            itemDesc = itemDocstring['documentation'] or ""
            itemCodeSample = itemDocstring['code_sample'] or ""
            if itemCodeSample != "":
                itemDesc += f"\n```Lua\n{itemCodeSample}\n```"
            if itemDesc != "":
                itemDescs.append(f"- {itemName}: {itemDesc}")
            else:
                itemDescs.append(f"- {itemName}")
    itemDescs = "\n".join(itemDescs)

    referenceDoc = prepare_document_for_ingest(f"""# {enumObj["Name"]}
{desc}{code_sample}

## Items
{itemDescs}
""")
    # write.write_text(referenceDoc, f"build/Enum_{enumObj['Name']}.md")
    return referenceDoc

def createClassReference(classObj, api_docstrings):
    classKey = "@roblox/globaltype/" + classObj["Name"]
    if not classKey in api_docstrings:
        # print("Couldn't find any docsstring for " + classObj["Name"])
        return ""

    api_docstring = api_docstrings[classKey]

    desc = api_docstring["documentation"] or ""

    code_sample = ""
    if 'code_sample' in api_docstring and api_docstring['code_sample'] != "":
        code_sample = f"\n```Lua\n{api_docstring['code_sample']}\n```"

    propertyDescs = []
    methodDescs = []
    eventDescs = []
    callbackDescs = []

    for member in classObj["Members"]:
        if "Tags" in member and "Deprecated" in member['Tags']:
            continue
        memberKey = classKey + "." + member["Name"]
        # print(memberKey)

        if member["MemberType"] == "Property":
            memberBaseDesc = f"### {member['Name']} ({member['ValueType']['Name']})"
            if not memberKey in api_docstrings:
                # print("Couldn't find any docsstring for property " + classObj['Name'] + "." + member["Name"])
                propertyDescs.append(memberBaseDesc)
                continue
            memberDocstring = api_docstrings[memberKey]
            memberDesc = memberDocstring['documentation'] or ""
            memberCodeSample = memberDocstring['code_sample'] or ""
            if memberCodeSample != "":
                memberDesc += f"\n```Lua\n{memberCodeSample}\n```"
            if memberDesc != "":
                propertyDescs.append(memberBaseDesc + "\n- " + memberDesc)
            else:
                propertyDescs.append(memberBaseDesc)

        elif member["MemberType"] == "Function":
            if not memberKey in api_docstrings:
                # print("Couldn't find any docsstring for method " + classObj['Name'] + "." + member["Name"])
                # Use fallback info from dump
                memberParamsDesc = []
                for param in member["Parameters"]:
                    memberParamsDesc.append(f"- {param['Name']}: {param['Type']['Name']}")
                memberParamsDesc = "\n".join(memberParamsDesc)

                memberReturnsDesc = f"- {member['ReturnType']['Name']}"

                memberDesc = f"### {member['Name']} ({member['ReturnType']['Name']})"
                if memberParamsDesc != "":
                    memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
                if memberReturnsDesc != "":
                    memberDesc += f"\n\n**Returns**\n{memberReturnsDesc}"
                memberDesc += "\n"

                methodDescs.append(memberDesc)
                continue

            memberDocstring = api_docstrings[memberKey]
            memberParamsDesc = []
            for i in range(len(memberDocstring['params'])):
                param = memberDocstring['params'][i]
                if param['name'] == "self":
                    continue
                dumpParam = member["Parameters"][i-1]

                memberParamsDesc.append(f"- {param['name']}: {dumpParam['Type']['Name']}\n{param['documentation']}")
            memberParamsDesc = "\n".join(memberParamsDesc)

            memberReturnsDesc = []
            if len(memberDocstring['returns']) == 0:
                memberReturnsDesc.append(f"- {member['ReturnType']['Name']}")
            else:
                for ret in memberDocstring['returns']:
                    memberReturnsDesc.append(f"- {member['ReturnType']['Name']}: {ret['documentation'] if isinstance(ret, dict) else ret}")
            memberReturnsDesc = "\n".join(memberReturnsDesc)

            memberSummary = memberDocstring['documentation'] or ""
            memberCodeSample = memberDocstring['code_sample'] or ""
            if memberCodeSample != "":
                memberCodeSample = f"\n```Lua\n{memberCodeSample}\n```"

            memberDesc = f"### {member['Name']} ({member['ReturnType']['Name']})"
            if memberSummary != "":
                memberDesc += f"\n{memberSummary}"
            if memberParamsDesc != "":
                memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
            if memberReturnsDesc != "":
                memberDesc += f"\n\n**Returns**\n{memberReturnsDesc}"
            if memberCodeSample != "":
                memberDesc += f"\n\n**Example**\n{memberCodeSample}"
            memberDesc += "\n"

            methodDescs.append(memberDesc)

        elif member["MemberType"] == "Callback":
            if not memberKey in api_docstrings:
                # print("Couldn't find any docsstring for method " + classObj['Name'] + "." + member["Name"])
                # Use fallback info from dump
                memberParamsDesc = []
                for param in member["Parameters"]:
                    memberParamsDesc.append(f"- {param['Name']}: {param['Type']['Name']}")
                memberParamsDesc = "\n".join(memberParamsDesc)

                memberReturnsDesc = f"- {member['ReturnType']['Name']}"

                memberDesc = f"### {member['Name']} ({member['ReturnType']['Name']})"
                if memberParamsDesc != "":
                    memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
                if memberReturnsDesc != "":
                    memberDesc += f"\n\n**Returns**\n{memberReturnsDesc}"
                memberDesc += "\n"

                callbackDescs.append(memberDesc)
                continue

            memberDocstring = api_docstrings[memberKey]
            memberParamsDesc = []
            for i in range(len(memberDocstring['params'])):
                param = memberDocstring['params'][i]
                dumpParam = member["Parameters"][i]

                memberParamsDesc.append(f"- {param['name']}: {dumpParam['Type']['Name']}\n{param['documentation']}")
            memberParamsDesc = "\n".join(memberParamsDesc)

            memberReturnsDesc = []
            if len(memberDocstring['returns']) == 0:
                memberReturnsDesc.append(f"- {member['ReturnType']['Name']}")
            else:
                for ret in memberDocstring['returns']:
                    memberReturnsDesc.append(f"- {member['ReturnType']['Name']}: {ret['documentation'] if isinstance(ret, dict) else ret}")
            memberReturnsDesc = "\n".join(memberReturnsDesc)

            memberSummary = memberDocstring['documentation'] or ""
            memberCodeSample = memberDocstring['code_sample'] or ""
            if memberCodeSample != "":
                memberCodeSample = f"\n```Lua\n{memberCodeSample}\n```"

            memberDesc = f"### {member['Name']} ({member['ReturnType']['Name']})"
            if memberSummary != "":
                memberDesc += f"\n{memberSummary}"
            if memberParamsDesc != "":
                memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
            if memberReturnsDesc != "":
                memberDesc += f"\n\n**Returns**\n{memberReturnsDesc}"
            if memberCodeSample != "":
                memberDesc += f"\n\n**Example**\n{memberCodeSample}"
            memberDesc += "\n"

            callbackDescs.append(memberDesc)

        elif member["MemberType"] == "Event":
            if not memberKey in api_docstrings:
                # print("Couldn't find any docsstring for event " + classObj['Name'] + "." + member["Name"])
                # Use fallback info from dump
                memberParamsDesc = []
                for param in member["Parameters"]:
                    memberParamsDesc.append(f"- {param['Name']}: {param['Type']['Name']}")
                memberParamsDesc = "\n".join(memberParamsDesc)

                memberDesc = f"### {member['Name']}"
                if memberParamsDesc != "":
                    memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
                memberDesc += "\n"

                eventDescs.append(memberDesc)
                continue

            memberDocstring = api_docstrings[memberKey]
            memberParamsDesc = []
            for param in member["Parameters"]:
                memberParamsDesc.append(f"- {param['Name']}: {param['Type']['Name']}")
            memberParamsDesc = "\n".join(memberParamsDesc)

            memberSummary = memberDocstring['documentation'] or ""
            memberCodeSample = memberDocstring['code_sample'] or ""
            if memberCodeSample != "":
                memberCodeSample = f"\n```Lua\n{memberCodeSample}\n```"

            memberDesc = f"### {member['Name']}"
            if memberSummary != "":
                memberDesc += f"\n{memberSummary}"
            if memberParamsDesc != "":
                memberDesc += f"\n\n**Parameters**\n{memberParamsDesc}"
            if memberCodeSample != "":
                memberDesc += f"\n\n**Example**\n{memberCodeSample}"
            memberDesc += "\n"

            eventDescs.append(memberDesc)
        else:
            print("Unhandled member type: " + classObj['Name'] + '.' + member['Name'] + '.' + member["MemberType"])

    propertyDescs = "\n\n".join(propertyDescs)
    methodDescs = "\n\n".join(methodDescs)
    eventDescs = "\n\n".join(eventDescs)
    callableDescs = "\n\n".join(callbackDescs)

    if desc == '' and code_sample == '' and propertyDescs == '' and methodDescs == '' and eventDescs == '' and callableDescs == '':
        return ""

    referenceDoc = f"# {classObj['Name']}\n{desc}{code_sample}"
    if propertyDescs != "":
        referenceDoc += f"\n\n## Properties\n\n{propertyDescs}"
    if methodDescs != "":
        referenceDoc += f"\n\n## Methods\n\n{methodDescs}"
    if eventDescs != "":
        referenceDoc += f"\n\n## Events\n\n{eventDescs}"
    if callableDescs != "":
        referenceDoc += f"\n\n## Callbacks\n\n{callableDescs}"
    referenceDoc += "\n"

    referenceDoc = prepare_document_for_ingest(referenceDoc)
    tokens = count_tokens(referenceDoc)
    if tokens > 8100:
        print("Skipping " + classObj['Name'] + " because it has too many tokens (" + str(tokens) + ")")
        return ""

    # write.write_text(referenceDoc, f"build/Class_{classObj['Name']}.md")
    return referenceDoc

def get_reference():
    api_dump = get_api_dump()
    api_docstrings = get_api_docstrings()

    api_reference = {}

    for enumObj in api_dump["Enums"]:
        enum_reference = createEnumReference(enumObj, api_docstrings)
        if enum_reference != "":
            api_reference['Enum.' + enumObj["Name"]] = enum_reference

    for classObj in api_dump["Classes"]:
        class_reference = createClassReference(classObj, api_docstrings)
        if class_reference != "":
            api_reference[classObj["Name"]] = class_reference

    print("Found " + str(len(api_reference)) + " api references")
    # print(json_dumps(api_reference, indent=2))
    return api_reference

def get_sha():
    data = fetch_tree_data()
    write.write_text(data["sha"], "build/api-source-commit.txt")
    return data["sha"]

indexer/api_reference.py
boatbomber/Roblox-Docs-AI-Search
import os
from dotenv import load_dotenv
load_dotenv()

# GitHub API token
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN")

# OpenAI API key for embedding model
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

GH_REQ_HEADERS = {
	"Authorization": f"Bearer {GITHUB_TOKEN}",
	"Accept": "application/vnd.github+json",
	"X-GitHub-Api-Version": "2022-11-28",
}

ALLOWLISTED_DOCUMENT_PATHS = {
	'en-us/', # All english documentation
}

indexer/config.py
boatbomber/Roblox-Docs-AI-Search
import os  # for managing paths
import requests  # for fetching the docs
import base64  # for decoding the doc file contents
import yaml  # for reading the metadata of doc files
import re  # for cutting metadata out of doc headers

import write
import config

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader


METADATA_PATTERN = re.compile("---(.+?)---", re.DOTALL)
SELFCLOSING_HTML_PATTERN = re.compile(r"<([^/]\w*)[^>]*/>", re.DOTALL)
HTML_PATTERN = re.compile(r"<([^/]\w*)[^>]*>.*?</\1>", re.DOTALL)
EXTRA_NEWLINES_PATTERN = re.compile(r"\n[\n\s]+", re.DOTALL)
INLINE_LINKS_PATTERN = re.compile(r"\[(.*?)\]\(.*?\)", re.DOTALL)
API_LINKING_WITH_NAME_PATTERN = re.compile(
    r"`[A-Z]\w*?\.[^\n]*?\|(.*?)`", re.DOTALL)
API_LINKING_PATTERN = re.compile(r"`[A-Z]\w*?\.([^\n]*?)`", re.DOTALL)


def fetch_tree_data():
    data_res = requests.get(
        "https://api.github.com/repos/Roblox/creator-docs/git/trees/main:content?recursive=true", headers=config.GH_REQ_HEADERS)
    data_res.raise_for_status()

    data = data_res.json()
    return data


def is_path_allowed(path):
    for allowed_path in config.ALLOWLISTED_DOCUMENT_PATHS:
        if path.startswith(allowed_path):
            return True

    return False


def get_documentation_filepaths(tree):
    documentation_filepaths = list(map(
        lambda item: "content/" + item["path"],
        filter(lambda item: item["type"] == "blob" and item["path"].endswith(
            ".md") and is_path_allowed(item["path"]), tree)
    ))

    return documentation_filepaths


def download_documents(documentation_filepaths):
    documents = {}

    total_files = len(documentation_filepaths)
    current_file = 0

    for filepath in documentation_filepaths:
        current_file += 1
        print(f"{current_file}/{total_files} Fetching " + filepath)

        res = requests.get(
            f"https://api.github.com/repos/Roblox/creator-docs/contents/{filepath}?ref=main", headers=config.GH_REQ_HEADERS)
        res.raise_for_status()

        content_encoded = res.json()["content"]
        content = base64.b64decode(
            content_encoded.encode("utf-8")).decode("utf-8")

        documents[filepath] = content

    return documents


def get_document_metadata(filepath, document):
    metadata_match = METADATA_PATTERN.match(document)

    metadata = {}
    metadata["path"] = filepath

    if metadata_match == None:
        print("> Failed to get metadata for " + filepath)
        # Use filename as title
        metadata["title"] = os.path.basename(filepath).replace(".md", "")
        metadata["description"] = ""
    else:
        metadata_str = metadata_match.group(1)
        metadata_dict = yaml.load(metadata_str, Loader=Loader)
        metadata["title"] = metadata_dict["title"] or os.path.basename(
            filepath).replace(".md", "")
        metadata["description"] = metadata_dict["description"] or ""

    return metadata


def prepare_document_for_ingest(document):
    # Remove metadata header (---...---)
    document = METADATA_PATTERN.sub("", document)
    # Remove html elements like <img /> and <video></video>
    document = HTML_PATTERN.sub("", document)
    document = SELFCLOSING_HTML_PATTERN.sub("", document)
    # Remove links but keep the text [text](url)
    document = INLINE_LINKS_PATTERN.sub(r"\1", document)
    # Remove extra newlines
    document = EXTRA_NEWLINES_PATTERN.sub("\n\n", document)
    # Remove Roblox's custom API linking
    document = API_LINKING_WITH_NAME_PATTERN.sub(r"`\1`", document)
    document = API_LINKING_PATTERN.sub(r"`\1`", document)

    return document


def get_document_sections(content, metadata):
    sections = {}
    split_content = content.split("\n## ")
    for section in split_content:
        sectionlines = section.split("\n", 1)
        if len(sectionlines) == 0:
            continue
        elif len(sectionlines) == 1:
            # With only one line, we just use it as content and use the title as the header
            sections[metadata["title"]] = sectionlines[0]
            continue
        else:
            # We have multiple lines, so line one is likely the header and the rest is the content
            section_header = sectionlines[0].strip()
            if section_header == "":
                section_header = metadata["title"]
            section_content = "\n".join(sectionlines[1:])
            sections[section_header] = section_content
            continue
    return sections


def get_documents():
    data = fetch_tree_data()

    tree = data["tree"]
    truncated = data["truncated"]

    # TODO: Do a recursive scan if truncated
    if truncated:
        print("Cannot handle truncated paths yet")
        exit(1)

    documentation_filepaths = get_documentation_filepaths(tree)
    print(len(documentation_filepaths), "documents found")

    documents = download_documents(documentation_filepaths)
    return documents


def get_sha():
    data = fetch_tree_data()
    write.write_text(data["sha"], "build/docs-source-commit.txt")
    return data["sha"]

indexer/creator_docs.py
boatbomber/Roblox-Docs-AI-Search
import os  # for creating the build directory
import openai  # for generating embeddings
import pandas as pd  # for DataFrames to store article sections and embeddings
import tiktoken  # for counting tokens
import json
from datetime import date  # for writing the date to the summary file

import write
import config
import creator_docs
import api_reference

# OpenAI's best embeddings as of Sept 2023
EMBEDDING_MODEL = "text-embedding-ada-002"

openai.api_key = config.OPENAI_API_KEY


def count_tokens(text: str, model: str = EMBEDDING_MODEL) -> int:
    """Return the number of tokens in a string."""
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))


def create_dataframe(documents, reference):
    data = {
        'type': [],
        'title': [],
        'content': [],
        'embedding': [],
    }

    count_files = len(documents) + len(reference)
    current_file = 0

    # Handle the documents
    for filepath in documents:
        document = documents[filepath]

        current_file += 1

        # First, we get the metadata
        metadata = creator_docs.get_document_metadata(
            filepath=filepath, document=document)

        # Then, we process it for ingestion
        content = creator_docs.prepare_document_for_ingest(document=document)

        # Then, we break it into sections using ## headers, so that we get a dict of header content -> section content
        sections = creator_docs.get_document_sections(content=content, metadata=metadata)

        print("[" + str(current_file) + "/" + str(count_files) + "] Generating data for the " +
              str(len(sections)) + " sections in " + metadata["title"])

        # Then we get the embeddings for each section
        embed_batch = []
        for header in sections:
            section_content = sections[header]

            data["type"].append("creator-docs")
            data["title"].append(metadata["title"] + " / " + header)
            data["content"].append(section_content)

            embeddable_content = "# " + \
                metadata["title"] + "\n## " + metadata["description"] + \
                "\n### " + header + "\n" + section_content
            if count_tokens(embeddable_content) > 8100:
                print("  Skipping " + header +
                      " content embedding because it has too many tokens")
                embed_batch.append(
                    "# " + metadata["title"] + "\n## " + metadata["description"] + "\n### " + header)
                continue

            embed_batch.append(embeddable_content)

        try:
            response = openai.Embedding.create(
                model=EMBEDDING_MODEL, input=embed_batch)
            data["embedding"].extend(
                list(map(lambda data: data['embedding'], response['data'])))
        except Exception as e:
            print("Failed to get embedding for " + metadata["title"], e)
            data["embedding"].extend([None] * len(embed_batch))

    # Handle the API reference
    reference_embed_batch, reference_batch_tokens = [], 0
    for key in reference:
        embeddable_content = reference[key]

        current_file += 1
        print("[" + str(current_file) + "/" + str(count_files) + "] Generating data for " + key)

        data["type"].append("api-reference")
        data["title"].append(key)
        data["content"].append(embeddable_content)

        embed_tokens = count_tokens(embeddable_content)
        if embed_tokens > 8100:
            print("  Skipping " + key +
                  " content embedding because it has too many tokens")
            data["embedding"].append(None)
            continue

        if reference_batch_tokens + embed_tokens > 8100:
            try:
                response = openai.Embedding.create(
                    model=EMBEDDING_MODEL, input=reference_embed_batch)
                data["embedding"].extend(
                    list(map(lambda data: data['embedding'], response['data'])))
            except Exception as e:
                print("Failed to get embedding for api-reference", e)
                data["embedding"].extend([None] * len(reference_embed_batch))

            reference_embed_batch, reference_batch_tokens = [], 0

        reference_batch_tokens += embed_tokens
        reference_embed_batch.append(embeddable_content)

    # Get the last batch of reference embeddings
    try:
        response = openai.Embedding.create(
            model=EMBEDDING_MODEL, input=reference_embed_batch)
        data["embedding"].extend(
            list(map(lambda data: data['embedding'], response['data'])))
    except Exception as e:
        print("Failed to get embedding for api-reference", e)
        data["embedding"].extend([None] * len(reference_embed_batch))

    # Build and return the DataFrame
    return pd.DataFrame(data)


if __name__ == "__main__":
    if not os.path.exists('build'):
        os.makedirs('build')

    documents = creator_docs.get_documents()
    reference = api_reference.get_reference()

    embeddings_dataframe = create_dataframe(
        documents=documents, reference=reference)

    embeddings_dataframe.to_json(
        "build/docs-embeddings.json", orient="records")

    write.write_text(f"""# Roblox Documentation Index

Generated on {date.today()} from:
- https://github.com/Roblox/creator-docs @ {creator_docs.get_sha()[:7]}
- https://github.com/MaximumADHD/Roblox-Client-Tracker/tree/roblox/api-docs @ {api_reference.get_sha()[:7]}

## Embeddings

With those files, {embeddings_dataframe['title'].count()} items were found and {embeddings_dataframe['embedding'].count()} embeddings were created. The embeddings, along with content and metadata, can be found in `docs-embeddings.json`.""", "build/summary.md")

    print("Documentation collection completed")

indexer/main.py
boatbomber/Roblox-Docs-AI-Search
import json

def write_text(data: str, filename: str):
	print(f"Writing {filename}")
	with open(filename, "w", encoding="utf-8") as f:
		f.write(data)


def write_json(data, filename: str):
	print(f"Writing {filename}")
	with open(filename, "w", encoding="utf-8") as f:
		json.dump(data, f, ensure_ascii=False, indent="\t")
		f.write("\n")

indexer/write.py
boatbomber/Roblox-Docs-AI-Search