From 1539f3dd62e054ee84d624b50d4a86da590c6820 Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Sat, 9 Nov 2024 16:31:47 +0300
Subject: [PATCH 1/7] added schemas and base classes for algos

---
 .../chatsky_llm_autoconfig/algorithms/base.py | 42 ++++++---
 .../algorithms/graph_generation.py            |  3 -
 .../algorithms/topic_graph_generation.py      | 87 +++++++++++++++++++
 .../autometrics/run_autometrics.py            |  3 -
 .../chatsky_llm_autoconfig/graph.py           | 30 +------
 .../chatsky_llm_autoconfig/pipeline.py        |  2 +-
 .../chatsky_llm_autoconfig/schemas.py         | 20 +++++
 7 files changed, 144 insertions(+), 43 deletions(-)
 delete mode 100644 dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/graph_generation.py
 create mode 100644 dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
 create mode 100644 dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
index b7185d4..03d7ed4 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
@@ -57,19 +57,41 @@ def invoke(self, dialogue: Dialogue, topic: str = "") -> Dialogue:
         raise NotImplementedError
 
 
+class GraphAugmentation(BaseAlgorithm):
+    """Graph generator that works only with topics."""
+
+    def invoke(self, topic: str, Gaph: BaseGraph) -> BaseGraph:
+        raise NotImplementedError
+
+    async def ainvoke(self, topic: str, Gaph: BaseGraph) -> BaseGraph:
+        raise NotImplementedError
+
+
+class TopicGraphGenerator(BaseAlgorithm):
+    """Graph generator that works only with topics."""
+
+    def invoke(self, topic: str) -> BaseGraph:
+        raise NotImplementedError
+
+    async def ainvoke(self, topic: str) -> BaseGraph:
+        raise NotImplementedError
+
+
 class GraphGenerator(BaseAlgorithm):
-    """
-    Base class for generating Graph objects.
+    """Graph generator that works only with topics."""
 
-    This class is used to create a Graph based on a Dialogue, a specified topic, or an existing Graph.
+    def invoke(self, dialogue: Dialogue) -> BaseGraph:
+        raise NotImplementedError
 
-    :param dialogue: The Dialogue object used for generating the Graph.
-    :param graph: An existing Graph object to base the generation on (optional).
-    :param topic: The topic to guide the Graph generation process (optional).
-    """
+    async def ainvoke(self, dialogue: Dialogue) -> BaseGraph:
+        raise NotImplementedError
 
-    def __init__(self):
-        super().__init__()
 
-    def invoke(self, dialogue: Dialogue = None, graph: BaseGraph = None, topic: str = "") -> BaseGraph:
+class GraphExtender(BaseAlgorithm):
+    """Graph generator that works only with topics."""
+
+    def invoke(self, dialogue: Dialogue, graph: BaseGraph) -> BaseGraph:
+        raise NotImplementedError
+
+    async def ainvoke(self, dialogue: Dialogue, graph: BaseGraph) -> BaseGraph:
         raise NotImplementedError
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/graph_generation.py
deleted file mode 100644
index 09efa8d..0000000
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/graph_generation.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from chatsky_llm_autoconfig.algorithms.base import GraphGenerator
-from chatsky_llm_autoconfig.graph import BaseGraph
-from chatsky_llm_autoconfig.dialogue import Dialogue
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
new file mode 100644
index 0000000..52e1ddb
--- /dev/null
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -0,0 +1,87 @@
+from chatsky_llm_autoconfig.algorithms.base import TopicGraphGenerator
+from chatsky_llm_autoconfig.schemas import DialogueGraph
+from langchain_openai import ChatOpenAI
+
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+
+from chatsky_llm_autoconfig.graph import BaseGraph, Graph
+
+from langchain.prompts import PromptTemplate
+import os
+
+cycle_graph_generation_prompt = PromptTemplate.from_template(
+    """
+    Create a cyclic dialogue graph where the conversation MUST return to an existing node.
+
+    **CRITICAL: Response Specificity**
+    Responses must acknowledge and build upon what the user has already specified:
+
+    INCORRECT flow:
+    - User: "I'd like to order a coffee"
+    - Staff: "What would you like to order?" (TOO GENERAL - ignores that they specified coffee)
+
+    CORRECT flow:
+    - User: "I'd like to order a coffee"
+    - Staff: "What kind of coffee would you like?" (GOOD - acknowledges they want coffee)
+
+    Example of a CORRECT cyclic graph for a coffee shop:
+    "edges": [
+        {{ "source": 1, "target": 2, "utterances": ["Hi, I'd like to order a coffee"] }},
+        {{ "source": 2, "target": 3, "utterances": ["A large latte please"] }},
+        {{ "source": 3, "target": 4, "utterances": ["Yes, that's correct"] }},
+        {{ "source": 4, "target": 5, "utterances": ["Here's my payment"] }},
+        {{ "source": 5, "target": 2, "utterances": ["I'd like to order another coffee"] }}
+    ],
+    "nodes": [
+        {{ "id": 1, "label": "welcome", "is_start": true, "utterances": ["Welcome! How can I help you today?"] }},
+        {{ "id": 2, "label": "ask_coffee_type", "is_start": false, "utterances": ["What kind of coffee would you like?"] }},
+        {{ "id": 3, "label": "confirm", "is_start": false, "utterances": ["That's a large latte. Is this correct?"] }},
+        {{ "id": 4, "label": "payment", "is_start": false, "utterances": ["Great! That'll be $5. Please proceed with payment."] }},
+        {{ "id": 5, "label": "completed", "is_start": false, "utterances": ["Thank you! Would you like another coffee?"] }}
+    ]
+
+    **Rules:**
+    1) Responses must acknowledge what the user has already specified
+    2) The final node MUST connect back to an existing node
+    3) Each node must have clear purpose
+    4) Return ONLY the JSON without commentary
+    5) Graph must be cyclic - no dead ends
+    6) All edges must connect to existing nodes
+    7) The cycle point should make logical sense
+
+    **Your task is to create a cyclic dialogue graph about the following topic:** {topic}.
+    """
+)
+
+
+class CycleGraphGenerator(TopicGraphGenerator):
+    """Generator specifically for topic-based cyclic graphs"""
+
+    def __init__(self):
+        super().__init__()
+
+    def invoke(self, topic: str) -> BaseGraph:
+        """
+        Generate a cyclic dialogue graph based on the topic input.
+
+        :param input_data: TopicInput containing the topic
+        :return: Generated Graph object with cyclic structure
+        """
+
+        prompt_template = cycle_graph_generation_prompt
+        parser = JsonOutputParser(pydantic_object=DialogueGraph)
+        model = ChatOpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
+        chain = prompt_template | model | parser
+
+        generated_graph = chain.invoke(topic)
+
+        return Graph(generated_graph)
+
+    async def ainvoke(self, *args, **kwargs):
+        pass
+
+
+cyc = CycleGraphGenerator()
+
+cyc.invoke(topic="cofee")
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/autometrics/run_autometrics.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/autometrics/run_autometrics.py
index 0c1ad88..d77417e 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/autometrics/run_autometrics.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/autometrics/run_autometrics.py
@@ -1,7 +1,4 @@
 from chatsky_llm_autoconfig.autometrics.registry import AlgorithmRegistry
-import chatsky_llm_autoconfig.algorithms.dialogue_generation
-import chatsky_llm_autoconfig.algorithms.dialogue_augmentation
-import chatsky_llm_autoconfig.algorithms.graph_generation
 
 import json
 from chatsky_llm_autoconfig.graph import Graph, BaseGraph
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/graph.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/graph.py
index d60be49..5cfdb00 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/graph.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/graph.py
@@ -1,6 +1,6 @@
 import networkx as nx
-from pydantic import Field, BaseModel
-from typing import Optional, Any, List
+from pydantic import BaseModel
+from typing import Optional, Any
 import matplotlib.pyplot as plt
 import abc
 import logging
@@ -28,7 +28,8 @@ def visualise(self, *args, **kwargs):
 class Graph(BaseGraph):
 
     def __init__(self, graph_dict: dict, **kwargs: Any):
-        super().__init__(graph_dict=graph_dict, **kwargs)  # Pass graph_dict to the parent class
+        # Pass graph_dict to the parent class
+        super().__init__(graph_dict=graph_dict, **kwargs)
         self.load_graph()
 
     def load_graph(self):
@@ -71,26 +72,3 @@ def visualise(self, *args, **kwargs):
         plt.title(__name__)
         plt.axis("off")
         plt.show()
-
-
-"""
-Pydantic models for Langchain structured output
-"""
-
-
-class Edge(BaseModel):
-    source: int = Field(description="ID of the source node")
-    target: int = Field(description="ID of the target node")
-    utterances: str = Field(description="User's utterance that triggers this transition")
-
-
-class Node(BaseModel):
-    id: int = Field(description="Unique identifier for the node")
-    label: str = Field(description="Label describing the node's purpose")
-    is_start: bool = Field(description="Whether this is the starting node")
-    utterances: List[str] = Field(description="Possible assistant responses at this node")
-
-
-class DialogueGraph(BaseModel):
-    edges: List[Edge] = Field(description="List of transitions between nodes")
-    nodes: List[Node] = Field(description="List of nodes representing assistant states")
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/pipeline.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/pipeline.py
index 3058669..8512698 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/pipeline.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/pipeline.py
@@ -1,6 +1,6 @@
+from chatsky_llm_autoconfig.algorithms.base import DialogAugmentation, DialogueGenerator, GraphGenerator
 from pydantic import BaseModel
 from typing import Union
-from chatsky_llm_autoconfig.algorithms import DialogueGenerator, DialogAugmentation, GraphGenerator
 
 
 class Pipeline(BaseModel):
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
new file mode 100644
index 0000000..badc470
--- /dev/null
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
@@ -0,0 +1,20 @@
+from typing import List
+from pydantic import BaseModel, Field
+
+
+class Edge(BaseModel):
+    source: int = Field(description="ID of the source node")
+    target: int = Field(description="ID of the target node")
+    utterances: List[str] = Field(description="User's utterances that trigger this transition")
+
+
+class Node(BaseModel):
+    id: int = Field(description="Unique identifier for the node")
+    label: str = Field(description="Label describing the node's purpose")
+    is_start: bool = Field(description="Whether this is the starting node")
+    utterances: List[str] = Field(description="Possible assistant responses at this node")
+
+
+class DialogueGraph(BaseModel):
+    edges: List[Edge] = Field(description="List of transitions between nodes")
+    nodes: List[Node] = Field(description="List of nodes representing assistant states")

From fde1e4c69eb133925250968f7dc0d51187c7d9e0 Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Sat, 9 Nov 2024 17:04:57 +0300
Subject: [PATCH 2/7] small fixes

---
 .../chatsky_llm_autoconfig/algorithms/base.py |  3 +-
 .../algorithms/dialogue_augmentation.py       | 73 ++++++++++++++++++-
 .../algorithms/dialogue_generation.py         | 21 ++++--
 .../algorithms/topic_graph_generation.py      | 14 ++--
 .../chatsky_llm_autoconfig/schemas.py         | 14 +++-
 test_classes.py                               |  2 +-
 6 files changed, 105 insertions(+), 22 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
index 03d7ed4..bb28ee0 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/base.py
@@ -1,3 +1,4 @@
+from typing import List
 from pydantic import BaseModel
 import abc
 from chatsky_llm_autoconfig.graph import BaseGraph
@@ -35,7 +36,7 @@ class DialogueGenerator(BaseAlgorithm):
     def __init__(self):
         super().__init__()
 
-    def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = 0, topic: str = "") -> Dialogue:
+    def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = 0, topic: str = "") -> List[Dialogue]:
         raise NotImplementedError
 
 
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
index afeb5e6..d00f8fe 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
@@ -1,3 +1,70 @@
-from chatsky_llm_autoconfig.algorithms.base import DialogAugmentation
-from chatsky_llm_autoconfig.graph import BaseGraph
-from chatsky_llm_autoconfig.dialogue import Dialogue
+
+
+from chatsky_llm_autoconfig.schemas import DialogueMessage
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from typing import List
+
+from langchain_openai import ChatOpenAI
+from pydantic import BaseModel
+
+augmentation_prompt = PromptTemplate.from_template("""
+You are tasked with augmenting a dialogue by adding variations to existing utterances while maintaining the original dialogue flow and intent.
+
+THEME: {topic}
+
+INPUT DIALOGUE:
+{dialogue}
+
+INSTRUCTIONS:
+1. For each message in the dialogue:
+   - Keep the same structure (participant, source, target if present)
+   - Create variation of the 'text' field that:
+     * Express the same meaning/intent
+     * Use different wording and phrasing
+     * Match the given theme
+     * Sound natural and conversational
+
+2. The output must be a list of dictionaries, where each dictionary has:
+   - 'text': string
+   - 'participant': either 'user' or 'assistant'
+   
+3. Ensure all utterance variations:
+   - Are appropriate for the theme
+   - Maintain consistency in tone and style
+   - Make sense in the conversation flow
+
+Return ONLY a valid JSON array containing the augmented dialogue messages. Each message should be in this exact format:
+For assistant messages: {{"text": "utterance text", "participant": "assistant"}}
+For user messages: {{"text": "utterance text", "participant": "user"}}
+
+Example format:
+[
+    {{"text": "How may I assist you today?", "participant": "assistant"}},
+    {{"text": "I need help with a package", "participant": "user"}},
+    {{"text": "What kind of package is it?", "participant": "assistant"}}
+]
+""")
+
+
+class DialogueSequence(BaseModel):
+    result: List[DialogueMessage]
+
+
+parser = JsonOutputParser(pydantic_object=DialogueSequence)
+
+# Usage
+model = ChatOpenAI(
+    model="gpt-4o-mini",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    base_url=os.getenv("OPENAI_BASE_URL"),
+    temperature=0.7
+)
+
+chain = augmentation_prompt | model | parser
+
+
+result = chain.invoke({
+    "topic": topic,
+    "dialogue": dialogue
+})
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
index 8328a56..cf5168c 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
@@ -14,7 +14,8 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
             end_node = list(nx_graph.nodes)[-1]
 
         all_dialogues = []
-        start_nodes = [n for n, attr in nx_graph.nodes(data=True) if attr.get("is_start", n == start_node)]
+        start_nodes = [n for n, attr in nx_graph.nodes(
+            data=True) if attr.get("is_start", n == start_node)]
 
         for start in start_nodes:
             # Stack contains: (current_node, path, visited_edges)
@@ -24,20 +25,24 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
                 current_node, path, visited_edges = stack.pop()
 
                 # Add assistant utterance
-                current_utterance = random.choice(nx_graph.nodes[current_node]["utterances"])
-                path.append({"text": current_utterance, "participant": "assistant"})
+                current_utterance = random.choice(
+                    nx_graph.nodes[current_node]["utterances"])
+                path.append({"text": current_utterance,
+                            "participant": "assistant"})
 
                 if current_node == end_node:
                     # Check if the last node has edges and add the last edge utterances
                     edges = list(nx_graph.edges(current_node, data=True))
                     if edges:
-                        last_edge_data = edges[-1][2]  # Get the last edge's data
+                        # Get the last edge's data
+                        last_edge_data = edges[-1][2]
                         last_edge_utterance = (
                             random.choice(last_edge_data["utterances"])
                             if isinstance(last_edge_data["utterances"], list)
                             else last_edge_data["utterances"]
                         )
-                        path.append({"text": last_edge_utterance, "participant": "user"})
+                        path.append(
+                            {"text": last_edge_utterance, "participant": "user"})
 
                     all_dialogues.append(Dialogue(dialogue=path.copy()))
                     path.pop()
@@ -55,11 +60,13 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
                     # if topic and edge_data.get("theme") != topic:
                     #     continue
 
-                    edge_utterance = random.choice(edge_data["utterances"]) if isinstance(edge_data["utterances"], list) else edge_data["utterances"]
+                    edge_utterance = random.choice(edge_data["utterances"]) if isinstance(
+                        edge_data["utterances"], list) else edge_data["utterances"]
 
                     # Create new path and visited_edges for this branch
                     new_path = path.copy()
-                    new_path.append({"text": edge_utterance, "participant": "user"})
+                    new_path.append(
+                        {"text": edge_utterance, "participant": "user"})
 
                     new_visited = visited_edges | {edge_key}
                     stack.append((target, new_path, new_visited))
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
index 52e1ddb..042d480 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -1,4 +1,5 @@
 from chatsky_llm_autoconfig.algorithms.base import TopicGraphGenerator
+from chatsky_llm_autoconfig.autometrics.registry import AlgorithmRegistry
 from chatsky_llm_autoconfig.schemas import DialogueGraph
 from langchain_openai import ChatOpenAI
 
@@ -10,6 +11,8 @@
 from langchain.prompts import PromptTemplate
 import os
 
+from pydantic import SecretStr
+
 cycle_graph_generation_prompt = PromptTemplate.from_template(
     """
     Create a cyclic dialogue graph where the conversation MUST return to an existing node.
@@ -55,6 +58,7 @@
 )
 
 
+@AlgorithmRegistry.register(input_type=str, output_type=Graph)
 class CycleGraphGenerator(TopicGraphGenerator):
     """Generator specifically for topic-based cyclic graphs"""
 
@@ -71,17 +75,13 @@ def invoke(self, topic: str) -> BaseGraph:
 
         prompt_template = cycle_graph_generation_prompt
         parser = JsonOutputParser(pydantic_object=DialogueGraph)
-        model = ChatOpenAI(model="gpt-4o", api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
+        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY") or ''),
+                           base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
         chain = prompt_template | model | parser
 
-        generated_graph = chain.invoke(topic)
+        generated_graph = chain.invoke({"topic": topic})
 
         return Graph(generated_graph)
 
     async def ainvoke(self, *args, **kwargs):
         pass
-
-
-cyc = CycleGraphGenerator()
-
-cyc.invoke(topic="cofee")
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
index badc470..5d2427b 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
@@ -5,16 +5,24 @@
 class Edge(BaseModel):
     source: int = Field(description="ID of the source node")
     target: int = Field(description="ID of the target node")
-    utterances: List[str] = Field(description="User's utterances that trigger this transition")
+    utterances: List[str] = Field(
+        description="User's utterances that trigger this transition")
 
 
 class Node(BaseModel):
     id: int = Field(description="Unique identifier for the node")
     label: str = Field(description="Label describing the node's purpose")
     is_start: bool = Field(description="Whether this is the starting node")
-    utterances: List[str] = Field(description="Possible assistant responses at this node")
+    utterances: List[str] = Field(
+        description="Possible assistant responses at this node")
 
 
 class DialogueGraph(BaseModel):
     edges: List[Edge] = Field(description="List of transitions between nodes")
-    nodes: List[Node] = Field(description="List of nodes representing assistant states")
+    nodes: List[Node] = Field(
+        description="List of nodes representing assistant states")
+
+
+class DialogueMessage(BaseModel):
+    text: str
+    participant: str
diff --git a/test_classes.py b/test_classes.py
index 0958027..b2769f3 100644
--- a/test_classes.py
+++ b/test_classes.py
@@ -1,5 +1,5 @@
+from chatsky_llm_autoconfig.algorithms.dialogue_generation import DialogueSampler
 from chatsky_llm_autoconfig.dialogue import Dialogue
-from chatsky_llm_autoconfig.sample_dialogue import DialogueSampler
 from chatsky_llm_autoconfig.graph import Graph
 from chatsky_llm_autoconfig.metrics.automatic_metrics import all_paths_sampled
 import json

From 37420205fe68d16d2e832f1251dddb207031f58b Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Sat, 9 Nov 2024 18:13:55 +0300
Subject: [PATCH 3/7] added dialogue augmentation

---
 .../algorithms/dialogue_augmentation.py       |  76 ++++++++---
 .../chatsky_llm_autoconfig/dialogue.py        | 129 +++++++++++++-----
 .../chatsky_llm_autoconfig/schemas.py         |   6 +
 3 files changed, 163 insertions(+), 48 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
index d00f8fe..3065ce7 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
@@ -1,5 +1,6 @@
 
 
+from chatsky_llm_autoconfig.dialogue import Dialogue
 from chatsky_llm_autoconfig.schemas import DialogueMessage
 from langchain.prompts import PromptTemplate
 from langchain_core.output_parsers import JsonOutputParser
@@ -7,6 +8,7 @@
 
 from langchain_openai import ChatOpenAI
 from pydantic import BaseModel
+import os
 
 augmentation_prompt = PromptTemplate.from_template("""
 You are tasked with augmenting a dialogue by adding variations to existing utterances while maintaining the original dialogue flow and intent.
@@ -51,20 +53,60 @@ class DialogueSequence(BaseModel):
     result: List[DialogueMessage]
 
 
-parser = JsonOutputParser(pydantic_object=DialogueSequence)
-
-# Usage
-model = ChatOpenAI(
-    model="gpt-4o-mini",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    base_url=os.getenv("OPENAI_BASE_URL"),
-    temperature=0.7
-)
-
-chain = augmentation_prompt | model | parser
-
-
-result = chain.invoke({
-    "topic": topic,
-    "dialogue": dialogue
-})
+class DialogAugmentation(BaseModel):
+    """Base class for augmenting Dialogues."""
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.parser = JsonOutputParser(pydantic_object=DialogueSequence)
+        self.model = ChatOpenAI(
+            model="gpt-4o-mini",
+            api_key=os.getenv("OPENAI_API_KEY"),
+            base_url=os.getenv("OPENAI_BASE_URL"),
+            temperature=0.7
+        )
+        self.chain = augmentation_prompt | self.model | self.parser
+
+    def invoke(self, *, dialogue: Dialogue, topic: str = "") -> Dialogue:
+        """
+        Augment the input dialogue with variations.
+
+        Args:
+            dialogue: The input Dialogue object to augment
+            topic: Optional topic to guide the augmentation
+
+        Returns:
+            Dialogue: Augmented dialogue object
+        """
+        # Convert dialogue to string format for prompt
+        # Предполагая, что у Dialogue есть str представление
+        dialogue_str = str(dialogue)
+
+        # Get augmented messages
+        result = self.chain.invoke({
+            "topic": topic,
+            "dialogue": dialogue_str
+        })
+
+        # Create new Dialogue object with augmented messages
+        return Dialogue(messages=result, topic=topic)
+
+    async def ainvoke(self, *, dialogue: Dialogue, topic: str = "") -> Dialogue:
+        """
+        Async version of dialogue augmentation.
+
+        Args:
+            dialogue: The input Dialogue object to augment
+            topic: Optional topic to guide the augmentation
+
+        Returns:
+            Dialogue: Augmented dialogue object
+        """
+        dialogue_str = str(dialogue)
+
+        result = await self.chain.ainvoke({
+            "topic": topic,
+            "dialogue": dialogue_str
+        })
+
+        return Dialogue(messages=result, topic=topic)
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
index 2eb1939..5c63a5a 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
@@ -1,42 +1,109 @@
-from pydantic import BaseModel
-from typing import Optional, Any, Union
-import json
-import matplotlib.pyplot as plt
-import abc
-import logging
+from typing import List, Union, Dict
+from chatsky_llm_autoconfig.schemas import DialogueMessage
+from pydantic import BaseModel, Field, ConfigDict
 
-logger = logging.getLogger(__name__)
 
+class Dialogue(BaseModel):
+    """Represents a complete dialogue consisting of multiple messages.
 
-class Dialogue(BaseModel, abc.ABC):
+    The class provides methods for creating dialogues from different formats
+    and converting dialogues to various representations.
     """
-    Base class for all of the dialogues.
+    messages: List[DialogueMessage] = Field(default_factory=list)
+    topic: str = ""
 
-    Attributes
-    ----------
-    dialogue : list of dict
-        A list containing dialogue entries, where each entry is a
-        dictionary with "text" and "participant" keys.
-        It is initialized with an empty dictionary.
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        frozen=False,  # Dialogue needs to be mutable to append messages
+    )
 
-    Examples
-    --------
-        Dialogue([{"text": "How can I help?", "participant": "assistant"}, {"text": "I need to make an order", "participant": "user"}])
+    @classmethod
+    def from_string(cls, string: str) -> "Dialogue":
+        """Creates a Dialogue from a tab-separated string format.
 
-    """
+        Args:
+            string: Tab-separated string with format: "participant\ttext\n"
 
-    dialogue: list[dict] = []
-    topic: str = ""
+        Returns:
+            Dialogue object with parsed messages
+        """
+        messages: List[DialogueMessage] = [
+            DialogueMessage(participant=line.split(
+                "\t")[0], text=line.split("\t")[1])
+            for line in string.strip().split("\n")
+        ]
+        return cls(messages=messages)
+
+    @classmethod
+    def from_list(cls, dialogue_list: List[Dict[str, str]]) -> "Dialogue":
+        """Creates a Dialogue from a list of message dictionaries.
+
+        Args:
+            dialogue_list: List of dicts with 'text' and 'participant' keys
+
+        Returns:
+            Dialogue object with parsed messages
+        """
+        messages = [DialogueMessage(**msg) for msg in dialogue_list]
+        return cls(messages=messages)
+
+    def to_list(self) -> List[Dict[str, str]]:
+        """Converts Dialogue to a list of message dictionaries."""
+        return [msg.model_dump() for msg in self.messages]
+
+    def __str__(self) -> str:
+        """Returns a readable string representation of the dialogue."""
+        return "\n".join(
+            f"{msg.participant}: {msg.text}"
+            for msg in self.messages
+        ).strip()
+
+    def append(self, text: str, participant: str) -> None:
+        """Adds a new message to the dialogue.
+
+        Args:
+            text: Content of the message
+            participant: Sender of the message
+        """
+        self.messages.append(DialogueMessage(
+            text=text, participant=participant))
+
+    def extend(self, messages: List[Union[DialogueMessage, Dict[str, str]]]) -> None:
+        """Adds multiple messages to the dialogue.
+
+        Args:
+            messages: List of DialogueMessage objects or dicts to add
+        """
+        new_messages = [
+            msg if isinstance(msg, DialogueMessage) else DialogueMessage(**msg)
+            for msg in messages
+        ]
+        self.messages.extend(new_messages)
+
+
+# Type-safe usage examples
+if __name__ == "__main__":
+    # Create from list of dicts
+    dialogue1 = Dialogue(messages=[
+        DialogueMessage(text="How can I help?", participant="assistant"),
+        DialogueMessage(text="I need coffee", participant="user")
+    ])
 
-    class Config:
-        arbitrary_types_allowed = True
+    # Create using from_list
+    dialogue2 = Dialogue.from_list([
+        {"text": "How can I help?", "participant": "assistant"},
+        {"text": "I need coffee", "participant": "user"}
+    ])
 
-    def parse_string(self, string):
-        self.dialogue = []
-        for utt in string.split("\n"):
-            participant, phrase = utt.split("\t")
-            self.dialogue.append({"text": phrase, "participant": participant})
+    # Create from string
+    dialogue3 = Dialogue.from_string("""
+        assistant\tHow can I help?
+        user\tI need coffee
+""".strip())
 
-    def __str__(self):
-        readable = "\n".join([utt["participant"] + ": " + utt["text"] for utt in self.dialogue])
-        return readable.strip()
+    # Append and extend
+    dialogue1.append("What kind of coffee?", "assistant")
+    dialogue1.extend([
+        {"text": "Espresso please", "participant": "user"},
+        DialogueMessage(text="Coming right up!", participant="assistant")
+    ])
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
index 5d2427b..f9ad79f 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
@@ -24,5 +24,11 @@ class DialogueGraph(BaseModel):
 
 
 class DialogueMessage(BaseModel):
+    """Represents a single message in a dialogue.
+
+    Attributes:
+        text: The content of the message
+        participant: The sender of the message (e.g. "user" or "assistant")
+    """
     text: str
     participant: str

From 6d896807896458a9c87211a205addd5b20c2d464 Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Sat, 9 Nov 2024 18:16:14 +0300
Subject: [PATCH 4/7] format

---
 .../algorithms/dialogue_augmentation.py       | 25 ++++-------
 .../algorithms/dialogue_generation.py         | 18 +++-----
 .../algorithms/topic_graph_generation.py      |  3 +-
 .../chatsky_llm_autoconfig/dialogue.py        | 41 +++++++------------
 .../chatsky_llm_autoconfig/schemas.py         | 10 ++---
 5 files changed, 32 insertions(+), 65 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
index 3065ce7..bd0fc9d 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_augmentation.py
@@ -1,5 +1,3 @@
-
-
 from chatsky_llm_autoconfig.dialogue import Dialogue
 from chatsky_llm_autoconfig.schemas import DialogueMessage
 from langchain.prompts import PromptTemplate
@@ -10,7 +8,8 @@
 from pydantic import BaseModel
 import os
 
-augmentation_prompt = PromptTemplate.from_template("""
+augmentation_prompt = PromptTemplate.from_template(
+    """
 You are tasked with augmenting a dialogue by adding variations to existing utterances while maintaining the original dialogue flow and intent.
 
 THEME: {topic}
@@ -46,7 +45,8 @@
     {{"text": "I need help with a package", "participant": "user"}},
     {{"text": "What kind of package is it?", "participant": "assistant"}}
 ]
-""")
+"""
+)
 
 
 class DialogueSequence(BaseModel):
@@ -59,12 +59,7 @@ class DialogAugmentation(BaseModel):
     def __init__(self, **data):
         super().__init__(**data)
         self.parser = JsonOutputParser(pydantic_object=DialogueSequence)
-        self.model = ChatOpenAI(
-            model="gpt-4o-mini",
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=os.getenv("OPENAI_BASE_URL"),
-            temperature=0.7
-        )
+        self.model = ChatOpenAI(model="gpt-4o-mini", api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0.7)
         self.chain = augmentation_prompt | self.model | self.parser
 
     def invoke(self, *, dialogue: Dialogue, topic: str = "") -> Dialogue:
@@ -83,10 +78,7 @@ def invoke(self, *, dialogue: Dialogue, topic: str = "") -> Dialogue:
         dialogue_str = str(dialogue)
 
         # Get augmented messages
-        result = self.chain.invoke({
-            "topic": topic,
-            "dialogue": dialogue_str
-        })
+        result = self.chain.invoke({"topic": topic, "dialogue": dialogue_str})
 
         # Create new Dialogue object with augmented messages
         return Dialogue(messages=result, topic=topic)
@@ -104,9 +96,6 @@ async def ainvoke(self, *, dialogue: Dialogue, topic: str = "") -> Dialogue:
         """
         dialogue_str = str(dialogue)
 
-        result = await self.chain.ainvoke({
-            "topic": topic,
-            "dialogue": dialogue_str
-        })
+        result = await self.chain.ainvoke({"topic": topic, "dialogue": dialogue_str})
 
         return Dialogue(messages=result, topic=topic)
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
index cf5168c..196c289 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/dialogue_generation.py
@@ -14,8 +14,7 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
             end_node = list(nx_graph.nodes)[-1]
 
         all_dialogues = []
-        start_nodes = [n for n, attr in nx_graph.nodes(
-            data=True) if attr.get("is_start", n == start_node)]
+        start_nodes = [n for n, attr in nx_graph.nodes(data=True) if attr.get("is_start", n == start_node)]
 
         for start in start_nodes:
             # Stack contains: (current_node, path, visited_edges)
@@ -25,10 +24,8 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
                 current_node, path, visited_edges = stack.pop()
 
                 # Add assistant utterance
-                current_utterance = random.choice(
-                    nx_graph.nodes[current_node]["utterances"])
-                path.append({"text": current_utterance,
-                            "participant": "assistant"})
+                current_utterance = random.choice(nx_graph.nodes[current_node]["utterances"])
+                path.append({"text": current_utterance, "participant": "assistant"})
 
                 if current_node == end_node:
                     # Check if the last node has edges and add the last edge utterances
@@ -41,8 +38,7 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
                             if isinstance(last_edge_data["utterances"], list)
                             else last_edge_data["utterances"]
                         )
-                        path.append(
-                            {"text": last_edge_utterance, "participant": "user"})
+                        path.append({"text": last_edge_utterance, "participant": "user"})
 
                     all_dialogues.append(Dialogue(dialogue=path.copy()))
                     path.pop()
@@ -60,13 +56,11 @@ def invoke(self, graph: BaseGraph, start_node: int = 1, end_node: int = -1, topi
                     # if topic and edge_data.get("theme") != topic:
                     #     continue
 
-                    edge_utterance = random.choice(edge_data["utterances"]) if isinstance(
-                        edge_data["utterances"], list) else edge_data["utterances"]
+                    edge_utterance = random.choice(edge_data["utterances"]) if isinstance(edge_data["utterances"], list) else edge_data["utterances"]
 
                     # Create new path and visited_edges for this branch
                     new_path = path.copy()
-                    new_path.append(
-                        {"text": edge_utterance, "participant": "user"})
+                    new_path.append({"text": edge_utterance, "participant": "user"})
 
                     new_visited = visited_edges | {edge_key}
                     stack.append((target, new_path, new_visited))
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
index 042d480..93f2bbe 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -75,8 +75,7 @@ def invoke(self, topic: str) -> BaseGraph:
 
         prompt_template = cycle_graph_generation_prompt
         parser = JsonOutputParser(pydantic_object=DialogueGraph)
-        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY") or ''),
-                           base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
+        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY") or ""), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
         chain = prompt_template | model | parser
 
         generated_graph = chain.invoke({"topic": topic})
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
index 5c63a5a..1d3801a 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/dialogue.py
@@ -9,6 +9,7 @@ class Dialogue(BaseModel):
     The class provides methods for creating dialogues from different formats
     and converting dialogues to various representations.
     """
+
     messages: List[DialogueMessage] = Field(default_factory=list)
     topic: str = ""
 
@@ -28,9 +29,7 @@ def from_string(cls, string: str) -> "Dialogue":
             Dialogue object with parsed messages
         """
         messages: List[DialogueMessage] = [
-            DialogueMessage(participant=line.split(
-                "\t")[0], text=line.split("\t")[1])
-            for line in string.strip().split("\n")
+            DialogueMessage(participant=line.split("\t")[0], text=line.split("\t")[1]) for line in string.strip().split("\n")
         ]
         return cls(messages=messages)
 
@@ -53,10 +52,7 @@ def to_list(self) -> List[Dict[str, str]]:
 
     def __str__(self) -> str:
         """Returns a readable string representation of the dialogue."""
-        return "\n".join(
-            f"{msg.participant}: {msg.text}"
-            for msg in self.messages
-        ).strip()
+        return "\n".join(f"{msg.participant}: {msg.text}" for msg in self.messages).strip()
 
     def append(self, text: str, participant: str) -> None:
         """Adds a new message to the dialogue.
@@ -65,8 +61,7 @@ def append(self, text: str, participant: str) -> None:
             text: Content of the message
             participant: Sender of the message
         """
-        self.messages.append(DialogueMessage(
-            text=text, participant=participant))
+        self.messages.append(DialogueMessage(text=text, participant=participant))
 
     def extend(self, messages: List[Union[DialogueMessage, Dict[str, str]]]) -> None:
         """Adds multiple messages to the dialogue.
@@ -74,36 +69,28 @@ def extend(self, messages: List[Union[DialogueMessage, Dict[str, str]]]) -> None
         Args:
             messages: List of DialogueMessage objects or dicts to add
         """
-        new_messages = [
-            msg if isinstance(msg, DialogueMessage) else DialogueMessage(**msg)
-            for msg in messages
-        ]
+        new_messages = [msg if isinstance(msg, DialogueMessage) else DialogueMessage(**msg) for msg in messages]
         self.messages.extend(new_messages)
 
 
 # Type-safe usage examples
 if __name__ == "__main__":
     # Create from list of dicts
-    dialogue1 = Dialogue(messages=[
-        DialogueMessage(text="How can I help?", participant="assistant"),
-        DialogueMessage(text="I need coffee", participant="user")
-    ])
+    dialogue1 = Dialogue(
+        messages=[DialogueMessage(text="How can I help?", participant="assistant"), DialogueMessage(text="I need coffee", participant="user")]
+    )
 
     # Create using from_list
-    dialogue2 = Dialogue.from_list([
-        {"text": "How can I help?", "participant": "assistant"},
-        {"text": "I need coffee", "participant": "user"}
-    ])
+    dialogue2 = Dialogue.from_list([{"text": "How can I help?", "participant": "assistant"}, {"text": "I need coffee", "participant": "user"}])
 
     # Create from string
-    dialogue3 = Dialogue.from_string("""
+    dialogue3 = Dialogue.from_string(
+        """
         assistant\tHow can I help?
         user\tI need coffee
-""".strip())
+""".strip()
+    )
 
     # Append and extend
     dialogue1.append("What kind of coffee?", "assistant")
-    dialogue1.extend([
-        {"text": "Espresso please", "participant": "user"},
-        DialogueMessage(text="Coming right up!", participant="assistant")
-    ])
+    dialogue1.extend([{"text": "Espresso please", "participant": "user"}, DialogueMessage(text="Coming right up!", participant="assistant")])
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
index f9ad79f..0fe1437 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/schemas.py
@@ -5,22 +5,19 @@
 class Edge(BaseModel):
     source: int = Field(description="ID of the source node")
     target: int = Field(description="ID of the target node")
-    utterances: List[str] = Field(
-        description="User's utterances that trigger this transition")
+    utterances: List[str] = Field(description="User's utterances that trigger this transition")
 
 
 class Node(BaseModel):
     id: int = Field(description="Unique identifier for the node")
     label: str = Field(description="Label describing the node's purpose")
     is_start: bool = Field(description="Whether this is the starting node")
-    utterances: List[str] = Field(
-        description="Possible assistant responses at this node")
+    utterances: List[str] = Field(description="Possible assistant responses at this node")
 
 
 class DialogueGraph(BaseModel):
     edges: List[Edge] = Field(description="List of transitions between nodes")
-    nodes: List[Node] = Field(
-        description="List of nodes representing assistant states")
+    nodes: List[Node] = Field(description="List of nodes representing assistant states")
 
 
 class DialogueMessage(BaseModel):
@@ -30,5 +27,6 @@ class DialogueMessage(BaseModel):
         text: The content of the message
         participant: The sender of the message (e.g. "user" or "assistant")
     """
+
     text: str
     participant: str

From 2e57a3cac2a38eaa38d730b8ca77672652f6b477 Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Sat, 9 Nov 2024 18:19:11 +0300
Subject: [PATCH 5/7] format

---
 .../algorithms/topic_graph_generation.py              |  2 --
 .../metrics/automatic_metrics.py                      | 11 ++++++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
index 93f2bbe..412f1dd 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -7,8 +7,6 @@
 from langchain_core.output_parsers import JsonOutputParser
 
 from chatsky_llm_autoconfig.graph import BaseGraph, Graph
-
-from langchain.prompts import PromptTemplate
 import os
 
 from pydantic import SecretStr
diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/metrics/automatic_metrics.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/metrics/automatic_metrics.py
index 57db881..c45cdf2 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/metrics/automatic_metrics.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/metrics/automatic_metrics.py
@@ -78,10 +78,12 @@ def triplet_match(G1: BaseGraph, G2: BaseGraph, change_to_original_ids=False):
                     if node1_trg_nx == node2_trg_nx:
                         node_mapping[node1_trg + 1] = node2_trg + 1
                     print(
-                        f'The nodes of edges {edges1[i]} and {edges2[j]} has something in common, but not complete match: Sources: {node1_src_nx["utterances"]}, {node2_src_nx["utterances"]}'
+                        f'The nodes of edges {edges1[i]} and {edges2[j]} has something in common, but not complete match: Sources: {
+                            node1_src_nx["utterances"]}, {node2_src_nx["utterances"]}'
                     )
                     print(
-                        f'The nodes of edges {edges1[i]} and {edges2[j]} has something in common, but not complete match: Targets: {node1_trg_nx["utterances"]}, {node2_trg_nx["utterances"]}'
+                        f'The nodes of edges {edges1[i]} and {edges2[j]} has something in common, but not complete match: Targets: {
+                            node1_trg_nx["utterances"]}, {node2_trg_nx["utterances"]}'
                     )
 
     if G1.node_mapping != {} and change_to_original_ids:
@@ -101,7 +103,10 @@ def triplet_match(G1: BaseGraph, G2: BaseGraph, change_to_original_ids=False):
 
         for edge1, edge2 in edge_mapping.items():
             src1, trg1 = edge1.split("->")
-            new_edge_mapping[f"{inverse_mapping[int(src1)]}->{inverse_mapping[int(trg1)]}"] = edge2
+            new_edge_mapping[
+                f"{
+                inverse_mapping[int(src1)]}->{inverse_mapping[int(trg1)]}"
+            ] = edge2
         return new_node_mapping, new_edge_mapping
 
     return node_mapping, edge_mapping

From 75ec0154b00455e99311aad8eb11b05cb0b86531 Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Mon, 11 Nov 2024 15:09:29 +0300
Subject: [PATCH 6/7] small fixes

---
 .../algorithms/topic_graph_generation.py      | 32 ++++++++++---------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
index 412f1dd..12850fa 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -11,8 +11,15 @@
 
 from pydantic import SecretStr
 
-cycle_graph_generation_prompt = PromptTemplate.from_template(
-    """
+
+@AlgorithmRegistry.register(input_type=str, output_type=Graph)
+class CycleGraphGenerator(TopicGraphGenerator):
+    """Generator specifically for topic-based cyclic graphs"""
+
+    def __init__(self):
+        super().__init__()
+        self.cycle_graph_generation_prompt = PromptTemplate.from_template(
+            """
     Create a cyclic dialogue graph where the conversation MUST return to an existing node.
 
     **CRITICAL: Response Specificity**
@@ -53,15 +60,7 @@
 
     **Your task is to create a cyclic dialogue graph about the following topic:** {topic}.
     """
-)
-
-
-@AlgorithmRegistry.register(input_type=str, output_type=Graph)
-class CycleGraphGenerator(TopicGraphGenerator):
-    """Generator specifically for topic-based cyclic graphs"""
-
-    def __init__(self):
-        super().__init__()
+        )
 
     def invoke(self, topic: str) -> BaseGraph:
         """
@@ -70,11 +69,10 @@ def invoke(self, topic: str) -> BaseGraph:
         :param input_data: TopicInput containing the topic
         :return: Generated Graph object with cyclic structure
         """
-
-        prompt_template = cycle_graph_generation_prompt
         parser = JsonOutputParser(pydantic_object=DialogueGraph)
-        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY") or ""), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
-        chain = prompt_template | model | parser
+        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY")
+                           or ""), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
+        chain = self.cycle_graph_generation_prompt | model | parser
 
         generated_graph = chain.invoke({"topic": topic})
 
@@ -82,3 +80,7 @@ def invoke(self, topic: str) -> BaseGraph:
 
     async def ainvoke(self, *args, **kwargs):
         pass
+
+
+if __name__ == "__main__":
+    cycle_graph_generator = CycleGraphGenerator()

From 4587993901de25f5700777b94dbc85afeefc90da Mon Sep 17 00:00:00 2001
From: Dmitrii Martynov <dmitrii.martynov@deptagency.com>
Date: Mon, 11 Nov 2024 15:10:45 +0300
Subject: [PATCH 7/7] small fixes

---
 .../algorithms/topic_graph_generation.py                       | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
index 12850fa..9075714 100644
--- a/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
+++ b/dev_packages/chatsky_llm_autoconfig/chatsky_llm_autoconfig/algorithms/topic_graph_generation.py
@@ -70,8 +70,7 @@ def invoke(self, topic: str) -> BaseGraph:
         :return: Generated Graph object with cyclic structure
         """
         parser = JsonOutputParser(pydantic_object=DialogueGraph)
-        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY")
-                           or ""), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
+        model = ChatOpenAI(model="gpt-4o", api_key=SecretStr(os.getenv("OPENAI_API_KEY") or ""), base_url=os.getenv("OPENAI_BASE_URL"), temperature=0)
         chain = self.cycle_graph_generation_prompt | model | parser
 
         generated_graph = chain.invoke({"topic": topic})