LlamaIndex

pip install llama-index llama-index-llms-openai llama-index-embeddings-openai

import os
from llama_index.llms.openai import OpenAI

chat_model = os.environ.get("RUNINFRA_CHAT_MODEL")
if not chat_model:
    raise RuntimeError("Set RUNINFRA_CHAT_MODEL from GET /v1/models.")

llm = OpenAI(
    model=chat_model,
    api_base="https://api.runinfra.ai/v1",
    api_key=os.environ["RUNINFRA_GATEWAY_KEY"],
)

response = llm.complete("What is RunInfra?")
print(response.text)

import os
from llama_index.embeddings.openai import OpenAIEmbedding

embedding_model = os.environ.get("RUNINFRA_EMBEDDING_MODEL")
if not embedding_model:
    raise RuntimeError("Set RUNINFRA_EMBEDDING_MODEL from GET /v1/models.")

embed = OpenAIEmbedding(
    model=embedding_model,
    api_base="https://api.runinfra.ai/v1",
    api_key=os.environ["RUNINFRA_GATEWAY_KEY"],
)

vector = embed.get_text_embedding("Hello world")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings

Settings.llm = llm
Settings.embed_model = embed

documents = SimpleDirectoryReader("./docs").load_data()
index = VectorStoreIndex.from_documents(documents)

query_engine = index.as_query_engine()
print(query_engine.query("How fast are RunInfra cold starts?"))

for chunk in llm.stream_complete("Tell me a short story"):
    print(chunk.delta, end="", flush=True)