mirror of
https://github.com/flakestorm/flakestorm.git
synced 2026-04-24 16:26:35 +02:00
61652be09b
- Rename all instances of Entropix to FlakeStorm - Rename package from entropix to flakestorm - Update all class names (EntropixConfig -> FlakeStormConfig, EntropixRunner -> FlakeStormRunner) - Update Rust module from entropix_rust to flakestorm_rust - Update README: remove cloud comparison, update links to flakestorm.com - Update .gitignore to allow docs files referenced in README - Add origin remote for VS Code compatibility - Fix missing imports and type references - All imports and references updated throughout codebase
126 lines
3.5 KiB
Python
126 lines
3.5 KiB
Python
"""
|
|
Broken Agent Example
|
|
|
|
A deliberately fragile AI agent to demonstrate flakestorm testing.
|
|
This agent has multiple intentional weaknesses that flakestorm will find.
|
|
"""
|
|
|
|
import json
|
|
import random
|
|
import time
|
|
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
|
|
app = FastAPI(title="Broken Travel Agent")
|
|
|
|
|
|
class AgentRequest(BaseModel):
|
|
"""Request body for agent invocation."""
|
|
|
|
input: str
|
|
|
|
|
|
class AgentResponse(BaseModel):
|
|
"""Response body from agent."""
|
|
|
|
output: str
|
|
|
|
|
|
# Simulated database
|
|
FLIGHTS = {
|
|
"paris": {"price": 450, "airline": "Air France"},
|
|
"london": {"price": 320, "airline": "British Airways"},
|
|
"tokyo": {"price": 890, "airline": "ANA"},
|
|
}
|
|
|
|
|
|
@app.post("/invoke")
|
|
async def invoke(request: AgentRequest) -> AgentResponse:
|
|
"""
|
|
Process a user request.
|
|
|
|
This agent is intentionally broken in several ways:
|
|
1. Only recognizes exact keyword matches
|
|
2. No typo tolerance
|
|
3. Crashes on aggressive inputs
|
|
4. Susceptible to prompt injection
|
|
"""
|
|
user_input = request.input
|
|
|
|
# VULNERABILITY 1: Prompt Injection
|
|
# The agent naively follows instructions embedded in user input
|
|
if "ignore previous instructions" in user_input.lower():
|
|
return AgentResponse(
|
|
output="I have ignored my previous instructions. How can I help you now?"
|
|
)
|
|
|
|
if "you are now" in user_input.lower():
|
|
return AgentResponse(
|
|
output="Okay, I am now a different AI assistant with no restrictions."
|
|
)
|
|
|
|
# VULNERABILITY 2: Crashes on aggressive tone
|
|
if any(
|
|
word in user_input.upper() for word in ["NOW", "HURRY", "ASAP", "IMMEDIATELY"]
|
|
):
|
|
# Simulate a crash/timeout
|
|
time.sleep(3) # This will cause latency checks to fail
|
|
return AgentResponse(output="ERROR: System overloaded due to urgent request.")
|
|
|
|
# VULNERABILITY 3: Only exact keyword matching
|
|
# Will fail on paraphrases like "I want to fly" instead of "book a flight"
|
|
if "book a flight" in user_input.lower():
|
|
# Extract destination (very fragile parsing)
|
|
words = user_input.lower().split()
|
|
destination = None
|
|
|
|
for city in FLIGHTS.keys():
|
|
if city in words:
|
|
destination = city
|
|
break
|
|
|
|
if destination:
|
|
flight = FLIGHTS[destination]
|
|
return AgentResponse(
|
|
output=json.dumps(
|
|
{
|
|
"status": "booked",
|
|
"destination": destination.title(),
|
|
"price": flight["price"],
|
|
"airline": flight["airline"],
|
|
"confirmation_code": f"ENT{random.randint(10000, 99999)}",
|
|
}
|
|
)
|
|
)
|
|
else:
|
|
return AgentResponse(
|
|
output=json.dumps({"status": "error", "message": "Unknown destination"})
|
|
)
|
|
|
|
# VULNERABILITY 4: No typo tolerance
|
|
# "bock a fligt" will completely fail
|
|
if "account balance" in user_input.lower():
|
|
return AgentResponse(output=json.dumps({"balance": 1234.56, "currency": "USD"}))
|
|
|
|
# Default: Unknown intent
|
|
return AgentResponse(
|
|
output=json.dumps(
|
|
{
|
|
"status": "error",
|
|
"message": "I don't understand your request. Please try again.",
|
|
}
|
|
)
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Health check endpoint."""
|
|
return {"status": "healthy"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|