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 mermaid diagram to explain a submited text. The diagram is returned in different formats(html and raw code)
    
    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: *default for agentic chatbots* full kroki html to embed with link with image and link to kroki page.
        mermaid_diag_type: type of 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