kllll

main.py

@@ -69,6 +69,24 @@ def format_response(response: str) -> str:
     formatted_response = formatted_response.replace('\\n', '\n')
     
     return formatted_response
+def remove_notes_from_sequence_diagram(mermaid_code: str) -> str:
+    lines = mermaid_code.split('\n')
+    filtered_lines = []
+    in_mermaid_block = False
+
+    for line in lines:
+        stripped_line = line.strip()
+        if stripped_line.startswith("```mermaid"):
+            in_mermaid_block = True
+        elif stripped_line.startswith("```") and in_mermaid_block:
+            in_mermaid_block = False
+        
+        if in_mermaid_block and stripped_line.startswith("note"):
+            continue
+
+        filtered_lines.append(line)
+    
+    return '\n'.join(filtered_lines)
 
 @app.post("/generate-response/")
 async def generate_response(request: PromptRequest):
@@ -76,7 +94,7 @@ async def generate_response(request: PromptRequest):
     model_choice = request.model.lower()
 
     response = get_model_response(model_choice , input_text)
-    return {"response": format_response(response)}
+    return {"response": remove_notes_from_sequence_diagram(format_response(response))}
 
 # To run the FastAPI app, use:
 # uvicorn main:app --reload

system_message.py

@@ -5,7 +5,7 @@ The user will ask you to explain some concept and your task is to explain concep
          
 Follow these guidelines for generating the correct types of diagrams:
     1. Flowcharts are suitable for: Algorithm visualization, process flows ,conditional logic, and control structures (if-else, loops).
-    2. Sequence Diagrams are suitable for: Interaction between objects or components, communication protocols, method calls in OOP, and client-server interactions. 
+    2. Sequence Diagrams are suitable for: Interaction between objects or components, communication protocols, method calls in OOP, and client-server interactions. You must not add notes to sequence diagrams.  
     3. Class Diagrams are suitable for: Object-Oriented Programming (OOP) concepts, class structures, inheritance and relationships, and system architecture. 
     4. State Diagrams are suitable for: State machines, lifecycle of objects, protocol states, and workflow states
     5. Entity-Relationship Diagrams (ERD) are suitable for: Database schema design, relationships between data entities, and data modeling. 
@@ -22,7 +22,7 @@ When generating diagrams using the Mermaid diagramming language, ensure to follo
         - A --> text>B. (incorrect)
         - A -->|text|> B (incorrect)
     
-    3. Never add notes or alts when creating sequence diagrams,additional comments or alternative scenarios are not supported directly. DON'T ADD NOTES TO SEQUENCE DIAGRAMS. 
+    3. In sequence diagrams, don't add notes to them. Never add notes or alts when creating sequence diagrams,additional comments or alternative scenarios are not supported directly. DON'T ADD NOTES TO SEQUENCE DIAGRAMS. 
 
 When responding, follow this format:
 

test.py

@@ -1,16 +0,0 @@
-import requests
-
-# Define the URL of the API endpoint
-url = "http://127.0.0.1:8000/generate-response/"
-
-# Define the payload
-payload = {
-    "input": "give me a roadmap for frontend development using a mindmap.",
-    "model": "llama"  # Change to "mixtral" to test the Mixtral model
-}
-
-# Send a POST request to the API
-response = requests.post(url, json=payload)
-
-# Print the response from the API
-print(response.json())