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