akankshar639's picture
Update main.py
44de4c0 verified
from state import TraderState
from agents import hunter_agent, skeptic_agent, calibrator_agent, fetch_data
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
import os
import re
import json
import pytz
from datetime import datetime, timedelta
import yfinance as yf
load_dotenv()
# os.environ["OPENAI_API_KEY"] = os.getenv("OPENROUTER_API_KEY", "")
# llm_parser = ChatOpenAI(model="openai/gpt-3.5-turbo", openai_api_base="https://openrouter.ai/api/v1", temperature=0.0)
llm_parser = ChatGroq(model = "moonshotai/kimi-k2-instruct-0905",api_key=os.getenv("GROQ_API_KEY"), temperature=0.0)
def resolve_symbol(user_input: str) -> str:
"""
Handles Indian stocks (.NS/.BO).
"""
prompt = f"""
Analyze this user query: "{user_input}".
- Extract the company/brand name.
- Infer the likely stock market.
- Suggest the stock ticker symbol with the appropriate exchange suffix
(e.g., "GOOGL" for US, "TCS.NS" for NSE India).
- If it is an Indian stock, it MUST end with ".NS" or ".BO".
- Respond only with the symbol. If unsure, say "UNKNOWN".
"""
try:
suggested_symbol = llm_parser.invoke(prompt).content.strip().upper()
if suggested_symbol == "UNKNOWN":
return None
if not (suggested_symbol.endswith(".NS") or suggested_symbol.endswith(".BO")):
is_india_prompt = f"Is the stock/company '{suggested_symbol}' from the query '{user_input}' traded on Indian exchanges? Answer only YES or NO."
is_india = llm_parser.invoke(is_india_prompt).content.strip().upper()
if "YES" in is_india:
suggested_symbol += ".NS"
# Try NSE, fallback to BSE
ticker = yf.Ticker(suggested_symbol)
if not ticker.history(period="1d").empty:
return suggested_symbol
# If NSE (.NS) failed, it might be a BSE-only stock
if suggested_symbol.endswith(".NS"):
bse_symbol = suggested_symbol.replace(".NS", ".BO")
if not yf.Ticker(bse_symbol).history(period="1d").empty:
return bse_symbol
# Fallback
refine_prompt = f"The symbol '{suggested_symbol}' didn't validate. Suggest the EXACT ticker for: '{user_input}'. Respond ONLY with the ticker."
refined_symbol = llm_parser.invoke(refine_prompt).content.strip().upper()
# The same suffix logic
if not (refined_symbol.endswith(".NS") or refined_symbol.endswith(".BO")):
if "YES" in is_india:
refined_symbol += ".NS"
if not yf.Ticker(refined_symbol).history(period="1d").empty:
return refined_symbol
except Exception as e:
print(f"Symbol resolution error: {e}")
return None
def parse_query(user_input: str) -> tuple:
"""FIXED: Proper horizon detection based on correct trading logic"""
prompt = f"""
Analyze this user query: "{user_input}".
TRADING HORIZONS (CRITICAL):
- "intraday": same day buy and same day sell (hours)
- "scalping": ultra-short (minutes)
- "swing": buy this week, sell next week (1-4 weeks)
- "momentum": trend-following over weeks (2-4 weeks)
- "long_term": months to years (position trading, e.g., "bought 1 month ago")
Extract:
- stock symbol (e.g., "HINDUNILVR.NS")
- action: "buy" or "sell" (default "buy")
- horizon: match to above definitions
- date: "today", "tomorrow", or future (e.g., "after 2 weeks")
- holding_period: If mentioned (e.g., "1 week ago", "a week ago"), provide as "X unit ago". Else, None.
CRITICAL: If user says "bought X time ago" → horizon = "long_term"
If user says "tomorrow" without history → horizon = "intraday" (same-day strategy for next day)
Respond ONLY in JSON: {{"symbol": "HINDUNILVR.NS", "action": "sell", "horizon": "long_term", "date": "2025-12-27", "holding_period": "1 week ago"}}. Respond sholud be clean.
"""
try:
response = llm_parser.invoke(prompt).content.strip()
parsed = json.loads(response)
symbol = parsed.get("symbol")
action = parsed.get("action", "buy")
horizon = parsed.get("horizon", "intraday")
date_str = parsed.get("date")
holding_period = parsed.get("holding_period")
# Override horizon detection with explicit logic
if "ago" in user_input.lower() or "bought" in user_input.lower() or "purchased" in user_input.lower():
horizon = "long_term"
elif "tomorrow" in user_input.lower() and "ago" not in user_input.lower():
horizon = "intraday"
elif "week" in user_input.lower() and "ago" not in user_input.lower():
horizon = "swing"
elif "minute" in user_input.lower() or "scalp" in user_input.lower():
horizon = "scalping"
elif "month" in user_input.lower() and "ago" in user_input.lower():
horizon = "long_term"
ist = pytz.timezone('Asia/Kolkata')
now_ist = datetime.now(ist)
date = None
if date_str == "today":
date = now_ist.strftime("%Y-%m-%d")
elif date_str == "tomorrow":
date = (now_ist + timedelta(days=1)).strftime("%Y-%m-%d")
elif date_str and "after" in user_input.lower():
match = re.search(r'after\s+(\d+)\s*(week|month)', user_input.lower())
if match:
num, unit = int(match.group(1)), match.group(2)
days = num * 7 if unit == "week" else num * 30
date = (now_ist + timedelta(days=days)).strftime("%Y-%m-%d")
return symbol, date, action, horizon, holding_period
except Exception as e:
print(f"Parsing error: {e}. Fallback.")
symbol = resolve_symbol(user_input)
action = "sell" if "sell" in user_input.lower() else "buy"
# CRITICAL FIX: Correct horizon fallback
if "ago" in user_input.lower() or "month" in user_input.lower() or "bought" in user_input.lower():
horizon = "long_term"
elif "week" in user_input.lower() and "ago" not in user_input.lower():
horizon = "swing"
elif "minute" in user_input.lower():
horizon = "scalping"
elif "tomorrow" in user_input.lower():
horizon = "intraday"
else:
horizon = "intraday"
ist = pytz.timezone('Asia/Kolkata')
now_ist = datetime.now(ist)
date = now_ist.strftime("%Y-%m-%d") if "today" in user_input.lower() else (now_ist + timedelta(days=1)).strftime("%Y-%m-%d") if "tomorrow" in user_input.lower() else None
match = re.search(r'(\d+)\s*(week|month|day)\s*ago', user_input.lower())
holding_period = f"{match.group(1)} {match.group(2)} ago" if match else None
return symbol, date, action, horizon, holding_period
def generate_alert_message(state: TraderState, symbol: str, action: str, horizon: str) -> str:
confidence = state['confidence']
raw_data = state.get('raw_data', {}) or {}
indicators = raw_data.get("indicators", {})
news = raw_data.get("news", [])
prediction = raw_data.get("prediction", "No prediction available.")
holding = raw_data.get("holding", {})
claim = state.get('claim', 'No claim.')
skepticism = state.get('skepticism', 'No skepticism.')
features = raw_data.get("features", "")
rsi = indicators.get('rsi', 50)
macd = indicators.get('macd', 0)
current_price = indicators.get('current_price', 1)
volatility = 0
if current_price > 0:
bb_upper = indicators.get('bb_upper', current_price * 1.1)
bb_lower = indicators.get('bb_lower', current_price * 0.9)
volatility = (bb_upper - bb_lower) / current_price * 100
risk_adjustment = min(20, volatility / 5) if volatility > 0 else 0
if holding:
pnl_adjustment = 10 if holding.get('holding_return', 0) > 5 else -10 if holding.get('holding_return', 0) < -5 else 0
allocation = max(0, min(100, confidence - risk_adjustment + pnl_adjustment))
else:
allocation = max(0, min(100, confidence - risk_adjustment))
if indicators.get("error") == "No data available":
return f" No Data Available\nSorry, we couldn't fetch reliable data for {symbol}. Check official sources like NSE or Yahoo Finance. Try searching for '{symbol}.NS' if it's an Indian stock."
# HORIZON-AWARE best date suggestion
horizon_desc = {
"intraday": "same-day trading (hours)",
"scalping": "ultra-short (minutes)",
"swing": "1-4 weeks",
"momentum": "2-4 weeks",
"long_term": "months to years"
}
best_date_prompt = f"""
Based on prediction: '{prediction}', indicators: RSI {rsi:.1f}, MACD {macd:.2f}, trend: {features.split(';')[0] if features else 'N/A'}.
User query involves '{action}' for '{horizon}' ({horizon_desc.get(horizon, horizon)}).
Suggest the SPECIFIC best date/time to {action} {symbol} for maximum benefit.
- For intraday/scalping: suggest time of day (e.g., "Sell tomorrow at 10:30 AM IST for 1.5% gain")
- For swing/momentum: suggest specific date within the horizon (e.g., "Buy on Jan 2, 2026")
- For long_term: suggest holding duration (e.g., "Hold until Feb 2026 for 8% gain")
Include % potential gain/loss. Be accurate and align with user intent.
"""
try:
best_date = llm_parser.invoke(best_date_prompt).content.strip()
except:
best_date = f"Based on trends, {action} within {horizon_desc.get(horizon, horizon)} for potential 1-3% {'gain' if action == 'buy' else 'exit'}."
why_prompt = f"""
Based on claim: '{claim}', skepticism: '{skepticism}', features: '{features}', prediction: '{prediction}', holding: '{holding}'.
RSI is {rsi:.1f} (low = oversold/buy signal, high = overbought/sell signal). MACD is {macd:.2f} (positive = bullish, negative = bearish).
Explain why the signal is Green/Yellow/Red for {action} {symbol} ({horizon} = {horizon_desc.get(horizon, horizon)}).
Use layman terms, reference indicators accurately, tie to data. Be genuine—align with RSI/MACD signals.
"""
try:
why = llm_parser.invoke(why_prompt).content.strip()
except:
why = f"Based on RSI ({rsi:.1f}) and MACD ({macd:.2f}), conditions are {'favorable' if confidence > 60 else 'mixed' if confidence > 40 else 'unfavorable'} for {action}."
if confidence >= 70:
color = " Green: 'Yes, Go Ahead!'"
should_action = f"Yes, {action} now. Allocate {allocation:.0f}% as signals are strong."
elif confidence >= 40:
color = " Yellow: 'Wait and Watch'"
should_action = f"Maybe, monitor closely. Allocate {allocation:.0f}% cautiously due to mixed signals and potential risk."
else:
color = " Red: 'No, Stop!'"
should_action = f"No, avoid {action}. Signals are weak; wait for better conditions."
holding_summary = ""
if holding:
h_return = holding['holding_return']
holding_summary = f"\n**Holding Summary**: Purchased {holding['days_held']} days ago at ₹{holding['purchase_price']:.2f}, current: ₹{holding.get('current_price', current_price):.2f}, return: {h_return:.1f}% ({holding['pnl']})."
if h_return > 0:
advice = "lock in gains"
elif h_return < 0:
advice = "wait for recovery or cut losses"
else:
advice = "exit at breakeven"
should_action += f" Based on {h_return:.1f}% {holding['pnl']}, {advice}."
report = f"""
--- AI Trader Analyzer Report ---
Symbol: {symbol}
**Current Price**: {current_price:,.4f}
Action: {action.capitalize()} ({horizon.replace('_', '-').capitalize()} = {horizon_desc.get(horizon, horizon)}) "/n"
Date: {state.get('query_date', 'N/A')} (IST)
Signal: {color}
Confidence: {confidence}%
Key Metrics: RSI {rsi:.1f}, MACD {macd:.2f}, Trend {features.split(';')[0] if features else 'N/A'}
{holding_summary}
Should I {action.capitalize()}?** {should_action}
Best Date/Time to {action.capitalize()}**: {best_date}
Why? {why}
Prediction: {prediction}
Recommendation: Allocate {allocation:.0f}% based on analysis. Not financial advice.
"""
return report.strip()
def main():
symbol, date, action, horizon, holding_period = None, None, None, None, None
while not symbol:
user_input = input("Describe the stock (e.g., 'I want to buy Apple stock for today' or 'I purchased TCS.NS stock 1 month ago, now sell today'): ").strip()
symbol, date, action, horizon, holding_period = parse_query(user_input)
if not symbol:
print("No symbol detected. Try again.")
print(f"Analyzing: {symbol} on {date or 'recent'} (IST) for {action} ({horizon})" + (f", holding: {holding_period}" if holding_period else ""))
# Fetch data ONCE before the loop
hist = fetch_data(symbol, horizon)
state = TraderState(
input_symbol=symbol,
query_date=date,
action=action,
horizon=horizon,
holding_period=holding_period,
# Pre-load data
raw_data={'hist': hist},
claim=None,
skepticism=None,
confidence=50,
iterations=0,
stop=False,
alert_message=None
)
max_iterations = 5
for _ in range(max_iterations):
state = hunter_agent(state)
state = skeptic_agent(state)
state = calibrator_agent(state)
if state['stop']:
break
state['alert_message'] = generate_alert_message(state, symbol, action, horizon)
result = {
"symbol": symbol,
"date": date,
"action": action,
"horizon": horizon,
"holding_period": holding_period,
"alert_message": state['alert_message'],
"iterations": state['iterations'],
"confidence": state['confidence']
}
print("\n--- AI Trader Analyzer Result ---")
print(f"Symbol: {result['symbol']}")
print(f"Date: {result['date']} (IST)")
print(f"Action: {result['action']}")
print(f"Horizon: {result['horizon']}")
print(f"Holding: {result['holding_period'] or 'None'}")
print(f"Message:\n{result['alert_message']}")
print(f"Iterations: {result['iterations']}")
print(f"Confidence: {result['confidence']}%")
if __name__ == "__main__":
main()