Spaces:
Sleeping
Sleeping
Sávio Santos
commited on
Commit
·
b5c4d16
1
Parent(s):
e432eb1
new version
Browse files- .gitattributes +0 -4
- .gitignore +1 -36
- PlantUML_Language_Reference_Guide_en.pdf +0 -3
- README.md +3 -6
- app.py +15 -30
- preview.png +0 -3
- requirements.txt +5 -5
- src/agent.py +35 -0
- src/documents.py +30 -18
- src/examples.py +0 -53
- src/interface.py +54 -48
- src/plantuml.py +0 -9
- src/preview.py +17 -0
- src/retriever.py +0 -24
- src/utils.py +0 -15
.gitattributes
DELETED
@@ -1,4 +0,0 @@
|
|
1 |
-
*.pdf filter=lfs diff=lfs merge=lfs -text
|
2 |
-
*.png filter=lfs diff=lfs merge=lfs -text
|
3 |
-
preview.png filter=lfs diff=lfs merge=lfs -text
|
4 |
-
PlantUML_Language_Reference_Guide_en.pdf filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
.gitignore
CHANGED
@@ -1,37 +1,2 @@
|
|
1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
23 |
-
*.pth filter=lfs diff=lfs merge=lfs -text
|
24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
36 |
__pycache__
|
37 |
-
.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
__pycache__
|
2 |
+
.env
|
PlantUML_Language_Reference_Guide_en.pdf
DELETED
@@ -1,3 +0,0 @@
|
|
1 |
-
version https://git-lfs.github.com/spec/v1
|
2 |
-
oid sha256:1455bd5f63603ede0545f66d67de2162cce7b7b80e03e98d20e674e6c272a8a6
|
3 |
-
size 10048681
|
|
|
|
|
|
|
|
README.md
CHANGED
@@ -37,7 +37,7 @@ The PlantUML Agent is an interactive tool that leverages Large Language Models (
|
|
37 |
### Agent Flow
|
38 |
|
39 |
```
|
40 |
-
User Input -> Retriever Tool ->
|
41 |
```
|
42 |
|
43 |
## Setup and Installation
|
@@ -54,7 +54,7 @@ To set up and run the PlantUML Agent locally, follow these steps:
|
|
54 |
2. **Create a conda environment:**
|
55 |
|
56 |
```bash
|
57 |
-
conda create -n plantuml-agent python=3.
|
58 |
conda activate plantuml-agent
|
59 |
```
|
60 |
|
@@ -76,7 +76,4 @@ To start the PlantUML Agent application, run the `app.py` file:
|
|
76 |
|
77 |
```bash
|
78 |
python app.py
|
79 |
-
```
|
80 |
-
|
81 |
-
## Preview
|
82 |
-

|
|
|
37 |
### Agent Flow
|
38 |
|
39 |
```
|
40 |
+
User Input -> Retriever Tool -> PlantUML Tool -> PlantUML Code + PlantUML Diagram
|
41 |
```
|
42 |
|
43 |
## Setup and Installation
|
|
|
54 |
2. **Create a conda environment:**
|
55 |
|
56 |
```bash
|
57 |
+
conda create -n plantuml-agent python=3.13 --y
|
58 |
conda activate plantuml-agent
|
59 |
```
|
60 |
|
|
|
76 |
|
77 |
```bash
|
78 |
python app.py
|
79 |
+
```
|
|
|
|
|
|
app.py
CHANGED
@@ -1,46 +1,31 @@
|
|
1 |
-
from src.
|
2 |
-
from src.
|
3 |
-
from src.plantuml import render_plantuml
|
4 |
-
from src.utils import get_prompt, get_processed_response
|
5 |
from src.interface import create_interface
|
|
|
6 |
|
7 |
-
from smolagents import InferenceClientModel, CodeAgent
|
8 |
from cachetools import LRUCache, cached
|
9 |
|
10 |
-
import os
|
11 |
from dotenv import load_dotenv
|
12 |
|
13 |
load_dotenv()
|
14 |
|
15 |
-
|
16 |
-
|
17 |
|
18 |
-
|
19 |
-
|
20 |
-
agent = CodeAgent(
|
21 |
-
tools=[retriever_tool],
|
22 |
-
model=model,
|
23 |
-
max_steps=4,
|
24 |
-
verbosity_level=2,
|
25 |
-
)
|
26 |
|
27 |
cache = LRUCache(maxsize=100)
|
28 |
|
29 |
-
def get_plantuml_diagram(message):
|
30 |
-
prompt = get_prompt(message)
|
31 |
-
response = agent.run(prompt)
|
32 |
-
|
33 |
-
return get_processed_response(response)
|
34 |
-
|
35 |
@cached(cache)
|
36 |
-
def respond(message):
|
37 |
-
|
38 |
-
|
39 |
-
|
|
|
|
|
|
|
40 |
|
41 |
if __name__ == "__main__":
|
42 |
cache.clear()
|
43 |
-
demo = create_interface(respond)
|
44 |
-
demo.launch()
|
45 |
-
|
46 |
-
demo.launch()
|
|
|
1 |
+
from src.agent import PlantUMLAgent
|
2 |
+
from src.documents import get_processed_documents, get_diagrams_names
|
|
|
|
|
3 |
from src.interface import create_interface
|
4 |
+
from src.preview import render_plantuml, sanitize_plantuml_code
|
5 |
|
|
|
6 |
from cachetools import LRUCache, cached
|
7 |
|
|
|
8 |
from dotenv import load_dotenv
|
9 |
|
10 |
load_dotenv()
|
11 |
|
12 |
+
diagrams_names = get_diagrams_names()
|
13 |
+
docs_processed = get_processed_documents(diagrams_names)
|
14 |
|
15 |
+
agent = PlantUMLAgent('mistral-small-latest', docs_processed)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
16 |
|
17 |
cache = LRUCache(maxsize=100)
|
18 |
|
|
|
|
|
|
|
|
|
|
|
|
|
19 |
@cached(cache)
|
20 |
+
def respond(message, diagram_name):
|
21 |
+
response = agent.predict(message, diagram_name)
|
22 |
+
|
23 |
+
plantuml_code = sanitize_plantuml_code(response)
|
24 |
+
plantuml_diagram = render_plantuml(plantuml_code)
|
25 |
+
|
26 |
+
return message, plantuml_code, plantuml_diagram
|
27 |
|
28 |
if __name__ == "__main__":
|
29 |
cache.clear()
|
30 |
+
demo = create_interface(diagrams_names, respond)
|
31 |
+
demo.launch()
|
|
|
|
preview.png
DELETED
Git LFS Details
|
requirements.txt
CHANGED
@@ -1,11 +1,11 @@
|
|
1 |
gradio
|
2 |
-
pymupdf
|
3 |
langchain
|
4 |
langchain_community
|
5 |
-
smolagents
|
6 |
-
smolagents[gradio]
|
7 |
-
plantweb
|
8 |
rank_bm25
|
9 |
numpy==1.26.4
|
10 |
python-dotenv
|
11 |
-
cachetools
|
|
|
|
|
|
|
|
|
|
1 |
gradio
|
|
|
2 |
langchain
|
3 |
langchain_community
|
|
|
|
|
|
|
4 |
rank_bm25
|
5 |
numpy==1.26.4
|
6 |
python-dotenv
|
7 |
+
cachetools
|
8 |
+
beautifulsoup4
|
9 |
+
docling
|
10 |
+
plantweb
|
11 |
+
mistralai
|
src/agent.py
ADDED
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from langchain_community.retrievers import BM25Retriever
|
2 |
+
from mistralai import Mistral
|
3 |
+
import os
|
4 |
+
|
5 |
+
class PlantUMLAgent:
|
6 |
+
def __init__(self, model, docs_processed):
|
7 |
+
self.model = model
|
8 |
+
self.docs_processed = docs_processed
|
9 |
+
self.client = Mistral(api_key=os.getenv("MISTRAL_API_KEY"))
|
10 |
+
|
11 |
+
def predict(self, message, diagram_name):
|
12 |
+
diagram_name = diagram_name.replace(" ", "-").lower()
|
13 |
+
|
14 |
+
filtered_docs = [doc for doc in self.docs_processed if diagram_name.lower() in doc.metadata["source"].lower()]
|
15 |
+
|
16 |
+
retriever = BM25Retriever.from_documents(filtered_docs, k=10)
|
17 |
+
|
18 |
+
documents = retriever.invoke(message)
|
19 |
+
documents_texts = [doc.page_content for doc in documents]
|
20 |
+
|
21 |
+
chat_response = self.client.chat.complete(
|
22 |
+
model=self.model,
|
23 |
+
messages=[
|
24 |
+
{
|
25 |
+
"role": "system",
|
26 |
+
"content": "You are a PlantUML code generator. Generate PlantUML diagram code based on the provided description."
|
27 |
+
},
|
28 |
+
{
|
29 |
+
"role": "user",
|
30 |
+
"content": f"Description:\n{message}\n\nUse the following documents as reference:\n{documents_texts}\n\nReturn the output between ```plantuml and ```."
|
31 |
+
}
|
32 |
+
]
|
33 |
+
)
|
34 |
+
response = chat_response.choices[0].message.content
|
35 |
+
return response.strip()
|
src/documents.py
CHANGED
@@ -1,9 +1,35 @@
|
|
1 |
-
import fitz
|
2 |
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
3 |
from langchain.docstore.document import Document
|
|
|
|
|
|
|
4 |
|
5 |
-
|
6 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
7 |
|
8 |
text_splitter = RecursiveCharacterTextSplitter(
|
9 |
chunk_size=500,
|
@@ -15,18 +41,4 @@ def get_processed_documents(pdf_path: str):
|
|
15 |
|
16 |
docs_processed = text_splitter.split_documents(source_docs)
|
17 |
|
18 |
-
return docs_processed
|
19 |
-
|
20 |
-
def load_pdf_as_documents(pdf_path: str):
|
21 |
-
doc = fitz.open(pdf_path)
|
22 |
-
documents = []
|
23 |
-
for i, page in enumerate(doc):
|
24 |
-
text = page.get_text().strip()
|
25 |
-
if text:
|
26 |
-
documents.append(
|
27 |
-
Document(
|
28 |
-
page_content=text,
|
29 |
-
metadata={"source": f"page_{i+1}"}
|
30 |
-
)
|
31 |
-
)
|
32 |
-
return documents
|
|
|
|
|
1 |
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
2 |
from langchain.docstore.document import Document
|
3 |
+
from docling.document_converter import DocumentConverter
|
4 |
+
from bs4 import BeautifulSoup
|
5 |
+
import requests
|
6 |
|
7 |
+
plantuml_url = "https://plantuml.com"
|
8 |
+
|
9 |
+
def get_diagrams_names():
|
10 |
+
response = requests.get(plantuml_url)
|
11 |
+
soup = BeautifulSoup(response.content, 'html.parser')
|
12 |
+
menu = soup.find('div', class_='menu1')
|
13 |
+
links = menu.find_all('a')
|
14 |
+
diagrams = [link['href'] for link in links if 'href' in link.attrs]
|
15 |
+
return diagrams
|
16 |
+
|
17 |
+
def load_documents(diagrams_names):
|
18 |
+
converter = DocumentConverter()
|
19 |
+
documents = []
|
20 |
+
for diagram_name in diagrams_names:
|
21 |
+
result = converter.convert(f"{plantuml_url}/{diagram_name}")
|
22 |
+
markdown = result.document.export_to_markdown()
|
23 |
+
documents.append(
|
24 |
+
Document(
|
25 |
+
page_content=markdown,
|
26 |
+
metadata={"source": diagram_name}
|
27 |
+
)
|
28 |
+
)
|
29 |
+
return documents
|
30 |
+
|
31 |
+
def get_processed_documents(diagrams_names):
|
32 |
+
source_docs = load_documents(diagrams_names)
|
33 |
|
34 |
text_splitter = RecursiveCharacterTextSplitter(
|
35 |
chunk_size=500,
|
|
|
41 |
|
42 |
docs_processed = text_splitter.split_documents(source_docs)
|
43 |
|
44 |
+
return docs_processed
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
src/examples.py
DELETED
@@ -1,53 +0,0 @@
|
|
1 |
-
def get_plantuml_examples():
|
2 |
-
return {
|
3 |
-
'Default example': {
|
4 |
-
"description": "Select a diagram type to see examples.",
|
5 |
-
"input": "Generate a simple use case diagram."
|
6 |
-
},
|
7 |
-
"Use Case Diagram": {
|
8 |
-
"description": "A use case diagram showing how users (actors) interact with a system and what functionalities (use cases) they use.",
|
9 |
-
"input": """Create a use case diagram for a library system.
|
10 |
-
Actors are: Reader, Librarian.
|
11 |
-
Use cases for Reader are: Search Book, Borrow Book, Return Book.
|
12 |
-
Use cases for Librarian are: Add Book, Remove Book, Manage Loans."""
|
13 |
-
},
|
14 |
-
"Class Diagram": {
|
15 |
-
"description": "A class diagram displaying the static structure of a system, showing its classes, attributes, operations, and relationships.",
|
16 |
-
"input": """Generate a class diagram for an e-commerce system.
|
17 |
-
Classes are: Customer, Product, Order, OrderItem.
|
18 |
-
Customer has attributes: name, email, address.
|
19 |
-
Product has attributes: name, price, description.
|
20 |
-
Order has attributes: date, status.
|
21 |
-
OrderItem has attributes: quantity, unitPrice.
|
22 |
-
Relationships:
|
23 |
-
- Customer can place 0 or many Orders.
|
24 |
-
- Order contains 1 or many OrderItem.
|
25 |
-
- OrderItem is associated with 1 Product."""
|
26 |
-
},
|
27 |
-
"Sequence Diagram": {
|
28 |
-
"description": "A sequence diagram illustrating the order of interactions between objects in a specific scenario, showing message exchange over time.",
|
29 |
-
"input": """Draw a sequence diagram for user login.
|
30 |
-
Participants are: User, Login System, Database.
|
31 |
-
1. User sends (enter credentials) to Login System.
|
32 |
-
2. Login System queries (verify credentials) Database.
|
33 |
-
3. Database responds (valid/invalid credentials) to Login System.
|
34 |
-
4. If valid credentials, Login System sends (successful login) to User.
|
35 |
-
5. If invalid credentials, Login System sends (login error) to User."""
|
36 |
-
},
|
37 |
-
"Activity Diagram": {
|
38 |
-
"description": "An activity diagram modeling the control flow of activities in a process or workflow, including decisions and branches.",
|
39 |
-
"input": """Create an activity diagram for the online purchasing process.
|
40 |
-
Start activity.
|
41 |
-
1. Browse Products.
|
42 |
-
2. Add to Cart.
|
43 |
-
3. Decision (Cart Full?).
|
44 |
-
- If yes: Go to Checkout.
|
45 |
-
- If no: Continue Browse.
|
46 |
-
4. Process Payment.
|
47 |
-
5. Decision (Payment Approved?).
|
48 |
-
- If yes: Confirm Order.
|
49 |
-
- If no: Notify Payment Error.
|
50 |
-
6. Ship Order.
|
51 |
-
End activity."""
|
52 |
-
}
|
53 |
-
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
src/interface.py
CHANGED
@@ -1,73 +1,79 @@
|
|
1 |
import gradio as gr
|
2 |
-
from src.examples import get_plantuml_examples
|
3 |
|
4 |
-
def
|
5 |
-
|
6 |
-
|
|
|
|
|
|
|
|
|
7 |
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
def set_example_input(example_name):
|
12 |
-
return PLANTUML_EXAMPLES.get(example_name, {}).get('input', "")
|
13 |
-
|
14 |
-
with gr.Blocks(title="PlantUML Agent") as demo:
|
15 |
gr.HTML(
|
16 |
"""
|
17 |
-
<h1 style="text-align: center;">PlantUML Agent</h1>
|
18 |
-
<div style="text-align: center;">
|
19 |
-
<img src="https://repository-images.githubusercontent.com/553868400/2f7375d8-c2ee-44fe-81d0-a8a29ec284da" alt="PlantUML Logo" style="height:
|
20 |
</div>
|
21 |
-
<p style="text-align: center; font-size: 1.1em;">
|
22 |
-
Generate UML diagrams using <a href="https://plantuml.com" target="_blank">PlantUML</a> syntax with the help of an AI agent.
|
23 |
</p>
|
|
|
24 |
"""
|
25 |
)
|
26 |
|
27 |
-
with gr.Row(
|
28 |
-
with gr.Column(scale=1
|
29 |
-
gr.Markdown("
|
30 |
msg = gr.Textbox(
|
31 |
-
label="UML
|
32 |
placeholder="e.g., create a basic use case class",
|
33 |
-
|
34 |
-
|
|
|
35 |
)
|
36 |
|
37 |
-
gr.Markdown("
|
38 |
-
|
39 |
-
choices=
|
40 |
-
value=
|
41 |
-
label="
|
42 |
interactive=True
|
43 |
)
|
44 |
|
45 |
-
|
46 |
-
fn=set_example_input,
|
47 |
-
inputs=example_dropdown,
|
48 |
-
outputs=msg
|
49 |
-
)
|
50 |
|
51 |
-
|
|
|
52 |
|
53 |
-
|
54 |
-
gr.
|
55 |
-
|
56 |
-
|
|
|
|
|
57 |
|
58 |
-
gr.Markdown("
|
59 |
-
diagram_output = gr.HTML(
|
|
|
|
|
|
|
60 |
|
61 |
-
gr.
|
62 |
-
"
|
63 |
-
"
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
|
|
|
|
|
|
68 |
)
|
69 |
|
70 |
-
submit.click(respond, [msg], [msg, diagram_text_output, diagram_output])
|
71 |
-
msg.submit(respond, [msg], [msg, diagram_text_output, diagram_output])
|
72 |
|
73 |
return demo
|
|
|
1 |
import gradio as gr
|
|
|
2 |
|
3 |
+
def beautify_diagram_types(diagram_types):
|
4 |
+
first_diagram = "use-case-diagram"
|
5 |
+
if first_diagram in diagram_types:
|
6 |
+
diagram_types.remove(first_diagram)
|
7 |
+
diagram_types.insert(0, first_diagram)
|
8 |
+
diagram_types.sort()
|
9 |
+
return [name.replace("-", " ").title() for name in diagram_types]
|
10 |
|
11 |
+
def create_interface(diagram_types, respond):
|
12 |
+
diagram_types = beautify_diagram_types(diagram_types)
|
13 |
+
with gr.Blocks(title="PlantUML Agent", theme=gr.themes.Soft()) as demo:
|
|
|
|
|
|
|
|
|
14 |
gr.HTML(
|
15 |
"""
|
16 |
+
<h1 style="text-align: center; color: #333;">PlantUML Agent</h1>
|
17 |
+
<div style="text-align: center; margin-bottom: 20px;">
|
18 |
+
<img src="https://repository-images.githubusercontent.com/553868400/2f7375d8-c2ee-44fe-81d0-a8a29ec284da" alt="PlantUML Logo" style="height: 60px;">
|
19 |
</div>
|
20 |
+
<p style="text-align: center; font-size: 1.1em; color: #555;">
|
21 |
+
Generate UML diagrams using <a href="https://plantuml.com" target="_blank" style="color: #1A73E8; text-decoration: none;">PlantUML</a> syntax with the help of an AI agent.
|
22 |
</p>
|
23 |
+
<hr style="border: 0; height: 1px; background: #eee; margin: 20px 0;">
|
24 |
"""
|
25 |
)
|
26 |
|
27 |
+
with gr.Row():
|
28 |
+
with gr.Column(scale=1):
|
29 |
+
gr.Markdown("## 📝 Input Description")
|
30 |
msg = gr.Textbox(
|
31 |
+
label="Describe your UML diagram",
|
32 |
placeholder="e.g., create a basic use case class",
|
33 |
+
value="Create a basic use case diagram",
|
34 |
+
lines=7,
|
35 |
+
interactive=True
|
36 |
)
|
37 |
|
38 |
+
gr.Markdown("## 📂 Select Diagram Type")
|
39 |
+
diagram_dropdown = gr.Dropdown(
|
40 |
+
choices=diagram_types,
|
41 |
+
value="Use Case Diagram",
|
42 |
+
label="Choose the type of UML diagram",
|
43 |
interactive=True
|
44 |
)
|
45 |
|
46 |
+
submit = gr.Button("🚀 Generate Diagram", variant="primary", size="lg")
|
|
|
|
|
|
|
|
|
47 |
|
48 |
+
with gr.Column(scale=2):
|
49 |
+
gr.Markdown("## 📊 Generated Diagram")
|
50 |
|
51 |
+
gr.Markdown("### PlantUML Code")
|
52 |
+
diagram_text_output = gr.Code(
|
53 |
+
label="The generated PlantUML code",
|
54 |
+
lines=15,
|
55 |
+
interactive=False,
|
56 |
+
)
|
57 |
|
58 |
+
gr.Markdown("### Diagram Preview")
|
59 |
+
diagram_output = gr.HTML(
|
60 |
+
label="Visual representation of your diagram",
|
61 |
+
elem_id="diagram-preview"
|
62 |
+
)
|
63 |
|
64 |
+
gr.HTML(
|
65 |
+
"""
|
66 |
+
<hr style="border: 0; height: 1px; background: #eee; margin: 30px 0 15px 0;">
|
67 |
+
<div style='text-align: center; font-size: 0.9em; color: #777;'>
|
68 |
+
Powered by Mistral, BM25Retriever, LangChain, Docling, LRUCache, Gradio, and PlantUML.
|
69 |
+
</div>
|
70 |
+
<div style='text-align: center; font-size: 0.9em; color: #777; margin-top: 5px;'>
|
71 |
+
Created by <a href='https://linkedin.com/in/savi8sant8s' target='_blank' style="color: #1A73E8; text-decoration: none;">Sávio Santos</a>
|
72 |
+
</div>
|
73 |
+
"""
|
74 |
)
|
75 |
|
76 |
+
submit.click(respond, [msg, diagram_dropdown], [msg, diagram_text_output, diagram_output])
|
77 |
+
msg.submit(respond, [msg, diagram_dropdown], [msg, diagram_text_output, diagram_output])
|
78 |
|
79 |
return demo
|
src/plantuml.py
DELETED
@@ -1,9 +0,0 @@
|
|
1 |
-
from plantweb import render
|
2 |
-
|
3 |
-
def render_plantuml(uml_code):
|
4 |
-
try:
|
5 |
-
svg_output, _, _, _ = render.render(uml_code, engine='plantuml', format='svg')
|
6 |
-
svg_output = svg_output.decode('utf-8')
|
7 |
-
return svg_output
|
8 |
-
except Exception as e:
|
9 |
-
return f"Error rendering PlantUML: {e}"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
src/preview.py
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from plantweb import render
|
2 |
+
import re
|
3 |
+
|
4 |
+
def sanitize_plantuml_code(plantuml_code):
|
5 |
+
match = re.search(r'```plantuml\s*(.*?)\s*```', plantuml_code, re.DOTALL)
|
6 |
+
if match:
|
7 |
+
return match.group(1).strip()
|
8 |
+
return "@startuml\n'Error\n@enduml"
|
9 |
+
|
10 |
+
def render_plantuml(plantuml_code):
|
11 |
+
try:
|
12 |
+
svg_output, _, _, _ = render.render(plantuml_code, engine='plantuml', format='svg')
|
13 |
+
|
14 |
+
return svg_output.decode('utf-8')
|
15 |
+
|
16 |
+
except Exception as e:
|
17 |
+
return f"Error rendering PlantUML: {e}"
|
src/retriever.py
DELETED
@@ -1,24 +0,0 @@
|
|
1 |
-
from langchain_community.retrievers import BM25Retriever
|
2 |
-
from smolagents import Tool
|
3 |
-
|
4 |
-
class RetrieverTool(Tool):
|
5 |
-
name = "retriever"
|
6 |
-
description = "Retrieves relevant chunks from the PlantUML reference guide."
|
7 |
-
inputs = {
|
8 |
-
"query": {
|
9 |
-
"type": "string",
|
10 |
-
"description": "The information to search for in the PlantUML documentation.",
|
11 |
-
}
|
12 |
-
}
|
13 |
-
output_type = "string"
|
14 |
-
|
15 |
-
def __init__(self, docs, **kwargs):
|
16 |
-
super().__init__(**kwargs)
|
17 |
-
self.retriever = BM25Retriever.from_documents(docs, k=10)
|
18 |
-
|
19 |
-
def forward(self, query: str) -> str:
|
20 |
-
assert isinstance(query, str), "Query must be a string"
|
21 |
-
docs = self.retriever.invoke(query)
|
22 |
-
return "\nRetrieved documents:\n" + "".join(
|
23 |
-
[f"\n\n===== Document {i} =====\n{doc.page_content}" for i, doc in enumerate(docs)]
|
24 |
-
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
src/utils.py
DELETED
@@ -1,15 +0,0 @@
|
|
1 |
-
import re
|
2 |
-
|
3 |
-
def get_prompt(message):
|
4 |
-
return ("Generate a PlantUML diagram based on the following message:\n"
|
5 |
-
f"{message}\n"
|
6 |
-
"Use the PlantUML syntax and include all necessary components.\n"
|
7 |
-
"Return the PlantUML code wrapped in @startuml and @enduml tags."
|
8 |
-
"Do not include any additional text outside these tags without commenting it out.")
|
9 |
-
|
10 |
-
def get_processed_response(response):
|
11 |
-
matches = re.findall(r"@startuml(.*?)@enduml", response, re.DOTALL)
|
12 |
-
if matches:
|
13 |
-
last_match = matches[-1].strip()
|
14 |
-
return f"@startuml\n{last_match}\n@enduml"
|
15 |
-
return "@startuml\n@enduml"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|