More updates to agentic framework

Author: Breta01

brain/agentic.py

@@ -16,14 +16,30 @@
 MAX_EXPLORATION_COUNT = 3
 
 
+def plateau_condition(state: GraphState) -> bool:
+    """Check if the exploration has plateaued."""
+    storage = state["storage"]
+    best_alpha = storage.best_alpha
+    old_best_alpha = state.get("old_best_alpha")
+
+    if best_alpha is None or old_best_alpha is None:
+        return False
+
+    # Check if the score has not improved significantly
+    score_diff = (
+        best_alpha.fitness - old_best_alpha.fitness + best_alpha.sharpe - old_best_alpha.sharpe
+    )
+    return score_diff < 0.1
+
+
 def seed_finder_node(state: GraphState) -> GraphState:
     """Find a seed alpha to start the exploration."""
     # Iterate databse till we find some decent alpha, or some other seed idea
-    alpha = invoke_seeder(state)
-    print(f"Seed alpha: {alpha.prompt_format()}")
+    alpha_idea, config = invoke_seeder(state)
+    print(f"Seed alpha: {alpha_idea}")
     return {
-        "alpha_idea": alpha.regular,
-        "default_config": get_config(alpha),
+        "alpha_idea": alpha_idea,
+        "default_config": config,
         "node": "plan",
         "state": "explore",
         "explore_count": 0,
@@ -35,10 +51,11 @@ def seed_finder_node(state: GraphState) -> GraphState:
 def planner_node(state: GraphState) -> GraphState:
     """Plan the next steps based on the current alpha and state."""
     state["explore_count"] += 1
+    state["old_best_alpha"] = state["storage"].best_alpha
 
     plan = invoke_planner(state)
     if not plan:
-        return {**state, "node": "seed", "state": "explore"}
+        return {**state, "node": "seed"}
     return {**state, "node": "execute", "plan": plan}
 
 
@@ -70,7 +87,7 @@ def explore_test_node(state: GraphState) -> GraphState:
             }
 
     # TODO: Test plateau condition, compare previous best, with current best alpha
-    if state["explore_count"] < MAX_EXPLORATION_COUNT:
+    if state["explore_count"] < MAX_EXPLORATION_COUNT and plateau_condition(state):
         return {**state, "node": "plan", "state": "explore"}
 
     return {**state, "node": "seed", "state": "explore"}
@@ -89,7 +106,7 @@ def submit_test_node(state: GraphState) -> GraphState:
         return {**state, "node": "seed", "state": "explore", "explore_count": 0}
 
     # TODO: Test plateau condition, compare previous best, with current best alpha
-    if state["explore_count"] < MAX_EXPLORATION_COUNT:
+    if state["explore_count"] < MAX_EXPLORATION_COUNT and plateau_condition(state):
         return {**state, "node": "plan", "state": "fine-tune"}
 
     return {**state, "node": "seed", "state": "explore"}

brain/agents/executor.py

@@ -18,7 +18,7 @@
 from brain.tools.operators import describe_operators, operators
 from brain.tools.simulation import StopException, submit_alpha
 
-STEPS_PER_INSTRUCTION = 4
+STEPS_PER_INSTRUCTION = 2
 
 
 def prompt(
@@ -70,10 +70,11 @@ class CustomState(AgentState):
 )
 
 
-def create_alpha_simulation(
-    storage: Storage, plan: list[str], config: dict = DEFAULT_CONFIG
-) -> tuple[Response, Alpha]:
+def create_alpha_simulation(storage: Storage, state: GraphState) -> tuple[Response, Alpha]:
     """Create a new alpha based on the given ID."""
+    plan = state.get("plan", [])
+    config = state.get("default_config", DEFAULT_CONFIG)
+
     top_k = 5
 
     alphas = {
@@ -93,6 +94,9 @@ def create_alpha_simulation(
 Create a new alpha based on the best alpha according to the following instructions:
 {plan[(storage.counter // STEPS_PER_INSTRUCTION) % len(plan)]}
 
+ALPHA IDEA
+----------
+{state.get("alpha_idea", "None")}
 
 BEST ALPHA
 ----------
@@ -126,7 +130,7 @@ def create_alpha_simulation(
                     ]
                 },
                 config={
-                    "recursion_limit": 50,
+                    "recursion_limit": 30,
                     "configurable": {
                         **config,
                         "alphas": alphas_store,
@@ -158,7 +162,7 @@ def invoke(state: GraphState) -> GraphState:
         config = get_config(storage.best_alpha)
     else:
         config = state.get("default_config", DEFAULT_CONFIG)
-    create_alpha = partial(create_alpha_simulation, plan=plan, config=config)
+    create_alpha = partial(create_alpha_simulation, state=state)
 
     storage.reset_counter()
 

brain/agents/planner.py

@@ -3,6 +3,8 @@
 
 from brain.graph_state import GraphState
 from brain.model import MODEL
+from brain.tools.datafields import get_random_datafields, search_datafields
+from brain.tools.ideas import get_random_idea
 from brain.tools.operators import operators
 
 
@@ -14,24 +16,25 @@ class ResponseFormat(BaseModel):
 
 agent = create_react_agent(
     model=MODEL,
-    tools=[],
+    tools=[get_random_idea, search_datafields, get_random_datafields],
     response_format=ResponseFormat,
     prompt=f"""
 You are a senior quantitative researcher working with the World Quant Brain platform to create alphas. Based on an initial alpha idea provided (either in written form or as an expression), outline specific steps to implement concrete changes to the alpha using the Fast Expression Language, which includes various operators.
 
 Allowed operators are: {operators.all()}
 
-1. Analyze the initial alpha idea to identify key components such as variables, conditions, and expected outcomes.
-2. Propose specific changes to the alpha expression, detailing how each change can enhance its performance. For example:
+1. Execute the original alpha to obtain a baseline performance.
+2. Analyze the initial alpha idea to identify key components such as variables, conditions, and expected outcomes.
+3. Propose specific changes to the alpha expression, detailing how each change can enhance its performance. For example:
    <example>
    - Change the parameter value of 'X' from 0.5 to 0.7 to test sensitivity.
    - Introduce a new operator, such as 'rank', to refine the conditions under which trades are executed.
    - Suggest using a different data field, like 'volume', instead of 'price' to capture market dynamics more accurately.
    - Suggest type of data field to search for, e.g. 'momentum', 'trend', 'news' etc.
    - Add a group neutralization to the alpha expression to reduce sector bias.
    </example>
-3. For each proposed change, provide a rationale explaining why it may improve the alpha's effectiveness.
-4. Ensure that each change is concrete and directly related to the alpha expression, focusing on parameter adjustments, operator modifications, or data field substitutions.
+4. For each proposed change, provide a rationale explaining why it may improve the alpha's effectiveness.
+5. Ensure that each change is concrete and directly related to the alpha expression, focusing on parameter adjustments, operator modifications, or data field substitutions.
 
 List only the steps related to changing the alpha expression. Do NOT include steps such as evaluation, backtesting, documentation, or communication.
 
@@ -49,7 +52,7 @@ def invoke(state: GraphState) -> ResponseFormat:
             "You are exploring a new alphas based on the new idea.\n"
             f"{state.get('alpha_idea', '')}\n"
             f"Best alpha so far:\n {best_alpha.prompt_format() if best_alpha else 'None'}\n"
-            "Create a plan to explore this idea, start with simple alpha expression.\n"
+            "Create a plan to explore this idea.\n"
         )
     elif state.get("state", "explore") == "fine_tune" and best_alpha is not None:
         message = (

brain/agents/seeder.py

@@ -1,13 +1,26 @@
+import json
 import random
+from pathlib import Path
 
+from brain.agent_config import get_config
 from brain.alpha_class import Alpha
 from brain.database import Database
 from brain.graph_state import GraphState
 
+TEMPLATES = json.load((Path(__file__).parent.parent / "tools/data/alpha_templates.json").open())
 
-def invoke(state: GraphState) -> Alpha:
+
+def invoke(state: GraphState) -> tuple[str, dict]:
     """Invoke seeder agent to generate alphas."""
-    page = random.randint(1, 100)
+    if random.random() < 0.3:
+        template = random.choice(TEMPLATES)["template"]
+        # TODO: Add random config
+        return (
+            f"Use following template, fix syntax and try replacing different datafields\n```\n{template}\n```",
+            {},
+        )
+
+    page = 0
     while True:
         database = Database()
         alphas = database.get_alphas_page(page)
@@ -17,7 +30,10 @@ def invoke(state: GraphState) -> Alpha:
             corr = alpha.update_self_correlation()
             database.upsert_alpha(alpha)
 
-            if score > 0 and corr < 0.5:
-                return alpha
+            if score > 0 and corr < 0.7 and alpha.sharpe > 1:
+                print(
+                    f"Found alpha: {alpha.regular} with score {score} and self-correlation {corr}"
+                )
+                return f"Use following alpha\n```\n{alpha.regular}\n```", get_config(alpha)
 
         page += 1

brain/agents/tester.py

@@ -8,7 +8,7 @@
 
 
 def invoke(state: GraphState) -> GraphState:
-    """Invoke fine-tuning agent."""
+    """Invoke tester agent, returns False if any test fails."""
     storage = state.get("storage")
 
     best_alpha = storage.best_alpha
@@ -31,8 +31,8 @@ def invoke(state: GraphState) -> GraphState:
 
         if (
             alpha.fitness < best_alpha.fitness * 0.5
-            and alpha.sharpe < best_alpha.sharpe * 0.5
-            and alpha.long_count + alpha.short_count < 400
+            or alpha.sharpe < best_alpha.sharpe * 0.5
+            or alpha.long_count + alpha.short_count < 400
         ):
             return False
 

brain/database.py

@@ -86,8 +86,10 @@ def get_alphas_page(
               FROM alphas
              WHERE self_correlation < %s
                AND (score IS NULL OR score > %s)
+               AND sharpe > 1.0
+               AND fitness > 0.9
                AND (long_count IS NULL OR long_count > 100)
-          ORDER BY sharpe DESC, fitness DESC
+          ORDER BY score DESC NULLS LAST, sharpe DESC, fitness DESC
              LIMIT %s
             OFFSET %s
         """

brain/graph_state.py

@@ -1,11 +1,13 @@
 from typing import TypedDict
 
+from brain.alpha_class import Alpha
 from brain.alpha_storage import Storage
 
 
 class GraphState(TypedDict, total=False):
     alpha_idea: str
     storage: Storage
+    old_best_alpha: Alpha | None
     # Number of times the alpha idea has been explored
     explore_count: int
     # Static fine-tuning