src/graph_definition.py

'''
The LangGraph Orchestrator that defines states, creates nodes (agents) and their edges (interdependencies and sequencing). It calls ticker_agent, sentiment_agent, options agent and ml_predictor_agent. This is the starting trigger point for all agents.

The final state for one ticker after calling all agents is like below.

{'tickers': ['AAPL'], 
 'sentiment': {'AAPL': 'positive'}, 
 'expiries': {'AAPL': '2026-12-18'}, 
 'top_strikes': {'AAPL': [{'strike': 70.0, 'openInterest': 2526, 'contractSymbol': 'AAPL261218C00070000'}]}, 
 'predicted_prices': {'AAPL': 206.58}, 
 'best_strike_prices': {'AAPL': {'strike': 70.0, 'openInterest': 2526, 'contractSymbol': 'AAPL261218C00070000'}} 
'''
import os
os.environ["HF_HOME"]="./hf_cache";

# graph_definition.py
from langgraph.graph import StateGraph, END
from langchain_core.runnables import RunnableLambda
from typing import List, Dict, Optional, Literal, TypedDict, Any

from ticker_agent import ticker_agent
from sentiment_agent_optimized import sentiment_agent
from options_agent import options_agent
from ml_predictor_agent_optimized import ml_predictor_agent

# Define shared state schema
class GraphState(TypedDict, total=False):
    tickers: List[str]
    sentiment: Dict[str, Literal["positive", "negative"]]
    expiries: Dict[str, Optional[str]]
    top_strikes: Dict[str, List[Dict]]  # ticker -> list of strikes
    predicted_prices: Dict[str, float]  # ticker -> predicted price
    best_strike_prices: Dict[str, Dict[str, Any]]  # ticker -> best matching option dict


def create_options_graph():
    builder = StateGraph(GraphState)

    # Register the nodes (all must be Runnables or callables)
    builder.add_node("TickerAgent", ticker_agent)
    builder.add_node("SentimentAgent", sentiment_agent)
    builder.add_node("OptionsAgent", options_agent)
    builder.add_node("MLPredictorAgent", ml_predictor_agent)
 
    # Set flow
    builder.set_entry_point("TickerAgent")
    builder.add_edge("TickerAgent", "SentimentAgent")
    builder.add_edge("SentimentAgent", "OptionsAgent")
    builder.add_edge("OptionsAgent", "MLPredictorAgent")
    builder.add_edge("MLPredictorAgent", END)

    # Call all agents as per the nodes and edges in graph configured above
    return builder.compile()