fix: migrating out conversations from Kura

ivanleomk · Jan 19, 2025 · dcb66c7 · dcb66c7
1 parent 8cbf4a1
commit dcb66c7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 5 deletions.
diff --git a/kura/cli/server.py b/kura/cli/server.py
@@ -71,9 +71,8 @@ async def analyse_conversations(conversation_data: ConversationData):
             checkpoint_dir=str(
                 Path(os.path.abspath(os.environ["KURA_CHECKPOINT_DIR"]))
             ),
-            conversations=conversations,
         )
-        await kura.cluster_conversations()
+        clusters = await kura.cluster_conversations(conversations)
 
     with open(clusters_file) as f:
         clusters_data = []

diff --git a/kura/kura.py b/kura/kura.py
@@ -25,7 +25,6 @@
 class Kura:
     def __init__(
         self,
-        conversations: list[Conversation] = [],
         embedding_model: BaseEmbeddingModel = OpenAIEmbeddingModel(),
         summarisation_model: BaseSummaryModel = SummaryModel(),
         cluster_model: BaseClusterModel = ClusterModel(),
@@ -43,7 +42,6 @@ def __init__(
         self.embedding_model = embedding_model
         self.embedding_model = embedding_model
         self.summarisation_model = summarisation_model
-        self.conversations = conversations
         self.max_clusters = max_clusters
         self.cluster_model = cluster_model
         self.meta_cluster_model = meta_cluster_model
@@ -69,7 +67,7 @@ def __init__(
             os.makedirs(self.checkpoint_dir)
 
     def load_checkpoint(
-        self, checkpoint_path: str, response_model: T
+        self, checkpoint_path: str, response_model: type[T]
     ) -> Union[list[T], None]:
         if not self.disable_checkpoints:
             if os.path.exists(checkpoint_path):