Open Source Experiential Router Uses Request Data to Fine-Tune Models

CREATE TABLE llm_traces (
    trace_id UUID PRIMARY KEY,
    prompt JSONB NOT NULL,
    response JSONB,
    model VARCHAR(50) NOT NULL,
    latency_ms INT,
    prompt_tokens INT,
    completion_tokens INT,
    feedback_score FLOAT,
    corrected_text TEXT,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);
 
CREATE INDEX idx_traces_feedback ON llm_traces (feedback_score) WHERE feedback_score IS NOT NULL;

import uuid
import httpx
from fastapi import FastAPI, Request, BackgroundTasks
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import json
 
app = FastAPI()
UPSTREAM_URL = "https://api.openai.com/v1/chat/completions"
OPENAI_API_KEY = "your-api-key-here"
 
db = {}
 
def log_transaction(trace_id: str, request_body: dict, response_body: dict):
    db[trace_id] = {
        "prompt": request_body.get("messages", []),
        "response": response_body.get("choices", [{}])[0].get("message", {}),
        "model": request_body.get("model"),
        "feedback_score": None,
        "corrected_text": None
    }
    print(f"Logged transaction {trace_id}")
 
def log_stream_transaction(trace_id: str, request_body: dict, chunks: list):
    full_content = ""
    role = "assistant"
    for chunk in chunks:
        try:
            data = json.loads(chunk)
            delta = data.get("choices", [{}])[0].get("delta", {})
            if "content" in delta:
                full_content += delta["content"]
            if "role" in delta:
                role = delta["role"]
        except json.JSONDecodeError:
            continue
    
    db[trace_id] = {
        "prompt": request_body.get("messages", []),
        "response": {"role": role, "content": full_content},
        "model": request_body.get("model"),
        "feedback_score": None,
        "corrected_text": None
    }
    print(f"Logged stream transaction {trace_id}")
 
@app.post("/v1/chat/completions")
async def route_request(request: Request, background_tasks: BackgroundTasks):
    body = await request.json()
    trace_id = str(uuid.uuid4())
    
    headers = {
        "Authorization": f"Bearer {OPENAI_API_KEY}",
        "Content-Type": "application/json"
    }
    
    is_stream = body.get("stream", False)
    
    if not is_stream:
        async with httpx.AsyncClient() as client:
            response = await client.post(UPSTREAM_URL, json=body, headers=headers, timeout=60.0)
        response_data = response.json()
        background_tasks.add_task(log_transaction, trace_id, body, response_data)
        return StreamingResponse(
            content=iter([json.dumps(response_data)]),
            media_type="application/json",
            headers={"x-trace-id": trace_id}
        )
    
    async def stream_generator():
        chunks = []
        async with httpx.AsyncClient() as client:
            async with client.stream("POST", UPSTREAM_URL, json=body, headers=headers, timeout=60.0) as r:
                async for line in r.aiter_lines():
                    if not line.strip():
                        continue
                    yield line + "\n"
                    if line.startswith("data: "):
                        data_str = line[6:].strip()
                        if data_str == "[DONE]":
                            break
                        chunks.append(data_str)
        background_tasks.add_task(log_stream_transaction, trace_id, body, chunks)
 
    return StreamingResponse(
        stream_generator(),
        media_type="text/event-stream",
        headers={"x-trace-id": trace_id}
    )

class FeedbackPayload(BaseModel):
    trace_id: str
    score: float
    corrected_text: str | None = None
 
@app.post("/v1/feedback")
async def receive_feedback(payload: FeedbackPayload):
    if payload.trace_id not in db:
        return {"status": "error", "message": "Trace ID not found"}
        
    db[payload.trace_id]["feedback_score"] = payload.score
    if payload.corrected_text:
        db[payload.trace_id]["corrected_text"] = payload.corrected_text
        
    return {"status": "success", "trace_id": payload.trace_id}

import sys
import io
 
def execute_and_evaluate(code: str) -> float:
    old_stdout = sys.stdout
    redirected_output = sys.stdout = io.StringIO()
    try:
        exec(code, {}, {})
        sys.stdout = old_stdout
        return 1.0
    except Exception:
        sys.stdout = old_stdout
        return 0.0

import re
 
def clean_text(text: str) -> str:
    text = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', '[EMAIL]', text)
    text = re.sub(r'sk-[a-zA-Z0-9]{32,}', '[API_KEY]', text)
    return text
 
def validate_chat_format(messages: list) -> bool:
    for msg in messages:
        if "role" not in msg or "content" not in msg:
            return False
        if msg["role"] not in ["system", "user", "assistant"]:
            return False
        if not isinstance(msg["content"], str) or len(msg["content"]) == 0:
            return False
    return True
 
def generate_training_data(database: dict):
    dataset = []
    for trace_id, entry in database.items():
        score = entry.get("feedback_score")
        if score is None or score < 0.8:
            continue
            
        target_output = entry["corrected_text"] if entry["corrected_text"] else entry["response"].get("content")
        if not target_output:
            continue
            
        system_msg = ""
        user_msg = ""
        
        for msg in entry["prompt"]:
            if msg["role"] == "system":
                system_msg = msg["content"]
            elif msg["role"] == "user":
                user_msg = msg["content"]
                
        messages = [
            {"role": "system", "content": clean_text(system_msg)},
            {"role": "user", "content": clean_text(user_msg)},
            {"role": "assistant", "content": clean_text(target_output)}
        ]
        
        if validate_chat_format(messages):
            dataset.append({"messages": messages})
            
    return dataset

{"messages": [{"role": "system", "content": "You convert natural language to SQL."}, {"role": "user", "content": "Find users registered last week."}, {"role": "assistant", "content": "SELECT * FROM users WHERE created_at >= DATE_SUB(NOW(), INTERVAL 7 DAY);"}]}

base_model: meta-llama/Meta-Llama-3-8B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer
 
load_in_8bit: false
load_in_4bit: true
strict: false
 
datasets:
  - path: /data/processed_feedback.jsonl
    type: chat_template
    chat_template: llama3
 
dataset_prepared_path:
val_set_size: 0.05
output_dir: ./outputs/llama3-custom
 
adapter: qlora
lora_model_dir:
 
sequence_len: 2048
sample_packing: true
pad_to_sequence_len: true
 
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05
lora_target_modules:
  - q_proj
  - v_proj
  - k_proj
  - o_proj
 
wandb_project: experiential-router-tuning
tf32: true
 
epochs: 3
micro_batch_size: 2
gradient_accumulation_steps: 4
optimizer: adamw_torch
learning_rate: 0.0002
lr_scheduler: cosine

import random
 
class RouterConfig:
    def __init__(self):
        self.weights = [0.0, 1.0] 
        self.custom_model_id = "llama3-custom"
        self.commercial_model_id = "gpt-4o"
 
    def select_model(self) -> str:
        choices = [self.custom_model_id, self.commercial_model_id]
        return random.choices(choices, weights=self.weights, k=1)[0]
 
router_config = RouterConfig()

def get_rolling_average(model_name: str, limit: int = 100) -> float:
    scores = [
        entry["feedback_score"]
        for entry in db.values()
        if entry.get("model") == model_name and entry.get("feedback_score") is not None
    ][-limit:]
    if not scores:
        return 0.0
    return sum(scores) / len(scores)
 
def adjust_routing_weights():
    custom_avg = get_rolling_average(router_config.custom_model_id)
    commercial_avg = get_rolling_average(router_config.commercial_model_id)
    
    if custom_avg >= (commercial_avg * 0.95):
        current_custom_weight = router_config.weights[0]
        new_custom_weight = min(current_custom_weight + 0.1, 1.0)
        router_config.weights = [new_custom_weight, 1.0 - new_custom_weight]
        print(f"Increased custom model weight to {new_custom_weight}")
    else:
        router_config.weights = [0.0, 1.0]
        print("Performance dropped. Rolled back to commercial model.")

def check_context_limit(prompt: list, limit: int = 8192) -> bool:
    total_chars = sum(len(msg.get("content", "")) for msg in prompt)
    estimated_tokens = total_chars // 4
    return estimated_tokens <= limit