agent/nodes.py

from getpass import getpass
import os
from typing import Literal, cast
from langchain_core.tools import BaseTool
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.runnables import Runnable
from langchain_core.messages import BaseMessage
from langchain_core.language_models.base import LanguageModelInput
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from langchain_deepseek import ChatDeepSeek
from pydantic import BaseModel, Field, SecretStr
from agent.prompts import get_system_prompt
from agent.state import State
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.prebuilt import ToolNode
from langchain_core.rate_limiters import InMemoryRateLimiter

from agent.config import API_BASE_URL, MODEL_NAME, API_KEY_ENV_VAR, MODEL_TEMPERATURE

if API_KEY_ENV_VAR not in os.environ:
    print(f"Please set the environment variable {API_KEY_ENV_VAR}.")
    os.environ[API_KEY_ENV_VAR] = getpass(f"Enter your {API_KEY_ENV_VAR} (will not be echoed): ")

# Global singleton rate limiter
_rate_limiter = InMemoryRateLimiter(
    requests_per_second=1,
    check_every_n_seconds=0.1,  # Wake up every 100 ms to check whether allowed to make a request,
    max_bucket_size=5,  # Controls the maximum burst size.
)

### Helper functions ###

def _get_model() -> BaseChatModel:

    # api_key = os.getenv("GOOGLE_API_KEY")
    # return ChatGoogleGenerativeAI(
    #     api_key=SecretStr(api_key) if api_key else None,
    #     model="gemini-2.5-pro"
    # )

    # api_key = os.getenv(API_KEY_ENV_VAR)

    # return ChatOpenAI(
    #     api_key=SecretStr(api_key) if api_key else None,
    #     base_url=API_BASE_URL,
    #     model=MODEL_NAME,
    #     temperature=MODEL_TEMPERATURE if MODEL_TEMPERATURE else 0.0,
    #     rate_limiter=_rate_limiter,
    #     metadata={
    #         "reasoning": {
    #             "effort": "high"  # Use high reasoning effort
    #         }
    #     }
    # )

    return ChatDeepSeek(
        model="deepseek-chat",
        temperature=MODEL_TEMPERATURE if MODEL_TEMPERATURE else 0.0,
        max_retries=2,
    )

def _get_tools() -> list[BaseTool]:
    from tools import get_all_tools
    return get_all_tools()

def _bind_model(model: BaseChatModel) -> Runnable[LanguageModelInput, BaseMessage]:
    return model.bind_tools(_get_tools())

### NODES ###
# Call model node
def call_model(state: State, config) -> dict[str, list[BaseMessage]]:
    messages = state["messages"]
    app_name = config.get('configurable', {}).get("app_name", "OracleBot")
    
    # Add system prompt if not already present
    if not messages or messages[0].type != "system":
        # Use dynamic system prompt if sports are mentioned
        system_prompt = get_system_prompt()
        system_message: BaseMessage = SystemMessage(content=system_prompt)
        messages = [system_message] + list(messages)
    
    model = _get_model()
    model = _bind_model(model)
    response = model.invoke(messages)
    
    return {"messages": [response]}

# Tool node
tool_node = ToolNode(tools=_get_tools())