WIP: New architecture - multi-simulation

Author: Breta01

brain/agent.py

@@ -1,13 +1,25 @@
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 from langchain_core.messages import AnyMessage
+from langchain_core.messages.utils import count_tokens_approximately
 from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.prebuilt import create_react_agent
 from langgraph.prebuilt.chat_agent_executor import AgentState
+from langmem.short_term import SummarizationNode
 
-from brain.agent_config import get_universe_config
+from brain.agent_config import DEFAULT_CONFIG, get_universe_config
+from brain.api import DEFAULT_CONFIG as API_DEFAULT_CONFIG
+from brain.api import BrainAPI
 from brain.model import MODEL
 from brain.tools.datafields import get_random_datafields, search_datafields
 from brain.tools.operators import describe_operators, operators
-from brain.tools.simulation import simulate_alpha
+from brain.tools.simulation import submit_alpha
+
+MAX_WORKERS = 3
+
+from langgraph.graph import END
+from langgraph.types import Command
 
 
 def prompt(
@@ -19,11 +31,24 @@ def prompt(
     #       - store best alphas so far
     conf = get_universe_config(config)
 
+    if state.get("end", False):
+        state["is_last_step"] = True
+        return "END (do not continue)"
+
     prompt = f"""
-You are a quant researcher using World Quant Brain platform.
-You trying to find best alpha which passes all the checks.
-You can simulate alpha expressions and get results.
-Return alpha passing all the checks.
+You are a quant researcher using World Quant Brain platform. Based on provided strategy, create
+a new alpha.
+
+Steps:
+1. Propose the alpha and run the alpha simulation.
+2. Check the resulting metrics of alpha (try to maximize fitness)
+3. Repeat the process until the alpha is satisfactory (passing all the checks).
+
+- Try applying different functions and changing parameter values.
+- 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.
+- Change alpha completely if you think it is not good enough.
 
 Working with equity in universe: {conf['universe']}, region: {conf['region']}
 and delay {conf['delay']}.
@@ -34,8 +59,118 @@ def prompt(
     return [{"role": "system", "content": prompt}] + state["messages"]
 
 
+summarization_node = SummarizationNode(
+    token_counter=count_tokens_approximately,
+    model=MODEL,
+    max_tokens=2000,
+    max_summary_tokens=1000,
+    output_messages_key="llm_input_messages",
+)
+
+
+class CustomState(AgentState):
+    end: bool = False
+
+
 agent = create_react_agent(
     model=MODEL,
-    tools=[simulate_alpha, describe_operators, search_datafields, get_random_datafields],
+    tools=[submit_alpha, describe_operators, search_datafields, get_random_datafields],
     prompt=prompt,
+    state_schema=CustomState,
 )
+
+
+def create_alpha_simulation(alphas_dict, alphas_categories):
+    """Create a new alpha based on the given ID."""
+    alphas = []
+    agent.invoke(
+        {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": """
+Your task is to create a new alpha. Below is possible suggestion for the alpha.
+
+/*
+ HYPOTHESIS: "Fair value accounting" aims to measure and report on-going valuations of assets and liabilities. A recent increase in the fair value of liabilities may indicate a higher cost to be borne by the company in the future. This may deteriorate the company's financial health, potentially leading to lower profitability or financial distress.
+
+ IMPLEMENTATION: Using ts_rank() operator, sell stocks when there is an increase in the fair value of liabilities within a year using fundamental data.
+
+ HINT TO IMPROVE THE ALPHA: Could observing the increase over a shorter period improve accuracy?
+ */
+
+ -ts_rank(fn_liab_fair_val_l1_a,252)
+
+""",
+                }
+            ]
+        },
+        config={
+            "configurable": {
+                **DEFAULT_CONFIG,
+                "alphas": alphas,
+            }
+        },
+    )
+
+    print("Alpha", alphas[-1])
+    return alphas[-1]
+
+
+def monitor_alpha(response, alpha_config):
+    """Monitor the alpha simulation."""
+    simulation_result = BrainAPI().simulation_progress(response)
+    if not simulation_result["completed"]:
+        return {
+            "alpha_id": None,
+            "simulate_data": alpha_config,
+            "error": simulation_result["error"],
+        }
+
+    BrainAPI().set_alpha_properties(simulation_result["result"]["id"])
+
+    return BrainAPI().get_specified_alpha_stats(
+        simulation_result["result"]["id"], alpha_config, **API_DEFAULT_CONFIG
+    )
+
+
+print("Creating alpha...")
+response, config = create_alpha_simulation({}, [])
+res = monitor_alpha(response, config)
+print("RES", res)
+quit()
+
+
+def main():
+    """Main function to run the agent."""
+    alphas_dict = {}
+    alphas_categories = {
+        "passing": [],
+        "failing": [],
+        "pending": [],
+    }
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
+        # map Future → URL so we know which request just finished
+        for _ in range(MAX_WORKERS):
+            response, alpha_config = generate_alpha()
+            live_jobs = {pool.submit(monitor_alpha, alpha_id): alpha_id}
+
+        while live_jobs:
+            for job in as_completed(live_jobs):
+                url = live_jobs.pop(job)  # remove from “running” set
+
+                try:
+                    body = job.result()
+                    print(f"[OK] {url}")
+
+                except Exception as exc:
+                    print(f"[ERR] {url} – {exc}")
+                    break  # go back to outer loop
+
+                # MAIN‑THREAD logic: decide next jobs based on *this* response
+                for nxt in extract_follow_up(body):
+                    if nxt not in seen:  # avoid loops / dupes
+                        seen.add(nxt)
+                        live_jobs[pool.submit(fetch, nxt)] = nxt
+                break

brain/api.py

@@ -564,7 +564,7 @@ def _request(self, method: str, endpoint: str, max_retries: int = 200, **kwargs)
             else:
                 return response
 
-        raise Exception(f"Max retries exceeded for {method} request to {endpoint}")
+        raise TimeoutError(f"Max retries exceeded for {method} request to {endpoint}")
 
     def _authenticate(self):
         """Authenticate user."""

brain/main.py

@@ -17,16 +17,27 @@ def main():
             "messages": [
                 {
                     "role": "user",
-                    "content": (
-                        "Come up with a new alpha expression, start by exploring possible data"
-                        "fields and then test it."
-                    ),
+                    "content": """
+Your task is to create a new alpha. Below is possible suggestion for the alpha.
+
+/*
+ HYPOTHESIS: "Fair value accounting" aims to measure and report on-going valuations of assets and liabilities. A recent increase in the fair value of liabilities may indicate a higher cost to be borne by the company in the future. This may deteriorate the company's financial health, potentially leading to lower profitability or financial distress.
+
+ IMPLEMENTATION: Using ts_rank() operator, sell stocks when there is an increase in the fair value of liabilities within a year using fundamental data.
+
+ HINT TO IMPROVE THE ALPHA: Could observing the increase over a shorter period improve accuracy?
+ */
+
+ -ts_rank(fn_liab_fair_val_l1_a,252)
+
+""",
                 }
             ]
         },
-        config={"configurable": DEFAULT_CONFIG},
+        config={"configurable": {**DEFAULT_CONFIG, "thread_id": "1"}},
     )
-    result.pretty_print()
+    for msg in result["messages"]:
+        print(f"{type(msg)}, {msg}")
 
 
 if __name__ == "__main__":

brain/tools/simulation.py

@@ -1,35 +1,53 @@
 """Agent tool for simulating alpha expressions and obtaining output."""
 
+import time
+from typing import Annotated
+
 import pandas as pd
+from langchain_core.messages import ToolMessage
 from langchain_core.runnables import RunnableConfig
+from langchain_core.tools import InjectedToolCallId, tool
+from langgraph.graph import END
+from langgraph.types import Command
 
 from brain.agent_config import get_universe_config
 from brain.alpha_helpers import generate_alpha
 from brain.api import BrainAPI
 
 
-def simulate_alpha(alpha: str, config: RunnableConfig) -> str:
+def submit_alpha(
+    alpha: str, tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig
+) -> Command:
     """Execute alpha expression."""
     conf = get_universe_config(config)
     alpha_config = generate_alpha(alpha, **conf)
 
-    print("Alpha configuration:")
-    print(alpha_config)
+    print(f"Alpha config: {alpha_config}")
 
+    result, response = {}, ""
     try:
-        result = BrainAPI().simulate_alpha_list([alpha_config])[0]
+        response = BrainAPI().start_simulation(alpha_config)
+        endpoint = response.headers["Location"]
+        time.sleep(5)
+        result = BrainAPI().s.request("GET", endpoint).json()
     except Exception as e:
         print(f"Error during simulation: {e}")
         quit()
 
-    if error := result.get("error"):
+    if result.get("status", "") == "ERROR":
+        error = result.get("message", "Unknown error")
         print(f"Error occured: {error}")
         # TODO: More instructions on how to fix (operators, etc.)
         return f"Incorrect alpha. Fix following error and simulate it again: {error}"
 
-    response = format_alpha_prompt(result)
-    print(response)
-    return response
+    config["configurable"]["alphas"].append((response, alpha_config))
+
+    return Command(
+        update={
+            "end": True,
+            "messages": [ToolMessage(content=END, tool_call_id=tool_call_id)],
+        },
+    )
 
 
 def format_alpha_prompt(alpha_data):