llll

main.py

@@ -69,32 +69,13 @@ 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):
     input_text = request.input
     model_choice = request.model.lower()
 
     response = get_model_response(model_choice , input_text)
-    return {"response": remove_notes_from_sequence_diagram(format_response(response))}
+    return {"response": format_response(response)}
 
 # To run the FastAPI app, use:
 # uvicorn main:app --reload

system_message.py

@@ -22,7 +22,7 @@ When generating diagrams using the Mermaid diagramming language, ensure to follo
         - A --> text>B. (incorrect)
         - A -->|text|> B (incorrect)
     
-    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. 
+    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. DON'T ADD NOTES TO SEQUENCE DIAGRAMS. 
 
 When responding, follow this format: