app.py

#!/usr/bin/env python3
import os
import json
import logging
import re
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

# --- Setup logging ---
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("code-assistant")

# --- 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")

# --- Flask app setup ---
app = Flask(__name__)
CORS(app)

# --- 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,
)

# --- Constants ---
LLM_PARSE_ERROR_MESSAGE = (
    "Sorry, I couldn't understand the last response due to formatting issues. "
    "Please try rephrasing or simplifying your query."
)

SYSTEM_PROMPT = """
You are an expert programming assistant. You help with code suggestions, bug fixes, explanations, and contextual help.

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 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:
        # 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 detect_language(text: str, default: str = "Python") -> str:
    """
    Detect programming language from user text.
    Returns detected language or default.
    """
    if not text:
        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
    else:
        # No user message found, add a dummy one with context
        messages.append({"role": "user", "content": f"[Context: Language={language}, Summary={conversation_summary}]"})

    return messages

@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}
    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")

    # 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)

    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,
            },
            "suggested_tags": parsed.get("suggested_tags", []),
        }
        return jsonify(response)

    except Exception as e:
        logger.exception("Error during LLM invocation or parsing")
        return jsonify({
            "assistant_reply": LLM_PARSE_ERROR_MESSAGE,
            "updated_state": {
                "conversationSummary": conversation_summary,
                "language": language,
            },
            "suggested_tags": [],
            "error": str(e),
        }), 500

@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)