Refactor API configuration and implement rate limiting in agent calls

agent/nodes.py

@@ -13,14 +13,23 @@ 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
 
-API_BASE_URL = "https://api.openrouter.ai/v1"
-MODEL_NAME = "qwen/qwen3-235b-a22b:free"
+
+API_BASE_URL = "https://openrouter.ai/api/v1"
+MODEL_NAME = "openai/gpt-oss-120b"
 API_KEY_ENV_VAR = "OPENROUTER_API_KEY"
 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:
@@ -32,10 +41,12 @@ def _get_model() -> BaseChatModel:
     # )
 
     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,
+        rate_limiter=_rate_limiter,
         metadata={
             "reasoning": {
                 "effort": "high"  # Use high reasoning effort

app.py

@@ -2,24 +2,12 @@ import os
 import gradio as gr
 import requests
 import pandas as pd
-from langchain_openai import ChatOpenAI
 from os import getenv
 from dotenv import load_dotenv
-from typing import Annotated
-from pydantic import SecretStr
-
-from typing_extensions import TypedDict
-
-from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import add_messages
-import asyncio  # Added for async processing
-import time  # Added for rate limiting
-
-from langchain_community.tools import WikipediaQueryRun
-from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
-from langgraph.prebuilt import tools_condition
-from langgraph.checkpoint.memory import MemorySaver
-from langgraph.prebuilt import create_react_agent
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import RunnableConfig
+import asyncio
+from typing import cast
 
 # Phoenix imports
 from phoenix.otel import register
@@ -59,30 +47,17 @@ start_phoenix()
 class BasicAgent:
     def __init__(self):
         self.agent = get_agent()
-        self._last_request_time = 0
-        self._request_lock = asyncio.Lock()
 
     async def __call__(self, question: str) -> str:
         print(f"Agent received question: {question}")
         
-        # Rate limiting: ensure at least 1 second between requests
-        async with self._request_lock:
-            current_time = time.time()
-            time_since_last_request = current_time - self._last_request_time
-            if time_since_last_request < 1.0:
-                sleep_time = 1.0 - time_since_last_request
-                print(f"Rate limiting: sleeping for {sleep_time:.2f} seconds")
-                await asyncio.sleep(sleep_time)
-            
-            self._last_request_time = time.time()
-        
         # Create configuration like in main.py
         config = create_agent_config(app_name=APP_NAME)
         
         # Call the agent with the question and config (like main.py)
         answer = await self.agent.ainvoke(
-            {"messages": [{"role": "user", "content": question}]},
-            config=config
+            {"messages": [HumanMessage(content=question)]},
+            cast(RunnableConfig, config)
         )
 
         print(f"Agent returning answer: {answer}")