Skip to content

Commit c3b0b6c

Browse files
authored
Continue implementation of experimental Redis vector search (#2430)
* Add important note to redis adapter * Implement search text fields * Add noop tagger * Add dedicated search class * Order results by redis simalarity * Update documentation * Simplify tagger logic * Update imports
1 parent bf812e8 commit c3b0b6c

8 files changed

Lines changed: 575 additions & 94 deletions

File tree

chatterbot/chatterbot.py

Lines changed: 62 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from typing import Union
33
from chatterbot.storage import StorageAdapter
44
from chatterbot.logic import LogicAdapter
5-
from chatterbot.search import TextSearch, IndexedTextSearch
5+
from chatterbot.search import TextSearch, IndexedTextSearch, SemanticVectorSearch
66
from chatterbot.tagging import PosLemmaTagger
77
from chatterbot.conversation import Statement
88
from chatterbot import languages
@@ -74,41 +74,60 @@ def __init__(self, name, stream=False, **kwargs):
7474

7575
tagger_language = kwargs.get('tagger_language', languages.ENG)
7676

77-
try:
78-
Tagger = kwargs.get('tagger', PosLemmaTagger)
79-
80-
# Allow instances to be provided for performance optimization
81-
# (Example: a pre-loaded model in a tagger when unit testing)
82-
if not isinstance(Tagger, type):
83-
self.tagger = Tagger
84-
else:
85-
self.tagger = Tagger(language=tagger_language)
86-
except IOError as io_error:
87-
# Return a more helpful error message if possible
88-
if "Can't find model" in str(io_error):
89-
model_name = utils.get_model_for_language(tagger_language)
90-
if hasattr(tagger_language, 'ENGLISH_NAME'):
91-
language_name = tagger_language.ENGLISH_NAME
77+
# Check if storage adapter has a preferred tagger
78+
PreferredTagger = self.storage.get_preferred_tagger()
79+
80+
if PreferredTagger is not None:
81+
# Storage adapter specifies its own tagger
82+
self.tagger = PreferredTagger(language=tagger_language)
83+
else:
84+
# Use default or user-specified tagger
85+
try:
86+
Tagger = kwargs.get('tagger', PosLemmaTagger)
87+
88+
# Allow instances to be provided for performance optimization
89+
# (Example: a pre-loaded model in a tagger when unit testing)
90+
if not isinstance(Tagger, type):
91+
self.tagger = Tagger
9292
else:
93-
language_name = tagger_language
94-
raise self.ChatBotException(
95-
'Setup error:\n'
96-
f'The Spacy model for "{language_name}" language is missing.\n'
97-
'Please install the model using the command:\n\n'
98-
f'python -m spacy download {model_name}\n\n'
99-
'See https://spacy.io/usage/models for more information about available models.'
100-
) from io_error
101-
else:
102-
raise io_error
93+
self.tagger = Tagger(language=tagger_language)
94+
except IOError as io_error:
95+
# Return a more helpful error message if possible
96+
if "Can't find model" in str(io_error):
97+
model_name = utils.get_model_for_language(tagger_language)
98+
if hasattr(tagger_language, 'ENGLISH_NAME'):
99+
language_name = tagger_language.ENGLISH_NAME
100+
else:
101+
language_name = tagger_language
102+
raise self.ChatBotException(
103+
'Setup error:\n'
104+
f'The Spacy model for "{language_name}" language is missing.\n'
105+
'Please install the model using the command:\n\n'
106+
f'python -m spacy download {model_name}\n\n'
107+
'See https://spacy.io/usage/models for more information about available models.'
108+
) from io_error
109+
else:
110+
raise io_error
103111

112+
# Initialize search algorithms
104113
primary_search_algorithm = IndexedTextSearch(self, **kwargs)
105114
text_search_algorithm = TextSearch(self, **kwargs)
115+
semantic_vector_search_algorithm = SemanticVectorSearch(self, **kwargs)
106116

107117
self.search_algorithms = {
108118
primary_search_algorithm.name: primary_search_algorithm,
109-
text_search_algorithm.name: text_search_algorithm
119+
text_search_algorithm.name: text_search_algorithm,
120+
semantic_vector_search_algorithm.name: semantic_vector_search_algorithm
110121
}
111122

123+
# Check if storage adapter has a preferred search algorithm
124+
preferred_search_algorithm = self.storage.get_preferred_search_algorithm()
125+
if preferred_search_algorithm and preferred_search_algorithm in self.search_algorithms:
126+
# Set as default for logic adapters that don't specify their own search algorithm
127+
# This ensures BestMatch and other adapters use the optimal search method
128+
self.logger.info(f'Storage adapter prefers search algorithm: {preferred_search_algorithm}')
129+
kwargs.setdefault('search_algorithm_name', preferred_search_algorithm)
130+
112131
for adapter in logic_adapters:
113132
utils.validate_adapter_class(adapter, LogicAdapter)
114133
logic_adapter = utils.initialize_class(adapter, self, **kwargs)
@@ -191,15 +210,22 @@ def get_response(self, statement: Union[Statement, str, dict] = None, **kwargs)
191210
input_statement.in_response_to = previous_statement.text
192211

193212
# Make sure the input statement has its search text saved
194-
195-
if not input_statement.search_text:
196-
_search_text = self.tagger.get_text_index_string(input_statement.text)
197-
input_statement.search_text = _search_text
198-
199-
if not input_statement.search_in_response_to and input_statement.in_response_to:
200-
input_statement.search_in_response_to = self.tagger.get_text_index_string(
201-
input_statement.in_response_to
202-
)
213+
if not self.tagger.needs_text_indexing():
214+
# Tagger doesn't transform text, use it directly
215+
if not input_statement.search_text:
216+
input_statement.search_text = input_statement.text
217+
if not input_statement.search_in_response_to and input_statement.in_response_to:
218+
input_statement.search_in_response_to = input_statement.in_response_to
219+
else:
220+
# Use tagger for text indexing or transformations
221+
if not input_statement.search_text:
222+
_search_text = self.tagger.get_text_index_string(input_statement.text)
223+
input_statement.search_text = _search_text
224+
225+
if not input_statement.search_in_response_to and input_statement.in_response_to:
226+
input_statement.search_in_response_to = self.tagger.get_text_index_string(
227+
input_statement.in_response_to
228+
)
203229

204230
response = self.generate_response(
205231
input_statement,

chatterbot/search.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,3 +157,73 @@ def search(self, input_statement, **additional_parameters):
157157
if confidence >= 1.0:
158158
self.chatbot.logger.info('Exact match found, stopping search')
159159
break
160+
161+
162+
class SemanticVectorSearch:
163+
"""
164+
Semantic vector search for storage adapters that use vector embeddings.
165+
Does not require a tagger or comparison function - relies on the storage
166+
adapter's native vector similarity search capabilities.
167+
168+
:param search_page_size:
169+
The maximum number of records to load into memory at a time when searching.
170+
Defaults to 1000
171+
"""
172+
173+
name = 'semantic_vector_search'
174+
175+
def __init__(self, chatbot, **kwargs):
176+
self.chatbot = chatbot
177+
178+
self.search_page_size = kwargs.get(
179+
'search_page_size', 1000
180+
)
181+
182+
def search(self, input_statement, **additional_parameters):
183+
"""
184+
Search for semantically similar statements using vector similarity.
185+
Confidence scores are calculated by the storage adapter based on
186+
vector distances and returned in the results.
187+
188+
:param input_statement: A statement.
189+
:type input_statement: chatterbot.conversation.Statement
190+
191+
:param **additional_parameters: Additional parameters to be passed
192+
to the ``filter`` method of the storage adapter when searching.
193+
194+
:rtype: Generator yielding one closest matching statement at a time.
195+
"""
196+
self.chatbot.logger.info('Beginning semantic vector search')
197+
198+
search_parameters = {
199+
'search_in_response_to_contains': input_statement.text,
200+
'persona_not_startswith': 'bot:',
201+
'page_size': self.search_page_size
202+
}
203+
204+
if additional_parameters:
205+
search_parameters.update(additional_parameters)
206+
207+
statement_list = self.chatbot.storage.filter(**search_parameters)
208+
209+
best_confidence_so_far = 0
210+
211+
self.chatbot.logger.info('Processing search results')
212+
213+
# Yield statements with confidence scores from vector similarity
214+
for statement in statement_list:
215+
# Confidence should already be set by the storage adapter
216+
confidence = getattr(statement, 'confidence', 0.0)
217+
218+
if confidence > best_confidence_so_far:
219+
best_confidence_so_far = confidence
220+
221+
self.chatbot.logger.info('Similar statement found: {} {}'.format(
222+
statement.in_response_to, confidence
223+
))
224+
225+
yield statement
226+
227+
if confidence >= 1.0:
228+
self.chatbot.logger.info('Exact match found, stopping search')
229+
break

0 commit comments

Comments
 (0)