chATLAS_Frontend 1.0.0__py3-none-any.whl

chATLAS_Frontend/frontend/html_generation.py

@@ -0,0 +1,361 @@
+import ast
+import json
+import logging
+import os
+import re
+
+import markdown
+from flask import render_template
+
+from chATLAS_Frontend.frontend.resources import get_db_logger
+
+logger = logging.getLogger(__name__)
+
+
+def get_file_name(file_name, CAPTION_SWITCH):
+    if ":" in file_name:
+        return file_name, None
+
+    name_out = os.path.basename(file_name).rsplit(".", 1)[0]
+    caption_out = None
+
+    return name_out, caption_out
+
+
+def generate_full_response(user_input, sources, answer, sources_used_by_model, assistant_mode=True):
+    """
+    Render all html for the response. Some pieces are not rendered when assistant_mode is False.
+
+    In assistant mode, citations (used docs + all sources) are placed in a hidden block
+    and a "Citations" button opens them in the right sidebar. In search mode, sources
+    are rendered inline.
+
+    Returns:
+        dict with keys:
+          - 'bot_html': the rendered bot reply fragment
+          - 'citations_html': the citations block HTML (assistant mode) or sources HTML (search mode)
+    """
+    _user_display, bot_display = generate_chat(answer, user_input, len(sources), assistant_mode=assistant_mode)
+
+    # Raw sources list (no show/hide button wrapper)
+    sources_display_raw = generate_sources_display(sources, assistant_mode=False)
+
+    if assistant_mode:
+        used_sources_html = generate_used_docs_display(sources, sources_used_by_model)
+        citations_content = used_sources_html + sources_display_raw
+        citations_block_html = render_template("response/citations_block.html", citations_content=citations_content)
+        # citations_block goes inside the bot reply so the "Citations" button
+        # lives in the chat row (sourcePanel.js reads it on click)
+        full_bot_html = bot_display + citations_block_html
+    else:
+        citations_content = sources_display_raw
+        full_bot_html = bot_display
+
+    return {"bot_html": full_bot_html, "citations_html": citations_content}
+
+
+def _used_doc_utility_for_source(doc_utility_map: dict, source_metadata_name: str) -> str | None:
+    """
+    Return the utility for a retrieved source if it is listed in used_docs (exact or suffix match).
+
+    Inputs:
+        doc_utility_map: Map from document name (as returned by LLM) to utility string.
+        source_metadata_name: The 'name' from the retrieved source's metadata.
+    Outputs:
+        Utility string if the source is in used_docs, else None.
+    """
+    stripped = (source_metadata_name or "").strip('"').strip("'")
+    if not stripped:
+        return None
+    if stripped in doc_utility_map:
+        return doc_utility_map[stripped]
+    for key, utility in doc_utility_map.items():
+        if key == stripped:
+            return utility
+        if key.endswith("/" + stripped) or (key and key.rstrip("/").endswith("/" + stripped)):
+            return utility
+    return None
+
+
+def generate_used_docs_display(sources, used_docs):
+    """
+    Generates html to display the sources used by the model in the response.
+
+    Parameters:
+    ----------
+    used_docs : list of dict
+        List of dictionaries containing document information used by the model.
+        Each dictionary has 'name' and 'utility' keys.
+
+    sources: list Document objects
+
+    Returns:
+    --------
+    str
+        HTML representation of the used documents
+    """
+    # Enforce that used_docs must be a list
+    if not isinstance(used_docs, list):
+        raise TypeError(f"used_docs must be a list of dictionaries, not {type(used_docs)}")
+
+    if not used_docs:
+        # If there are no used docs, render a template with empty data
+        return render_template("response/used_docs.html", has_sources=False)
+
+    # Get source ids for feedback buttons and sources data
+    db_logger = get_db_logger()
+    source_ids = db_logger.get_most_recent_source_ids()
+
+    # Create a mapping of document names to their utility values (exact keys from LLM)
+    doc_utility_map = {}
+    for doc in used_docs:
+        if isinstance(doc, dict) and "name" in doc and "utility" in doc:
+            name = doc["name"]
+            if isinstance(name, str):
+                doc_utility_map[name.strip('"').strip("'")] = doc["utility"]
+            else:
+                doc_utility_map[str(name).strip('"').strip("'")] = doc["utility"]
+
+    # Prepare the source data: include a source if it matches used_docs (exact or suffix match)
+    source_entries = []
+    for source in sources:
+        source_name = source.metadata.get("name", "Unknown")
+        utility = _used_doc_utility_for_source(doc_utility_map, source_name)
+        if utility is None:
+            continue
+        score_metric = "rrf_score" if "rrf_score" in source.metadata else "similarity"
+        source_data = extract_source_data(source, source_ids, score_metric=score_metric)
+        source_data["utility"] = utility
+        source_entries.append(source_data)
+
+    # Render the template with the prepared data
+    return render_template("response/used_docs.html", has_sources=True, sources=source_entries)
+
+
+def extract_source_data(source, source_ids, score_metric="similarity"):
+    """Format the source data to a dictionary for html generation"""
+
+    importance_thresholds = {"similarity": 0.5, "rrf_score": 0.015}
+
+    if score_metric not in importance_thresholds:
+        raise ValueError(f"Unknown score metric: {score_metric}")
+
+    source_name = source.metadata.get("name", "Unknown")
+    source_name_key = source.metadata.get("name")
+    source_id = source_ids.get(source_name_key)
+    score = round(source.metadata.get(score_metric, 0), 3)
+    url = source.metadata.get("url", "#")
+    is_important = score > importance_thresholds[score_metric]
+    NAME_OF_SOURCE, summary = get_file_name(source_name, True)
+
+    return {
+        "url": url,
+        "display_name": NAME_OF_SOURCE,
+        "summary": summary,
+        "similarity_score": score,
+        "similarity_formatted": f"{score:.3f}",
+        "is_important": is_important,
+        # Add variables needed by feedback_buttons.html
+        "button_type": "source",
+        "feedback_id": source_id,
+    }
+
+
+def generate_sources_display(sources, assistant_mode=True):
+    """
+    Generates html for the used sources, grouped by Type and SearchType
+    with scrollable container.
+
+    In assistant mode, the sources are wrapped in a button to show/hide them.
+    """
+
+    if not sources:
+        all_sources_html = render_template("response/all_sources.html", all_sources_data=[])
+        return (
+            render_template("response/source_display_button.html", sources_display=all_sources_html)
+            if assistant_mode
+            else all_sources_html
+        )
+
+    # Get the ids of the sources
+    db_logger = get_db_logger()
+    source_ids = db_logger.get_most_recent_source_ids()
+
+    if not source_ids:
+        source_ids = {}
+
+    # Group sources by type (twiki, indico, CDS ...) and search type (vector, text)
+    grouped_by_source_and_search_type = {}
+
+    if "rrf_score" in sources[0].metadata:
+        score_metric = "rrf_score"
+        for source in sources:
+            source.metadata["search_type"] = "RRF"
+
+    else:
+        score_metric = "similarity"
+
+    for source in sources:
+        source_type = source.metadata.get("source", "Unknown Type")
+        search_type = source.metadata.get("search_type", "Unknown Search")
+        grouped_by_source_and_search_type.setdefault(source_type, {}).setdefault(search_type, []).append(source)
+
+    # Prepare data structure for template context
+    all_sources_data = []
+
+    # For each source type (twiki, indico, CDS ...)
+    for source_type, sources_by_search_type in grouped_by_source_and_search_type.items():
+        search_types_data = []
+
+        # For each search type (vector, text)
+        for search_type, sources_list in sources_by_search_type.items():
+            source_entries = []
+
+            # For each individual source
+            for source in sources_list:
+                source_data = extract_source_data(source, source_ids, score_metric=score_metric)
+
+                source_entries.append(source_data)
+
+            # Add this search type's data
+            search_types_data.append({"search_type": search_type, "sources": source_entries})
+
+        # Add this source type's data
+        all_sources_data.append({"source_type": source_type, "search_types": search_types_data})
+
+    # Render everything at once
+    all_sources_html = render_template("response/all_sources.html", all_sources_data=all_sources_data)
+
+    # wrap source display behind a "Show Full Sources" button
+    if assistant_mode:
+        all_sources_html = render_template("response/source_display_button.html", sources_display=all_sources_html)
+
+    return all_sources_html
+
+
+def fix_markdown_links(text):
+    """
+    Converts links from [[url][label]] format to standard [label](url) Markdown,
+    but only if the url looks like a valid web address.
+    """
+    pattern = r"\[\[(https?://[^\[\]]+|www\.[^\[\]]+)\]\[([^\[\]]+)\]\]"
+    return re.sub(pattern, r"[\2](\1)", text)
+
+
+def parse_answer(chatbot_answer):
+    """
+    Parse the LLM output to extract the actual answer and JSON.
+    """
+    # Clean up the input from markdown/code block markers
+    cleaned_answer = chatbot_answer.strip()
+    cleaned_answer = cleaned_answer.removeprefix("```json")
+    cleaned_answer = cleaned_answer.removeprefix("```")
+    cleaned_answer = cleaned_answer.removesuffix("```")
+
+    try:
+        # Handle cases where JSON is escaped
+        if cleaned_answer.startswith("{") and '"' in cleaned_answer:
+            cleaned_answer = json.loads(cleaned_answer)
+        # If cleaned_answer is still a string, try parsing as JSON
+        parsed_ans = json.loads(cleaned_answer) if isinstance(cleaned_answer, str) else cleaned_answer
+
+        answer = parsed_ans.get("text", "")
+        # Fix malformed markdown links if there are any
+        answer = fix_markdown_links(answer)
+
+        used_docs = parsed_ans.get("sources", [])
+        return answer, used_docs
+    except json.JSONDecodeError:
+        try:
+            # Fallback: attempt parsing using ast.literal_eval
+            parsed_ans = ast.literal_eval(cleaned_answer)
+            if isinstance(parsed_ans, dict):
+                return parsed_ans.get("text", ""), parsed_ans.get("sources", [])
+        except (ValueError, SyntaxError):
+            pass
+    except Exception as e:
+        print(f"Parsing error: {e}")
+
+    # Fallback: return raw text if parsing fails
+    logger.warning(
+        "parse_answer fallback: returning raw text. chatbot_answer=%r",
+        chatbot_answer[:500] if chatbot_answer else chatbot_answer,
+    )
+    return str(chatbot_answer), []
+
+
+def post_process_markdown(html_content):
+    """
+    Clean up markdown-generated HTML to fix:
+    - Code blocks with language identifiers
+    - Code blocks that don't have <pre> tags
+    """
+    # Fix code blocks with language identifiers
+    html_content = re.sub(r"<code>(\w+)\n", r"<code>", html_content)
+
+    # Fix code blocks that don't have <pre> tags
+    html_content = re.sub(r"<br />\s*<code>", r"<pre><code>", html_content)
+    html_content = re.sub(r"</code>\s*<br />", r"</code></pre>", html_content)
+
+    return html_content
+
+
+def generate_chat(answer, user_input, num_sources, assistant_mode=True):
+    """
+    Renders the HTML templates for the bot response and feedback button.
+    Returns (user_output, bot_display) — user_output is kept for compatibility
+    but is no longer included in the JSON response blob (the client renders the
+    user message itself before sending).
+    """
+    user_output = render_template("response/user_query.html", user_input=user_input, result_count=num_sources)
+
+    if assistant_mode:
+        if answer:
+            # --- Preprocess for code blocks ---
+            answer = re.sub(r"```(\w*)", r"```\1", answer)
+
+            code_blocks = re.findall(r"```(?:\w*)\n(.*?)```", answer, re.DOTALL)
+            for block in code_blocks:
+                lines = block.split("\n")
+                non_empty_lines = [line for line in lines if line.strip()]
+                if non_empty_lines:
+                    indentation = min(len(line) - len(line.lstrip()) for line in non_empty_lines if line.strip())
+                    normalized_block = "\n".join(line[indentation:] if line.strip() else line for line in lines)
+                    answer = answer.replace(block, normalized_block)
+
+            # Ensure a blank line before any table
+            answer = re.sub(r"(?m)(?<!\|)\n(?=\|)", r"\n\n", answer)
+
+            # --- Markdown conversion ---
+            # Note: nl2br is intentionally excluded — it converts bare newlines to <br>
+            # before the list/heading parsers run, which breaks bullet list rendering.
+            extensions = [
+                "markdown.extensions.fenced_code",
+                "markdown.extensions.codehilite",
+                "markdown.extensions.tables",
+                "markdown.extensions.extra",
+            ]
+            extension_configs = {
+                "markdown.extensions.codehilite": {
+                    "css_class": "codehilite",
+                    "linenums": False,
+                    "guess_lang": False,
+                    "use_pygments": False,
+                },
+                "markdown.extensions.fenced_code": {"lang_prefix": ""},
+            }
+            processed_answer = markdown.markdown(answer, extensions=extensions, extension_configs=extension_configs)
+            processed_answer = post_process_markdown(processed_answer)
+        else:
+            processed_answer = ""
+
+        bot_display = render_template(
+            "response/bot_response.html",
+            bot_response=processed_answer,
+            button_type="response",
+            feedback_id=get_db_logger().get_most_recent_query_id(),
+        )
+    else:
+        bot_display = ""
+
+    return user_output, bot_display