File size: 1,413 Bytes
b5c4d16 b0fac8f b5c4d16 f897f8d 98fd32e b0fac8f 98fd32e b5c4d16 f6662f9 37f4efb e432eb1 99e2b34 f897f8d 99e2b34 b5c4d16 f6662f9 b0fac8f 37f4efb b5c4d16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
from src.agent import PlantUMLAgent
from src.documents import get_processed_documents, get_diagrams_names
from src.interface import create_interface
from src.preview import render_plantuml, sanitize_plantuml_code
import gradio as gr
from cachetools import LRUCache, cached
diagrams_names = get_diagrams_names()
docs_processed = get_processed_documents(diagrams_names)
cache = LRUCache(maxsize=100)
@cached(cache)
def respond(message, diagram_name, image_input, api_key):
if not api_key:
gr.Warning("Please provide a valid Mistral API key to generate diagrams.")
return message, "", ""
if not message:
gr.Warning("Please provide a description for the UML diagram.")
return "", "", ""
agent = PlantUMLAgent(docs_processed, api_key)
input = "Description:\n" + message.strip() + "\n\nDiagram Type: " + diagram_name.strip() + "\n\n"
if image_input:
image_analysis = agent.recognize_image(diagram_name, image_input)
input += f"\n\nImage Analysis:\n{image_analysis}\n\n" if image_analysis else ""
response = agent.predict(diagram_name, input)
plantuml_code = sanitize_plantuml_code(response)
plantuml_diagram = render_plantuml(plantuml_code)
return message, plantuml_code, plantuml_diagram
if __name__ == "__main__":
cache.clear()
demo = create_interface(diagrams_names, respond)
demo.launch() |