StrategyGeneratorV001 / strategy_generator.py
JuanFuriaz's picture
Upload folder using huggingface_hub
80ac665 verified
Raw
History Blame Contribute Delete
6.39 kB
import anthropic
import os
from dotenv import load_dotenv
import yaml
from openai import OpenAI
from huggingface_hub import InferenceClient
load_dotenv(override=True)
os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY')
os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY')
google_api_key = os.getenv('GOOGLE_API_KEY')
deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')
grok_api_key = os.getenv("XAI_API_KEY")
with open('config/var_dev.yaml', 'r') as f:
config = yaml.safe_load(f)
# Logging stuff
if config["local"]:
hf_token = os.getenv('HF_TOKEN') # Sign in to HuggingFace Hub
else:
hf_token = None
#huggingface_hub.login(hf_token)
# Initialize clients
openai = OpenAI()
deepseek_api= OpenAI(
api_key=deepseek_api_key,
base_url="https://api.deepseek.com"
)
gemini_api = OpenAI(
api_key=google_api_key,
base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
)
grok_api = OpenAI(api_key=grok_api_key, base_url="https://api.x.ai/v1")
claude = anthropic.Anthropic()
client = InferenceClient() # For HuggingFace Inference API
LLM_MODEL = None
def user_prompt_for(user_msg):
return f"""
Trading strategy description:
\"\"\"{user_msg}\"\"\"
Task:
- Convert the description into executable Python code.
- Use only the library {config["api_fin"]}.
- Respond only with valid Python code, following Python best practices
"""
example_1= f'''
# User prompt:
# "Go long when the 10-period SMA crosses above the 100-period SMA,
# and exit when the 10-period SMA crosses below the 100-period SMA."
# Generated Python code:
import backtrader as bt
class SmaCross(bt.Strategy):
"""
Simple moving average crossover strategy.
Buy when fast SMA crosses above slow SMA.
Sell when fast SMA crosses below slow SMA.
"""
params = dict(pfast=10, pslow=100)
def __init__(self):
self.sma_fast = bt.ind.SMA(period=self.p.pfast)
self.sma_slow = bt.ind.SMA(period=self.p.pslow)
self.crossover = bt.ind.CrossOver(self.sma_fast, self.sma_slow)
def next(self):
if not self.position:
if self.crossover > 0: # Golden cross
self.buy()
elif self.crossover < 0: # Death cross
self.close()
# Initialize Cerebro
cerebro = bt.Cerebro()
cerebro.addstrategy(SmaCross, pfast=10, pslow=100)
'''
list_of_pyclasses = [example_1]
system_message = f'''
You are a financial assistant specialized in transforming natural language descriptions of trading strategies into clean, production-ready Python code.
Guidelines:
- Use only the library {config["api_fin"]}.
- Always create a class with the abreviation of the strategy with the form `NameOfStrategy(bt.Strategy)`.
- Implement strategy logic in `__init__` (indicators/signals) and `next()` (trade execution).
- Implement the strategy for this intervall of time {config["interval"]}
- Finish with initializing the strategy in Cerebro:
cerebro = bt.Cerebro()
cerebro.addstrategy(MyStrategy, param1=value, param2=value)
- Keep code minimal, clear, and follow Java best practices (PEP8, clear naming, modularity).
- If a strategy cannot be implemented with {config["api_fin"]}, respond with: "Unable to implement with {config["api_fin"]}."
- If used any addional libraries, add it in the code: import MyUsedLibrary
- If you don't know the answer, just say that you don't know, don't try to make up an answer.
Example(s) of transformation from user prompt of Python code: \n
'''
for pyclass in list_of_pyclasses:
system_message += pyclass
# Messages in Openai format
def messages_for(user_msg):
return [
{"role": "system", "content": system_message},
{"role": "user", "content": user_prompt_for(user_msg)}
]
def stream_llms(user_msg, typ_llm="gpt"):
""" Stream responses from different LLMs based on user message and selected model. """
global LLM_MODEL
llm_key = typ_llm.lower()
model_overrides = {
"deepseek": ("DEEPSEEK_MODEL", config.get("deepseek_model")),
"gimini": ("GEMINI_MODEL", config.get("gemini_model")),
"grok4": ("GROK4_MODEL", config.get("grok4_model")),
"claude": ("CLAUDE_MODEL", config.get("claude_model")),
"gpt": ("OPENAI_MODEL", config.get("openai_model")),
}
if llm_key not in model_overrides:
raise ValueError("Unknown model")
env_var, default_model = model_overrides[llm_key]
if not default_model:
raise KeyError(f"Missing default model configuration for '{llm_key}'")
selected_model = os.getenv(env_var) or default_model
LLM_MODEL = selected_model
messages = messages_for(user_msg)
if llm_key == "deepseek":
stream = deepseek_api.chat.completions.create(
model=selected_model,
messages=messages,
stream=True
)
elif llm_key == "gimini":
stream = gemini_api.chat.completions.create(
model=selected_model,
messages=messages,
stream=True
)
elif llm_key == "grok4":
stream = grok_api.chat.completions.create(
model=selected_model,
messages=messages,
stream= True
)
elif llm_key == "claude":
stream = claude.messages.stream(
model=selected_model,
max_tokens=2000,
system=messages[0]['content'],
messages=[messages[1]],
)
elif llm_key == "gpt":
stream = openai.chat.completions.create(model=selected_model, messages=messages, stream=True)
else:
raise ValueError("Unknown model")
reply = f"# Model {LLM_MODEL}\n"
if typ_llm.lower() == "claude":
with stream as stream_clde:
for fragment in stream_clde.text_stream:
reply += fragment
yield reply.replace("```python\n","").replace("```","")
else:
for chunk in stream:
if chunk and chunk.choices:
fragment = chunk.choices[0].delta.content or ""
reply += fragment
yield reply.replace("```python\n","").replace("```","")
def stream_manager(user_msg, model):
""" Streaming manager for different LLMs based on user message and selected model. """
result = stream_llms(user_msg, model)
for stream_so_far in result:
yield stream_so_far