Update app.py

app.py

@@ -1,290 +1,297 @@
 #!/usr/bin/env python3
+# filename: app_refactored.py
 import os
 import json
 import logging
 import re
-from typing import Dict, Any, List, Optional
 from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from dotenv import load_dotenv
-from werkzeug.utils import secure_filename
-from langchain_groq import ChatGroq
-from typing_extensions import TypedDict
-
-# --- Type Definitions for State Management ---
-class TaggedReply(TypedDict):
-    reply: str
-    tags: List[str]
-
-class AssistantState(TypedDict):
-    conversationSummary: str
-    language: str 
-    taggedReplies: List[TaggedReply]
-    # Note: lastUserMessage is calculated on request, not stored in state
-
-# --- Logging ---
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger("code-assistant")
 
-# --- Load environment ---
+# Replace with your LLM client import; kept generic here.
+# from langchain_groq import ChatGroq
+
+# === Config ===
 load_dotenv()
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 if not GROQ_API_KEY:
-    logger.error("GROQ_API_KEY not set in environment")
-    # For deployment, consider raising an exception instead of exiting:
-    # raise ValueError("GROQ_API_KEY not set in environment") 
-    exit(1)
+    raise RuntimeError("GROQ_API_KEY not set in environment")
 
-# --- Flask app setup ---
-BASE_DIR = Path(__file__).resolve().parent
-static_folder = BASE_DIR / "static"
+LLM_MODEL = os.getenv("LLM_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
+LLM_TIMEOUT_SECONDS = float(os.getenv("LLM_TIMEOUT_SECONDS", "20"))
+MAX_HISTORY_MESSAGES = int(os.getenv("MAX_HISTORY_MESSAGES", "12"))
+VALID_LANGUAGES = {"python", "javascript", "java", "c++", "c#", "go", "ruby", "php", "typescript", "swift"}
 
-app = Flask(__name__, static_folder=str(static_folder), static_url_path="/static")
-CORS(app)
+# === Logging ===
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("code-assistant")
 
-# --- LLM setup ---
-llm = ChatGroq(
-    model=os.getenv("LLM_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct"),
-    temperature=0.1, # Set a lower, deterministic temperature
-    max_tokens=2048, # Ensure max_tokens is set to avoid truncation
-    api_key=GROQ_API_KEY,
+# === LLM client (example) ===
+# NOTE: adapt this block to match your SDK. Keep a tolerant accessor for response text.
+class DummyLLM:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+    def invoke(self, messages: List[Dict[str, str]], timeout: Optional[float] = None):
+        # stub: replace with real client call
+        class Resp: pass
+        r = Resp()
+        r.content = json.dumps({
+            "assistant_reply": "This is a dummy reply. Replace with real LLM client.",
+            "code_snippet": "",
+            "state_updates": {"conversationSummary": "dummy", "language": "Python"},
+            "suggested_tags": ["example"]
+        })
+        return r
+
+# llm = ChatGroq(model=LLM_MODEL, api_key=GROQ_API_KEY, temperature=0.1, max_tokens=2048)
+llm = DummyLLM(model=LLM_MODEL, api_key=GROQ_API_KEY)  # replace with real client
+
+# === Prompt ===
+SYSTEM_PROMPT = (
+    "You are an expert programming assistant. Prefer to return a JSON object with keys: "
+    "assistant_reply (string), code_snippet (string, optional, can be multiline), "
+    "state_updates (object), suggested_tags (array). If including code, put it in triple backticks. "
+    "Do NOT escape newlines in code_snippet; return natural multi-line strings."
 )
 
-# --- Define the standard error message for failed parsing
-LLM_PARSE_ERROR_MESSAGE = "I'm sorry, I couldn't process the last response correctly due to a formatting issue. Could you please rephrase or try a simpler query?"
-
-PROGRAMMING_ASSISTANT_PROMPT = f"""
-You are an expert programming assistant. Your role is to provide code suggestions, fix bugs, explain programming concepts, and offer contextual help based on the user's query and preferred programming language.
+# === Utilities ===
+def clamp_summary(s: str, max_len: int = 1200) -> str:
+    s = (s or "").strip()
+    return s if len(s) <= max_len else s[:max_len-3] + "..."
 
-**CONTEXT HANDLING RULES (Follow these strictly):**
-- **Conversation Summary:** At the end of every response, you MUST provide an updated, concise `conversationSummary` based on the entire chat history provided. This summary helps you maintain context.
-- **Language Adaptation:** Adjust your suggestions, code, and explanations to the programming language specified in the 'language' field of the 'AssistantState'.
-
-STRICT OUTPUT FORMAT (JSON ONLY):
-Return a single JSON object with the following keys. **The JSON object MUST be enclosed in a single ```json block.**
-- assistant_reply: string  // A natural language reply to the user (short and helpful). Do NOT include code blocks here.
-- code_snippet: string  // If suggesting code, provide it here in a markdown code block. **CRITICALLY, you must escape all internal newlines as '\\n' and backslashes as '\\\\'** to keep the string value valid JSON. If no code is required, use an empty string: "".
-- state_updates: object  // updates to the internal state, must include: language, conversationSummary
-- suggested_tags: array of strings // a list of 1-3 relevant tags for the assistant_reply
+def canonicalize_language(text: Optional[str]) -> Optional[str]:
+    if not text:
+        return None
+    t = text.strip().lower()
+    # quick membership test
+    for lang in VALID_LANGUAGES:
+        if lang in t or t == lang:
+            return lang
+    return None
 
-Rules:
-- ALWAYS include all four top-level keys: `assistant_reply`, `code_snippet`, `state_updates`, and `suggested_tags`.
-- ALWAYS include `assistant_reply` as a non-empty string.
-- Do NOT produce any text outside the JSON block.
-"""
+def try_parse_json(s: str) -> Optional[Dict[str, Any]]:
+    try:
+        return json.loads(s)
+    except Exception:
+        return None
 
-def extract_json_from_llm_response(raw_response: str) -> dict:
-    # The default object is only used if parsing fails, providing a clean error message.
-    # The actual state preservation logic is in the /chat route.
+def extract_code_fence(text: str) -> Optional[str]:
+    m = re.search(r"```(?:[a-zA-Z0-9_+\-]*)\n([\s\S]*?)```", text)
+    return m.group(1).strip() if m else None
+
+def parse_llm_output(raw: str) -> Dict[str, Any]:
+    """
+    Tolerant multi-strategy parser:
+      1) Direct JSON
+      2) JSON inside a ```json``` fence
+      3) Heuristic extraction: assistant_reply lines, code fences for code_snippet, simple state_updates line (json)
+    """
     default = {
-        "assistant_reply": LLM_PARSE_ERROR_MESSAGE,
+        "assistant_reply": "I couldn't parse the model response. Please rephrase or simplify the request.",
         "code_snippet": "",
-        "state_updates": {"conversationSummary": "", "language": "Python"},
+        "state_updates": {"conversationSummary": "", "language": "python"},
         "suggested_tags": [],
+        "parse_ok": False,
     }
-    
-    if not raw_response or not isinstance(raw_response, str):
+    if not raw or not isinstance(raw, str):
         return default
-    
-    # 1. Use regex to find the JSON content inside the first code block (```json)
-    m = re.search(r"```json\s*([\s\S]*?)\s*```", raw_response)
-    json_string = m.group(1).strip() if m else raw_response
-    
-    # 2. Refine candidate to just the JSON object content
-    first = json_string.find('{')
-    last = json_string.rfind('}')
-    candidate = json_string[first:last+1] if first != -1 and last != -1 and first < last else json_string
-    
-    # 3. Remove trailing commas which can break JSON parsing
-    candidate = re.sub(r',\s*(?=[}\]])', '', candidate)
-    
-    try:
-        parsed = json.loads(candidate)
-    except Exception as e:
-        logger.warning("Failed to parse JSON from LLM output: %s. Candidate: %s", e, candidate[:200]) # Truncate candidate for cleaner logs
-        return default
-    
-    # 4. Validate and clean up the parsed dictionary
-    if isinstance(parsed, dict) and "assistant_reply" in parsed:
+
+    raw = raw.strip()
+
+    # 1) direct JSON
+    parsed = try_parse_json(raw)
+    if parsed and isinstance(parsed, dict) and "assistant_reply" in parsed:
         parsed.setdefault("code_snippet", "")
         parsed.setdefault("state_updates", {})
-        parsed["state_updates"].setdefault("conversationSummary", "")
-        parsed["state_updates"].setdefault("language", "Python")
         parsed.setdefault("suggested_tags", [])
-        
-        # Ensure reply is not empty
-        if not parsed["assistant_reply"].strip():
-             parsed["assistant_reply"] = "I need a clearer instruction to provide a reply."
-        
+        parsed["parse_ok"] = True
         return parsed
+
+    # 2) JSON inside any code fence (```json ... ```)
+    m_json_fence = re.search(r"```json\s*([\s\S]*?)```", raw, re.IGNORECASE)
+    if m_json_fence:
+        candidate = m_json_fence.group(1)
+        parsed = try_parse_json(candidate)
+        if parsed and "assistant_reply" in parsed:
+            parsed.setdefault("code_snippet", "")
+            parsed.setdefault("state_updates", {})
+            parsed.setdefault("suggested_tags", [])
+            parsed["parse_ok"] = True
+            return parsed
+
+    # 3) Heuristics: find assistant_reply: ...; code fence for code; state_updates as inline JSON
+    assistant_reply = ""
+    code_snippet = ""
+    state_updates = {}
+    suggested_tags = []
+
+    # a) extract code fence (first code block)
+    code_snippet = extract_code_fence(raw) or ""
+
+    # b) extract assistant_reply by looking for lines like "assistant_reply:" or markdown bold
+    m = re.search(r'assistant_reply\s*[:\-]\s*(["\']?)([\s\S]*?)(?=\n[a-z_]+[\s\-:]{1}|$)', raw, re.IGNORECASE)
+    if m:
+        assistant_reply = m.group(2).strip()
     else:
-        logger.warning("Parsed JSON missing required keys or invalid format. Returning default.")
-        return default
+        # fallback: take everything up to the first code fence or up to "state_updates"
+        cut_idx = raw.find("```")
+        state_idx = raw.lower().find("state_updates")
+        end = min([i for i in (cut_idx if cut_idx>=0 else len(raw), state_idx if state_idx>=0 else len(raw))])
+        assistant_reply = raw[:end].strip()
+        # strip any leading labels like "**assistant_reply**:" or similar
+        assistant_reply = re.sub(r'^\**\s*assistant_reply\**\s*[:\-]?\s*', '', assistant_reply, flags=re.IGNORECASE).strip()
+
+    # c) find state_updates JSON block if present
+    m_state = re.search(r"state_updates\s*[:\-]?\s*(\{[\s\S]*?\})", raw, re.IGNORECASE)
+    if m_state:
+        try:
+            state_updates = json.loads(m_state.group(1))
+        except Exception:
+            state_updates = {}
+
+    # d) suggested_tags simple extract
+    m_tags = re.search(r"suggested_tags\s*[:\-]?\s*(\[[^\]]*\])", raw, re.IGNORECASE)
+    if m_tags:
+        try:
+            suggested_tags = json.loads(m_tags.group(1))
+        except Exception:
+            suggested_tags = []
+
+    result = {
+        "assistant_reply": assistant_reply or default["assistant_reply"],
+        "code_snippet": code_snippet or "",
+        "state_updates": state_updates or {"conversationSummary": "", "language": "python"},
+        "suggested_tags": suggested_tags or [],
+        "parse_ok": bool(assistant_reply or code_snippet),
+    }
+    return result
 
-def detect_language_from_text(text: str) -> Optional[str]:
-    """Simple check for common programming languages."""
-    if not text:
-        return None
-    lower = text.lower()
-    known_languages = ["python", "javascript", "java", "c++", "c#", "go", "ruby", "php", "typescript", "swift"]
-    
-    lang_match = re.search(r'\b(in|using|for)\s+(' + '|'.join(known_languages) + r')\b', lower)
-    if lang_match:
-        return lang_match.group(2).capitalize()
-    return None
+# === Flask app ===
+BASE_DIR = Path(__file__).resolve().parent
+app = Flask(__name__, static_folder=str(BASE_DIR / "static"), static_url_path="/static")
+CORS(app)
 
-# --- Flask routes ---
 @app.route("/", methods=["GET"])
 def serve_frontend():
     try:
-        return app.send_static_file("frontend.html") 
+        return app.send_static_file("frontend.html")
     except Exception:
         return "<h3>frontend.html not found in static/ — please add your frontend.html there.</h3>", 404
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    data = request.get_json(force=True)
-    if not isinstance(data, dict):
+    payload = request.get_json(force=True, silent=True)
+    if not isinstance(payload, dict):
         return jsonify({"error": "invalid request body"}), 400
 
-    chat_history: List[Dict[str, str]] = data.get("chat_history") or []
-    assistant_state: AssistantState = data.get("assistant_state") or {}
+    chat_history = payload.get("chat_history", [])
+    assistant_state = payload.get("assistant_state", {})
 
-    # Initialize state from input. This is the "safe" state.
-    state: AssistantState = {
-        "conversationSummary": assistant_state.get("conversationSummary", ""),
-        "language": assistant_state.get("language", "Python"),
+    # validate/normalize assistant_state
+    state = {
+        "conversationSummary": assistant_state.get("conversationSummary", "").strip(),
+        "language": assistant_state.get("language", "python").strip().lower(),
         "taggedReplies": assistant_state.get("taggedReplies", []),
     }
-    
-    # 1. Prepare LLM Messages from Full History (same as before)
-    llm_messages = [{"role": "system", "content": PROGRAMMING_ASSISTANT_PROMPT}]
-    
-    last_user_message = ""
-    for msg in chat_history:
-        role = msg.get("role")
-        content = msg.get("content")
-        if role in ["user", "assistant"] and content:
-            llm_messages.append({"role": role, "content": content})
-            if role == "user":
-                last_user_message = content
-
-    # 2. Language Detection & State Update (same as before)
-    detected_lang = detect_language_from_text(last_user_message)
-    if detected_lang and detected_lang.lower() != state["language"].lower():
-        logger.info("Detected new language: %s", detected_lang)
-        state["language"] = detected_lang
-        
-    # 3. Inject Contextual Hint and State into the LAST user message (same as before)
-    context_hint = f"Current Language: {state['language']}. Conversation Summary so far: {state['conversationSummary']}"
-    if llm_messages and llm_messages[-1]["role"] == "user":
-        llm_messages[-1]["content"] = f"USER MESSAGE: {last_user_message}\n\n[CONTEXT HINT: {context_hint}]"
-    elif last_user_message:
-        llm_messages.append({"role": "user", "content": f"USER MESSAGE: {last_user_message}\n\n[CONTEXT HINT: {context_hint}]"})
-    
-    
-    try:
-        logger.info("Invoking LLM with full history and prepared prompt...")
-        llm_response = llm.invoke(llm_messages)
-        raw_response = llm_response.content if hasattr(llm_response, "content") else str(llm_response)
-        
-        logger.info(f"Raw LLM response: {raw_response}")
-        parsed_result = extract_json_from_llm_response(raw_response)
 
+    # limit history length to recent messages to control token usage
+    if isinstance(chat_history, list) and len(chat_history) > MAX_HISTORY_MESSAGES:
+        chat_history = chat_history[-MAX_HISTORY_MESSAGES:]
+
+    # build messages for LLM (do not mutate user's last message)
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for m in chat_history:
+        if not isinstance(m, dict):
+            continue
+        role = m.get("role")
+        content = m.get("content")
+        if role in ("user", "assistant") and content:
+            messages.append({"role": role, "content": content})
+
+    # append a supplemental context message (do not overwrite)
+    context_hint = f"[CONTEXT] language={state['language']} summary={clamp_summary(state['conversationSummary'], 300)}"
+    messages.append({"role": "system", "content": context_hint})
+
+    # call LLM (wrap in try/except)
+    try:
+        raw_resp = llm.invoke(messages, timeout=LLM_TIMEOUT_SECONDS)
+        # tolerate different shapes
+        raw_text = getattr(raw_resp, "content", None) or getattr(raw_resp, "text", None) or str(raw_resp)
+        logger.info("LLM raw text: %.300s", raw_text.replace('\n', ' ')[:300])
     except Exception as e:
-        logger.exception("LLM invocation failed")
-        error_detail = str(e)
-        if 'decommissioned' in error_detail:
-             error_detail = "LLM Model Error: The model is likely decommissioned. Please check the 'LLM_MODEL' environment variable or the default model in app.py."
-        return jsonify({"error": "LLM invocation failed", "detail": error_detail}), 500
-
-    # 4. State Update from LLM (NEW ROBUST LOGIC)
-    
-    # Check if parsing failed (by comparing the reply to the known error message)
-    if parsed_result.get("assistant_reply") == LLM_PARSE_ERROR_MESSAGE:
-        final_reply_content = LLM_PARSE_ERROR_MESSAGE
-        
-        # State and tags remain as initialized (from the input assistant_state), fulfilling the user request.
-        response_payload = {
-            "assistant_reply": final_reply_content,
-            "updated_state": state, # Keep the original input state
+        logger.exception("LLM invocation error")
+        return jsonify({"error": "LLM invocation failed", "detail": str(e)}), 500
+
+    parsed = parse_llm_output(raw_text)
+
+    # If parse failed, don't overwrite the existing state; give helpful message.
+    if not parsed.get("parse_ok"):
+        logger.warning("Parse failure. Returning fallback message.")
+        return jsonify({
+            "assistant_reply": parsed["assistant_reply"],
+            "code_snippet": "",
+            "updated_state": state,
             "suggested_tags": [],
-        }
-        return jsonify(response_payload)
-
-    # Parsing was successful. Safely update the state.
-    updated_state_from_llm = parsed_result.get("state_updates", {})
-    
-    if 'conversationSummary' in updated_state_from_llm:
-        state["conversationSummary"] = updated_state_from_llm["conversationSummary"]
-    # We allow the language to be updated only if it's explicitly set by the LLM AND it's a valid change.
-    if 'language' in updated_state_from_llm and updated_state_from_llm['language'].strip():
-        state["language"] = updated_state_from_llm["language"]
-
-
-    assistant_reply = parsed_result.get("assistant_reply")
-    code_snippet = parsed_result.get("code_snippet")
-    
-    # 5. Final Response Payload: Combine the reply and the code snippet
-    final_reply_content = assistant_reply
-    if code_snippet and code_snippet.strip():
-        if final_reply_content.strip():
-            final_reply_content += "\n\n"
-        final_reply_content += code_snippet
-        
-    if not final_reply_content.strip():
-        final_reply_content = "I'm here to help with your code! What programming language are you using?"
-
-    response_payload = {
-        "assistant_reply": final_reply_content, 
-        "updated_state": state,
-        "suggested_tags": parsed_result.get("suggested_tags", []),
-    }
+            "parse_ok": False,
+        }), 200
+
+    # Validate and apply state_updates conservatively
+    updates = parsed.get("state_updates", {}) or {}
+    if isinstance(updates, dict):
+        if "conversationSummary" in updates:
+            state["conversationSummary"] = clamp_summary(str(updates["conversationSummary"]))
+        if "language" in updates:
+            lang = canonicalize_language(str(updates["language"]))
+            if lang:
+                state["language"] = lang
+
+    # limit suggested tags
+    tags = parsed.get("suggested_tags", []) or []
+    if isinstance(tags, list):
+        tags = [str(t).strip() for t in tags if t and isinstance(t, (str,))]
+        tags = tags[:3]
 
-    return jsonify(response_payload)
+    return jsonify({
+        "assistant_reply": parsed.get("assistant_reply", ""),
+        "code_snippet": parsed.get("code_snippet", ""),
+        "updated_state": state,
+        "suggested_tags": tags,
+        "parse_ok": True,
+    }), 200
 
 @app.route("/tag_reply", methods=["POST"])
 def tag_reply():
-    data = request.get_json(force=True)
+    data = request.get_json(force=True, silent=True)
     if not isinstance(data, dict):
         return jsonify({"error": "invalid request body"}), 400
 
     reply_content = data.get("reply")
-    tags = data.get("tags")
-    assistant_state: AssistantState = data.get("assistant_state") or {}
-
-    if not reply_content or not tags:
+    tags = data.get("tags", [])
+    if not reply_content or not tags or not isinstance(tags, list):
         return jsonify({"error": "Missing 'reply' or 'tags' in request"}), 400
-    
-    tags = [str(t).strip() for t in tags if str(t).strip()]
-    if not tags:
-        return jsonify({"error": "Tags list cannot be empty"}), 400
 
-    state: AssistantState = {
+    tags_clean = [str(t).strip().lower() for t in tags if re.match(r'^[\w\-]{1,30}$', str(t).strip())]
+    if not tags_clean:
+        return jsonify({"error": "No valid tags provided"}), 400
+
+    assistant_state = data.get("assistant_state", {})
+    state = {
         "conversationSummary": assistant_state.get("conversationSummary", ""),
-        "language": assistant_state.get("language", "Python"),
+        "language": assistant_state.get("language", "python"),
         "taggedReplies": assistant_state.get("taggedReplies", []),
     }
 
-    new_tagged_reply: TaggedReply = {
-        "reply": reply_content,
-        "tags": tags,
-    }
-
-    state["taggedReplies"].append(new_tagged_reply)
-    
-    logger.info("Reply tagged with: %s", tags)
-
-    return jsonify({
-        "message": "Reply saved and tagged successfully.",
-        "updated_state": state,
-    }), 200
+    state["taggedReplies"].append({"reply": reply_content, "tags": tags_clean})
+    logger.info("Tagged reply saved: %s", tags_clean)
+    return jsonify({"message": "Reply saved", "updated_state": state}), 200
 
 @app.route("/ping", methods=["GET"])
 def ping():
     return jsonify({"status": "ok"})
 
 if __name__ == "__main__":
-    port = int(os.getenv("PORT", 7860))
-    app.run(host="0.0.0.0", port=port, debug=True)
+    port = int(os.getenv("PORT", "7860"))
+    app.run(host="0.0.0.0", port=port, debug=True)