tools/diagrams.py

import zlib
import base64
import os
from dotenv import load_dotenv
from huggingface_hub import InferenceClient

load_dotenv()
os.environ["SAMBANOVA_API_KEY"]

PROMPT="""
###PERSONA : Tu es un analyste métier sénior, expert en vulgarisation et en diagrammes formels. 
###INSTRUCTION : En fonction d’un sujet donné, sélectionne le type de diagramme le plus adapté parmi ceux listés, puis génère un diagramme complet, lisible, et détaillé. Syntaxe Mermaid uniquement, sans explication hors-code.
Génère le code mermaid, et rien que le code mermaid expliquant le mieux la description qui suit.
### Sélection du type de diagramme (`mermaid_diag_type`) :
- `Graph TB` : **choix par défaut**. Regroupe les éléments par blocs. Applique des couleurs si possible.
- `flowchart TD` : pour les **processus** et **arbres de décision**.
- `sequenceDiagram` : pour les **enchaînements complexes** entre **acteurs ou systèmes**.
- `classDiagram` : si **demande explicite** ou si **entrants techniques** (ex. schéma SQL).
- `stateDiagram-v2` : si **demande explicite** ou **entrants techniques** comme une user story.
- `erDiagram` : autre **choix par défaut** pour exposer les relations entités-attributs. Regroupe et colore les blocs.
- `gantt` : si **thématique projet ou travaux**, avec des **tâches identifiables**.
- `timeline` : pour toute **chronologie**, **biographie**, ou **historique**. Regroupe les événements par **décennie**, ou **siècle** si nécessaire. Exemple de syntaxe : 
timeline
title History of Social Media Platform
2002 : LinkedIn
2004 : Facebook
: Google
2005 : YouTube
2006 : Twitter

### Règles de syntaxe :
- Pas de parenthèses dans les libellés, c'est réservé syntaxe mermaid, exemple JeuxDeDonnees[450 000 Jeux de Données (2024)] devient JeuxDeDonnees[450 000 Jeux de Données en 2024]
- Les noms d’entités doivent être **concis**, s'ils ne sont pas affichés, en sécuriser la syntax avec des underscores, pas d'accent etc.
- Évite aussi les caractères spéciaux non compatibles Mermaid (`@`, `{}`, `[]`, `(`, `)` etc.).
- Utilise des noms explicites pour les nœuds (ex. `AgentDemandeur` plutôt que `A`)
- Utilise `%%` pour commenter ou annoter hors rendu visuel.
- Segmente clairement, évite les diagrammes illisibles ou trop denses.
- Rédige dans la **langue du demandeur** ou celle du **texte d’origine**.

Objectif : produire un diagramme Mermaid clair, fidèle, et directement exploitable dans un contexte explicatif, technique ou décisionnel.
"""
import os



def getModelResponseFromContent(content, provider_name="sambanova", provider_env_key="SAMBANOVA_API_KEY", model="meta-llama/Llama-4-Scout-17B-16E-Instruct"):
    client = InferenceClient(
        provider=provider_name,
        api_key=os.environ[provider_env_key],
    )

    completion = client.chat.completions.create(
        model=model,
        messages=[
            {
                "role": "user",
                f"content": content
            }
        ],
        temperature=0.5
    )

    raw_response = completion.choices[0].message.content
    mermaid_code_cleaned = raw_response.replace("```mermaid", "").replace("```", "").strip()
    mermaid_diagram_type = mermaid_code_cleaned.split('\n', 1)[0].strip()
    
    return mermaid_code_cleaned, mermaid_diagram_type

def kroki_img_from_mermaid(mermaid_code: str, output_format: str = "svg") -> str:
    compressed = zlib.compress(mermaid_code.encode('utf-8'))
    b64 = base64.urlsafe_b64encode(compressed).decode('utf-8')
    image_url = f"https://kroki.io/mermaid/{output_format}/{b64}"
    html_img = f"""<a href="{image_url}" target="_blank"> <img src="{image_url}" alt="Diagramme Mermaid via Kroki" > </a>"""
    return html_img


def generateDiagram(input_text="I broke my car yesterday, and the mechanics will give it back to me after insurance payment.", input_precisions="Choose the best diagram paradigm"):
    """Generates a diagram to explain a submitted text. The diagram is returned in different formats (HTML and raw code). It handles various types like Graph TB (default, block grouping with colors), flowchart TD (decision processes), sequenceDiagram (actor/system interactions), classDiagram (technical models like SQL schemas), stateDiagram-v2 (state logic or user stories), erDiagram (data entity relationships), gantt (projects and task timelines along dates), and timeline (historical/biographical sequences grouped by decade or century).
    
    Args:
        input_text: the text to explain with a diagram.
        input_precisions: diagram settings, by default the tool has the ability to choose the best type of diagram.
        
    Returns:
        kroki_html_result: **Ready to display for chatbots** html diagram.
        mermaid_diag_type: type of the diagram selected by the tool
        mermaid_code: mermaid code
    """
    full_content_prompt = f"{PROMPT}. ###INPUT : {input_text}. ###PRECISIONS {input_precisions}"
    mermaid_code, mermaid_diag_type = getModelResponseFromContent(content=full_content_prompt )
    kroki_html_link = kroki_img_from_mermaid(mermaid_code)
    #kroki_html_img =  kroki_html_link[kroki_html_link.find('<img '):kroki_html_link.find('>', kroki_html_link.find('<img ')) + 1]
    return kroki_html_link, mermaid_diag_type, mermaid_code


#La solution tire a son avantage la rapidité du serveur d'inférence SAmbanova couplée à l'ingéniosité de Scout.