import json
import tempfile
from pathlib import Path
from mlflow.gateway.app import create_app_from_path
# This HTML was obtained by sending a request to the `/docs` route and saving the response.
# To hide the "try it out" button, we set `supportedSubmitMethods` to an empty list.
# The url was changed to "./openapi.json" from "/openapi.json" because `api.html` and `openapi.json`
# are served from the same directory.
API_HTML = """
MLflow AI Gateway - Swagger UI
"""
def main():
config = """
endpoints:
- name: chat
endpoint_type: llm/v1/chat
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_key: key
- name: completions
endpoint_type: llm/v1/completions
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_key: key
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: openai
name: text-embedding-ada-002
config:
openai_api_key: key
"""
with tempfile.TemporaryDirectory() as tmpdir:
config_path = Path(tmpdir).joinpath("config.yaml")
config_path.write_text(config)
app = create_app_from_path(config_path)
docs_build = Path("build/html/llms/deployments")
docs_build.mkdir(parents=True, exist_ok=True)
with docs_build.joinpath("openapi.json").open("w") as f:
json.dump(app.openapi(), f)
with docs_build.joinpath("api.html").open("w") as f:
f.write(API_HTML)
if __name__ == "__main__":
main()