cdshelat/MorphAI
0
1"""2chat.py — Multi-provider AI chat interface for PhysicsSim.3 4Uses litellm as a universal gateway so the user can choose any LLM provider5(OpenAI, Anthropic Claude, Gemini, Ollama, Azure, etc.) by supplying their own6API key and model string in the UI. No key is ever written to disk.7"""8 9import json10import os11import re12 13# Provider presets shown in the UI dropdown14PROVIDER_PRESETS = {15 "Anthropic Claude": {16 "models": ["claude-sonnet-4-6", "claude-opus-4-6", "claude-haiku-4-5-20251001"],17 "key_env": "ANTHROPIC_API_KEY",18 "key_prefix": "sk-ant-",19 "litellm_prefix": "", # litellm uses bare model name for Claude20 },21 "OpenAI": {22 "models": ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "o3-mini"],23 "key_env": "OPENAI_API_KEY",24 "key_prefix": "sk-",25 "litellm_prefix": "",26 },27 "Google Gemini": {28 "models": ["gemini/gemini-1.5-pro", "gemini/gemini-1.5-flash"],29 "key_env": "GEMINI_API_KEY",30 "key_prefix": "",31 "litellm_prefix": "gemini/",32 },33 "Ollama (local)": {34 "models": ["ollama/llama3.2", "ollama/llama3.1", "ollama/mistral", "ollama/phi3"],35 "key_env": "",36 "key_prefix": "",37 "litellm_prefix": "ollama/",38 },39}40 41# System prompt that seeds any LLM with the task context42SYSTEM_PROMPT = """You are a structural engineering assistant for a topology optimization tool called PhysicsSim.43Your job is to extract simulation parameters from a user's natural language description of a mechanical part.44 45The design space is a rectangular bounding box with exactly 6 named faces:46 left, right, top, bottom, front, back47 48Rules:49- fixed_face = the face that is bolted/welded/attached to a rigid surface50- load_face = the face where the external force is applied51- force_direction must be one of: -X, +X, -Y, +Y, -Z, +Z52 (-Y = downward/gravity, +Y = upward, -X/+X = horizontal left/right, -Z/+Z = depth)53- Mass → Force: multiply kg × 9.81 to get Newtons (e.g. "20 kg" → 196.2 N)54- If no material is mentioned, default to "PLA (Bioplastic)"55- If no safety factor is mentioned, default to 2.056- If no volume fraction / fill percentage is mentioned, default to 0.457 58Available materials (use EXACT name from this list):59 PLA (Bioplastic), PETG (Engineering Plastic), ABS (Acrylonitrile Butadiene Styrene),60 Nylon PA12, TPU (Flexible), Carbon Fiber PETG (Composite),61 Titanium Ti-6Al-4V (Reference), Aluminum 6061 (Reference)62 63Respond ONLY with a valid JSON object — no markdown, no prose, no code fences. Example:64{65 "fixed_face": "left",66 "load_face": "right",67 "force_direction": "-Y",68 "applied_force_n": 196.2,69 "material": "PLA (Bioplastic)",70 "safety_factor": 2.0,71 "volume_fraction": 0.4,72 "load_scenario": "Downward only",73 "confidence_notes": "Assumed downward gravity for 20kg load."74}75 76load_scenario must be one of:77 Downward only, Lateral only, Down + Lateral, Down + Upward (top),78 Down + Downward (bot), Symmetric (top+bot), Torsion (top+bot opp)79"""80 81REQUIRED_KEYS = [82 "fixed_face", "load_face", "force_direction",83 "applied_force_n", "material",84]85 86VALID_FACES = {"left", "right", "top", "bottom", "front", "back"}87VALID_DIRECTIONS = {"-X", "+X", "-Y", "+Y", "-Z", "+Z"}88VALID_SCENARIOS = {89 "Downward only", "Lateral only", "Down + Lateral",90 "Down + Upward (top)", "Down + Downward (bot)",91 "Symmetric (top+bot)", "Torsion (top+bot opp)",92}93 94# Map from chat-extracted face names (no suffix) → app FACES list values95FACE_MAP = {96 "left": "Left (X=0)",97 "right": "Right (X=W)",98 "bottom": "Bottom (Y=0)",99 "top": "Top (Y=H)",100 "front": "Front (Z=0)",101 "back": "Back (Z=D)",102}103 104 105def get_api_key_from_env(provider_name: str) -> str:106 """Check environment / Streamlit secrets for a pre-configured API key."""107 preset = PROVIDER_PRESETS.get(provider_name, {})108 env_var = preset.get("key_env", "")109 if not env_var:110 return ""111 # Try st.secrets first (Streamlit Cloud / local secrets.toml)112 try:113 import streamlit as st114 return st.secrets.get(env_var, "")115 except Exception:116 pass117 return os.environ.get(env_var, "")118 119 120def extract_params(user_message: str, model: str, api_key: str,121 provider_name: str = "") -> dict:122 """Call the chosen LLM and extract simulation parameters as a dict.123 124 Returns a validated dict with keys matching REQUIRED_KEYS plus optional keys.125 Raises ValueError with a human-readable message on failure.126 """127 try:128 import litellm129 except ImportError:130 raise ImportError(131 "litellm is not installed. Run: pip install litellm"132 )133 134 # Build litellm kwargs135 kwargs = {136 "model": model,137 "messages": [138 {"role": "system", "content": SYSTEM_PROMPT},139 {"role": "user", "content": user_message},140 ],141 "max_tokens": 512,142 "temperature": 0.1,143 }144 145 # Set API key via environment variable so litellm picks it up146 preset = PROVIDER_PRESETS.get(provider_name, {})147 env_var = preset.get("key_env", "")148 original_val = None149 if env_var and api_key:150 original_val = os.environ.get(env_var)151 os.environ[env_var] = api_key152 153 # Ollama doesn't need a key154 if provider_name == "Ollama (local)":155 kwargs.pop("temperature", None)156 157 try:158 response = litellm.completion(**kwargs)159 raw = response.choices[0].message.content.strip()160 finally:161 # Restore env var162 if env_var and original_val is not None:163 os.environ[env_var] = original_val164 elif env_var and api_key:165 os.environ.pop(env_var, None)166 167 # Parse JSON — strip any accidental markdown fences168 raw_clean = re.sub(r"^```[a-z]*\n?|```$", "", raw.strip(), flags=re.MULTILINE).strip()169 try:170 params = json.loads(raw_clean)171 except json.JSONDecodeError as e:172 raise ValueError(173 f"Model returned invalid JSON: {e}\n\nRaw response:\n{raw[:400]}"174 )175 176 # Validate required keys177 missing = [k for k in REQUIRED_KEYS if k not in params]178 if missing:179 raise ValueError(180 f"Model response is missing required fields: {missing}\n"181 f"Try rephrasing your description to include: "182 f"which face is fixed, which face receives the load, "183 f"force direction, force magnitude, and material."184 )185 186 # Validate values187 if params.get("fixed_face") not in VALID_FACES:188 raise ValueError(189 f"fixed_face '{params.get('fixed_face')}' is not valid. "190 f"Must be one of: {sorted(VALID_FACES)}"191 )192 if params.get("load_face") not in VALID_FACES:193 raise ValueError(194 f"load_face '{params.get('load_face')}' is not valid. "195 f"Must be one of: {sorted(VALID_FACES)}"196 )197 if params.get("force_direction") not in VALID_DIRECTIONS:198 raise ValueError(199 f"force_direction '{params.get('force_direction')}' is not valid. "200 f"Must be one of: {sorted(VALID_DIRECTIONS)}"201 )202 203 # Apply defaults for optional keys204 params.setdefault("safety_factor", 2.0)205 params.setdefault("volume_fraction", 0.4)206 params.setdefault("load_scenario", "Downward only")207 params.setdefault("confidence_notes", "")208 209 # Clamp numbers to safe ranges210 params["applied_force_n"] = max(1.0, float(params["applied_force_n"]))211 params["safety_factor"] = max(1.0, min(5.0, float(params["safety_factor"])))212 params["volume_fraction"] = max(0.1, min(0.9, float(params["volume_fraction"])))213 214 # Map bare face name → app-style label with suffix215 params["fixed_face"] = FACE_MAP.get(params["fixed_face"], params["fixed_face"])216 params["load_face"] = FACE_MAP.get(params["load_face"], params["load_face"])217 218 # Ensure load_scenario is valid219 if params["load_scenario"] not in VALID_SCENARIOS:220 params["load_scenario"] = "Downward only"221 222 return params223 224 225def call_llm(model: str, api_key: str, prompt: str,226 provider_name: str = "") -> str:227 """Call the LLM with a plain-text prompt and return the response as a string.228 229 Unlike extract_params, this does NOT enforce JSON — it is used for free-form230 commentary (e.g. topology narration).231 """232 try:233 import litellm234 except ImportError:235 raise ImportError("litellm is not installed. Run: pip install litellm")236 237 preset = PROVIDER_PRESETS.get(provider_name, {})238 env_var = preset.get("key_env", "")239 # Auto-detect provider from model string when provider_name is empty240 if not env_var:241 if model.startswith("claude"):242 env_var = "ANTHROPIC_API_KEY"243 elif model.startswith("gpt") or model.startswith("o3") or model.startswith("o1"):244 env_var = "OPENAI_API_KEY"245 elif model.startswith("gemini"):246 env_var = "GEMINI_API_KEY"247 248 original_val = None249 if env_var and api_key:250 original_val = os.environ.get(env_var)251 os.environ[env_var] = api_key252 253 try:254 resp = litellm.completion(255 model=model,256 messages=[{"role": "user", "content": prompt}],257 max_tokens=400,258 temperature=0.4,259 )260 return resp.choices[0].message.content.strip()261 finally:262 if env_var and original_val is not None:263 os.environ[env_var] = original_val264 elif env_var and api_key:265 os.environ.pop(env_var, None)266 267 268def test_connection(model: str, api_key: str, provider_name: str) -> tuple[bool, str]:269 """Quick connectivity test — sends a tiny message to verify the key works.270 271 Returns (success: bool, message: str).272 """273 try:274 import litellm275 except ImportError:276 return False, "litellm not installed. Run: pip install litellm"277 278 preset = PROVIDER_PRESETS.get(provider_name, {})279 env_var = preset.get("key_env", "")280 original_val = None281 if env_var and api_key:282 original_val = os.environ.get(env_var)283 os.environ[env_var] = api_key284 285 try:286 resp = litellm.completion(287 model=model,288 messages=[{"role": "user", "content": "Reply with the single word: OK"}],289 max_tokens=5,290 temperature=0,291 )292 text = resp.choices[0].message.content.strip()293 return True, f"Connected — model replied: \"{text}\""294 except Exception as e:295 return False, f"Connection failed: {str(e)[:200]}"296 finally:297 if env_var and original_val is not None:298 os.environ[env_var] = original_val299 elif env_var and api_key:300 os.environ.pop(env_var, None)301