LangChain

pip install langchain-openai

import os
from langchain_openai import ChatOpenAI

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

llm = ChatOpenAI(
    model=chat_model,
    openai_api_base="https://api.runinfra.ai/v1",
    openai_api_key=os.environ["RUNINFRA_GATEWAY_KEY"],
)

response = llm.invoke("What is RunInfra?")
print(response.content)

import os
from langchain_openai import OpenAIEmbeddings

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

embeddings = OpenAIEmbeddings(
    model=embedding_model,
    openai_api_base="https://api.runinfra.ai/v1",
    openai_api_key=os.environ["RUNINFRA_GATEWAY_KEY"],
)

vectors = embeddings.embed_documents(["Hello", "World"])

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

from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"{city}: 21C, partly cloudy"

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful weather assistant."),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_openai_tools_agent(llm, [get_weather], prompt)
executor = AgentExecutor(agent=agent, tools=[get_weather])
result = executor.invoke({"input": "What's the weather in Paris?"})
print(result["output"])

from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA

store = FAISS.from_texts(
    ["RunInfra cold starts under 2 seconds.", "RunInfra serves Llama, Qwen, Mistral..."],
    embeddings,
)
qa = RetrievalQA.from_chain_type(llm=llm, retriever=store.as_retriever(k=2))
print(qa.invoke({"query": "How fast are cold starts?"})["result"])