File size: 5,287 Bytes
b2c86fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 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()
|