Completed vector context search tool.

functions/rag.py

@@ -5,14 +5,14 @@ import logging
 import queue
 from semantic_text_splitter import TextSplitter
 from tokenizers import Tokenizer
-from upstash_vector import Index, Vector
+from upstash_vector import Index
 
 
 def ingest(rag_ingest_queue: queue.Queue) -> None:
     '''Semantically chunks article and upsert to Upstash vector db
     using article title as namespace.'''
 
-    logger = logging.getLevelName(__name__ + '.ingest()')
+    logger = logging.getLogger(__name__ + '.ingest()')
 
     index = Index(
         url='https://living-whale-89944-us1-vector.upstash.io',
@@ -24,27 +24,37 @@ def ingest(rag_ingest_queue: queue.Queue) -> None:
         namespaces = index.list_namespaces()
 
         item = rag_ingest_queue.get()
+        logger.info(item)
         title = item['title']
-        text = item['content']
-        logger.info('Got %s from RAG ingest queue', title)
 
         if title not in namespaces:
+            text = item['content']
+            logger.info('Got "%s" from RAG ingest queue', title)
 
             tokenizer=Tokenizer.from_pretrained('bert-base-uncased')
             splitter=TextSplitter.from_huggingface_tokenizer(tokenizer, 256)
             chunks=splitter.chunks(text)
 
             for i, chunk in enumerate(chunks):
+                # index.upsert(
+                #     vectors=[
+                #         Vector(
+                #             id=hash(f'{title}-{i}'),
+                #             data=chunk,
+                #         )
+                #     ],
+                #     namespace=title
+                # )
+
                 index.upsert(
-                    vectors=[
-                        Vector(
-                            id=hash(f'{title}-{i}'),
-                            data=chunk,
+                    [
+                        (
+                            hash(f'{title}-{i}'),
+                            chunk,
+                            {'namespace': title}
                         )
                     ],
-                    namespace=title
                 )
-
             logger.info('Ingested %s chunks into vector DB', i + 1)
 
         else:

functions/tools.py

@@ -59,7 +59,7 @@ def get_feed(website: str) -> list:
         if item['content'] is not None:
 
             RAG_INGEST_QUEUE.put(item)
-            logger.info('%s sent to RAG ingest', item['title'])
+            logger.info('"%s" sent to RAG ingest', item['title'])
 
             summary = summarization_funcs.summarize_content(
                 item['title'],
@@ -67,7 +67,7 @@ def get_feed(website: str) -> list:
             )
 
             content[i]['summary'] = summary
-            logger.info('Summary of %s generated', item['title'])
+            logger.info('Summary of "%s" generated', item['title'])
 
         content[i].pop('content', None)
 
@@ -96,8 +96,10 @@ def context_search(query: str, article_title: str = None) -> str:
     results = None
 
     results = index.query(
-        [query],
+        data=query,
         top_k=3,
+        include_metadata=True,
+        include_data=True,
         namespace=article_title
     )
 

rss_server.py

@@ -40,7 +40,7 @@ with gr.Blocks() as demo:
     gr.HTML(html.DESCRIPTION)
 
     # Log output
-    dialog_output = gr.Textbox(label='Server logs', lines=10, max_lines=100)
+    dialog_output = gr.Textbox(label='Server logs', lines=10, max_lines=10)
     timer = gr.Timer(0.5, active=True)
 
     timer.tick( # pylint: disable=no-member
@@ -51,16 +51,28 @@ with gr.Blocks() as demo:
 
     # Get feed tool
     website_url = gr.Textbox('hackernews.com', label='Website')
-    output = gr.Textbox(label='RSS entries', lines=10)
+    feed_output = gr.Textbox(label='RSS entries', lines=10, max_lines=10)
     submit_button = gr.Button('Submit')
 
     submit_button.click( # pylint: disable=no-member
         fn=tool_funcs.get_feed,
         inputs=website_url,
-        outputs=output,
+        outputs=feed_output,
         api_name='Get RSS feed content'
     )
 
+    # Vector search tool
+    search_query = gr.Textbox('Does apple offer parental controls?', label='Vector search query')
+    search_output = gr.Textbox(label='Vector search results', lines=10, max_lines=10)
+    submit_button = gr.Button('Submit')
+
+    submit_button.click( # pylint: disable=no-member
+        fn=tool_funcs.context_search,
+        inputs=search_query,
+        outputs=search_output,
+        api_name='Context vector search'
+    )
+
 
 if __name__ == '__main__':