michael-c-137/gemini-mythos

1import json2import os3from http.server import HTTPServer, BaseHTTPRequestHandler4import urllib.request5from dotenv import load_dotenv6 7# Load .env only if it exists (for local development)8if os.path.exists(".env"):9    load_dotenv()10 11BACKEND_URL = "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"12 13def get_api_key():14    """Retrieves the API Key from environment variables."""15    return os.environ.get("GOOGLE_API_KEY")16 17class LobsterTrapMock(BaseHTTPRequestHandler):18    def do_POST(self):19        content_length = int(self.headers['Content-Length'])20        post_data = self.rfile.read(content_length)21        payload = json.loads(post_data.decode('utf-8'))22 23        # Log Lobster Trap Metadata24        lt_metadata = payload.get("extra_body", {}).get("_lobstertrap") or payload.get("_lobstertrap")25        print(f"\n[LOBSTER TRAP PROXY] Intercepted Request from: {self.client_address}")26        if lt_metadata:27            print(f"[LOBSTER TRAP PROXY] Governance Metadata: {json.dumps(lt_metadata, indent=2)}")28        else:29            print("[LOBSTER TRAP PROXY] WARNING: No _lobstertrap metadata found!")30 31        # Clean payload for Gemini (remove extra_body and _lobstertrap)32        if "extra_body" in payload:33            del payload["extra_body"]34        if "_lobstertrap" in payload:35            del payload["_lobstertrap"]36 37        # Get API Key38        api_key = get_api_key()39        40        # Prepare headers41        headers = {"Content-Type": "application/json"}42        43        if api_key:44            headers["Authorization"] = f"Bearer {api_key}"45            print(f"[LOBSTER TRAP PROXY] Forwarding with API Key")46        else:47            print("[LOBSTER TRAP PROXY] ERROR: GOOGLE_API_KEY missing in environment!")48            self.send_response(500)49            self.end_headers()50            self.wfile.write(b"GOOGLE_API_KEY missing")51            return52 53        req = urllib.request.Request(54            BACKEND_URL,55            data=json.dumps(payload).encode('utf-8'),56            headers=headers,57            method="POST"58        )59 60        try:61            with urllib.request.urlopen(req) as response:62                res_data = response.read()63                self.send_response(response.status)64                self.send_header("Content-Type", "application/json")65                self.end_headers()66                self.wfile.write(res_data)67                print(f"[LOBSTER TRAP PROXY] Request Forwarded Successfully (Status: {response.status})")68        except urllib.error.HTTPError as e:69            print(f"[LOBSTER TRAP PROXY] Error from Backend: {e.code} {e.reason}")70            self.send_response(e.code)71            self.send_header("Content-Type", "application/json")72            self.end_headers()73            self.wfile.write(e.read())74 75def run(server_class=HTTPServer, handler_class=LobsterTrapMock, port=8081):76    server_address = ('', port)77    httpd = server_class(server_address, handler_class)78    print(f"Starting Lobster Trap Mock on port {port}...")79    httpd.serve_forever()80 81if __name__ == "__main__":82    run()83