Update app.py

app.py

@@ -3,17 +3,29 @@ import os
 import json
 import logging
 import re
-from typing import Dict, Any, List, Optional
+from typing import Dict, List, Optional
+from pathlib import Path
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from dotenv import load_dotenv
 from langchain_groq import ChatGroq
+from typing_extensions import TypedDict
 
-# --- Setup logging ---
+# --- Type Definitions ---
+class TaggedReply(TypedDict):
+    reply: str
+    tags: List[str]
+
+class AssistantState(TypedDict):
+    conversationSummary: str
+    language: str 
+    taggedReplies: List[TaggedReply]
+
+# --- Logging ---
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 logger = logging.getLogger("code-assistant")
 
-# --- Load environment variables ---
+# --- Load environment ---
 load_dotenv()
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 if not GROQ_API_KEY:
@@ -21,7 +33,10 @@ if not GROQ_API_KEY:
     raise RuntimeError("GROQ_API_KEY not set in environment")
 
 # --- Flask app setup ---
-app = Flask(__name__)
+BASE_DIR = Path(__file__).resolve().parent
+static_folder = BASE_DIR / "static"
+
+app = Flask(__name__, static_folder=str(static_folder), static_url_path="/static")
 CORS(app)
 
 # --- LLM setup ---
@@ -33,203 +48,170 @@ llm = ChatGroq(
 )
 
 # --- Constants ---
-LLM_PARSE_ERROR_MESSAGE = (
-    "Sorry, I couldn't understand the last response due to formatting issues. "
-    "Please try rephrasing or simplifying your query."
-)
+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 = """
+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.
+
+**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'.
 
-SYSTEM_PROMPT = """
-You are an expert programming assistant. You help with code suggestions, bug fixes, explanations, and contextual help.
+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
 
 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.
+- 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 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.
-    """
+def extract_json_from_llm_response(raw_response: str) -> dict:
+    default = {
+        "assistant_reply": LLM_PARSE_ERROR_MESSAGE,
+        "code_snippet": "",
+        "state_updates": {"conversationSummary": "", "language": "Python"},
+        "suggested_tags": [],
+    }
+    if not raw_response or not isinstance(raw_response, str):
+        return default
+    m = re.search(r"```json\s*([\s\S]*?)\s*```", raw_response)
+    json_string = m.group(1).strip() if m else raw_response
+    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
+    candidate = re.sub(r',\s*(?=[}\]])', '', candidate)
     try:
-        # 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
+        parsed = json.loads(candidate)
     except Exception as e:
-        logger.warning(f"Failed to parse JSON from LLM response: {e}")
-        return None
-
-def detect_language(text: str, default: str = "Python") -> str:
-    """
-    Detect programming language from user text.
-    Returns detected language or default.
-    """
-    if not text:
+        logger.warning("Failed to parse JSON from LLM output: %s. Candidate: %s", e, candidate[:200])
         return default
-    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
+    if 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", [])
+        if not parsed["assistant_reply"].strip():
+             parsed["assistant_reply"] = "I need a clearer instruction to provide a reply."
+        return parsed
     else:
-        # No user message found, add a dummy one with context
-        messages.append({"role": "user", "content": f"[Context: Language={language}, Summary={conversation_summary}]"})
+        logger.warning("Parsed JSON missing required keys or invalid format. Returning default.")
+        return default
+
+def detect_language_from_text(text: str) -> Optional[str]:
+    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
 
-    return messages
+# --- Routes ---
+
+@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
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    """
-    Main chat endpoint.
-    Expects JSON with keys:
-    - chat_history: list of messages {role: "user"/"assistant", content: str}
-    - assistant_state: {conversationSummary: str, language: str, taggedReplies: list}
-    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
-
-    chat_history = data.get("chat_history", [])
-    assistant_state = data.get("assistant_state", {})
-
-    # Initialize state with defaults
-    conversation_summary = assistant_state.get("conversationSummary", "")
-    language = assistant_state.get("language", "Python")
-    tagged_replies = assistant_state.get("taggedReplies", [])
-
-    # 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
-
-    # Build messages for LLM
-    messages = build_llm_messages(SYSTEM_PROMPT, chat_history, conversation_summary, language)
+        return jsonify({"error": "invalid request body"}), 400
 
-    try:
-        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,
-                "taggedReplies": tagged_replies,
-            },
-            "suggested_tags": parsed.get("suggested_tags", []),
-        }
-        return jsonify(response)
+    chat_history: List[Dict[str, str]] = data.get("chat_history") or []
+    assistant_state: AssistantState = data.get("assistant_state") or {}
+
+    state: AssistantState = {
+        "conversationSummary": assistant_state.get("conversationSummary", ""),
+        "language": assistant_state.get("language", "Python"),
+        "taggedReplies": assistant_state.get("taggedReplies", []),
+    }
 
+    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
+
+    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
+
+    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)
     except Exception as e:
-        logger.exception("Error during LLM invocation or parsing")
+        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
+
+    if parsed_result.get("assistant_reply") == LLM_PARSE_ERROR_MESSAGE:
         return jsonify({
             "assistant_reply": LLM_PARSE_ERROR_MESSAGE,
-            "updated_state": {
-                "conversationSummary": conversation_summary,
-                "language": language,
-                "taggedReplies": tagged_replies,
-            },
+            "updated_state": state,
             "suggested_tags": [],
-            "error": str(e),
-        }), 500
+        })
+
+    updated_state_from_llm = parsed_result.get("state_updates", {})
+    if 'conversationSummary' in updated_state_from_llm:
+        state["conversationSummary"] = updated_state_from_llm["conversationSummary"]
+    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")
+
+    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?"
+
+    return jsonify({
+        "assistant_reply": final_reply_content,
+        "updated_state": state,
+        "suggested_tags": parsed_result.get("suggested_tags", []),
+    })
 
 @app.route("/tag_reply", methods=["POST"])
 def tag_reply():
-    """
-    Endpoint to save/bookmark a reply with tags.
-    Expects JSON with keys:
-    - reply: string
-    - tags: list of strings
-    - assistant_state: current state dict
-    Returns updated state with new tagged reply appended.
-    """
     data = request.get_json(force=True)
     if not isinstance(data, dict):
         return jsonify({"error": "invalid request body"}), 400
 
     reply_content = data.get("reply")
     tags = data.get("tags")
-    assistant_state = data.get("assistant_state") or {}
+    assistant_state: AssistantState = data.get("assistant_state") or {}
 
     if not reply_content or not tags:
         return jsonify({"error": "Missing 'reply' or 'tags' in request"}), 400
@@ -238,20 +220,24 @@ def tag_reply():
     if not tags:
         return jsonify({"error": "Tags list cannot be empty"}), 400
 
-    tagged_replies = assistant_state.get("taggedReplies", [])
-    tagged_replies.append({"reply": reply_content, "tags": tags})
-
-    updated_state = {
+    state: AssistantState = {
         "conversationSummary": assistant_state.get("conversationSummary", ""),
         "language": assistant_state.get("language", "Python"),
-        "taggedReplies": tagged_replies,
+        "taggedReplies": assistant_state.get("taggedReplies", []),
     }
 
-    logger.info(f"Tagged reply saved with tags: {tags}")
+    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": updated_state,
+        "updated_state": state,
     }), 200
 
 @app.route("/ping", methods=["GET"])