Support Custom Note Max Length

Author: coolbeevip

README.md

@@ -79,6 +79,12 @@ Sequence Diagram
 callback_handler = diagram.sequence_diagram_callback()
 ```
 
+Support Custom Note Max Length(default 1000)
+
+```python
+callback_handler = diagram.activity_diagram_callback(note_max_length=2000)
+```
+
 ## Exporting PlantUML to PNG
 
 You can download [plantuml.1.2023.10.jar](https://github.com/plantuml/plantuml/releases/download/v1.2023.10/plantuml-1.2023.10.jar)

langchain_plantuml/activity_diagram_beta/plantuml_activity_diagram_beta_callback_handler.py

@@ -33,13 +33,17 @@ class PlantUMLActivityDiagramCallbackHandler(BasePlantUMLCallbackHandler):
     _runs_metrics: dict = {}
 
     def __init__(
-        self, color: Optional[str] = None, skin_param: List[str] = DEFAULT_SKIN_PARAM
+        self,
+        color: Optional[str] = None,
+        skin_param: List[str] = DEFAULT_SKIN_PARAM,
+        note_max_length: int = 1000,
     ) -> None:
         super().__init__()
         for param in skin_param:
             self.uml_content.append(param)
         self.uml_content.append("start")
         self.color = color
+        self.note_max_length = note_max_length
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any

langchain_plantuml/activity_diagram_beta/plantuml_sequence_diagram_callback_handler.py

@@ -34,7 +34,10 @@ class PlantUMLSequenceDiagramCallbackHandler(BasePlantUMLCallbackHandler):
     _runs_metrics: dict = {}
 
     def __init__(
-        self, color: Optional[str] = None, skin_param: List[str] = DEFAULT_SKIN_PARAM
+        self,
+        color: Optional[str] = None,
+        skin_param: List[str] = DEFAULT_SKIN_PARAM,
+        note_max_length: int = 1000,
     ) -> None:
         super().__init__()
         for param in skin_param:
@@ -43,6 +46,7 @@ def __init__(
         self.participants = {}
         self.participant_name_indexes = []
         self.color = color
+        self.note_max_length = note_max_length
 
     def on_llm_start(
         self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any

langchain_plantuml/diagram.py

@@ -20,9 +20,9 @@
     PlantUMLSequenceDiagramCallbackHandler
 
 
-def activity_diagram_callback() -> BaseCallbackHandler:
-    return PlantUMLActivityDiagramCallbackHandler()
+def activity_diagram_callback(note_max_length: int = 1000) -> BaseCallbackHandler:
+    return PlantUMLActivityDiagramCallbackHandler(note_max_length=note_max_length)
 
 
-def sequence_diagram_callback() -> BaseCallbackHandler:
-    return PlantUMLSequenceDiagramCallbackHandler()
+def sequence_diagram_callback(note_max_length: int = 1000) -> BaseCallbackHandler:
+    return PlantUMLSequenceDiagramCallbackHandler(note_max_length=note_max_length)

pyproject.toml

@@ -18,7 +18,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "langchain-plantuml"
-version = "0.0.3"
+version = "0.0.4"
 description = "Subscribe to events using a callback and store them in PlantUML format. You can easily subscribe to events and keep them in a form that is easy to visualize and analyze."
 authors = [{ name = "Lei Zhang", email = "[email protected]" }]
 license = {file = "LICENSE"}

tests/activity_diagram_beta/test_scene_agent.py

@@ -28,8 +28,8 @@ def setUpClass(cls):
         cls._agent = SceneAgent()
 
     def test(self):
-        activity_diagram_callback = diagram.activity_diagram_callback()
-        sequence_diagram_callback = diagram.sequence_diagram_callback()
+        activity_diagram_callback = diagram.activity_diagram_callback(note_max_length=2000)
+        sequence_diagram_callback = diagram.sequence_diagram_callback(note_max_length=2000)
         question = "What did biden say about ketanji brown jackson in the state of the union address?"
         try:
             self._agent.run(question=question, callbacks=[activity_diagram_callback, sequence_diagram_callback])