OpenAI SDK (Python)

The official OpenAI Python SDK works with PAIS out of the box by setting base_url and api_key. This is the lowest-overhead approach and gives full control over prompts, parameters, and output parsing.

Installation

pip install openai python-dotenv

Setup

import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    base_url=os.environ["PAIS_API_BASE"],
    api_key=os.environ["PAIS_API_KEY"],
)

Structured output

Llama 3.1 supports JSON mode for structured extraction:

from pydantic import BaseModel
import json

class PaperSummary(BaseModel):
    title: str
    key_contributions: list[str]
    methods: list[str]
    limitations: list[str]

abstract = """
We present SciBench, a benchmark evaluating LLMs on university-level science problems
across physics, chemistry, and mathematics. We find that chain-of-thought prompting
with self-critique improves accuracy by 12% over zero-shot baselines...
"""

response = client.chat.completions.create(
    model="llama-3-1-8b",
    messages=[
        {
            "role": "system",
            "content": "Extract structured information from research abstracts. "
                       "Respond with valid JSON only.",
        },
        {
            "role": "user",
            "content": f"Abstract:\n{abstract}\n\nExtract into this schema: "
                       f"{PaperSummary.model_json_schema()}",
        },
    ],
    response_format={"type": "json_object"},
    temperature=0.1,
)

summary = PaperSummary.model_validate_json(response.choices[0].message.content)
print(summary.key_contributions)

Function / tool calling

Build a simple agent that decides which tool to call:

import json

tools = [
    {
        "type": "function",
        "function": {
            "name": "search_arxiv",
            "description": "Search arXiv for papers by keyword",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search query"},
                    "max_results": {"type": "integer", "default": 5},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "Execute a Python expression and return the result",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"},
                },
                "required": ["code"],
            },
        },
    },
]

def search_arxiv(query: str, max_results: int = 5) -> str:
    # Real implementation would use the arxiv Python package
    return f"[Found {max_results} papers for '{query}']"

def run_python(code: str) -> str:
    result = {}
    exec(code, {}, result)
    return str(result)

def run_agent(user_message: str) -> str:
    messages = [
        {"role": "system", "content": "You are a research assistant. Use tools when helpful."},
        {"role": "user", "content": user_message},
    ]

    while True:
        response = client.chat.completions.create(
            model="llama-3-1-8b",
            messages=messages,
            tools=tools,
            tool_choice="auto",
        )
        msg = response.choices[0].message
        messages.append(msg)

        if msg.tool_calls:
            for tool_call in msg.tool_calls:
                fn = tool_call.function.name
                args = json.loads(tool_call.function.arguments)
                result = search_arxiv(**args) if fn == "search_arxiv" else run_python(**args)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": result,
                })
        else:
            return msg.content

print(run_agent("Find recent papers on LoRA fine-tuning, then compute 2**10."))

Async usage

For parallel requests (e.g., embedding a large document corpus):

import asyncio
from openai import AsyncOpenAI

async_client = AsyncOpenAI(
    base_url=os.environ["PAIS_API_BASE"],
    api_key=os.environ["PAIS_API_KEY"],
)

async def embed_batch(texts: list[str]) -> list[list[float]]:
    response = await async_client.embeddings.create(
        model="qwen3-vl-embedding-8b",
        input=texts,
    )
    return [d.embedding for d in response.data]

async def process_corpus(chunks: list[str], batch_size: int = 32):
    batches = [chunks[i:i+batch_size] for i in range(0, len(chunks), batch_size)]
    tasks = [embed_batch(b) for b in batches]
    results = await asyncio.gather(*tasks)
    return [emb for batch in results for emb in batch]

Rate limits with async

Async requests can hit the 60 req/min rate limit quickly. Add a semaphore to limit concurrent requests:

sem = asyncio.Semaphore(4)
async def rate_limited_embed(texts):
    async with sem:
        return await embed_batch(texts)