Added tool to find title of article being refered to by query with vector search of content.

functions/tools.py

@@ -6,6 +6,7 @@ import time
 import json
 import logging
 import queue
+from typing import Tuple
 from upstash_vector import Index
 
 import functions.feed_extraction as extraction_funcs
@@ -85,7 +86,7 @@ def get_feed(website: str, n: int = 3) -> list:
     return json.dumps(articles)
 
 
-def context_search(query: str, article_title: str = None) -> str:
+def context_search(query: str, article_title: str = None) -> list[Tuple[float, str]]:
     '''Searches for context relevant to query in article vector store. 
     Use this Function to search for additional information before 
     answering the user's question about an article. If article_title is
@@ -99,7 +100,7 @@ def context_search(query: str, article_title: str = None) -> str:
         context from a specific context, defaults to None
             
     Returns:
-        Context information which bests matches the query.
+        List of tuples with the following format: [(relevance score, 'context string')]
     '''
 
     logger = logging.getLogger(__name__ + 'context_search')
@@ -114,11 +115,54 @@ def context_search(query: str, article_title: str = None) -> str:
     results = index.query(
         data=query,
         top_k=3,
-        include_metadata=True,
         include_data=True,
         namespace=article_title
     )
 
     logger.info('Retrieved %s chunks for "%s"', len(results), query)
 
-    return results
+    contexts = []
+
+    for result in results:
+        contexts.append((result.score, result.data))
+
+    return contexts
+
+
+def find_article(query: str) -> list[Tuple[float, str]]:
+    '''Uses vector search to find the most likely title of the article 
+    referred to by query. Use this function if the user is asking about
+    and article, but it is unclear which article.
+    
+    Args:
+        query: query to to find source article for
+        
+    Returns:
+        List of tuples of most likely context article titles in the following format:
+        [(relevance score, 'article title')]
+    '''
+
+    logger = logging.getLogger(__name__ + 'context_search')
+
+    index = Index(
+        url='https://living-whale-89944-us1-vector.upstash.io',
+        token=os.environ['UPSTASH_VECTOR_KEY']
+    )
+
+    results = None
+
+    results = index.query(
+        data=query,
+        top_k=3,
+        include_metadata=True,
+        include_data=True
+    )
+
+    logger.info('Retrieved %s chunks for "%s"', len(results), query)
+
+    contexts = []
+
+    for result in results:
+        contexts.append((result.score, result.metadata['namespace']))
+
+    return contexts

rss_server.py

@@ -71,20 +71,56 @@ with gr.Blocks() as demo:
 
     # Vector search tool
     gr.Markdown('### 2. `context_search()`')
-    search_query = gr.Textbox('Does apple offer parental controls?', label='Vector search query')
-    search_output = gr.Textbox(label='Vector search results', lines=7, max_lines=7)
+
+    context_search_query = gr.Textbox(
+        'Does apple offer parental controls?',
+        label='Context search query'
+    )
+    context_search_output = gr.Textbox(
+        label='Context search results',
+        lines=7,
+        max_lines=7
+    )
 
     with gr.Row():
-        search_submit_button = gr.Button('Submit query')
-        search_clear_button = gr.ClearButton(components=[search_query, search_output])
+        context_search_submit_button = gr.Button('Submit query')
+        context_search_clear_button = gr.ClearButton(
+            components=[context_search_query, context_search_output]
+        )
 
-    search_submit_button.click( # pylint: disable=no-member
+    context_search_submit_button.click( # pylint: disable=no-member
         fn=tool_funcs.context_search,
-        inputs=search_query,
-        outputs=search_output,
+        inputs=context_search_query,
+        outputs=context_search_output,
         api_name='Context vector search'
     )
 
+    # Find article tool
+    gr.Markdown('### 3. `find_article()`')
+
+    article_search_query = gr.Textbox(
+        'Does apple offer parental controls?',
+        label='Article search query'
+    )
+    article_search_output = gr.Textbox(
+        label='Article search results',
+        lines=7,
+        max_lines=7
+    )
+
+    with gr.Row():
+        article_search_submit_button = gr.Button('Submit query')
+        article_search_clear_button = gr.ClearButton(
+            components=[article_search_query, article_search_output]
+        )
+
+    article_search_submit_button.click( # pylint: disable=no-member
+        fn=tool_funcs.find_article,
+        inputs=article_search_query,
+        outputs=article_search_output,
+        api_name='Article vector search'
+    )
+
 
 if __name__ == '__main__':