AI Refinery Runtime

version: 1
runtime: "ai_refinery"
start:
  command: ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7070"]
  port: 7070
  configSchema:
    properties:
      apiKey:
        type: "string"
        env_variable_name: "API_KEY"
        description: "Your AI Refinery API key"
        required: true
        mask: true

# Use a slim Python base image for a smaller final image size
FROM python:3.12-slim-bookworm

# Set the working directory inside the container
WORKDIR /app

# Copy the rest of your application's code into the container
COPY . .

# Install the Python dependencies
RUN pip install --no-cache-dir -r requirements.txt

# Expose the port that the application will run on
EXPOSE 7070

# The command to run the uvicorn server for your main.py file
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7070"]

import os
import uuid
import json
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from dotenv import load_dotenv

# --- Custom Agent Imports ---
from air import DistillerClient
from custom import executor_dict # Import the agent logic

# --- Initial Setup ---
load_dotenv()
api_key = str(os.getenv("API_KEY"))
project_name = "recommender_project"

# Initialize the Distiller Client and create the project on startup
print(f"Initializing DistillerClient for project: {project_name}")
distiller_client = DistillerClient(api_key=api_key)
# It's better to ensure the project exists rather than creating it on every startup
# In a production scenario, project creation would be a separate, one-time step.
try:
    distiller_client.create_project(config_path="config.yaml", project=project_name)
    print(f"Project '{project_name}' created or already exists.")
except Exception as e:
    # Handle cases where the project might already exist gracefully
    if "already exists" in str(e):
        print(f"Project '{project_name}' already exists.")
    else:
        print(f"An error occurred during project setup: {e}")
        raise

# --- FastAPI Application ---
app = FastAPI(
    title="Custom Agent Server",
    description="A FastAPI server to interact with the Recommender AIR agent via SSE.",
)

# --- Pydantic Models for Request Bodies ---
class QueryRequest(BaseModel):
    """Defines the body for sending a query to the agent."""
    query: str

# --- API Endpoints ---
@app.post("/agent/query")
async def agent_query_stream(request: Request, body: QueryRequest):
    """
    Accepts a query and streams responses from the Recommender agent via SSE.
    
    - **query**: The natural language query to send to the agent.
    """
    
    async def event_generator():
        """
        The async generator that handles the agent query and yields SSE events.
        """
        user_id = f"user-{uuid.uuid4()}"
        
        try:
            # Yield status updates in SSE format
            yield f"event: status\ndata: {json.dumps(f'Connecting to project {project_name} for user {user_id}...')}\n\n"
            
            # Use the async context manager for the DistillerClient
            async with distiller_client(
                project=project_name,
                uuid=user_id,
                executor_dict=executor_dict
            ) as dc:
                
                yield f"event: status\ndata: {json.dumps('Connection successful. Sending query...')}\n\n"
                
                # Send the query to the agent
                responses = await dc.query(query=body.query)
                
                # Stream responses back to the client
                async for response in responses:
                    if await request.is_disconnected():
                        print("Client disconnected, stopping stream.")
                        break
                    
                    # Serialize the Pydantic object to JSON and format as a valid SSE message
                    yield f"event: message\ndata: {response.model_dump_json()}\n\n"
            
            yield f"event: close\ndata: {json.dumps('Stream finished.')}\n\n"
        except Exception as e:
            print(f"An error occurred during query stream: {e}")
            import traceback
            traceback.print_exc()
            error_message = f"An unexpected server error occurred: {type(e).__name__}"
            # Yield the error message in SSE format
            yield f"event: error\ndata: {json.dumps(error_message)}\n\n"

    return StreamingResponse(event_generator(), media_type="text/event-stream")

@app.get("/")
def read_root():
    """A simple root endpoint to confirm the server is running."""
    return {
        "message": f"Custom Agent Server is running. Loaded project: '{project_name}'. POST to /agent/query to interact."
    }

import os
from air import AsyncAIRefinery
from dotenv import load_dotenv

# Load environment variables from a .env file
load_dotenv()
api_key = str(os.getenv("API_KEY"))

async def recommender_agent(query: str) -> str:
    """
    A custom agent that provides recommendations based on a user query.

    Args:
        query: The user's request for a recommendation.

    Returns:
        A string containing the recommendation and a justification.
    """
    prompt = """Given the query below, your task is to provide the user with a useful and cool
       recommendation followed by a one-sentence justification.\n\nQUERY: {query}"""

    formatted_prompt = prompt.format(query=query)

    # Initialize the AIRefinery client to interact with the LLM
    airefinery_client = AsyncAIRefinery(api_key=api_key)

    # Call the chat completions API
    response = await airefinery_client.chat.completions.create(
        messages=[{"role": "user", "content": formatted_prompt}],
        model="meta-llama/Llama-3.1-70B-Instruct",
    )

    return response.choices[0].message.content

# This dictionary maps the agent name from config.yaml to its function
executor_dict = {"Recommender Agent": recommender_agent}

utility_agents:
  - agent_class: CustomAgent
    agent_name: "Recommender Agent"
    agent_description: |
      The Recommender Agent is a specialist in item recommendations. For instance,
      it can provide users with costume recommendations, items to purchase, food,
      decorations, and so on. 
    config: {}

super_agents: []

orchestrator:
  agent_list:
    - agent_name: "Recommender Agent"

airefinery-sdk==1.21.0
fastapi==0.119.1