Update app.py

app.py

@@ -1,297 +1,219 @@
 #!/usr/bin/env python3
-# filename: app_refactored.py
 import os
 import json
 import logging
 import re
-from pathlib import Path
-from typing import Dict, Any, List, Optional, Tuple
-
+from typing import Dict, Any, List, Optional
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from dotenv import load_dotenv
+from langchain_groq import ChatGroq
 
-# Replace with your LLM client import; kept generic here.
-# from langchain_groq import ChatGroq
+# --- Setup logging ---
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("code-assistant")
 
-# === Config ===
+# --- Load environment variables ---
 load_dotenv()
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 if not GROQ_API_KEY:
+    logger.error("GROQ_API_KEY not set in environment")
     raise RuntimeError("GROQ_API_KEY not set in environment")
 
-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"}
-
-# === Logging ===
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-logger = logging.getLogger("code-assistant")
-
-# === 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
+# --- Flask app setup ---
+app = Flask(__name__)
+CORS(app)
 
-# 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
+# --- LLM setup ---
+llm = ChatGroq(
+    model=os.getenv("LLM_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct"),
+    temperature=0.1,
+    max_tokens=2048,
+    api_key=GROQ_API_KEY,
+)
 
-# === 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."
+# --- Constants ---
+LLM_PARSE_ERROR_MESSAGE = (
+    "Sorry, I couldn't understand the last response due to formatting issues. "
+    "Please try rephrasing or simplifying your query."
 )
 
-# === 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] + "..."
+SYSTEM_PROMPT = """
+You are an expert programming assistant. You help with code suggestions, bug fixes, explanations, and contextual help.
 
-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 respond with a single JSON object enclosed in a ```json ... ``` code block.
+- The JSON must have these keys:
+  - assistant_reply: string (short, helpful natural language reply, no code blocks)
+  - code_snippet: string (code in markdown code block, with newlines escaped as \\n and backslashes as \\\\; empty string if none)
+  - state_updates: object with keys:
+      - conversationSummary: string (concise summary of the conversation so far)
+      - language: string (programming language context)
+  - suggested_tags: array of strings (1-3 relevant tags)
+
+- Always include all keys.
+- Adapt code and explanations to the language in state_updates.language.
+"""
 
-def try_parse_json(s: str) -> Optional[Dict[str, Any]]:
+def extract_json_from_response(text: str) -> Optional[Dict[str, Any]]:
+    """
+    Extract JSON object from LLM response text inside a ```json ... ``` block.
+    Return None if parsing fails.
+    """
     try:
-        return json.loads(s)
-    except Exception:
+        # Extract JSON code block content
+        match = re.search(r"```json\s*([\s\S]*?)\s*```", text)
+        json_text = match.group(1) if match else text
+
+        # Find first and last braces to isolate JSON object
+        first = json_text.find('{')
+        last = json_text.rfind('}')
+        if first == -1 or last == -1 or last < first:
+            logger.warning("No valid JSON braces found in LLM response")
+            return None
+        json_str = json_text[first:last+1]
+
+        # Remove trailing commas before } or ]
+        json_str = re.sub(r",\s*(?=[}\]])", "", json_str)
+
+        parsed = json.loads(json_str)
+        return parsed
+    except Exception as e:
+        logger.warning(f"Failed to parse JSON from LLM response: {e}")
         return None
 
-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]:
+def detect_language(text: str, default: str = "Python") -> str:
     """
-    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)
+    Detect programming language from user text.
+    Returns detected language or default.
     """
-    default = {
-        "assistant_reply": "I couldn't parse the model response. Please rephrase or simplify the request.",
-        "code_snippet": "",
-        "state_updates": {"conversationSummary": "", "language": "python"},
-        "suggested_tags": [],
-        "parse_ok": False,
-    }
-    if not raw or not isinstance(raw, str):
+    if not text:
         return default
-
-    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.setdefault("suggested_tags", [])
-        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()
+    text_lower = text.lower()
+    languages = ["python", "javascript", "java", "c++", "c#", "go", "ruby", "php", "typescript", "swift"]
+    for lang in languages:
+        if re.search(rf"\b(in|using|for)\s+{lang}\b", text_lower):
+            return lang.capitalize()
+    return default
+
+def build_llm_messages(
+    system_prompt: str,
+    chat_history: List[Dict[str, str]],
+    conversation_summary: str,
+    language: str,
+) -> List[Dict[str, str]]:
+    """
+    Build messages list for LLM invocation.
+    Inject conversation summary and language context into the last user message.
+    """
+    messages = [{"role": "system", "content": system_prompt}]
+    for msg in chat_history:
+        if msg.get("role") in ["user", "assistant"] and msg.get("content"):
+            messages.append({"role": msg["role"], "content": msg["content"]})
+
+    # Inject context hint into last user message
+    for i in reversed(range(len(messages))):
+        if messages[i]["role"] == "user":
+            messages[i]["content"] += f"\n\n[Context: Language={language}, Summary={conversation_summary}]"
+            break
     else:
-        # 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
+        # No user message found, add a dummy one with context
+        messages.append({"role": "user", "content": f"[Context: Language={language}, Summary={conversation_summary}]"})
 
-# === Flask app ===
-BASE_DIR = Path(__file__).resolve().parent
-app = Flask(__name__, static_folder=str(BASE_DIR / "static"), static_url_path="/static")
-CORS(app)
-
-@app.route("/", methods=["GET"])
-def serve_frontend():
-    try:
-        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
+    return messages
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    payload = request.get_json(force=True, silent=True)
-    if not isinstance(payload, dict):
-        return jsonify({"error": "invalid request body"}), 400
-
-    chat_history = payload.get("chat_history", [])
-    assistant_state = payload.get("assistant_state", {})
+    """
+    Main chat endpoint.
+    Expects JSON with keys:
+    - chat_history: list of messages {role: "user"/"assistant", content: str}
+    - assistant_state: {conversationSummary: str, language: str}
+    Returns JSON with:
+    - assistant_reply: str
+    - updated_state: dict
+    - suggested_tags: list
+    """
+    data = request.get_json(force=True)
+    if not isinstance(data, dict):
+        return jsonify({"error": "Invalid request body"}), 400
 
-    # validate/normalize assistant_state
-    state = {
-        "conversationSummary": assistant_state.get("conversationSummary", "").strip(),
-        "language": assistant_state.get("language", "python").strip().lower(),
-        "taggedReplies": assistant_state.get("taggedReplies", []),
-    }
+    chat_history = data.get("chat_history", [])
+    assistant_state = data.get("assistant_state", {})
 
-    # 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:]
+    # Initialize state with defaults
+    conversation_summary = assistant_state.get("conversationSummary", "")
+    language = assistant_state.get("language", "Python")
 
-    # 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})
+    # Detect language from last user message if possible
+    last_user_msg = ""
+    for msg in reversed(chat_history):
+        if msg.get("role") == "user" and msg.get("content"):
+            last_user_msg = msg["content"]
+            break
+    detected_lang = detect_language(last_user_msg, default=language)
+    if detected_lang.lower() != language.lower():
+        logger.info(f"Language changed from {language} to {detected_lang}")
+        language = detected_lang
 
-    # 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})
+    # Build messages for LLM
+    messages = build_llm_messages(SYSTEM_PROMPT, chat_history, conversation_summary, language)
 
-    # 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 error")
-        return jsonify({"error": "LLM invocation failed", "detail": str(e)}), 500
+        logger.info("Invoking LLM...")
+        llm_response = llm.invoke(messages)
+        raw_text = getattr(llm_response, "content", str(llm_response))
+        logger.info(f"LLM raw response: {raw_text}")
+
+        parsed = extract_json_from_response(raw_text)
+        if not parsed:
+            raise ValueError("Failed to parse JSON from LLM response")
+
+        # Validate keys
+        required_keys = {"assistant_reply", "code_snippet", "state_updates", "suggested_tags"}
+        if not required_keys.issubset(parsed.keys()):
+            raise ValueError(f"Missing keys in LLM response JSON: {required_keys - parsed.keys()}")
+
+        # Update state
+        state_updates = parsed.get("state_updates", {})
+        conversation_summary = state_updates.get("conversationSummary", conversation_summary)
+        language = state_updates.get("language", language)
+
+        # Compose final assistant reply with optional code snippet
+        assistant_reply = parsed["assistant_reply"].strip()
+        code_snippet = parsed["code_snippet"].strip()
+        if code_snippet:
+            # Unescape newlines and backslashes for display
+            code_snippet_display = code_snippet.replace("\\n", "\n").replace("\\\\", "\\")
+            assistant_reply += f"\n\n```{language.lower()}\n{code_snippet_display}\n```"
+
+        # Prepare response
+        response = {
+            "assistant_reply": assistant_reply,
+            "updated_state": {
+                "conversationSummary": conversation_summary,
+                "language": language,
+            },
+            "suggested_tags": parsed.get("suggested_tags", []),
+        }
+        return jsonify(response)
 
-    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.")
+    except Exception as e:
+        logger.exception("Error during LLM invocation or parsing")
         return jsonify({
-            "assistant_reply": parsed["assistant_reply"],
-            "code_snippet": "",
-            "updated_state": state,
+            "assistant_reply": LLM_PARSE_ERROR_MESSAGE,
+            "updated_state": {
+                "conversationSummary": conversation_summary,
+                "language": language,
+            },
             "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({
-        "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, silent=True)
-    if not isinstance(data, dict):
-        return jsonify({"error": "invalid request body"}), 400
-
-    reply_content = data.get("reply")
-    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_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"),
-        "taggedReplies": assistant_state.get("taggedReplies", []),
-    }
-
-    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
+            "error": str(e),
+        }), 500
 
 @app.route("/ping", methods=["GET"])
 def ping():
     return jsonify({"status": "ok"})
 
 if __name__ == "__main__":
-    port = int(os.getenv("PORT", "7860"))
+    port = int(os.getenv("PORT", 7860))
     app.run(host="0.0.0.0", port=port, debug=True)