diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 017c23c02503e45e36b8580c0b5c6664d59f016f..7e3ad27128f3cdd09c40ce555b0882dd32d78618 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -54,4 +54,4 @@ jobs:
CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
with:
file: ./coverage.xml
- fail_ci_if_error: true
+ fail_ci_if_error: false
diff --git a/docs/examples/conversation-topic-splitter.ipynb b/docs/examples/conversation-topic-splitter.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..5ac130b0fcd8277f559606817a353e1b261f12f6
--- /dev/null
+++ b/docs/examples/conversation-topic-splitter.ipynb
@@ -0,0 +1,257 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Split Conversations by Topic"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Topics Splitters have been implemented in the code in `semantic-router/splitters`.\n",
+ "\n",
+ "These allow a set of utterances to be automatically grouped/clustered into (un-labelled) topics. \n",
+ "\n",
+ "Additionally, splitters have been integrated with `Conversation` objects allowing conversations to be progressively spit by topic as they evolve. This is beneficial to routing, as earlier messages in a conversation topic might provide useful context when determining routes. By using all utterances in the latest conversation this additional context allows for correct routes to be more reliably chosen."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Example: IT Support Dialogue"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Setup"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "First, we import the necessary classes and initialize the conversation with dialogue."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/jamesbriggs/opt/anaconda3/envs/decision-layer/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+ " from .autonotebook import tqdm as notebook_tqdm\n"
+ ]
+ }
+ ],
+ "source": [
+ "from semantic_router.text import Conversation\n",
+ "from semantic_router.schema import Message\n",
+ "\n",
+ "# Initialize the Conversation\n",
+ "conversation = Conversation()\n",
+ "\n",
+ "# Define the IT support dialogue\n",
+ "messages = [\n",
+ " Message(role=\"user\", content=\"Hi, there, please can you confirm your full name\"),\n",
+ " Message(role=\"user\", content=\"Hi, my name is John Doe.\"),\n",
+ " Message(role=\"bot\", content=\"Okay, how can I help you today?\"),\n",
+ " Message(role=\"user\", content=\"My computer keeps crashing\"),\n",
+ " Message(\n",
+ " role=\"bot\", content=\"Okay, is our software running when the computer crashes.\"\n",
+ " ),\n",
+ " Message(role=\"user\", content=\"Yeah, v.3.11.2.\"),\n",
+ "]\n",
+ "\n",
+ "# Add messages to the conversation\n",
+ "conversation.add_new_messages(messages)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Initialize an Encoder"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from semantic_router.encoders.openai import OpenAIEncoder\n",
+ "\n",
+ "encoder = OpenAIEncoder(openai_api_key=\"sk-...\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Split Conversation by Topic"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[31muser: Hi, there, please can you confirm your full name\u001b[0m\n",
+ "\u001b[31muser: Hi, my name is John Doe.\u001b[0m\n",
+ "\u001b[31mbot: Okay, how can I help you today?\u001b[0m\n",
+ "\u001b[32muser: My computer keeps crashing\u001b[0m\n",
+ "\u001b[32mbot: Okay, is our software running when the computer crashes.\u001b[0m\n",
+ "\u001b[32muser: Yeah, v.3.11.2.\u001b[0m\n"
+ ]
+ }
+ ],
+ "source": [
+ "conversation.configure_splitter(\n",
+ " encoder=encoder, threshold=0.78, split_method=\"cumulative_similarity\"\n",
+ ")\n",
+ "\n",
+ "all_topics, new_topics = conversation.split_by_topic()\n",
+ "\n",
+ "# Display all topics\n",
+ "print(conversation)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Notice that the last message says \"Yeah, it crashes right after I start the software\".\n",
+ "\n",
+ "This might be correctly routed by the semantic-router, particularly if the route is quite generic, intended for \"software\" and/or \"crashes\".\n",
+ "\n",
+ "However, as an illustrative example, what if the routes were \n",
+ "\n",
+ "Route A: \"Sotware Crashes - v3.11\"\n",
+ "\n",
+ "Route B: \"Computer Crashes - v3.11\"\n",
+ "\n",
+ "If just the last utterance was used, then Route A would likely be chosen. However, if instead every utterance from the last topic (Topic 4), concatenated together, were sent to the semantic-router, then this important additional context would most likely result in Route A being chosen.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Topic Splitting After Topic Continuation"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Note that topics can be continued even after `conversation.split_by_topic()` has already been run. "
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Add some new messages."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Define the IT support dialogue\n",
+ "messages = [\n",
+ " Message(\n",
+ " role=\"user\", content=\"What do the system logs say, right before the crash?\"\n",
+ " ),\n",
+ " Message(role=\"user\", content=\"I'll check soon, but first let's talk refund.\"),\n",
+ " Message(role=\"bot\", content=\"Okay let me sort out a refund.\"),\n",
+ "]\n",
+ "\n",
+ "# Add messages to the conversation\n",
+ "conversation.add_new_messages(messages)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[31muser: Hi, there, please can you confirm your full name\u001b[0m\n",
+ "\u001b[31muser: Hi, my name is John Doe.\u001b[0m\n",
+ "\u001b[31mbot: Okay, how can I help you today?\u001b[0m\n",
+ "\u001b[32muser: My computer keeps crashing\u001b[0m\n",
+ "\u001b[32mbot: Okay, is our software running when the computer crashes.\u001b[0m\n",
+ "\u001b[32muser: Yeah, v.3.11.2.\u001b[0m\n",
+ "\u001b[32muser: What do the system logs say, right before the crash?\u001b[0m\n",
+ "\u001b[33muser: I'll check soon, but first let's talk refund.\u001b[0m\n",
+ "\u001b[33mbot: Okay let me sort out a refund.\u001b[0m\n"
+ ]
+ }
+ ],
+ "source": [
+ "conversation.configure_splitter(\n",
+ " encoder=encoder, threshold=0.78, split_method=\"cumulative_similarity\"\n",
+ ")\n",
+ "\n",
+ "all_topics, new_topics = conversation.split_by_topic(force=True)\n",
+ "\n",
+ "print(conversation)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "As you can see, we:\n",
+ "\n",
+ "1) Added the first six messages, as seen above, to the `Conversation`.\n",
+ "2) Ran the Topic Splitter.\n",
+ "3) Added the last two messages to the `Conversation`.\n",
+ "4) Ran the Topic Splitter again.\n",
+ "\n",
+ "Despite \"user: Yeah, v.3.11.2 is running when it crashes\" and \"user: What do the system logs say, right before the crash?\" being added and separately, and despite the conversation splitter being run twice (once before user: What do the system logs say, right before the crash?\" was added, and once after), both these utterances were successfully assigned the same Topic - `Topic 4`.\n"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "semantic_splitter_1",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.5"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/poetry.lock b/poetry.lock
index e7c7c5957164751b22db0c863e49ac67a390bd99..2498a9a6c5327741198c1a5e1ac6ed2986abe8a9 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3388,4 +3388,4 @@ local = ["llama-cpp-python", "torch", "transformers"]
[metadata]
lock-version = "2.0"
python-versions = "^3.9"
-content-hash = "49fd469a4cf8a0a31d2e4df9e9a2b40d1f6fefeba87ec26d76d1cb716f4f51ca"
+content-hash = "10453196b0249ab854bdcd965ce8631ad9e0db33ae4423d35b250cb8b07b9898"
diff --git a/pyproject.toml b/pyproject.toml
index fef3a6fc4c00573c54ec0afcb2d43eea0f314927..b396601d8ca30dd802c5116f5493b0571a585f17 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "semantic-router"
-version = "0.0.20"
+version = "0.0.21"
description = "Super fast semantic router for AI decision making"
authors = [
"James Briggs <james@aurelio.ai>",
@@ -28,6 +28,7 @@ torch = {version = "^2.1.0", optional = true}
transformers = {version = "^4.36.2", optional = true}
llama-cpp-python = {version = "^0.2.28", optional = true}
black = "^23.12.1"
+colorama = "^0.4.6"
[tool.poetry.extras]
hybrid = ["pinecone-text"]
diff --git a/semantic_router/__init__.py b/semantic_router/__init__.py
index d8f5459e575f8db8d0ec21cd33dbf545a7025882..3b7582d2d8e2f842ac8904d7d5edcf2843e4f2ef 100644
--- a/semantic_router/__init__.py
+++ b/semantic_router/__init__.py
@@ -4,4 +4,4 @@ from semantic_router.route import Route
__all__ = ["RouteLayer", "HybridRouteLayer", "Route", "LayerConfig"]
-__version__ = "0.0.20"
+__version__ = "0.0.21"
diff --git a/semantic_router/encoders/cohere.py b/semantic_router/encoders/cohere.py
index f534fa51cb3524b047416260ff688fd9b9769bee..6bc7d661b1794fc0c23443e06dab78a977fa9d28 100644
--- a/semantic_router/encoders/cohere.py
+++ b/semantic_router/encoders/cohere.py
@@ -9,16 +9,23 @@ from semantic_router.encoders import BaseEncoder
class CohereEncoder(BaseEncoder):
client: Optional[cohere.Client] = None
type: str = "cohere"
+ input_type: Optional[str] = "search_query"
def __init__(
self,
name: Optional[str] = None,
cohere_api_key: Optional[str] = None,
score_threshold: float = 0.3,
+ input_type: Optional[str] = "search_query",
):
if name is None:
name = os.getenv("COHERE_MODEL_NAME", "embed-english-v3.0")
- super().__init__(name=name, score_threshold=score_threshold)
+ super().__init__(
+ name=name,
+ score_threshold=score_threshold,
+ input_type=input_type, # type: ignore
+ )
+ self.input_type = input_type
cohere_api_key = cohere_api_key or os.getenv("COHERE_API_KEY")
if cohere_api_key is None:
raise ValueError("Cohere API key cannot be 'None'.")
@@ -33,7 +40,9 @@ class CohereEncoder(BaseEncoder):
if self.client is None:
raise ValueError("Cohere client is not initialized.")
try:
- embeds = self.client.embed(docs, input_type="search_query", model=self.name)
+ embeds = self.client.embed(
+ docs, input_type=self.input_type, model=self.name
+ )
return embeds.embeddings
except Exception as e:
raise ValueError(f"Cohere API call failed. Error: {e}") from e
diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index 64f2cd1d7ca910a5165518af31209f60b78b280a..46ee7f590c3275b314b9ecb3e52adc54f16012fe 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -1,5 +1,5 @@
from enum import Enum
-from typing import List, Literal, Optional
+from typing import List, Optional
from pydantic.v1 import BaseModel
from pydantic.v1.dataclasses import dataclass
@@ -10,7 +10,6 @@ from semantic_router.encoders import (
FastEmbedEncoder,
OpenAIEncoder,
)
-from semantic_router.utils.splitters import DocumentSplit, semantic_splitter
class EncoderType(Enum):
@@ -66,19 +65,11 @@ class Message(BaseModel):
def to_llamacpp(self):
return {"role": self.role, "content": self.content}
+ def __str__(self):
+ return f"{self.role}: {self.content}"
-class Conversation(BaseModel):
- messages: List[Message]
-
- def split_by_topic(
- self,
- encoder: BaseEncoder,
- threshold: float = 0.5,
- split_method: Literal[
- "consecutive_similarity_drop", "cumulative_similarity_drop"
- ] = "consecutive_similarity_drop",
- ) -> list[DocumentSplit]:
- docs = [f"{m.role}: {m.content}" for m in self.messages]
- return semantic_splitter(
- encoder=encoder, docs=docs, threshold=threshold, split_method=split_method
- )
+
+class DocumentSplit(BaseModel):
+ docs: List[str]
+ is_triggered: bool = False
+ triggered_score: Optional[float] = None
diff --git a/semantic_router/splitters/__init__.py b/semantic_router/splitters/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/semantic_router/splitters/base.py b/semantic_router/splitters/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..c7ca37c2ed2660448b4549181e354aa0e5b8727a
--- /dev/null
+++ b/semantic_router/splitters/base.py
@@ -0,0 +1,13 @@
+from typing import List
+
+from pydantic.v1 import BaseModel
+from semantic_router.encoders import BaseEncoder
+
+
+class BaseSplitter(BaseModel):
+ name: str
+ encoder: BaseEncoder
+ score_threshold: float
+
+ def __call__(self, docs: List[str]) -> List[List[float]]:
+ raise NotImplementedError("Subclasses must implement this method")
diff --git a/semantic_router/splitters/consecutive_sim.py b/semantic_router/splitters/consecutive_sim.py
new file mode 100644
index 0000000000000000000000000000000000000000..6bd08845e0d3a45706bc99dd0740ab1355ae3840
--- /dev/null
+++ b/semantic_router/splitters/consecutive_sim.py
@@ -0,0 +1,51 @@
+from typing import List
+from semantic_router.splitters.base import BaseSplitter
+from semantic_router.encoders import BaseEncoder
+import numpy as np
+from semantic_router.schema import DocumentSplit
+
+
+class ConsecutiveSimSplitter(BaseSplitter):
+
+ """
+ Called "consecutive sim splitter" because we check the similarities of consecutive document embeddings (compare ith to i+1th document embedding).
+ """
+
+ def __init__(
+ self,
+ encoder: BaseEncoder,
+ name: str = "consecutive_similarity_splitter",
+ score_threshold: float = 0.45,
+ ):
+ super().__init__(name=name, score_threshold=score_threshold, encoder=encoder)
+ encoder.score_threshold = score_threshold
+
+ def __call__(self, docs: List[str]):
+ # Check if there's only a single document
+ if len(docs) == 1:
+ raise ValueError(
+ "There is only one document provided; at least two are required to determine topics based on similarity."
+ )
+
+ doc_embeds = self.encoder(docs)
+ norm_embeds = doc_embeds / np.linalg.norm(doc_embeds, axis=1, keepdims=True)
+ sim_matrix = np.matmul(norm_embeds, norm_embeds.T)
+ total_docs = len(docs)
+ splits = []
+ curr_split_start_idx = 0
+ curr_split_num = 1
+
+ for idx in range(1, total_docs):
+ curr_sim_score = sim_matrix[idx - 1][idx]
+ if idx < len(sim_matrix) and curr_sim_score < self.score_threshold:
+ splits.append(
+ DocumentSplit(
+ docs=list(docs[curr_split_start_idx:idx]),
+ is_triggered=True,
+ triggered_score=curr_sim_score,
+ )
+ )
+ curr_split_start_idx = idx
+ curr_split_num += 1
+ splits.append(DocumentSplit(docs=list(docs[curr_split_start_idx:])))
+ return splits
diff --git a/semantic_router/splitters/cumulative_sim.py b/semantic_router/splitters/cumulative_sim.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba8f4bd31faa15d268acb729aaedb4147d9c97e0
--- /dev/null
+++ b/semantic_router/splitters/cumulative_sim.py
@@ -0,0 +1,67 @@
+from typing import List
+from semantic_router.splitters.base import BaseSplitter
+import numpy as np
+from semantic_router.schema import DocumentSplit
+from semantic_router.encoders import BaseEncoder
+
+
+class CumulativeSimSplitter(BaseSplitter):
+
+ """
+ Called "cumulative sim" because we check the similarities of the embeddings of cumulative concatenated documents with the next document.
+ """
+
+ def __init__(
+ self,
+ encoder: BaseEncoder,
+ name: str = "cumulative_similarity_splitter",
+ score_threshold: float = 0.45,
+ ):
+ super().__init__(name=name, score_threshold=score_threshold, encoder=encoder)
+ encoder.score_threshold = score_threshold
+
+ def __call__(self, docs: List[str]):
+ total_docs = len(docs)
+ # Check if there's only a single document
+ if total_docs == 1:
+ raise ValueError(
+ "There is only one document provided; at least two are required to determine topics based on similarity."
+ )
+ splits = []
+ curr_split_start_idx = 0
+
+ for idx in range(0, total_docs):
+ if idx + 1 < total_docs: # Ensure there is a next document to compare with.
+ if idx == 0:
+ # On the first iteration, compare the first document directly to the second.
+ curr_split_docs = docs[idx]
+ else:
+ # For subsequent iterations, compare cumulative documents up to the current one with the next.
+ curr_split_docs = "\n".join(docs[curr_split_start_idx : idx + 1])
+ next_doc = docs[idx + 1]
+
+ # Embedding and similarity calculation remains the same.
+ curr_split_docs_embed = self.encoder([curr_split_docs])[0]
+ next_doc_embed = self.encoder([next_doc])[0]
+ curr_sim_score = np.dot(curr_split_docs_embed, next_doc_embed) / (
+ np.linalg.norm(curr_split_docs_embed)
+ * np.linalg.norm(next_doc_embed)
+ )
+ # Decision to split based on similarity score.
+ if curr_sim_score < self.score_threshold:
+ splits.append(
+ DocumentSplit(
+ docs=list(docs[curr_split_start_idx : idx + 1]),
+ is_triggered=True,
+ triggered_score=curr_sim_score,
+ )
+ )
+ curr_split_start_idx = (
+ idx + 1
+ ) # Update the start index for the next segment.
+
+ # Add the last segment after the loop.
+ if curr_split_start_idx < total_docs:
+ splits.append(DocumentSplit(docs=list(docs[curr_split_start_idx:])))
+
+ return splits
diff --git a/semantic_router/text.py b/semantic_router/text.py
new file mode 100644
index 0000000000000000000000000000000000000000..978da135eb71f3912c8d49318c743f3c91fc2b47
--- /dev/null
+++ b/semantic_router/text.py
@@ -0,0 +1,206 @@
+from colorama import Fore
+from colorama import Style
+
+from pydantic.v1 import BaseModel, Field
+from typing import Union, List, Literal, Tuple
+from semantic_router.splitters.consecutive_sim import ConsecutiveSimSplitter
+from semantic_router.splitters.cumulative_sim import CumulativeSimSplitter
+from semantic_router.encoders import BaseEncoder
+from semantic_router.schema import Message
+from semantic_router.schema import DocumentSplit
+
+# Define a type alias for the splitter to simplify the annotation
+SplitterType = Union[ConsecutiveSimSplitter, CumulativeSimSplitter, None]
+
+colors = [
+ Fore.WHITE,
+ Fore.RED,
+ Fore.GREEN,
+ Fore.YELLOW,
+ Fore.BLUE,
+ Fore.MAGENTA,
+ Fore.CYAN,
+]
+
+
+class Conversation(BaseModel):
+ messages: List[Message] = Field(
+ default_factory=list
+ ) # Ensure this is initialized as an empty list
+ topics: List[Tuple[int, str]] = []
+ splitter: SplitterType = None
+
+ def __str__(self):
+ if not self.messages:
+ return ""
+ if not self.topics:
+ return "\n".join([str(message) for message in self.messages])
+ else:
+ # we print each topic a different color
+ return_str_list = []
+ current_topic_id = None
+ color_idx = 0
+ for topic_id, message in self.topics:
+ if topic_id != current_topic_id:
+ # change color
+ color_idx = (color_idx + 1) % len(colors)
+ current_topic_id = topic_id
+ return_str_list.append(f"{colors[color_idx]}{message}{Style.RESET_ALL}")
+ return "\n".join(return_str_list)
+
+ def add_new_messages(self, new_messages: List[Message]):
+ """Adds new messages to the conversation.
+
+ :param messages: The new messages to be added to the conversation.
+ :type messages: List[Message]
+ """
+ self.messages.extend(new_messages)
+
+ def remove_topics(self):
+ self.topics = []
+
+ def configure_splitter(
+ self,
+ encoder: BaseEncoder,
+ threshold: float = 0.5,
+ split_method: Literal[
+ "consecutive_similarity", "cumulative_similarity"
+ ] = "consecutive_similarity",
+ ):
+ """
+ Configures the splitter for the conversation based on the specified method.
+
+ This method sets the splitter attribute of the Conversation class to an instance of the appropriate splitter class, based on the `split_method` parameter. It uses the provided encoder and similarity threshold to initialize the splitter.
+
+ :param encoder: The encoder to be used by the splitter for encoding messages.
+ :type encoder: BaseEncoder
+ :param threshold: The similarity threshold to be used by the splitter. Defaults to 0.5.
+ :type threshold: float
+ :param split_method: The method to be used for splitting the conversation into topics. Can be one of "consecutive_similarity" or "cumulative_similarity". Defaults to "consecutive_similarity".
+ :type split_method: Literal["consecutive_similarity", "cumulative_similarity"]
+ :raises ValueError: If an invalid split method is provided.
+ """
+
+ if split_method == "consecutive_similarity":
+ self.splitter = ConsecutiveSimSplitter(
+ encoder=encoder, score_threshold=threshold
+ )
+ elif split_method == "cumulative_similarity":
+ self.splitter = CumulativeSimSplitter(
+ encoder=encoder, score_threshold=threshold
+ )
+ else:
+ raise ValueError(f"Invalid split method: {split_method}")
+
+ def get_last_message_and_topic_id(self):
+ """
+ Retrieves the last message and its corresponding topic ID from the list of topics.
+
+ This method scans the list of topics, if any, and returns the topic ID and message of the last entry. If there are no topics, it returns None for both the topic ID and message.
+
+ The last message from a previous spiltting is useful because it can be passed to the splitter along with new messages, and if the first new message is assigned the same topic as the last message, then we know that the new message should continue with the same topic ID as the last message.
+
+ :return: A tuple containing the topic ID (int) and message (str) of the last topic, or (None, None) if there are no topics.
+ :rtype: tuple[int | None, str | None]
+ """
+
+ if self.topics:
+ return self.topics[-1]
+ else:
+ return None, None
+
+ def determine_topic_start_index(self, new_topics, last_topic_id, last_message):
+ """
+ Determines the starting index for new topics based on existing topics and the last message.
+
+ :param new_topics: The list of new topics generated by the splitter.
+ :type new_topics: List[DocumentSplit]
+ :param last_topic_id: The topic ID of the last message from the previous splitting.
+ :type last_topic_id: int, optional
+ :param last_message: The last message from the previous splitting.
+ :type last_message: str, optional
+ :return: The starting index for new topics.
+ :rtype: int
+ """
+ if not self.topics or not new_topics:
+ return 1
+ if (
+ last_topic_id is not None
+ and last_message
+ and last_message in new_topics[0].docs
+ ):
+ return last_topic_id
+ return self.topics[-1][0] + 1
+
+ def append_new_topics(self, new_topics, start) -> None:
+ """
+ Appends new topics to the list of topics with unique IDs.
+
+ This method takes a list of new topics generated by the splitter and appends them to the existing list of topics, ensuring each topic is assigned a unique ID starting from the specified starting index.
+
+ :param new_topics: The list of new topics generated by the splitter.
+ :type new_topics: List[DocumentSplit]
+ :param start: The starting index for new topics.
+ :type start: int
+ """
+ for i, topic in enumerate(new_topics, start=start):
+ for message in topic.docs:
+ self.topics.append((i, message))
+
+ def split_by_topic(
+ self, force: bool = False
+ ) -> Tuple[List[Tuple[int, str]], List[DocumentSplit]]:
+ """
+ Splits the messages into topics based on their semantic similarity.
+
+ This method processes unclustered messages, splits them into topics using the configured splitter, and appends the new topics to the existing list of topics with unique IDs. It ensures that messages belonging to the same topic are grouped together, even if they were not processed in the same batch.
+
+ :raises ValueError: If the splitter is not configured before calling this method.
+
+ :return: A tuple containing the updated list of topics and the list of new topics generated in this call.
+ :rtype: tuple[list[tuple[int, str]], list[DocumentSplit]]
+ """
+
+ if self.splitter is None:
+ raise ValueError(
+ "Splitter is not configured. Please call configure_splitter first."
+ )
+ new_topics: List[DocumentSplit] = []
+
+ if self.topics:
+ # reset self.topics
+ self.topics = []
+
+ # Get unclusteed messages.
+ unclustered_messages = self.messages[len(self.topics) :]
+ if not unclustered_messages:
+ print("No unclustered messages to process.")
+ return self.topics, new_topics
+
+ # Extract the last topic ID and message from the previous splitting, if they exist.
+ last_topic_id, last_message = self.get_last_message_and_topic_id()
+
+ # Initialize docs with the last message from the last topic if it exists, and with unclustered messages.
+ # TODO: Currenlty only getting last message from last topic in previous splitting. Should we get more for more reliable continuation of topic ids?
+ docs = [last_message] if last_message else []
+ docs.extend([f"{m.role}: {m.content}" for m in unclustered_messages])
+
+ new_topics = self.splitter(docs)
+
+ # Ensure there are new topics before proceeding
+ if not new_topics:
+ return self.topics, []
+
+ # If last_message and the first new message are assigned the same topic ID, then we know the new message should take last_message's place original topic id.
+ start = self.determine_topic_start_index(
+ new_topics, last_topic_id, last_message
+ )
+
+ # If the last message from the previous splitting is found in the first new topic, remove it
+ if self.topics and new_topics[0].docs[0] == self.topics[-1][1]:
+ new_topics[0].docs.pop(0)
+
+ self.append_new_topics(new_topics, start)
+
+ # TODO: Instead of self.topics as list of tuples should it also be a list of DocumentSplit objects?
+ return self.topics, new_topics
diff --git a/semantic_router/utils/splitters.py b/semantic_router/utils/splitters.py
deleted file mode 100644
index 0f3dc3410c2e6a0cb640b22dd7c111febd0f95cf..0000000000000000000000000000000000000000
--- a/semantic_router/utils/splitters.py
+++ /dev/null
@@ -1,97 +0,0 @@
-from typing import List, Literal, Optional
-
-import numpy as np
-from pydantic.v1 import BaseModel
-
-from semantic_router.encoders import BaseEncoder
-
-
-class DocumentSplit(BaseModel):
- docs: List[str]
- is_triggered: bool = False
- triggered_score: Optional[float] = None
-
-
-def semantic_splitter(
- encoder: BaseEncoder,
- docs: List[str],
- threshold: float,
- split_method: Literal[
- "consecutive_similarity_drop", "cumulative_similarity_drop"
- ] = "consecutive_similarity_drop",
-) -> List[DocumentSplit]:
- """
- Splits a list of documents base on semantic similarity changes.
-
- Method 1: "consecutive_similarity_drop" - This method splits documents based on
- the changes in similarity scores between consecutive documents.
- Method 2: "cumulative_similarity_drop" - This method segments the documents based
- on the changes in cumulative similarity score of the documents within the same
- split.
-
- Args:
- encoder (BaseEncoder): Encoder for document embeddings.
- docs (List[str]): Documents to split.
- threshold (float): The similarity drop value that will trigger a new document
- split.
- split_method (str): The method to use for splitting.
-
- Returns:
- Dict[str, List[str]]: Splits with corresponding documents.
- """
- total_docs = len(docs)
- splits = []
- curr_split_start_idx = 0
- curr_split_num = 1
-
- if split_method == "consecutive_similarity_drop":
- doc_embeds = encoder(docs)
- norm_embeds = doc_embeds / np.linalg.norm(doc_embeds, axis=1, keepdims=True)
- sim_matrix = np.matmul(norm_embeds, norm_embeds.T)
-
- for idx in range(1, total_docs):
- curr_sim_score = sim_matrix[idx - 1][idx]
- if idx < len(sim_matrix) and curr_sim_score < threshold:
- splits.append(
- DocumentSplit(
- docs=docs[curr_split_start_idx:idx],
- is_triggered=True,
- triggered_score=curr_sim_score,
- )
- )
- curr_split_start_idx = idx
- curr_split_num += 1
-
- elif split_method == "cumulative_similarity_drop":
- for idx in range(1, total_docs):
- if idx + 1 < total_docs:
- curr_split_docs = "\n".join(docs[curr_split_start_idx : idx + 1])
- next_doc = docs[idx + 1]
-
- curr_split_docs_embed = encoder([curr_split_docs])[0]
- next_doc_embed = encoder([next_doc])[0]
-
- curr_sim_score = np.dot(curr_split_docs_embed, next_doc_embed) / (
- np.linalg.norm(curr_split_docs_embed)
- * np.linalg.norm(next_doc_embed)
- )
-
- if curr_sim_score < threshold:
- splits.append(
- DocumentSplit(
- docs=docs[curr_split_start_idx : idx + 1],
- is_triggered=True,
- triggered_score=curr_sim_score,
- )
- )
- curr_split_start_idx = idx + 1
- curr_split_num += 1
-
- else:
- raise ValueError(
- "Invalid 'split_method'. Choose either 'consecutive_similarity_drop' or"
- " 'cumulative_similarity_drop'."
- )
-
- splits.append(DocumentSplit(docs=docs[curr_split_start_idx:]))
- return splits
diff --git a/tests/unit/test_splitters.py b/tests/unit/test_splitters.py
index f0e8e3f32734f988e2d54a878c26f7f5f1994228..165fc2dd4327320f12ed3ea3a68c440b650af1dd 100644
--- a/tests/unit/test_splitters.py
+++ b/tests/unit/test_splitters.py
@@ -1,68 +1,199 @@
-from unittest.mock import Mock
+from unittest.mock import Mock, create_autospec
import pytest
+import numpy as np
-from semantic_router.schema import Conversation, Message
-from semantic_router.utils.splitters import semantic_splitter
+from semantic_router.text import Conversation
+from semantic_router.schema import Message
+from semantic_router.splitters.consecutive_sim import ConsecutiveSimSplitter
+from semantic_router.splitters.cumulative_sim import CumulativeSimSplitter
+from semantic_router.encoders.base import BaseEncoder
+from semantic_router.encoders.cohere import CohereEncoder
+from semantic_router.splitters.base import BaseSplitter
-def test_semantic_splitter_consecutive_similarity_drop():
- # Mock the BaseEncoder
+def test_consecutive_sim_splitter():
+ # Create a Mock object for the encoder
mock_encoder = Mock()
- mock_encoder.return_value = [[0.5, 0], [0.5, 0], [0.5, 0], [0, 0.5], [0, 0.5]]
+ mock_encoder.return_value = np.array([[1, 0], [1, 0.1], [0, 1]])
- docs = ["doc1", "doc2", "doc3", "doc4", "doc5"]
- threshold = 0.5
- split_method = "consecutive_similarity_drop"
+ cohere_encoder = CohereEncoder(
+ name="",
+ cohere_api_key="a",
+ input_type="",
+ )
+ # Instantiate the ConsecutiveSimSplitter with the mock encoder
+ splitter = ConsecutiveSimSplitter(encoder=cohere_encoder, score_threshold=0.9)
+ splitter.encoder = mock_encoder
- result = semantic_splitter(mock_encoder, docs, threshold, split_method)
+ # Define some documents
+ docs = ["doc1", "doc2", "doc3"]
- assert result[0].docs == ["doc1", "doc2", "doc3"]
- assert result[1].docs == ["doc4", "doc5"]
+ # Use the splitter to split the documents
+ splits = splitter(docs)
+ # Verify the splits
+ assert len(splits) == 2, "Expected two splits based on the similarity threshold"
+ assert splits[0].docs == [
+ "doc1",
+ "doc2",
+ ], "First split does not match expected documents"
+ assert splits[1].docs == ["doc3"], "Second split does not match expected documents"
-def test_semantic_splitter_cumulative_similarity_drop():
+
+def test_cumulative_sim_splitter():
# Mock the BaseEncoder
mock_encoder = Mock()
+ # Adjust the side_effect to simulate the encoder's behavior for cumulative document comparisons
+ # This simplistic simulation assumes binary embeddings for demonstration purposes
+ # Define a side_effect function for the mock encoder
mock_encoder.side_effect = (
- lambda x: [[0.5, 0]] if "doc1" in x or "doc1\ndoc2" in x else [[0, 0.5]]
+ lambda x: [[0.5, 0]]
+ if "doc1" in x or "doc1\ndoc2" in x or "doc2" in x
+ else [[0, 0.5]]
+ )
+
+ # Instantiate the CumulativeSimSplitter with the mock encoder
+ cohere_encoder = CohereEncoder(
+ name="",
+ cohere_api_key="a",
+ input_type="",
)
+ splitter = CumulativeSimSplitter(encoder=cohere_encoder, score_threshold=0.9)
+ splitter.encoder = mock_encoder
+ # Define some documents
docs = ["doc1", "doc2", "doc3", "doc4", "doc5"]
- threshold = 0.5
- split_method = "cumulative_similarity_drop"
- result = semantic_splitter(mock_encoder, docs, threshold, split_method)
+ # Use the splitter to split the documents
+ splits = splitter(docs)
- assert result[0].docs == ["doc1", "doc2"]
- assert result[1].docs == ["doc3", "doc4", "doc5"]
+ # Verify the splits
+ # The expected outcome needs to match the logic defined in your mock_encoder's side_effect
+ assert len(splits) == 2, f"{len(splits)}"
+ assert splits[0].docs == [
+ "doc1",
+ "doc2",
+ ], "First split does not match expected documents"
+ assert splits[1].docs == [
+ "doc3",
+ "doc4",
+ "doc5",
+ ], "Second split does not match expected documents"
-def test_semantic_splitter_invalid_method():
- # Mock the BaseEncoder
+def test_split_by_topic_consecutive_similarity():
mock_encoder = Mock()
+ mock_encoder.return_value = [[0.5, 0], [0, 0.5]]
- docs = ["doc1", "doc2", "doc3", "doc4", "doc5"]
- threshold = 0.5
- split_method = "invalid_method"
+ messages = [
+ Message(role="User", content="What is the latest news?"),
+ Message(role="Assistant", content="How is the weather today?"),
+ ]
+ conversation = Conversation(messages=messages)
- with pytest.raises(ValueError):
- semantic_splitter(mock_encoder, docs, threshold, split_method)
+ cohere_encoder = CohereEncoder(
+ name="",
+ cohere_api_key="a",
+ input_type="",
+ )
+
+ conversation.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="consecutive_similarity"
+ )
+ conversation.splitter.encoder = mock_encoder
+
+ topics, new_topics = conversation.split_by_topic()
+ assert len(new_topics) == 2
+ assert new_topics[0].docs == ["User: What is the latest news?"]
+ assert new_topics[1].docs == ["Assistant: How is the weather today?"]
-def test_split_by_topic():
+
+def test_split_by_topic_cumulative_similarity():
mock_encoder = Mock()
- mock_encoder.return_value = [[0.5, 0], [0, 0.5]]
+ mock_encoder.side_effect = (
+ lambda x: [[0.5, 0]] if "User: What is the latest news?" in x else [[0, 0.5]]
+ )
messages = [
Message(role="User", content="What is the latest news?"),
- Message(role="Bot", content="How is the weather today?"),
+ Message(role="Assistant", content="How is the weather today?"),
]
conversation = Conversation(messages=messages)
- result = conversation.split_by_topic(
- encoder=mock_encoder, threshold=0.5, split_method="consecutive_similarity_drop"
+ cohere_encoder = CohereEncoder(
+ name="",
+ cohere_api_key="a",
+ input_type="",
+ )
+
+ conversation.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="cumulative_similarity"
+ )
+ conversation.splitter.encoder = mock_encoder
+
+ topics, new_topics = conversation.split_by_topic()
+
+ # Assertions may need to be adjusted based on the expected behavior of the cumulative similarity splitter
+ assert len(new_topics) == 2
+
+
+def test_split_by_topic_no_messages():
+ mock_encoder = create_autospec(BaseEncoder)
+ conversation = Conversation()
+ conversation.configure_splitter(
+ encoder=mock_encoder, threshold=0.5, split_method="consecutive_similarity"
)
- assert result[0].docs == ["User: What is the latest news?"]
- assert result[1].docs == ["Bot: How is the weather today?"]
+ topics, new_topics = conversation.split_by_topic()
+
+ assert len(new_topics) == 0
+ assert len(topics) == 0
+
+
+def test_split_by_topic_without_configuring_splitter():
+ conversation = Conversation(messages=[Message(role="User", content="Hello")])
+
+ with pytest.raises(ValueError):
+ conversation.split_by_topic()
+
+
+def test_consecutive_similarity_splitter_single_doc():
+ mock_encoder = create_autospec(BaseEncoder)
+ # Assuming any return value since it should not reach the point of using the encoder
+ mock_encoder.return_value = np.array([[0.5, 0]])
+
+ splitter = ConsecutiveSimSplitter(encoder=mock_encoder, score_threshold=0.5)
+
+ docs = ["doc1"]
+ with pytest.raises(ValueError) as excinfo:
+ splitter(docs)
+ assert "at least two are required" in str(excinfo.value)
+
+
+def test_cumulative_similarity_splitter_single_doc():
+ mock_encoder = create_autospec(BaseEncoder)
+ # Assuming any return value since it should not reach the point of using the encoder
+ mock_encoder.return_value = np.array([[0.5, 0]])
+
+ splitter = CumulativeSimSplitter(encoder=mock_encoder, score_threshold=0.5)
+
+ docs = ["doc1"]
+ with pytest.raises(ValueError) as excinfo:
+ splitter(docs)
+ assert "at least two are required" in str(excinfo.value)
+
+
+@pytest.fixture
+def base_splitter_instance():
+ # Now MockEncoder includes default values for required fields
+ mock_encoder = Mock(spec=BaseEncoder)
+ mock_encoder.name = "mock_encoder"
+ mock_encoder.score_threshold = 0.5
+ return BaseSplitter(name="test_splitter", encoder=mock_encoder, score_threshold=0.5)
+
+
+def test_base_splitter_call_not_implemented(base_splitter_instance):
+ with pytest.raises(NotImplementedError):
+ base_splitter_instance(["document"])
diff --git a/tests/unit/test_text.py b/tests/unit/test_text.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7490ec14d43593efa608261e4e6ef1d76915347
--- /dev/null
+++ b/tests/unit/test_text.py
@@ -0,0 +1,162 @@
+import pytest
+from unittest.mock import Mock
+from semantic_router.text import Conversation, Message
+from semantic_router.splitters.consecutive_sim import ConsecutiveSimSplitter
+from semantic_router.splitters.cumulative_sim import CumulativeSimSplitter
+from semantic_router.encoders.cohere import (
+ CohereEncoder,
+) # Adjust this import based on your project structure
+from semantic_router.schema import DocumentSplit
+
+
+@pytest.fixture
+def conversation_instance():
+ return Conversation()
+
+
+@pytest.fixture
+def cohere_encoder():
+ # Initialize CohereEncoder with necessary arguments
+ encoder = CohereEncoder(
+ name="cohere_encoder", cohere_api_key="dummy_key", input_type="text"
+ )
+ return encoder
+
+
+def test_add_new_messages(conversation_instance):
+ initial_len = len(conversation_instance.messages)
+ conversation_instance.add_new_messages([Message(role="user", content="Hello")])
+ assert len(conversation_instance.messages) == initial_len + 1
+
+
+def test_remove_topics(conversation_instance):
+ conversation_instance.topics.append((1, "Sample Topic"))
+ conversation_instance.remove_topics()
+ assert len(conversation_instance.topics) == 0
+
+
+def test_configure_splitter_consecutive_similarity(
+ conversation_instance, cohere_encoder
+):
+ conversation_instance.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="consecutive_similarity"
+ )
+ assert isinstance(conversation_instance.splitter, ConsecutiveSimSplitter)
+
+
+def test_configure_splitter_cumulative_similarity(
+ conversation_instance, cohere_encoder
+):
+ conversation_instance.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="cumulative_similarity"
+ )
+ assert isinstance(conversation_instance.splitter, CumulativeSimSplitter)
+
+
+def test_configure_splitter_invalid_method(conversation_instance, cohere_encoder):
+ with pytest.raises(ValueError):
+ conversation_instance.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="invalid_method"
+ )
+
+
+def test_split_by_topic_without_configuring_splitter(conversation_instance):
+ with pytest.raises(ValueError):
+ conversation_instance.split_by_topic()
+
+
+def test_split_by_topic_with_no_unclustered_messages(
+ conversation_instance, cohere_encoder, capsys
+):
+ conversation_instance.configure_splitter(
+ encoder=cohere_encoder, threshold=0.5, split_method="consecutive_similarity"
+ )
+ conversation_instance.splitter = Mock()
+ conversation_instance.split_by_topic()
+ captured = capsys.readouterr()
+ assert "No unclustered messages to process." in captured.out
+
+
+def test_get_last_message_and_topic_id_with_no_topics(conversation_instance):
+ # Test the method when there are no topics in the conversation
+ last_topic_id, last_message = conversation_instance.get_last_message_and_topic_id()
+ assert (
+ last_topic_id is None and last_message is None
+ ), "Expected None for both topic ID and message when there are no topics"
+
+
+def test_get_last_message_and_topic_id_with_topics(conversation_instance):
+ # Add some topics to the conversation instance
+ conversation_instance.topics.append((0, "First message"))
+ conversation_instance.topics.append((1, "Second message"))
+ conversation_instance.topics.append((2, "Third message"))
+
+ # Test the method when there are topics in the conversation
+ last_topic_id, last_message = conversation_instance.get_last_message_and_topic_id()
+ assert (
+ last_topic_id == 2 and last_message == "Third message"
+ ), "Expected last topic ID and message to match the last topic added"
+
+
+def test_determine_topic_start_index_no_existing_topics(conversation_instance):
+ # Scenario where there are no existing topics
+ new_topics = [
+ DocumentSplit(docs=["User: Hello!"], is_triggered=True, triggered_score=0.4)
+ ]
+ start_index = conversation_instance.determine_topic_start_index(
+ new_topics, None, None
+ )
+ assert (
+ start_index == 1
+ ), "Expected start index to be 1 when there are no existing topics"
+
+
+def test_determine_topic_start_index_with_existing_topics_not_including_last_message(
+ conversation_instance,
+):
+ # Scenario where existing topics do not include the last message
+ conversation_instance.topics.append((0, "First message"))
+ new_topics = [
+ DocumentSplit(docs=["User: Hello!"], is_triggered=True, triggered_score=0.4)
+ ]
+ start_index = conversation_instance.determine_topic_start_index(
+ new_topics, 0, "Non-existent last message"
+ )
+ assert (
+ start_index == 1
+ ), "Expected start index to increment when last message is not in new topics"
+
+
+def test_determine_topic_start_index_with_existing_topics_including_last_message(
+ conversation_instance,
+):
+ # Scenario where the first new topic includes the last message
+ conversation_instance.topics.append((0, "First message"))
+ new_topics = [
+ DocumentSplit(
+ docs=["First message", "Another message"],
+ is_triggered=True,
+ triggered_score=0.4,
+ )
+ ]
+ start_index = conversation_instance.determine_topic_start_index(
+ new_topics, 0, "First message"
+ )
+ assert (
+ start_index == 0
+ ), "Expected start index to be the same as last topic ID when last message is included in new topics"
+
+
+def test_determine_topic_start_index_increment_from_last_topic_id(
+ conversation_instance,
+):
+ # Scenario to test increment from the last topic ID when last message is not in new topics
+ conversation_instance.topics.append((1, "First message"))
+ conversation_instance.topics.append((2, "Second message"))
+ new_topics = [
+ DocumentSplit(docs=["User: Hello!"], is_triggered=True, triggered_score=0.4)
+ ]
+ start_index = conversation_instance.determine_topic_start_index(
+ new_topics, 2, "Non-existent last message"
+ )
+ assert start_index == 3, "Expected start index to be last topic ID + 1"