Upload folder using huggingface_hub

openenv.yaml

@@ -0,0 +1,5 @@
+name: opensleuth
+version: 0.1.0
+description: An OpenEnv environment for training LLMs to reverse-engineer black-box functions.
+author: Gemini
+contact: gemini@google.com

opensleuth_env/__init__.py

@@ -0,0 +1 @@
+# This file makes the 'opensleuth_env' directory a Python package.

opensleuth_env/black_box.py

@@ -0,0 +1,31 @@
+def fibonacci(n: int) -> int:
+    """
+    Calculates the nth Fibonacci number.
+    - Handles positive integers up to 90 to avoid large numbers.
+    - Raises ValueError for non-positive inputs or large inputs.
+    """
+    if not isinstance(n, int) or n <= 0 or n > 90:
+        raise ValueError("Input must be a positive integer less than or equal to 90.")
+    if n == 1:
+        return 1
+    a, b = 0, 1
+    for _ in range(n - 1):
+        a, b = b, a + b
+    return b
+
+# --- Add more black-box functions for later stages ---
+
+def reverse_string(s: str) -> str:
+    """
+    Reverses a string.
+    - Raises TypeError for non-string inputs.
+    """
+    if not isinstance(s, str):
+        raise TypeError("Input must be a string.")
+    return s[::-1]
+
+# --- Dictionary to hold all available black-box functions ---
+BLACK_BOX_FUNCTIONS = {
+    "fibonacci": fibonacci,
+    "reverse_string": reverse_string,
+}

opensleuth_env/client.py

@@ -0,0 +1,3 @@
+# This file is intentionally left blank.
+# It is used to create a client package if we were to publish this.
+# For the hackathon, we will interact with the server directly.

opensleuth_env/env.py

@@ -0,0 +1,93 @@
+from opensleuth_env.models import Action, Observation, State, ProbeAction, SubmitAction
+from opensleuth_env.black_box import BLACK_BOX_FUNCTIONS
+from opensleuth_env.verifier import verify_submission
+import random
+import traceback
+
+class OpenSleuthEnv:
+    def __init__(self):
+        self.state = None
+        # The verifier is now a static function, so no need to init it
+
+    def reset(self, target_name: str = "fibonacci") -> Observation:
+        """
+        Resets the environment to a new episode.
+        Selects a black-box function and clears the history.
+        """
+        if target_name not in BLACK_BOX_FUNCTIONS:
+            raise ValueError(f"Unknown target function: {target_name}")
+
+        self.state = State(
+            target_function_name=target_name,
+            probe_history=[],
+            seen_outputs=set(),
+            seen_error_types=set(),
+        )
+        return Observation(probe_history=[], last_error="")
+
+    def step(self, action: Action) -> tuple[Observation, float, bool]:
+        """
+        Takes a step in the environment.
+        """
+        if self.state is None:
+            # If reset() was not called, do it now.
+            self.reset()
+
+        # The Pydantic model binding in FastAPI should handle the conversion.
+        # This check is for robustness.
+        if not isinstance(action, (ProbeAction, SubmitAction)):
+            try:
+                if action.get("action_type") == "probe":
+                    action = ProbeAction(**action)
+                elif action.get("action_type") == "submit":
+                    action = SubmitAction(**action)
+                else:
+                    raise ValueError("Invalid action_type")
+            except Exception as e:
+                obs = Observation(probe_history=self.state.probe_history, last_error=f"Invalid action format: {e}")
+                return obs, -20.0, True
+
+
+        if action.action_type == "probe":
+            return self._handle_probe(action)
+        elif action.action_type == "submit":
+            return self._handle_submit(action)
+        else:
+            obs = Observation(probe_history=self.state.probe_history, last_error=f"Invalid action type: {action.action_type}")
+            return obs, -20.0, True
+
+    def _handle_probe(self, action: ProbeAction) -> tuple[Observation, float, bool]:
+        target_func = BLACK_BOX_FUNCTIONS[self.state.target_function_name]
+        intrinsic_reward = 0.0
+        last_error = ""
+
+        try:
+            eval_input = action.input
+            output = target_func(eval_input)
+            self.state.probe_history.append((eval_input, output))
+            if str(output) not in self.state.seen_outputs:
+                intrinsic_reward += 2.0
+                self.state.seen_outputs.add(str(output))
+
+        except Exception as e:
+            error_type = type(e).__name__
+            error_str = traceback.format_exc()
+            self.state.probe_history.append((action.input, error_str))
+            last_error = error_str
+            if error_type not in self.state.seen_error_types:
+                intrinsic_reward += 5.0
+                self.state.seen_error_types.add(error_type)
+
+        reward = intrinsic_reward - 1.0
+        obs = Observation(probe_history=self.state.probe_history, last_error=last_error)
+        return obs, reward, False
+
+    def _handle_submit(self, action: SubmitAction) -> tuple[Observation, float, bool]:
+        target_func = BLACK_BOX_FUNCTIONS[self.state.target_function_name]
+        execution_reward, complexity_penalty = verify_submission(action.code, target_func)
+        total_reward = execution_reward - complexity_penalty
+        if execution_reward == 100.0:
+            total_reward += 50.0
+
+        obs = Observation(probe_history=self.state.probe_history, last_error="")
+        return obs, total_reward, True

opensleuth_env/models.py

@@ -0,0 +1,29 @@
+from typing import Union, List, Tuple, Any, Literal
+from pydantic import BaseModel, Field
+
+class ProbeAction(BaseModel):
+    action_type: Literal["probe"] = "probe"
+    input: Any
+
+class SubmitAction(BaseModel):
+    action_type: Literal["submit"] = "submit"
+    code: str
+
+Action = Union[ProbeAction, SubmitAction]
+
+class Observation(BaseModel):
+    probe_history: List[Tuple[Any, Any]] = Field(
+        ...,
+        description="A list of (input, output) pairs from previous probes. Output can be a value or an error string."
+    )
+    last_error: str = Field(
+        "",
+        description="The error message from the last action, if any."
+    )
+
+class State(BaseModel):
+    target_function_name: str
+    probe_history: List[Tuple[Any, Any]]
+    # Store unique outputs and error types to calculate intrinsic reward
+    seen_outputs: set
+    seen_error_types: set

opensleuth_env/verifier.py

@@ -0,0 +1,68 @@
+import ast
+import random
+import string
+import math
+
+class ComplexityVisitor(ast.NodeVisitor):
+    def __init__(self):
+        self.complexity = 1
+    def visit_If(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+    def visit_For(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+    def visit_While(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+    def visit_And(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+    def visit_Or(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+    def visit_ExceptHandler(self, node):
+        self.complexity += 1
+        self.generic_visit(node)
+
+def _calculate_cyclomatic_complexity(code: str) -> int:
+    try:
+        tree = ast.parse(code)
+        visitor = ComplexityVisitor()
+        visitor.visit(tree)
+        return math.log(visitor.complexity)
+    except SyntaxError:
+        return 50
+
+def _generate_fuzz_inputs(target_func, count=100):
+    inputs = []
+    if target_func.__name__ == "fibonacci":
+        inputs = [random.randint(1, 90) for _ in range(count)]
+    elif target_func.__name__ == "reverse_string":
+        inputs = [''.join(random.choices(string.ascii_letters + string.digits, k=random.randint(1, 20))) for _ in range(count)]
+    return inputs
+
+def verify_submission(submitted_code: str, target_function: callable, fuzz_count: int = 100) -> tuple[float, float]:
+    try:
+        local_scope = {}
+        exec(submitted_code, {}, local_scope)
+        submitted_func = local_scope.get(target_function.__name__)
+        if not callable(submitted_func):
+            return 0.0, 50.0
+    except Exception:
+        return 0.0, 50.0
+
+    fuzz_inputs = _generate_fuzz_inputs(target_function, fuzz_count)
+    matches = 0
+    for inp in fuzz_inputs:
+        try:
+            target_output = target_function(inp)
+            submitted_output = submitted_func(inp)
+            if target_output == submitted_output:
+                matches += 1
+        except Exception:
+            continue
+    
+    execution_reward = 100.0 * (matches / fuzz_count)
+    complexity_penalty = _calculate_cyclomatic_complexity(submitted_code)
+    return execution_reward, complexity_penalty

requirements.txt

@@ -0,0 +1,3 @@
+fastapi
+uvicorn
+pydantic

server.py

@@ -0,0 +1,27 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from opensleuth_env.env import OpenSleuthEnv
+from opensleuth_env.models import Action, Observation
+
+app = FastAPI()
+env = OpenSleuthEnv()
+
+class ResetBody(BaseModel):
+    target_name: str = "fibonacci"
+
+@app.post("/reset", response_model=Observation)
+def reset_env(body: ResetBody):
+    # Ensure the environment is reset for a new session
+    return env.reset(target_name=body.target_name)
+
+@app.post("/step")
+def step_env(action: Action):
+    # The environment now handles the case where it's not reset
+    obs, reward, done = env.step(action)
+    return {"observation": obs, "reward": reward, "done": done}
+
+@app.get("/state")
+def get_state():
+    if env.state is None:
+        return {}
+    return env.get_state()

test_client.py

@@ -0,0 +1,29 @@
+import requests
+import json
+
+# The exact code to be submitted, without shell escaping issues
+code_to_submit = """
+def fibonacci(n: int) -> int:
+    if not isinstance(n, int) or n <= 0 or n > 90:
+        raise ValueError("Input must be a positive integer less than or equal to 90.")
+    if n == 1:
+        return 1
+    a, b = 0, 1
+    for _ in range(n - 1):
+        a, b = b, a + b
+    return b
+"""
+
+action = {
+    "action_type": "submit",
+    "code": code_to_submit
+}
+
+# Reset the environment first
+requests.post("http://127.0.0.1:8000/reset", json={"target_name": "fibonacci"})
+
+# Now send the step action
+response = requests.post("http://127.0.0.1:8000/step", json=action)
+
+print(response.status_code)
+print(response.json())

train.py

@@ -0,0 +1,157 @@
+import torch
+import requests
+from transformers import AutoTokenizer
+from unsloth import FastLanguageModel
+from trl import GPPOTrainer, PPOConfig
+import json
+import re
+
+# == 1. Constants ==
+MAX_STEPS_PER_EPISODE = 15
+ENV_URL = "http://127.0.0.1:8000"
+MODEL_NAME = "unsloth/qwen2-0.5b-instruct-sft-bnb-4bit"
+
+# == 2. Prompt Engineering ==
+def build_prompt(probe_history):
+    """
+    Creates the prompt for the LLM based on the probe history.
+    """
+    prompt = "You are a reverse-engineering AI. Your goal is to understand a hidden black-box function by probing it and then writing a Python replica.\\n\\n"
+    prompt += "== Probe History ==\\n"
+    if not probe_history:
+        prompt += "No probes yet. Your first action should be a probe.\\n"
+    else:
+        for i, (inp, out) in enumerate(probe_history):
+            prompt += f"{i+1}. IN: {inp} -> OUT: {out}\\n"
+    
+    prompt += "\\n== Your Action ==\\n"
+    prompt += "You can either PROBE or SUBMIT.\\n"
+    prompt += "To probe, respond with: PROBE(input)\\n"
+    prompt += "To submit your code, respond with: SUBMIT\\n```python\\n[your code here]\\n```\\n"
+    prompt += "Your decision: "
+    return prompt
+
+# == 3. Action Parsing ==
+def parse_action_from_response(response_text):
+    """
+    Parses the model's text response to determine the action.
+    """
+    probe_match = re.search(r"PROBE\\((.*)\\)", response_text)
+    if probe_match:
+        inp = probe_match.group(1).strip()
+        return {"action_type": "probe", "input": inp}
+
+    submit_match = re.search(r"SUBMIT\\s*```python\\n(.*)```", response_text, re.DOTALL)
+    if submit_match:
+        code = submit_match.group(1).strip()
+        return {"action_type": "submit", "code": code}
+
+    # Default to a probe if parsing fails
+    return {"action_type": "probe", "input": "1"}
+
+
+# == 4. Main Training Script ==
+def main():
+    # --- Initialize Model ---
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = MODEL_NAME,
+        max_seq_length = 2048,
+        dtype = None,
+        load_in_4bit = True,
+    )
+    # LoRA configuration
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = 16,
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        lora_alpha = 16,
+        lora_dropout = 0,
+        bias = "none",
+        use_gradient_checkpointing = True,
+        random_state = 3407,
+        use_rslora = False,
+        loftq_config = None,
+    )
+
+    # --- Initialize GPPO Trainer ---
+    # Note: GPPO is a new trainer in TRL and might require specific config.
+    # This is a placeholder configuration.
+    ppo_config = PPOConfig(
+        batch_size=4,
+        mini_batch_size=1,
+        learning_rate=1.41e-5,
+        adap_kl_ctrl=False,
+        log_with="tensorboard",
+        project_kwargs={"logging_dir": "./logs"}
+    )
+
+    # We need a dataset for the trainer, even if it's just a dummy one for initialization
+    # In a real RL loop, we provide the experiences directly to the `step` method.
+    dummy_dataset = [{"query": "dummy"}]
+    gppo_trainer = GPPOTrainer(
+        config=ppo_config,
+        model=model,
+        tokenizer=tokenizer,
+        dataset=dummy_dataset,
+    )
+
+    # --- Training Loop ---
+    for episode in range(10): # Run for 10 episodes for demonstration
+        print(f"--- Episode {episode+1} ---")
+        
+        # Reset environment
+        try:
+            resp = requests.post(f"{ENV_URL}/reset", json={"target_name": "fibonacci"})
+            obs = resp.json()
+        except requests.exceptions.ConnectionError as e:
+            print(f"ERROR: Could not connect to environment at {ENV_URL}. Is it running?")
+            print("Please run 'uvicorn server:app --host 0.0.0.0 --port 8000' in the 'opensleuth_env' directory.")
+            return
+
+        queries, responses, rewards = [], [], []
+
+        for step in range(MAX_STEPS_PER_EPISODE):
+            # Build prompt and generate action
+            prompt = build_prompt(obs.get("probe_history", []))
+            query_tensor = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+            
+            # Generate a response from the model
+            generation_kwargs = {"min_new_tokens": -1, "top_k": 0.0, "top_p": 1.0, "do_sample": True, "pad_token_id": tokenizer.eos_token_id, "max_new_tokens": 150}
+            response_tensor = gppo_trainer.generate(query_tensor, **generation_kwargs)
+            response_text = tokenizer.decode(response_tensor[0])
+            
+            # Parse action and execute in environment
+            action = parse_action_from_response(response_text)
+            step_resp = requests.post(f"{ENV_URL}/step", json=action)
+            step_data = step_resp.json()
+            
+            reward = torch.tensor(step_data["reward"], dtype=torch.float32)
+            obs = step_data["observation"]
+            done = step_data["done"]
+
+            # Store experience
+            queries.append(query_tensor.squeeze())
+            responses.append(response_tensor.squeeze())
+            rewards.append(reward)
+
+            print(f"Step {step+1}: Action: {action['action_type']}, Reward: {reward.item():.2f}")
+
+            if done:
+                break
+        
+        # --- Perform PPO Step ---
+        # This is a simplified view. The actual step requires careful handling of tensors.
+        # The `queries`, `responses`, `rewards` lists need to be formatted correctly.
+        try:
+            stats = gppo_trainer.step(queries, responses, rewards)
+            gppo_trainer.log_stats(stats, {}, rewards)
+            print(f"  PPO Step done. Mean reward: {stats['ppo/returns/mean']:.2f}")
+        except Exception as e:
+            print(f"ERROR during trainer.step: {e}")
+            print("  Skipping PPO step for this episode. This might happen if all trajectories are truncated.")
+
+
+if __name__ == "__main__":
+    # Ensure the server is running before starting training.
+    # We will run the server in the background from the CLI.
+    main()

verifier_log.txt

@@ -0,0 +1 @@
+\n--- Verifier Fuzzing ---\nInput: 88, Target: 1100087778366101931, Submitted: 1100087778366101931\nInput: 24, Target: 46368, Submitted: 46368\nInput: 14, Target: 377, Submitted: 377\nInput: 67, Target: 44945570212853, Submitted: 44945570212853\nInput: 35, Target: 9227465, Submitted: 9227465\nInput: 82, Target: 61305790721611591, Submitted: 61305790721611591\nInput: 82, Target: 61305790721611591, Submitted: 61305790721611591\nInput: 25, Target: 75025, Submitted: 75025\nInput: 1, Target: 1, Submitted: 1\nInput: 86, Target: 420196140727489673, Submitted: 420196140727489673\n--- End Verifier Fuzzing ---\nExecution Reward: 100.0, Complexity Penalty: 1.6094379124341003\n