Added RAG functions.

functions/rag.py

@@ -0,0 +1,51 @@
+'''Collection of function for RAG on article texts.'''
+
+import os
+import logging
+import queue
+from semantic_text_splitter import TextSplitter
+from tokenizers import Tokenizer
+from upstash_vector import Index, Vector
+
+
+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()')
+
+    index = Index(
+        url='https://living-whale-89944-us1-vector.upstash.io',
+        token=os.environ['UPSTASH_VECTOR_KEY']
+    )
+
+    while True:
+
+        namespaces = index.list_namespaces()
+
+        item = rag_ingest_queue.get()
+        title = item['title']
+        text = item['content']
+        logger.info('Got %s from RAG ingest queue', title)
+
+        if title not in namespaces:
+
+            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
+                )
+
+            logger.info('Ingested %s chunks into vector DB', i + 1)
+
+        else:
+            logger.info('%s already in RAG namespace', title)

functions/tools.py

@@ -1,10 +1,25 @@
 '''Tool functions for MCP server'''
 
+import os
+import threading
 import time
 import json
 import logging
+import queue
+from upstash_vector import Index, Vector
+
 import functions.feed_extraction as extraction_funcs
 import functions.summarization as summarization_funcs
+import functions.rag as rag_funcs
+
+RAG_INGEST_QUEUE = queue.Queue()
+
+rag_ingest_thread = threading.Thread(
+    target=rag_funcs.ingest,
+    args=(RAG_INGEST_QUEUE,)
+)
+
+rag_ingest_thread.start()
 
 
 def get_feed(website: str) -> list:
@@ -38,18 +53,52 @@ def get_feed(website: str) -> list:
     content = extraction_funcs.parse_feed(feed_uri)
     logger.info('parse_feed() returned %s entries', len(list(content.keys())))
 
-    # Summarize each post in the feed
+    # Summarize each post in the feed and submit full text for RAG ingest
     for i, item in content.items():
 
         if item['content'] is not None:
+
+            RAG_INGEST_QUEUE.put(item)
+            logger.info('%s sent to RAG ingest', item['title'])
+
             summary = summarization_funcs.summarize_content(
                 item['title'],
                 item['content']
             )
+
             content[i]['summary'] = summary
+            logger.info('Summary of %s generated', item['title'])
 
         content[i].pop('content', None)
 
     logger.info('Completed in %s seconds', round(time.time()-start_time, 2))
 
     return json.dumps(content)
+
+
+def context_search(query: str, article_title: str = None) -> str:
+    '''Searches for context relevant to query in article vector store.
+    
+    Ags:
+        query: user query to find context for
+        article_title: optional, use this argument to search only for context
+            from a specific context
+            
+    Returns:
+        Context which bests matches query as string.
+    '''
+
+    index = Index(
+        url='https://living-whale-89944-us1-vector.upstash.io',
+        token=os.environ['UPSTASH_VECTOR_KEY']
+    )
+
+    results = None
+
+    results = index.query(
+        [query],
+        top_k=3,
+        namespace=article_title
+    )
+
+    return results