Upload llm.py

llm.py

@@ -180,86 +180,188 @@ def build_extraction_template(interviewee_type: str) -> str:
 
 def parse_structured_response(text: str, interviewee_type: str) -> Dict:
     """Extract structured data from LLM response"""
-    
+
+    log(f"Parsing response ({len(text)} chars) for type: {interviewee_type}")
+    log(f"Response preview: {text[:500]}...")
+
     # Try to find JSON block
     json_match = re.search(r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', text, re.DOTALL)
-    
+
     if json_match:
+        log(f"Found JSON match: {json_match.group()[:200]}...")
         try:
             data = json.loads(json_match.group())
-            log(f"Successfully extracted JSON: {data}")
+            log(f"✅ Successfully extracted JSON with {len(data)} fields: {list(data.keys())}")
             return data
-        except json.JSONDecodeError:
-            log("Failed to parse JSON from response")
-    
+        except json.JSONDecodeError as e:
+            log(f"❌ JSON parsing failed: {e}")
+            log(f"Attempted to parse: {json_match.group()[:300]}")
+    else:
+        log("⚠️ No JSON block found in response, using regex fallback")
+
     # Fallback: Extract from text using patterns
     data = {}
-    
+
     if interviewee_type == "HCP":
+        log("Using HCP extraction patterns...")
         # Extract diagnoses
         diag_pattern = r'(?:diagnos[ei]s|condition):\s*([^\n]+)'
         data["diagnoses"] = re.findall(diag_pattern, text, re.IGNORECASE)
-        
+
         # Extract prescriptions
         rx_pattern = r'(?:prescri[bp]\w*|medication):\s*([^\n]+)'
         data["prescriptions"] = re.findall(rx_pattern, text, re.IGNORECASE)
-        
+
         # Extract treatment rationale
         treat_pattern = r'(?:treatment|therapy|rationale):\s*([^\n]+)'
         data["treatment_rationale"] = re.findall(treat_pattern, text, re.IGNORECASE)
-        
+
     elif interviewee_type == "Patient":
+        log("Using Patient extraction patterns...")
         # Extract symptoms
         symptom_pattern = r'(?:symptom|complaint|experienc\w*):\s*([^\n]+)'
         data["symptoms"] = re.findall(symptom_pattern, text, re.IGNORECASE)
-        
+
         # Extract concerns
         concern_pattern = r'(?:concern|worry|question|anxious):\s*([^\n]+)'
         data["concerns"] = re.findall(concern_pattern, text, re.IGNORECASE)
-        
+
         # Extract side effects
         se_pattern = r'(?:side effect|adverse|reaction):\s*([^\n]+)'
         data["side_effects"] = re.findall(se_pattern, text, re.IGNORECASE)
-    
+
     # Clean and deduplicate
     for key in data:
         data[key] = list(set([item.strip() for item in data[key] if item.strip()]))
-    
-    log(f"Extracted data from text: {data}")
+
+    log(f"Fallback extraction result: {len(data)} fields, {sum(len(v) for v in data.values())} total items")
+    log(f"Extracted fields: {data}")
     return data
 
 
-def query_llm_hf_api(prompt: str, max_tokens: int = 500) -> str:
-    """Use Hugging Face Inference API for better quality"""
+def query_llm_hf_api(prompt: str, max_tokens: int = 1500) -> str:
+    """Use Hugging Face Inference API with proper authentication"""
+    import requests
+    import json
+
+    hf_token = os.getenv("HUGGINGFACE_TOKEN", "")
+
+    if not hf_token:
+        error_msg = "[Error] HUGGINGFACE_TOKEN not set in environment!"
+        print(f"❌ {error_msg}")
+        return error_msg
+
+    print(f"[HF API] Using token for authentication: {hf_token[:20]}...")
+
     try:
-        from huggingface_hub import InferenceClient
-        
-        client = InferenceClient(token=HF_TOKEN)
-        
-        # Use chat completions instead
-        messages = [
-            {"role": "system", "content": "You are an expert transcript analyzer. Provide detailed, structured analysis."},
-            {"role": "user", "content": prompt}
-        ]
-        
-        response = client.chat_completion(
-            messages=messages,
-            model="microsoft/Phi-3-mini-4k-instruct",
-            max_tokens=max_tokens,
-            temperature=0.3
-        )
-        
-        return response.choices[0].message.content.strip()
-        
+        # Get model from environment variable (default to Phi-3 if not set)
+        hf_model = os.getenv("HF_MODEL", "microsoft/Phi-3-mini-4k-instruct")
+        API_URL = f"https://api-inference.huggingface.co/models/{hf_model}"
+
+        # Use Bearer token in Authorization header
+        headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "Content-Type": "application/json"
+        }
+
+        # Get temperature from environment
+        temperature = float(os.getenv("LLM_TEMPERATURE", "0.5"))
+
+        # Use the FULL prompt (don't truncate - the model can handle it)
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": max_tokens,  # Use parameter passed to function
+                "temperature": temperature,
+                "return_full_text": False
+            }
+        }
+
+        # Get timeout from environment
+        timeout = int(os.getenv("LLM_TIMEOUT", "60"))
+
+        print(f"[HF API] Calling {hf_model} ({max_tokens} tokens, temp={temperature})...")
+        response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
+
+        print(f"[HF API] Status code: {response.status_code}")
+
+        if response.status_code == 200:
+            result = response.json()
+            if isinstance(result, list) and len(result) > 0:
+                generated_text = result[0].get("generated_text", "")
+                print(f"[HF API] ✅ Response: {len(generated_text)} characters")
+                print(f"[HF API] First 200 chars: {generated_text[:200]}")
+                return generated_text
+            else:
+                print(f"[HF API] Unexpected response format: {result}")
+                return "[Error] Unexpected API response format"
+        elif response.status_code == 401:
+            print(f"[HF API] ❌ 401 Unauthorized - Token invalid or expired")
+            print(f"[HF API] Token used: {hf_token}")
+            print(f"[HF API] Response: {response.text[:500]}")
+            return "[Error] Invalid HuggingFace token - create a new one at https://huggingface.co/settings/tokens"
+        else:
+            print(f"[HF API] Failed with status {response.status_code}")
+            print(f"[HF API] Response: {response.text[:500]}")
+            return f"[Error] API returned status {response.status_code}"
+
     except Exception as e:
         import traceback
         full_error = traceback.format_exc()
-        log(f"HF API error: {e}\n{full_error}")
-        print(f"[HF API Full Error]\n{full_error}")  # Print to console
+        print(f"[HF API] Error:\n{full_error}")
         return f"[Error] HF API failed: {e}"
 
 
-def query_llm_local(prompt: str, max_tokens: int = 500) -> str:
+def query_llm_lmstudio(prompt: str, max_tokens: int = 1500) -> str:
+    """Query LM Studio local server (OpenAI-compatible API)"""
+    import requests
+    import json
+
+    lmstudio_url = os.getenv("LMSTUDIO_URL", "http://localhost:1234/v1/chat/completions")
+
+    print(f"[LM Studio] Calling {lmstudio_url}...")
+
+    try:
+        payload = {
+            "messages": [
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ],
+            "temperature": float(os.getenv("LLM_TEMPERATURE", "0.7")),
+            "max_tokens": max_tokens,
+            "stream": False
+        }
+
+        response = requests.post(lmstudio_url, json=payload, timeout=120)
+
+        print(f"[LM Studio] Status code: {response.status_code}")
+
+        if response.status_code == 200:
+            result = response.json()
+            generated_text = result["choices"][0]["message"]["content"]
+            print(f"[LM Studio] ✓ Response: {len(generated_text)} characters")
+            print(f"[LM Studio] First 300 chars: {generated_text[:300]}")
+            return generated_text
+        else:
+            error_msg = f"[Error] LM Studio returned status {response.status_code}: {response.text[:200]}"
+            print(f"[LM Studio] {error_msg}")
+            return error_msg
+
+    except requests.exceptions.ConnectionError:
+        error_msg = "[Error] Cannot connect to LM Studio. Make sure:\n1. LM Studio is running\n2. Server is started (in LM Studio's Server tab)\n3. A model is loaded\n4. Server is on http://localhost:1234"
+        print(f"[LM Studio] {error_msg}")
+        return error_msg
+    except Exception as e:
+        error_msg = f"[Error] LM Studio failed: {e}"
+        print(f"[LM Studio] {error_msg}")
+        import traceback
+        traceback.print_exc()
+        return error_msg
+
+
+def query_llm_local(prompt: str, max_tokens: int = 1500) -> str:
     """Local model optimized for L4 GPU"""
     try:
         from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -346,12 +448,17 @@ Be specific and include relevant details (dosages, durations, severity levels, e
         log(f"Prompt truncated to {len(full_prompt)} characters")
     
     def generate():
-        if os.getenv("USE_LMSTUDIO", "False").lower() == "true":
-            return query_llm_lmstudio(full_prompt, max_tokens=600)
-        elif USE_HF_API and HF_TOKEN:
-            return query_llm_hf_api(full_prompt, max_tokens=600)
+        # Check environment variables dynamically (not using module-level USE_HF_API)
+        use_lmstudio = os.getenv("USE_LMSTUDIO", "False").lower() == "true"
+        use_hf_api = os.getenv("USE_HF_API", "False").lower() == "true"
+        hf_token = os.getenv("HUGGINGFACE_TOKEN", "")
+
+        if use_lmstudio:
+            return query_llm_lmstudio(full_prompt, max_tokens=2000)
+        elif use_hf_api and hf_token:
+            return query_llm_hf_api(full_prompt, max_tokens=1500)
         else:
-            return query_llm_local(full_prompt, max_tokens=600)
+            return query_llm_local(full_prompt, max_tokens=1500)
     
     # Execute with timeout
     with ThreadPoolExecutor(max_workers=1) as executor:
@@ -359,13 +466,30 @@ Be specific and include relevant details (dosages, durations, severity levels, e
         try:
             response = future.result(timeout=timeout)
             log(f"LLM response received ({len(response)} chars)")
-            
+
             # Extract structured data if requested
             structured_data = {}
+            clean_response = response
             if extract_structured:
                 structured_data = parse_structured_response(response, interviewee_type)
-            
-            return response, structured_data
+
+                # Remove JSON blocks from the narrative text (handle nested braces)
+                # Remove all {....} blocks repeatedly until none remain
+                prev_response = ""
+                while prev_response != clean_response:
+                    prev_response = clean_response
+                    clean_response = re.sub(r'\{[^{}]*\}', '', clean_response, flags=re.DOTALL)
+
+                # Also remove common JSON artifacts
+                clean_response = re.sub(r'###\s*JSON\s*Structure:', '', clean_response, flags=re.IGNORECASE)
+                clean_response = re.sub(r'###\s*Analysis:', '', clean_response, flags=re.IGNORECASE)
+                clean_response = re.sub(r'###\s*Response:', '', clean_response, flags=re.IGNORECASE)
+                clean_response = re.sub(r'Please provide.*?structured JSON.*', '', clean_response, flags=re.IGNORECASE|re.DOTALL)
+
+                clean_response = clean_response.strip()
+                log(f"Cleaned response: {len(clean_response)} chars (removed JSON)")
+
+            return clean_response, structured_data
             
         except ThreadTimeout:
             log("LLM generation timed out")