handle duplicate citations

api/src/stampy_chat/get_blocks.py

@@ -1,13 +1,12 @@
 import dataclasses
 import datetime
-import itertools
 import numpy as np
 import openai
 import regex as re
 import requests
 import time
 from itertools import groupby
-from typing import Iterable, List, Tuple
+from typing import Iterable, List
 from stampy_chat.env import PINECONE_NAMESPACE, REMOTE_CHAT_INSTANCE, EMBEDDING_MODEL
 from stampy_chat import logging
 
@@ -54,12 +53,12 @@ def parse_block(match) -> Block:
 
     date = metadata.get('date_published') or metadata.get('date')
 
-    if isinstance(date, datetime.date):
-        date = date.isoformat()
-    elif isinstance(date, datetime.datetime):
+    if isinstance(date, datetime.datetime):
         date = date.date().isoformat()
+    elif isinstance(date, datetime.date):
+        date = date.isoformat()
     elif isinstance(date, (int, float)):
-        date = datetime.datetime.fromtimestamp(date).isoformat()
+        date = datetime.datetime.fromtimestamp(date).date().isoformat()
 
     authors = metadata.get('authors')
     if not authors and metadata.get('author'):
@@ -82,16 +81,22 @@ def join_blocks(blocks: Iterable[Block]) -> List[Block]:
     # that the combined block has the minimum index of the blocks combined.
 
     def to_tuple(block):
-        return (block.id, block.title or "", block.authors or [], block.date or "", block.url or "", block.tags or "")
+        return (block.title or "", block.authors or [], block.date or "", block.url or "", block.tags or "")
 
     def merge_texts(blocks):
         return "\n.....\n".join(sorted(block.text for block in blocks))
 
-    unified_blocks = [
-        Block(*key, merge_texts(group))
-        for key, group in groupby(blocks, key=to_tuple)
-    ]
-    return sorted(unified_blocks, key=to_tuple)
+    # There are sometimes duplicates in the dataset, but which have different ids, so the id
+    # is ignored when making sorting the blocks.
+    def make_block(key, group):
+        group = list(group)
+        # Just use the id of the first item - it doesn't matter that much in this case, as the other data points
+        # will be the same
+        return Block(group[0].id, *key, merge_texts(group))
+
+    blocks = sorted(blocks, key=to_tuple)
+    blocks = [make_block(key, group) for key, group in groupby(blocks, key=to_tuple)]
+    return blocks
 
 
 # Get the k blocks most semantically similar to the query using Pinecone.

api/tests/stampy_chat/test_get_blocks.py

@@ -8,22 +8,22 @@ from stampy_chat.get_blocks import Block, get_top_k_blocks, parse_block, join_bl
     ({}, {}),
 
     # Check dates
-    ({'date_published': '2023-01-02T03:04:05'}, {'date': '2023-01-02T03:04:05'}),
+    ({'date_published': '2023-01-01T03:04:05'}, {'date': '2023-01-01T03:04:05'}),
     (
         {'date_published': datetime.fromisoformat('2023-01-02T03:04:05')},
-        {'date': '2023-01-02T03:04:05'}
-    ),
-    (
-        {'date_published': datetime.fromisoformat('2023-01-02T03:04:05').date()},
         {'date': '2023-01-02'}
     ),
     (
-        {'date_published': datetime.fromisoformat('2023-01-02T03:04:05').timestamp()},
-        {'date': '2023-01-02T03:04:05'}
+        {'date_published': datetime.fromisoformat('2023-01-03T03:04:05').date()},
+        {'date': '2023-01-03'}
     ),
     (
-        {'date_published': int(datetime.fromisoformat('2023-01-02T03:04:05').timestamp())},
-        {'date': '2023-01-02T03:04:05'}
+        {'date_published': datetime.fromisoformat('2023-01-04T03:04:05').timestamp()},
+        {'date': '2023-01-04'}
+    ),
+    (
+        {'date_published': int(datetime.fromisoformat('2023-01-05T03:04:05').timestamp())},
+        {'date': '2023-01-05'}
     ),
 
     # Check authors
@@ -84,6 +84,20 @@ def test_parse_block(match_override, block_override):
             Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'),
         ]
     ),
+    (
+        [
+            Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-1'),
+            Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'),
+            Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-2'),
+            Block('id2', 'title2', ['author2'], 'date2', 'url2', 'tags2', 'text2'),
+            Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-3'),
+        ],
+        [
+            Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-1\n.....\ntext1-2\n.....\ntext1-3'),
+            Block('id2', 'title2', ['author2'], 'date2', 'url2', 'tags2', 'text2'),
+            Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'),
+        ]
+    ),
 ])
 def test_join_blocks(blocks, expected):
     assert list(join_blocks(blocks)) == expected