chore: Perform codebase cleanup and refactor App state management

This commit is contained in:
Yunxiao Xu
2026-02-17 02:34:47 -08:00
parent ec6760b5a7
commit a94cbc7f6d
9 changed files with 90 additions and 46 deletions

View File

@@ -9,6 +9,7 @@ from ea_chatbot.graph.checkpoint import get_checkpointer
from ea_chatbot.history.models import User as UserDB, Conversation
from ea_chatbot.history.utils import map_db_messages_to_langchain
from ea_chatbot.api.schemas import ChatRequest
from ea_chatbot.utils.plots import fig_to_bytes
import io
import base64
from langchain_core.runnables.config import RunnableConfig
@@ -90,9 +91,7 @@ async def stream_agent_events(
plots = output["plots"]
encoded_plots: list[str] = []
for fig in plots:
buf = io.BytesIO()
fig.savefig(buf, format="png")
plot_bytes = buf.getvalue()
plot_bytes = fig_to_bytes(fig)
assistant_plots.append(plot_bytes)
encoded_plots.append(base64.b64encode(plot_bytes).decode('utf-8'))
output_event["data"]["encoded_plots"] = encoded_plots

View File

@@ -6,6 +6,7 @@ from ea_chatbot.graph.workflow import app
from ea_chatbot.graph.state import AgentState
from ea_chatbot.utils.logging import get_logger
from ea_chatbot.utils.helpers import merge_agent_state
from ea_chatbot.utils.plots import fig_to_bytes
from ea_chatbot.config import Settings
from ea_chatbot.history.manager import HistoryManager
from ea_chatbot.auth import OIDCClient, AuthType, get_user_auth_type
@@ -419,9 +420,7 @@ def main():
for fig in final_state["plots"]:
st.pyplot(fig)
# Convert fig to bytes
buf = io.BytesIO()
fig.savefig(buf, format="png")
plot_bytes_list.append(buf.getvalue())
plot_bytes_list.append(fig_to_bytes(fig))
if final_state.get("dfs"):
for df_name, df in final_state["dfs"].items():

View File

@@ -162,16 +162,6 @@ class OIDCClient:
except JWTError as e:
raise ValueError(f"ID Token validation failed: {str(e)}")
def get_login_url(self) -> str:
"""Legacy method for generating simple authorization URL."""
metadata = self.fetch_metadata()
authorization_endpoint = metadata.get("authorization_endpoint")
if not authorization_endpoint:
raise ValueError("authorization_endpoint not found in OIDC metadata")
uri, state = self.oauth_session.create_authorization_url(authorization_endpoint)
return uri
def exchange_code_for_token(self, code: str, code_verifier: Optional[str] = None) -> Dict[str, Any]:
"""Exchange the authorization code for an access token, optionally using PKCE verifier."""
metadata = self.fetch_metadata()

View File

@@ -19,7 +19,7 @@ class Settings(BaseSettings):
data_dir: str = "data"
data_state: str = "new_jersey"
log_level: str = Field(default="INFO", alias="LOG_LEVEL")
dev_mode: bool = Field(default=True, alias="DEV_MODE")
dev_mode: bool = Field(default=False, alias="DEV_MODE")
frontend_url: str = Field(default="http://localhost:5173", alias="FRONTEND_URL")
# Voter Database configuration

View File

@@ -0,0 +1,17 @@
import io
import matplotlib.pyplot as plt
def fig_to_bytes(fig: plt.Figure, format: str = "png") -> bytes:
"""
Convert a Matplotlib figure to bytes.
Args:
fig: The Matplotlib figure to convert.
format: The image format to use (default: "png").
Returns:
bytes: The image data.
"""
buf = io.BytesIO()
fig.savefig(buf, format=format)
return buf.getvalue()

View File

@@ -55,15 +55,6 @@ def test_oidc_fetch_jwks(oidc_config, mock_metadata):
client.fetch_jwks()
assert mock_get.call_count == 1
def test_oidc_get_login_url_legacy(oidc_config, mock_metadata):
client = OIDCClient(**oidc_config)
client.metadata = mock_metadata
with patch.object(client.oauth_session, "create_authorization_url") as mock_create:
mock_create.return_value = ("https://url", "state")
url = client.get_login_url()
assert url == "https://url"
def test_oidc_get_user_info(oidc_config, mock_metadata):
client = OIDCClient(**oidc_config)
client.metadata = mock_metadata