app.py

#!/usr/bin/env python3
import os
import json
import logging
import re
from typing import Dict, Any, List
from pathlib import Path
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
    lastUserMessage: str
    language: str # New field to track the programming language
    taggedReplies: List[TaggedReply] # New field for saving/bookmarking replies

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

# --- Load environment ---
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
if not GROQ_API_KEY:
    logger.error("GROQ_API_KEY not set in environment")
    exit(1)

# --- Flask app setup ---
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 ---
# Using a model that's good for coding tasks
llm = ChatGroq(
    model=os.getenv("LLM_MODEL", "mixtral-8x7b-32768"), # Changed to a coding-friendly model
    temperature=0.1, # Slightly less creative than general chat
    max_tokens=2048, # Increased token limit for code
    api_key=GROQ_API_KEY,
)

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.

Behavior rules (follow these strictly):
- Contextual Help: Always aim to provide the most helpful, clear, and accurate information.
- Code Suggestions: When suggesting code, always enclose it in appropriate markdown code blocks (e.g., ```python\n...\n```).
- Error Explanation: When an error is provided, explain the root cause and provide a corrected code snippet if possible.
- Conceptual Questions: For questions like "What is a loop?", provide a clear, concise explanation with a simple, illustrative code example in the user's current language (if known, otherwise Python/JavaScript).
- Language Adaptation: Adjust your suggestions, code, and explanations to the programming language specified in the 'language' field of the 'AssistantState'. If 'language' is not set, ask the user what language they are working in.

STRICT OUTPUT FORMAT (JSON ONLY):
Return a single JSON object with the following keys:
- assistant_reply: string  // a natural language reply to the user (short, helpful, always present)
- state_updates: object  // updates to the internal state, keys may include: language, conversationSummary
- suggested_tags: array of strings // a list of 1-3 relevant tags for the assistant_reply (e.g., "Python", "Debugging", "Loop Concept")

Rules:
- ALWAYS include `assistant_reply` as a non-empty string.
- Do NOT produce any text outside the JSON object.
- Be concise in `assistant_reply`, but ensure the information is complete.
- Do not make up information.
"""

def extract_json_from_llm_response(raw_response: str) -> dict:
    # Helper function remains largely the same, adapted for new keys
    default = {
        "assistant_reply": "I'm sorry — I couldn't understand that. Could you please rephrase?",
        "state_updates": {},
        "suggested_tags": [],
    }
    # ... [JSON parsing logic remains similar] ...
    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('}')
    if first == -1 or last == -1 or first >= last:
        try:
            return json.loads(json_string)
        except Exception:
            logger.warning("Could not locate JSON braces in LLM output. Falling back to default.")
            return default
    candidate = json_string[first:last+1]
    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", e)
        return default
    
    # Validation for new keys
    if isinstance(parsed, dict) and "assistant_reply" in parsed and isinstance(parsed["assistant_reply"], str) and parsed["assistant_reply"].strip():
        parsed.setdefault("state_updates", {})
        parsed.setdefault("suggested_tags", [])
        return parsed
    else:
        logger.warning("Parsed JSON missing 'assistant_reply' or invalid format. Returning default.")
        return default

# --- Flask routes ---
@app.route("/", methods=["GET"])
def serve_frontend():
    try:
        # Assuming you will update frontend.html for the new assistant
        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

# UPLOAD routes are removed as they are no longer needed.

@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json(force=True)
    if not isinstance(data, dict):
        return jsonify({"error": "invalid request body"}), 400

    chat_history: List[Dict[str, str]] = data.get("chat_history") or []
    # Using 'assistant_state' to clearly separate from old patient_state
    assistant_state: AssistantState = data.get("assistant_state") or {}

    # Initialize/Clean up state
    state: AssistantState = {
        "conversationSummary": assistant_state.get("conversationSummary", ""),
        "lastUserMessage": "",
        "language": assistant_state.get("language", "Python"), # Default to Python
        "taggedReplies": assistant_state.get("taggedReplies", []),
    }
    
    # Find the last user message
    for msg in reversed(chat_history):
        if msg.get("role") == "user" and msg.get("content"):
            state["lastUserMessage"] = msg["content"]
            break

    # --- Language Detection (Simple check for common programming languages) ---
    last_msg_lower = state["lastUserMessage"].lower()
    known_languages = ["python", "javascript", "java", "c++", "c#", "go", "ruby", "php", "typescript", "swift"]
    
    # A simple regex to detect a language mention in the last message
    lang_match = re.search(r'\b(in|using|for)\s+(' + '|'.join(known_languages) + r')\b', last_msg_lower)
    if lang_match:
        detected_lang = lang_match.group(2).capitalize()
        if detected_lang != state["language"]:
            logger.info("Detected new language: %s", detected_lang)
            state["language"] = detected_lang
    
    # --- LLM Prompt Construction ---
    action_hint = ""
    if state["language"]:
        action_hint = f"Focus your answer on the {state['language']} programming language. If the user asks a conceptual question, use {state['language']} for examples."
    else:
        action_hint = "The current language is unknown. Please ask the user to specify the programming language they are working in."

    user_prompt = f"""
Current State: {json.dumps({"language": state["language"], "summary": state["conversationSummary"]})}
Last user message: {state["lastUserMessage"]}

SYSTEM_HINT: {action_hint}

Return ONLY valid JSON with keys: assistant_reply, state_updates, suggested_tags.
"""

    messages = [
        {"role": "system", "content": PROGRAMMING_ASSISTANT_PROMPT},
        {"role": "user", "content": user_prompt}
    ]

    try:
        logger.info("Invoking LLM for code assistant...")
        llm_response = llm.invoke(messages)
        raw_response = llm_response.content if hasattr(llm_response, "content") else str(llm_response)

        logger.info(f"Raw LLM response: {raw_response[:200]}...")
        parsed_result = extract_json_from_llm_response(raw_response)

    except Exception as e:
        logger.exception("LLM invocation failed")
        return jsonify({"error": "LLM invocation failed", "detail": str(e)}), 500

    # --- State Update from LLM ---
    updated_state_from_llm = parsed_result.get("state_updates", {})
    
    # Update state fields that the LLM is allowed to modify
    if 'conversationSummary' in updated_state_from_llm:
        state["conversationSummary"] = updated_state_from_llm["conversationSummary"]
    if 'language' in updated_state_from_llm:
        state["language"] = updated_state_from_llm["language"]

    assistant_reply = parsed_result.get("assistant_reply")
    if not assistant_reply or not isinstance(assistant_reply, str) or not assistant_reply.strip():
        assistant_reply = "I'm here to help with your code! What programming language are you using?"

    # --- Final Response Payload ---
    response_payload = {
        "assistant_reply": assistant_reply,
        "updated_state": state,
        "suggested_tags": parsed_result.get("suggested_tags", []), # Pass tags to frontend
    }

    return jsonify(response_payload)

# --- New Route for Tagging/Bookmarking Replies ---
@app.route("/tag_reply", methods=["POST"])
def tag_reply():
    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: AssistantState = data.get("assistant_state") or {}

    if not reply_content or not tags:
        return jsonify({"error": "Missing 'reply' or 'tags' in request"}), 400
    
    # Ensure tags is a list of strings
    tags = [str(t).strip() for t in tags if str(t).strip()]
    if not tags:
        return jsonify({"error": "Tags list cannot be empty"}), 400

    # Clean up state dictionary
    state: AssistantState = {
        "conversationSummary": assistant_state.get("conversationSummary", ""),
        "lastUserMessage": "",
        "language": assistant_state.get("language", "Python"),
        "taggedReplies": assistant_state.get("taggedReplies", []),
    }

    new_tagged_reply: TaggedReply = {
        "reply": reply_content,
        "tags": tags,
    }

    # Add the new tagged reply
    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

# --- Filtering/Search Route for Bookmarked Replies ---
@app.route("/search_tags", methods=["GET"])
def search_tags():
    tag_query = request.args.get("tag")
    # Using POST for /chat, so we'll pass state in the body
    # For a simple GET search, we'd need the state to be sent here, 
    # but for simplicity, let's assume the state is passed in a POST body
    # or fetched/maintained on the frontend and this route is just for logic.
    
    # Assuming the frontend sends the current state via a POST request for search
    if request.method == "GET":
        return jsonify({"error": "Please use POST and include 'assistant_state' in the body for tag search."}), 405
    
    # If using POST, you'd process request.get_json() here to get assistant_state
    # For now, let's stick to the simpler GET and assume the frontend handles the state.
    # To demonstrate the filtering logic:
    
    # --- DUMMY STATE FOR DEMO ---
    dummy_state: AssistantState = {
        "conversationSummary": "",
        "lastUserMessage": "",
        "language": "Python",
        "taggedReplies": [
            {"reply": "A Python loop example.", "tags": ["Python", "Loop Concept"]},
            {"reply": "Fix for 'undefined' error in JS.", "tags": ["JavaScript", "Debugging"]},
            {"reply": "Explanation of Polymorphism.", "tags": ["Java", "OOP"]},
        ],
    }
    
    if not tag_query:
        # Return all tagged replies if no query
        return jsonify({"tag_query": "", "results": dummy_state["taggedReplies"]}), 200

    tag_query_lower = tag_query.lower()
    
    filtered_results = [
        reply for reply in dummy_state["taggedReplies"] 
        if any(tag_query_lower in tag.lower() for tag in reply["tags"])
    ]
    
    return jsonify({
        "tag_query": tag_query,
        "results": filtered_results
    }), 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)