Migrate to sparse_vector query

ask_candid/retrieval/elastic.py

@@ -32,10 +32,10 @@ class RetrieverInput(BaseModel):
     user_input: str = Field(description="query to look up in retriever")
 
 
-def build_text_expansion_query(
+def build_sparse_vector_query(
     query: str,
     fields: Tuple[str],
-    model_id: str = ".elser_model_2_linux-x86_64"
+    inference_id: str = ".elser-2-elasticsearch"
 ) -> Dict[str, Any]:
     """Builds a valid Elasticsearch text expansion query payload
 
@@ -45,8 +45,8 @@ def build_text_expansion_query(
         Search context string
     fields : Tuple[str]
         Semantic text field names
-    model_id : str, optional
-        ID of model deployed in Elasticsearch, by default ".elser_model_2_linux-x86_64"
+    inference_id : str, optional
+        ID of model deployed in Elasticsearch, by default ".elser-2-elasticsearch"
 
     Returns
     -------
@@ -60,15 +60,14 @@ def build_text_expansion_query(
             "nested": {
                 "path": f"embeddings.{f}.chunks",
                 "query": {
-                    "text_expansion": {
-                        f"embeddings.{f}.chunks.vector": {
-                            "model_id": model_id,
-                            "model_text": query,
-                            "boost": 1 / len(fields)
-                        }
+                    "sparse_vector": {
+                        "field": f"embeddings.{f}.chunks.vector",
+                        "inference_id": inference_id,
+                        "query": query,
+                        "boost": 1 / len(fields)
                     }
                 },
-                 "inner_hits": {
+                "inner_hits": {
                     "_source": False,
                     "size": 2,
                     "fields": [f"embeddings.{f}.chunks.chunk"]
@@ -99,7 +98,7 @@ def query_builder(query: str, indices: List[str]) -> List[Dict[str, Any]]:
 
     for index in indices:
         if index == "issuelab":
-            q = build_text_expansion_query(
+            q = build_sparse_vector_query(
                 query=query,
                 fields=("description", "content", "combined_issuelab_findings", "combined_item_description")
             )
@@ -107,7 +106,7 @@ def query_builder(query: str, indices: List[str]) -> List[Dict[str, Any]]:
             q["size"] = 1
             queries.extend([{"index": ElasticIndexMapping.ISSUELAB_INDEX_ELSER}, q])
         elif index == "youtube":
-            q = build_text_expansion_query(
+            q = build_sparse_vector_query(
                 query=query,
                 fields=("captions_cleaned", "description_cleaned", "title")
             )
@@ -116,15 +115,15 @@ def query_builder(query: str, indices: List[str]) -> List[Dict[str, Any]]:
             q["size"] = 2
             queries.extend([{"index": ElasticIndexMapping.YOUTUBE_INDEX_ELSER}, q])
         elif index == "candid_blog":
-            q = build_text_expansion_query(
+            q = build_sparse_vector_query(
                 query=query,
-                fields=("content", "title")
+                fields=("content", "authors_text", "title_summary_tags")
             )
             q["_source"] = {"excludes": ["embeddings"]}
             q["size"] = 2
             queries.extend([{"index": ElasticIndexMapping.CANDID_BLOG_INDEX_ELSER}, q])
         elif index == "candid_learning":
-            q = build_text_expansion_query(
+            q = build_sparse_vector_query(
                 query=query,
                 fields=("content", "title", "training_topics", "staff_recommendations")
             )
@@ -132,7 +131,7 @@ def query_builder(query: str, indices: List[str]) -> List[Dict[str, Any]]:
             q["size"] = 2
             queries.extend([{"index": ElasticIndexMapping.CANDID_LEARNING_INDEX_ELSER}, q])
         elif index == "candid_help":
-            q = build_text_expansion_query(
+            q = build_sparse_vector_query(
                 query=query,
                 fields=("content", "combined_article_description")
             )
@@ -311,7 +310,8 @@ def get_results(user_input: str, indices: List[str]) -> Tuple[str, List[Document
     return content, output
 
 
-def get_context(field_name: str, hit: ElasticHitsResult, context_length: int = 1024) -> str:
+# TODO make it better!
+def get_context(field_name: str, hit: ElasticHitsResult, context_length: int = 1024, add_context: bool = True) -> str:
     """Pads the relevant chunk of text with context before and after
 
     Parameters
@@ -328,8 +328,10 @@ def get_context(field_name: str, hit: ElasticHitsResult, context_length: int = 1
         longer chunks stuffed together
     """
 
-    chunks_with_context = []
+    chunks = []
+    # TODO chunks have tokens, but long text is a normal text, but may contain html that also gets weird after tokenization
     long_text = hit.source.get(f"{field_name}", "")
+    long_text = long_text.lower()
     inner_hits_field = f"embeddings.{field_name}.chunks"
     found_chunks = hit.inner_hits.get(inner_hits_field, {})
     if found_chunks:
@@ -340,15 +342,17 @@ def get_context(field_name: str, hit: ElasticHitsResult, context_length: int = 1
             # cutting the middle because we may have tokenizing artifacts there
             chunk = chunk[3: -3]
 
-            # Find the start and end indices of the chunk in the large text
-            start_index = long_text.find(chunk)
-            if start_index != -1: # Chunk is found
-                end_index = start_index + len(chunk)
-                pre_start_index = max(0, start_index - context_length)
-                post_end_index = min(len(long_text), end_index + context_length)
-                chunks_with_context.append(long_text[pre_start_index:post_end_index])
-
-    return '\n\n'.join(chunks_with_context)
+            if add_context:
+                # Find the start and end indices of the chunk in the large text
+                start_index = long_text.find(chunk[:20])
+                if start_index != -1: # Chunk is found
+                    end_index = start_index + len(chunk)
+                    pre_start_index = max(0, start_index - context_length)
+                    post_end_index = min(len(long_text), end_index + context_length)
+                    chunks.append(long_text[pre_start_index:post_end_index])
+            else:
+                chunks.append(chunk)
+    return '\n\n'.join(chunks)
 
 
 def process_hit(hit: ElasticHitsResult) -> Union[Document, None]:
@@ -400,10 +404,12 @@ def process_hit(hit: ElasticHitsResult) -> Union[Document, None]:
     elif "candid-blog" in hit.index:
         excerpt = hit.source.get("excerpt", "")
         title = hit.source.get("title", "")
-        # we only need to process long texts
-        content_with_context_txt = get_context("content", hit, context_length=12)
+        # we only need to process long text
+        content_with_context_txt = get_context("content", hit, context_length=12, add_context=False)
+        authors = get_context("authors_text", hit, context_length=12, add_context=False)
+        tags = hit.source.get("title_summary_tags", "")
         doc = Document(
-            page_content='\n\n'.join([title, excerpt, content_with_context_txt]),
+            page_content='\n\n'.join([title, excerpt, content_with_context_txt, authors, tags]),
             metadata={
                 "title": title,
                 "source": "Candid Blog",