Biomanus_upload/demo_web/backend/session.py

import os
import uuid
from dataclasses import dataclass, field
from threading import Lock

from dotenv import load_dotenv

load_dotenv()

from biomni_agent import BiomniWebAgent

DEFAULT_LLM_PROVIDER = os.getenv("BIOMNI_LLM_PROVIDER", "deepseek").strip().lower()


@dataclass
class SessionState:
    agent: BiomniWebAgent | None = None
    uploaded_files: list[str] = field(default_factory=list)
    metabolomics_enabled: bool = False
    metabolomics_tools: list[str] = field(default_factory=list)
    paper_repro_enabled: bool = False
    paper_profile: str = "foxj1"
    stop_requested: bool = False
    llm_provider: str = field(default_factory=lambda: DEFAULT_LLM_PROVIDER)


def _create_agent(
    session_id: str,
    expected_files: list[str] | None = None,
    llm_provider: str = DEFAULT_LLM_PROVIDER,
) -> BiomniWebAgent:
    return BiomniWebAgent(
        session_id=session_id,
        expected_data_lake_files=expected_files or [],
        llm_provider=llm_provider,
    )


class SessionManager:
    def __init__(self):
        self._sessions: dict[str, SessionState] = {}
        self._lock = Lock()

    def create_session(self, llm_provider: str = DEFAULT_LLM_PROVIDER) -> str:
        sid = str(uuid.uuid4())
        with self._lock:
            # Lazy init agent to avoid blocking /session/new.
            self._sessions[sid] = SessionState(
                llm_provider=(llm_provider or DEFAULT_LLM_PROVIDER).strip().lower()
            )
        return sid

    def has_session(self, sid: str) -> bool:
        return sid in self._sessions

    def get_agent(self, sid: str) -> BiomniWebAgent:
        # Fast path: already initialized
        with self._lock:
            state = self._sessions[sid]
            if state.agent is not None:
                return state.agent
            expected_files = list(state.uploaded_files)
            llm_provider = state.llm_provider

        # Slow path: initialize outside lock
        new_agent = _create_agent(sid, expected_files=expected_files, llm_provider=llm_provider)
        if expected_files:
            new_agent.set_uploaded_files(expected_files)
        new_agent.set_metabolomics_settings(state.metabolomics_enabled, state.metabolomics_tools)
        new_agent.set_paper_reproduction_settings(state.paper_repro_enabled, state.paper_profile)

        # Publish initialized agent
        with self._lock:
            state = self._sessions.get(sid)
            if state is None:
                raise KeyError(f"Session not found: {sid}")
            if state.agent is None:
                state.agent = new_agent
            return state.agent

    def add_uploaded_file(self, sid: str, relative_path: str):
        with self._lock:
            state = self._sessions[sid]
            if relative_path not in state.uploaded_files:
                state.uploaded_files.append(relative_path)
            if state.agent is not None:
                state.agent.set_uploaded_files(state.uploaded_files)

    def set_metabolomics_settings(self, sid: str, enabled: bool, tools: list[str]):
        with self._lock:
            state = self._sessions[sid]
            state.metabolomics_enabled = bool(enabled)
            state.metabolomics_tools = sorted(set(tools))
            if state.agent is not None:
                state.agent.set_metabolomics_settings(state.metabolomics_enabled, state.metabolomics_tools)

    def set_llm_provider(self, sid: str, llm_provider: str):
        normalized = (llm_provider or DEFAULT_LLM_PROVIDER).strip().lower()
        with self._lock:
            state = self._sessions[sid]
            if state.llm_provider == normalized:
                return
            state.llm_provider = normalized
            # Recreate the agent on next access so the new backend takes effect.
            state.agent = None

    def request_stop(self, sid: str) -> None:
        with self._lock:
            state = self._sessions[sid]
            state.stop_requested = True
            if state.agent is not None:
                state.agent.request_cancel()

    def clear_stop(self, sid: str) -> None:
        with self._lock:
            state = self._sessions[sid]
            state.stop_requested = False
            if state.agent is not None:
                state.agent.clear_cancel_request()

    def is_stop_requested(self, sid: str) -> bool:
        with self._lock:
            state = self._sessions[sid]
            return bool(state.stop_requested)

    def set_paper_reproduction_settings(self, sid: str, enabled: bool, profile: str = "foxj1"):
        normalized_profile = (profile or "foxj1").strip().lower()
        with self._lock:
            state = self._sessions[sid]
            state.paper_repro_enabled = bool(enabled)
            state.paper_profile = normalized_profile
            if state.agent is not None:
                state.agent.set_paper_reproduction_settings(state.paper_repro_enabled, state.paper_profile)

    def get_uploaded_files(self, sid: str) -> list[str]:
        with self._lock:
            return list(self._sessions[sid].uploaded_files)

    def delete_session(self, sid: str):
        with self._lock:
            self._sessions.pop(sid, None)

    def cleanup_all(self):
        with self._lock:
            self._sessions.clear()