Agent improvements

Author: Breta01

brain/agent.py

@@ -36,13 +36,13 @@ def prompt(
 4. Propose the alpha and run the alpha simulation.
 
 - Try applying different functions and changing parameter values.
+- For d parameter use one of the following values: 2, 5, 10, 20, 30, 50, 100, 120, 250
 - Do NOT nest functions unless necessary, create simple alphas that are easy to understand
   and explain using basic operations like +, -, *, /.
-- Use only one operator in the alpha expression.
 - Try combining multiple data fields to create a new alpha.
 - Apply vec_avg(x) or vec_sum(x) if data field is a type VECTOR.
-- You can apply - to the alpha expression to negate fitness.
 - Change alpha completely if you think it is not good enough.
+- Self-correlation must be less than 0.5
 
 Working with equity in universe: {conf['universe']}, region: {conf['region']}
 and delay {conf['delay']}.
@@ -59,7 +59,7 @@ class CustomState(AgentState):
     tools=[
         submit_alpha,
         describe_operators,
-        search_datafields,
+        # search_datafields,
         get_random_datafields,
         get_random_idea,
     ],

brain/alpha_class.py

@@ -110,4 +110,4 @@ def from_stats(cls, stats: dict) -> "Alpha":
     @classmethod
     def from_row(cls, row: sqlite3.Row) -> "Alpha":
         """Build an Alpha from a sqlite3.Row."""
-        return cls(**{k: row[k] for k in cls.__dataclass_fields__.keys() if k in row.keys()})
+        return cls(**{k: row[k] for k in cls.__dataclass_fields__ if k in row.keys()})

brain/api.py

@@ -591,7 +591,7 @@ def _request(self, method: str, endpoint: str, max_retries: int = 200, **kwargs)
                 continue
 
             if float(response.headers.get("Retry-After", 0)) != 0:
-                time.sleep(float(response.headers["Retry-After"]))
+                time.sleep(float(response.headers["Retry-After"]) * 2)
                 max_retries -= 1
             else:
                 return response

brain/search_algorithm.py

@@ -114,6 +114,8 @@ def update_alphas_dict(alphas_dict, alphas_categories, stats, temp_id):
     else:
         alphas_categories["failing"].append(alpha_id)
 
+    return alphas_dict[alpha_id]
+
 
 def main():
     """Main function to run the agent."""
@@ -140,13 +142,17 @@ def main():
         while live_jobs:
             for job in as_completed(live_jobs):
                 # Update alphas_dict with the results
-                temp_id, _ = live_jobs.pop(job)  # remove from “running” set
+                temp_id, alpha_config = live_jobs.pop(job)  # remove from “running” set
                 stats = job.result()
                 print(f"Stats: {stats}")
-                update_alphas_dict(alphas_dict, alphas_categories, stats, temp_id)
+                alpha = update_alphas_dict(alphas_dict, alphas_categories, stats, temp_id)
 
                 # Start a new alpha simulation
-                response, alpha_config = create_alpha_simulation(alphas_dict, alphas_categories)
+                if alpha is not None and alpha.alpha_id is not None and alpha.fitness < -0.5:
+                    alpha_config = {**alpha_config, "regular": f"-{alpha.regular}"}
+                    response = BrainAPI.start_simulation(alpha_config)
+                else:
+                    response, alpha_config = create_alpha_simulation(alphas_dict, alphas_categories)
                 # Generate a unique ID for the alpha
                 temp_id = str(uuid.uuid4())
                 alphas_categories["pending"].append(temp_id)

brain/tools/datafields.py

@@ -17,6 +17,8 @@
 DELAY_0 = pd.read_csv(data_folder / "delay_0.csv", index_col=0)
 DELAY_1 = pd.read_csv(data_folder / "delay_1.csv", index_col=0)
 
+INDEX = len(DELAY_0) + len(DELAY_1)
+
 
 def format_datafields(data: pd.DataFrame) -> str:
     """Format the datafields DataFrame into a string."""
@@ -35,12 +37,19 @@ def search_datafields(datafield: str, config: RunnableConfig) -> str:
 
 def get_random_datafields(config: RunnableConfig) -> str:
     """Get a random datafields."""
+    global INDEX, DELAY_0, DELAY_1
     conf = get_universe_config(config)
 
+    if INDEX + 20 > len(DELAY_0 if conf["delay"] == 0 else DELAY_1):
+        DELAY_0 = DELAY_0.sample(frac=1).reset_index(drop=True)
+        DELAY_1 = DELAY_1.sample(frac=1).reset_index(drop=True)
+        INDEX = 0
+
     if conf["delay"] == 0:
         data = DELAY_0
     else:
         data = DELAY_1
 
-    data = data.sample(20)
+    data = data.iloc[INDEX : INDEX + 20]
+    INDEX += 20
     return format_datafields(data)

brain/tools/ideas.py

@@ -15,4 +15,5 @@ def get_random_idea():
     ```
 
     Think of the key features of this alpha and how you can use them to create a new alpha.
+    Create a new alpha using the same features but with different data fields.
     """