🤖 Flaxon AI

AI/LLM integration plugin for Flaxon framework with support for Google Gemini, OpenAI, and local Flax/JAX models.

Table of Contents

• Features

• Installation

• Quick Start

• Configuration

  • Environment Variables
  • With Flaxon Config

• Usage Examples

  • Basic Generation
  • Streaming Response
  • Using Decorators
  • Local Flax Model
  • Chat Completion
  • Embeddings

• Building a Flaxon Bot

• Pre-built Routes

• Security Best Practices

• Roadmap

• License

Features

• 🤖 Multiple AI Providers — Google Gemini, OpenAI, Flax/JAX local models
• ⚡ Async Generation — Non-blocking AI completions
• 📡 Streaming Support — Server-Sent Events (SSE) for real-time responses
• 🧠 Flax/JAX Integration — Local model inference with GPU/TPU acceleration
• 🎯 Route Decorators — Easy AI integration into Flaxon routes
• 📊 GraphQL Helpers — AI field resolvers for GraphQL
• 🚀 Pre-built Endpoints — /ai/generate, /ai/stream, /ai/chat
• 💾 Model Management — Load and cache local models

Installation

# Basic installation
pip install flaxon-ai

# With Google Gemini support
pip install flaxon-ai[gemini]

# With OpenAI support
pip install flaxon-ai[openai]

# With local Flax/JAX support
pip install flaxon-ai[flax]

# With all providers
pip install flaxon-ai[all]

Quick Start

from flaxon import Flaxon
from flaxon_ai import FlaxonAIPlugin
import os

app = Flaxon("my-app")

# Load AI plugin with Google Gemini
await app.plugins.load_plugin(FlaxonAIPlugin(
    provider="gemini",
    api_key=os.environ.get("GEMINI_API_KEY"),
    default_model="gemini-2.5-flash",
))

# Use in a route
@app.post("/api/generate")
async def generate(request):
    data = await request.json()
    result = await app.state.ai.generate(data["prompt"])
    return {"result": result}

Configuration

Environment Variables

# Google Gemini
GEMINI_API_KEY=your-api-key

# OpenAI
OPENAI_API_KEY=your-api-key

# Flax/JAX (no API key required)

With Flaxon Config

app = Flaxon("my-app", config={
    "AI_PROVIDER": "gemini",
    "AI_API_KEY": os.environ.get("GEMINI_API_KEY"),
    "AI_DEFAULT_MODEL": "gemini-2.5-flash",
    "AI_MAX_TOKENS": 100,
    "AI_TEMPERATURE": 0.7,
})

plugin = FlaxonAIPlugin.from_config(app.config)
await app.plugins.load_plugin(plugin)

Usage Examples

Basic Generation

@app.post("/api/summarize")
async def summarize(request):
    data = await request.json()
    text = data.get("text", "")
    
    summary = await app.state.ai.generate(
        f"Summarize this text in 2 sentences:\n{text}",
        max_tokens=100
    )
    return {"summary": summary}

Streaming Response

from flaxon_ai.streaming import StreamResponse

@app.post("/api/stream")
async def stream_response(request):
    data = await request.json()
    prompt = data.get("prompt", "Write a short story about a robot")

    return StreamResponse(
        app.state.ai.stream(prompt, model="gemini-2.5-flash"),
        metadata={"prompt": prompt},
    )

Using Decorators

from flaxon_ai import ai_prompt, stream_ai

@app.get("/api/smart-bio")
@ai_prompt("Write a professional bio based on: {data}")
async def get_user_data(request):
    user = await get_user(request.session.get("user_id"))
    return {"data": f"Name: {user.name}, Skills: {user.skills}"}

Local Flax Model

# Load local Flax model
app.plugins.load_plugin(FlaxonAIPlugin(
    provider="flax",
    default_model="gpt2",
    cache_models=True,
))

# Use it just like a cloud provider
result = await app.state.ai.generate(
    "Write a poem about Python",
    max_tokens=60
)

Chat Completion

@app.post("/api/chat")
async def chat(request):
    data = await request.json()
    messages = data.get("messages", [])
    
    response = await app.state.ai.chat(
        messages=messages,
        model="gemini-2.5-flash"
    )
    return {"response": response}

Embeddings

@app.post("/api/embed")
async def embed(request):
    data = await request.json()
    text = data.get("text", "")
    
    embedding = await app.state.ai.embed(text)
    return {"embedding": embedding}

Building a Flaxon Bot

You can use Flaxon AI to create a simple AI bot by connecting a chat route to the AI service.

The following example creates a small bot that accepts a user's message and returns an AI-generated response:

from flaxon import Flaxon
from flaxon_ai import FlaxonAIPlugin
import os

app = Flaxon("flaxon-bot")

await app.plugins.load_plugin(FlaxonAIPlugin(
    provider="gemini",
    api_key=os.environ.get("GEMINI_API_KEY"),
    default_model="gemini-2.5-flash",
))

@app.post("/bot/chat")
async def bot_chat(request):
    data = await request.json()
    message = data.get("message", "")

    response = await app.state.ai.chat(
        messages=[
            {
                "role": "system",
                "content": "You are a helpful Flaxon bot."
            },
            {
                "role": "user",
                "content": message
            }
        ],
        model="gemini-2.5-flash"
    )

    return {
        "message": message,
        "response": response
    }

You can then send a request to:

POST /bot/chat

with:

{
  "message": "What is Flaxon?"
}

The bot can be extended with authentication, conversation history, streaming responses, tools, database access, and custom application logic.

Pre-built Routes

Route	Method	Description
/ai/generate	POST	Generate text completion
/ai/stream	POST	Stream text via SSE
/ai/chat	POST	Chat completion
/ai/embed	POST	Generate embeddings
/ai/models	GET	List available models
/ai/health	GET	Health check

Security Best Practices

✅ Never hardcode API keys

✅ Use environment variables or secrets manager

✅ Validate and sanitize user prompts

✅ Implement rate limiting for AI endpoints

✅ Limit streaming duration and token count

✅ Use HTTPS in production

✅ Verify local model integrity before loading

Roadmap

Version	Features
0.1.0	Basic AI plugin, Gemini provider
0.2.0	OpenAI provider, streaming support
0.3.0	Flax/JAX local models
0.4.0	Embeddings, model management
0.5.0	Function calling, tool use
0.6.0	Fine-tuning support

License

MIT License - See LICENSE file for details.