Lightweight demo using the Anthropic Python SDK to experiment with Claude's Search and Retrieval capabilities over a variety of knowledge bases (Elasticsearch, vector databases, web search, and Wikipedia).
from .client import ClientWithRetrieval
claude_retriever/__init__.py
anthropics/anthropic-retrieval-demo
from typing import Optional, Tuple
from anthropic import Anthropic, HUMAN_PROMPT, AI_PROMPT
from .searcher.types import SearchTool, SearchResult, Tool
import logging
import re
from utils import format_results_full

logger = logging.getLogger(__name__)

RETRIEVAL_PROMPT = """
You will be given a query by a human user. Your job is solely to gather information from an external knowledge base that would help the user answer the query. To gather this information, you have been equipped with a search engine tool that you can use to query the external knowledge base. Here is a description of the search engine tool: <tool_description>{description}</tool_description>

You can make a call to the search engine tool by inserting a query within <search_query> tags like so: <search_query>query</search_query>. You'll then get results back within <search_result></search_result> tags. After these results back within, reflect briefly inside <search_quality></search_quality> tags about whether all the results together provide enough information to help the user answer the query, or whether more information is needed.

Before beginning to research the query, first think for a moment inside <thinking></thinking> tags about what information is necessary to gather to create a well-informed answer. 

If the query is complex, you may need to decompose the query into multiple subqueries and execute them individually. Sometimes the search engine will return empty search results, or the search results may not contain the information you need. In such cases, feel free to search again with a different query.

Do not try to answer the query. Your only job is to gather relevant search results that will help the user answer the query.

Here is the query: <query>{query}</query> 
"""

ANSWER_PROMPT = """
<search_results>{results}</search_results> Using the search results provided within the <search_results></search_results> tags, please answer the following query <query>{query}</query>.
"""

class ClientWithRetrieval(Anthropic):

    def __init__(self, search_tool: Optional[SearchTool] = None, verbose: bool = True, *args, **kwargs):
        """
        Initializes the ClientWithRetrieval class.
        
        Parameters:
            search_tool (SearchTool): SearchTool object to handle searching
            verbose (bool): Whether to print verbose logging
            *args, **kwargs: Passed to superclass init
        """
        super().__init__(*args, **kwargs)
        self.search_tool = search_tool
        self.verbose = verbose
    
    def retrieve(self,
                       query: str,
                       model: str,
                       n_search_results_to_use: int = 3,
                       stop_sequences: list[str] = [HUMAN_PROMPT],
                       max_tokens_to_sample: int = 1000,
                       max_searches_to_try: int = 5,
                       temperature: float = 1.0) -> list[SearchResult]:
        """
        Main method to retrieve relevant search results for a query with a provided search tool.
        
        Constructs RETRIEVAL prompt with query and search tool description. 
        Keeps sampling Claude completions until stop sequence hit.
        Extracts search results and accumulates all raw results.
        
        Returns:
            list[SearchResult]: List of all raw search results
        """
        assert self.search_tool is not None, "SearchTool must be provided to use .retrieve()"

        description = self.search_tool.tool_description
        prompt = f"{HUMAN_PROMPT} {RETRIEVAL_PROMPT.format(query=query, description=description)}{AI_PROMPT}"
        token_budget = max_tokens_to_sample
        all_raw_search_results: list[SearchResult] = []
        for tries in range(max_searches_to_try):
            partial_completion = self.completions.create(prompt = prompt,
                                                     stop_sequences=stop_sequences + ['</search_query>'],
                                                     model=model,
                                                     max_tokens_to_sample = token_budget,
                                                     temperature = temperature)
            partial_completion, stop_reason, stop_seq = partial_completion.completion, partial_completion.stop_reason, partial_completion.stop # type: ignore
            logger.info(partial_completion)
            token_budget -= self.count_tokens(partial_completion)
            prompt += partial_completion
            if stop_reason == 'stop_sequence' and stop_seq == '</search_query>':
                logger.info(f'Attempting search number {tries}.')
                raw_search_results, formatted_search_results = self._search_query_stop(partial_completion, n_search_results_to_use)
                prompt += '</search_query>' + formatted_search_results
                all_raw_search_results += raw_search_results
            else:
                break
        return all_raw_search_results
    
    def answer_with_results(self, raw_search_results: list[str]|list[SearchResult], query: str, model: str, temperature: float, format_results: bool =False):
        """Generates an RAG response based on search results and a query. If format_results is True,
           formats the raw search results first. Set format_results to True if you are using this method standalone without retrieve().
        
        Returns:
            str: Claude's answer to the query
        """
        if isinstance(raw_search_results[0], str):
            search_results  = [SearchResult(content=s) for s in raw_search_results] # type: ignore

        if format_results:
            processed_search_results = [search_result.content.strip() for search_result in search_results] # type: ignore
            formatted_search_results = format_results_full(processed_search_results)
        else:
            formatted_search_results = raw_search_results
        
        prompt = f"{HUMAN_PROMPT} {ANSWER_PROMPT.format(query=query, results=formatted_search_results)}{AI_PROMPT}"
        
        answer = self.completions.create(
            prompt=prompt, 
            model=model, 
            temperature=temperature, 
            max_tokens_to_sample=1000
        ).completion
        
        return answer
    
    def completion_with_retrieval(self,
                                        query: str,
                                        model: str,
                                        n_search_results_to_use: int = 3,
                                        stop_sequences: list[str] = [HUMAN_PROMPT],
                                        max_tokens_to_sample: int = 1000,
                                        max_searches_to_try: int = 5,
                                        temperature: float = 1.0) -> str:
        """
        Gets a final completion from retrieval results        
        
        Calls retrieve() to get search results.
        Calls answer_with_results() with search results and query.
        
        Returns:
            str: Claude's answer to the query
        """
        search_results = self.retrieve(query, model=model,
                                                 n_search_results_to_use=n_search_results_to_use, stop_sequences=stop_sequences,
                                                 max_tokens_to_sample=max_tokens_to_sample,
                                                 max_searches_to_try=max_searches_to_try,
                                                 temperature=temperature)
        answer = self.answer_with_results(search_results, query, model, temperature)
        return answer
    
    # Helper methods
    def _search_query_stop(self, partial_completion: str, n_search_results_to_use: int) -> Tuple[list[SearchResult], str]:
        """
        Helper to handle search query stop case.
        
        Extracts search query from completion text.
        Runs search using SearchTool. 
        Formats search results.
        
        Returns:
            tuple: 
                list[SearchResult]: Raw search results
                str: Formatted search result text
        """
        assert self.search_tool is not None, "SearchTool was not provided for client"

        search_query = self.extract_between_tags('search_query', partial_completion + '</search_query>') 
        if search_query is None:
            raise Exception(f'Completion with retrieval failed as partial completion returned mismatched <search_query> tags.')
        if self.verbose:
            logger.info('\n'+'-'*20 + f'\nPausing stream because Claude has issued a query in <search_query> tags: <search_query>{search_query}</search_query>\n' + '-'*20)
        logger.info(f'Running search query against SearchTool: {search_query}')
        search_results = self.search_tool.raw_search(search_query, n_search_results_to_use)
        extracted_search_results = self.search_tool.process_raw_search_results(search_results)
        formatted_search_results = format_results_full(extracted_search_results)

        if self.verbose:
            logger.info('\n' + '-'*20 + f'\nThe SearchTool has returned the following search results:\n\n{formatted_search_results}\n\n' + '-'*20 + '\n')
        return search_results, formatted_search_results
    
    def extract_between_tags(self, tag, string, strip=True):
        """
        Helper to extract text between XML tags.
        
        Finds last match of specified tags in string.
        Handles edge cases and stripping.
        
        Returns:
            str: Extracted string between tags
        """
        ext_list = re.findall(f"<{tag}\\s?>(.+?)</{tag}\\s?>", string, re.DOTALL)
        if strip:
            ext_list = [e.strip() for e in ext_list]
        
        if ext_list:
            return ext_list[-1]
        else:
            return None

claude_retriever/client.py
anthropics/anthropic-retrieval-demo
DEFAULT_EMBEDDER = 'sentence-transformers/all-mpnet-base-v2'
DEFAULT_SPARSE_EMBEDDER = 'naver/splade-cocondenser-ensembledistil'
claude_retriever/constants.py
anthropics/anthropic-retrieval-demo
from claude_retriever.searcher.types import Embedding, Embedder
from tenacity import retry, wait_exponential, stop_after_attempt
import requests
import json

class HuggingFaceEmbedder(Embedder):

    def __init__(self, api_key: str, model_name: str):

        self.api_key = api_key
        self.model_name = model_name

        self.url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{self.model_name}"
        self.headers = {"Authorization": f"Bearer {self.api_key}"}
        
        config_url = f'https://huggingface.co/{model_name}/resolve/main/config.json'
        response = requests.get(config_url)
        if response.status_code == 200:
            config = json.loads(response.text)
            embedding_dimension = config["hidden_size"]
        else:
            raise RuntimeError(
                f"Could not get config for model {model_name}"
            )

        self.dim = embedding_dimension

    def embed(self, text: str) -> Embedding:
        emb = self.embed_batch([text])
        return emb[0]
    
    @retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(10))
    def embed_batch(self, texts: list[str]) -> list[Embedding]:
        response = requests.post(
            self.url,
            headers=self.headers,
            json={"inputs": texts})
        result = response.json()
        embeddings = []
        if isinstance(result, list):
            embeddings = result
        elif list(result.keys())[0] == "error":
            raise RuntimeError(
                "The model is currently loading, please re-run the text."
            )
        return [Embedding(embedding=embedding, text=text) for embedding, text in zip(embeddings, texts)]

claude_retriever/searcher/embedders/huggingface.py
anthropics/anthropic-retrieval-demo
from claude_retriever.searcher.types import Embedding, SparseEmbeddingData, HybridEmbedding, Embedder
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForMaskedLM, AutoTokenizer
import numpy as np
import torch
from typing import cast


class LocalEmbedder(Embedder):

    def __init__(self, model_name: str):

        self.model_name = model_name

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = SentenceTransformer(model_name, device=self.device)
        
        dim = self.model.get_sentence_embedding_dimension()
        self.dim = cast(int, dim)

    def embed(self, text: str) -> Embedding:
        emb = self.embed_batch([text])
        return emb[0]

    def embed_batch(self, texts: list[str]) -> list[Embedding]:
        embeddings = self.model.encode(texts, show_progress_bar=False)
        assert isinstance(embeddings, np.ndarray)
        embeddings = [embedding.tolist() for embedding in embeddings]
        return [Embedding(embedding=embedding, text=text) for embedding, text in zip(embeddings, texts)]

class LocalHybridEmbedder(Embedder):

    def __init__(self, dense_model_name: str, sparse_model_name: str):

        self.dense_model_name = dense_model_name
        self.sparse_model_name = sparse_model_name

        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        self.dense_model = SentenceTransformer(dense_model_name, device=self.device)
        dense_dim = self.dense_model.get_sentence_embedding_dimension()
        self.dense_dim = cast(int, dense_dim)
        
        self.sparse_tokenizer = AutoTokenizer.from_pretrained(self.sparse_model_name)
        self.sparse_model = AutoModelForMaskedLM.from_pretrained(self.sparse_model_name)
        self.sparse_model.to(self.device)
        self.sparse_dim = self.sparse_model.config.vocab_size

    def _sparse_encode(self, texts: list[str]) -> list[SparseEmbeddingData]:
        with torch.no_grad():
            tokens = self.sparse_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512).to(self.device)
            vec = torch.max(
            torch.log(
                1 + torch.relu(self.sparse_model(**tokens).logits)
            ) * tokens.attention_mask.unsqueeze(-1),
            dim=1).values.cpu()
            indices = [v.nonzero().flatten() for v in vec]
            weights = [v[idx] for v, idx in zip(vec, indices)]
            sparse_embeddings = [SparseEmbeddingData(
                indices=c.tolist(),
                values=w.tolist(),
                max_index=self.sparse_dim
            ) for c, w in zip(indices, weights)]
        return sparse_embeddings

    def embed(self, text: str) -> HybridEmbedding:
        emb = self.embed_batch([text])
        return emb[0]

    def embed_batch(self, texts: list[str]) -> list[HybridEmbedding]:
        # Get dense embeddings
        dense_embeddings = self.dense_model.encode(texts, show_progress_bar=False)
        assert isinstance(dense_embeddings, np.ndarray)
        dense_embeddings = [embedding.tolist() for embedding in dense_embeddings]
        # Get sparse embeddings
        sparse_embeddings = self._sparse_encode(texts)
        assert isinstance(sparse_embeddings[0], SparseEmbeddingData)
        # Combine dense and sparse embeddings
        return [HybridEmbedding(embedding=embedding,
                                sparse_embedding=sparse_embedding,
                                text=text)
                                for embedding, sparse_embedding, text
                                in zip(dense_embeddings, sparse_embeddings, texts)]
claude_retriever/searcher/embedders/local.py
anthropics/anthropic-retrieval-demo
from typing import Optional
from claude_retriever.searcher.types import SearchResult, SearchTool
from anthropic import Anthropic
from elasticsearch import Elasticsearch

import logging
logger = logging.getLogger(__name__)

# Elasticsearch Cloud Searcher

class ElasticsearchCloudSearchTool(SearchTool):

    def __init__(self,
                tool_description: str,
                elasticsearch_cloud_id: str,
                elasticsearch_api_key_id: str,
                elasticsearch_api_key: str,
                elasticsearch_index: str,
                truncate_to_n_tokens: Optional[int] = 5000):
        
        self.index = elasticsearch_index
        self.cloud_id = elasticsearch_cloud_id
        self.api_key_id = elasticsearch_api_key_id
        self.api_key = elasticsearch_api_key
        self._connect_to_elasticsearch()

        self.tool_description = tool_description
        self.truncate_to_n_tokens = truncate_to_n_tokens
        if truncate_to_n_tokens is not None:
            self.tokenizer = Anthropic().get_tokenizer() 
    
    def _connect_to_elasticsearch(self):
        self.client = Elasticsearch(
            cloud_id=self.cloud_id,
            api_key=(self.api_key_id, self.api_key)
        )
        if not self.client.indices.exists(index=self.index):
            raise ValueError(f"Elasticsearch Index {self.index} does not exist.")
        index_mapping = self.client.indices.get_mapping(index=self.index)
        if "text" not in index_mapping.body[self.index]["mappings"]["properties"].keys():
            raise ValueError(f"Index {self.index} does not have a field called 'text'.")
    
    def truncate_page_content(self, page_content: str) -> str:
        if self.truncate_to_n_tokens is None:
            return page_content.strip()
        else:
            return self.tokenizer.decode(self.tokenizer.encode(page_content).ids[:self.truncate_to_n_tokens]).strip()

    def raw_search(self, query: str, n_search_results_to_use: int) -> list[SearchResult]:

        results = self.client.search(index=self.index,
                                     query={"match": {"text": query}})
        search_results: list[SearchResult] = []
        for result in results["hits"]["hits"]:
            if len(search_results) >= n_search_results_to_use:
                break
            content = result["_source"]["text"]
            search_results.append(SearchResult(content=content))

        return search_results
    
    def process_raw_search_results(self, results: list[SearchResult]) -> list[str]:
        processed_search_results = [self.truncate_page_content(result.content) for result in results]
        return processed_search_results
    
claude_retriever/searcher/searchtools/elasticsearch.py
anthropics/anthropic-retrieval-demo
from typing import Optional
from claude_retriever.searcher.types import Embedder, SearchResult, VectorStore, SearchTool
from claude_retriever.searcher.embedders.local import LocalEmbedder
from claude_retriever.constants import DEFAULT_EMBEDDER

import logging
logger = logging.getLogger(__name__)

# Embedding Searcher

class EmbeddingSearchTool(SearchTool):

    def __init__(self, tool_description: str, vector_store: VectorStore, embedder: Optional[Embedder] = None):
        self.tool_description = tool_description
        if embedder is None:
            logger.info(f"Using default embedder: {DEFAULT_EMBEDDER}")
            embedder = LocalEmbedder(DEFAULT_EMBEDDER)
        self.embedder = embedder
        self.vector_store = vector_store

    def raw_search(self, query: str, n_search_results_to_use: int) -> list[SearchResult]:
        query_embedding = self.embedder.embed(query)
        search_results = self.vector_store.query(query_embedding, n_search_results_to_use=n_search_results_to_use)
        return search_results
    
    def process_raw_search_results(self, results: list[SearchResult]) -> list[str]:
        processed_search_results = [result.content for result in results]
        return processed_search_results
    
claude_retriever/searcher/searchtools/embeddings.py
anthropics/anthropic-retrieval-demo
import os
from typing import Optional
from claude_retriever.searcher.types import SearchResult, SearchTool
from claude_retriever.utils import scrape_url
from dataclasses import dataclass
import requests
import asyncio
from tenacity import retry, wait_exponential, stop_after_attempt

import logging
logger = logging.getLogger(__name__)

@dataclass
class WebSearchResult(SearchResult):
    url: str

# Brave Searcher

BRAVE_DESCRIPTION = """Brave Search Engine Tool: The search engine will search using the Brave search engine for web pages with keywords similar to your query. It returns for each page its title, a summary and potentially the full page content. Use this tool if you want to get up-to-date and comprehensive information on a topic."""

class BraveAPI:
    def __init__(self, api_key: str):
        self.api_key = api_key

    @retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(10))
    def search(self, query: str) -> dict:
        headers = {"Accept": "application/json", "X-Subscription-Token": self.api_key}
        resp = requests.get(
            "https://api.search.brave.com/res/v1/web/search",
            params={"q": query,
                    "count": 20 # Max number of results to return, can filter down later
                    },
            headers=headers,
            timeout=60
        )
        if resp.status_code != 200:
            logger.error(f"Search request failed: {resp.text}")
            return {}
        return resp.json()

class BraveSearchTool(SearchTool):

    def __init__(self, brave_api_key: str,
                 tool_description: str = BRAVE_DESCRIPTION,
                 summarize_with_claude: bool = False,
                 anthropic_api_key: Optional[str] = None):
        """
        :param brave_api_key: The Brave API key to use for searching.
        :param tool_description: The description of the tool.
        :param summarize_with_claude: Whether to summarize the scraped web pages with Claude.
        :param anthropic_api_key: The anthropic API key to use for summarizing with Claude.
        """

        self.api = BraveAPI(brave_api_key)
        self.tool_description = tool_description
        self.summarize_with_claude = summarize_with_claude
        if summarize_with_claude and anthropic_api_key is None:
            try:
                anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]
            except KeyError:
                raise ValueError("If you want to summarize with Claude, you must provide an anthropic_api_key.")
        self.anthropic_api_key = anthropic_api_key

    def parse_faq(self, faq: dict) -> WebSearchResult:
        """
        https://api.search.brave.com/app/documentation/responses#FAQ
        """
        snippet = f"""FAQ Title: {faq.get('title', "Unknown")}
Question: {faq.get('question', "Unknown")}
Answer: {faq.get('answer', "Unknown")}"""
        
        return WebSearchResult(
            url=faq.get("url", ""),
            content=snippet
        )
    
    def parse_news(self, news_item: dict) -> Optional[WebSearchResult]:
        """
        https://api.search.brave.com/app/documentation/responses#News
        """
        article_description: str = news_item.get("description", "")

        # Throw out items where the description is tiny or doesn't exist.
        if len(article_description) < 5:
            return None

        snippet = f"""News Article Title: {news_item.get('title', "Unknown")}
News Article Description: {article_description}
News Article Age: {news_item.get("age", "Unknown")}
News Article Source: {news_item.get("meta_url", {}).get('hostname', "Unknown")}"""
        
        return WebSearchResult(
            url=news_item.get("url", ""),
            content=snippet
        )

    @staticmethod
    def remove_strong(web_description: str):
        # this is for cleaning up the brave web descriptions
        return (
            web_description.replace("<strong>", "")
            .replace("</strong>", "")
            .replace("&#x27;", "'")
        )
    
    async def parse_web(self, web_item: dict, query: str) -> WebSearchResult:
        """
        https://api.search.brave.com/app/documentation/responses#Search
        """
        url = web_item.get("url", "")
        title = web_item.get("title", "")
        description = self.remove_strong(web_item.get("description", ""))
        snippet = f"""Web Page Title: {title}
Web Page URL: {url}
Web Page Description: {description}"""

        try:
            # Currently there is no retry logic here, so if the scrape fails, we just skip it and return the snippet.
            if self.summarize_with_claude:
                content = await scrape_url(url, summarize_with_claude=True,
                                           anthropic_api_key=self.anthropic_api_key, query=query)
            else:
                content = await scrape_url(url)

            if content.startswith('<summary>'):
                snippet+="\nWeb Page Summary: "+content
            else:
                snippet+="\nWeb Page Content: "+content
        except:
            logger.warning(f"Failed to scrape {url}")
        return WebSearchResult(
            url=url,
            content=snippet
        )


    def raw_search(self, query: str, n_search_results_to_use: int) -> list[WebSearchResult]:
        """
        Run a search using the BraveAPI and return search results. Here are some details on the Brave API:

        Each search call to the Brave API returns the following fields:
         - faq: Frequently asked questions that are relevant to the search query (only on paid Brave tier).
         - news: News results relevant to the query.
         - web: Web search results relevant to the query.
         - [Thrown Out] videos: Videos relevant to the query.
         - [Thrown Out] locations: Places of interest (POIs) relevant to location sensitive queries.
         - [Thrown Out] infobox: Aggregated information on an entity showable as an infobox.
         - [Thrown Out] discussions: Discussions clusters aggregated from forum posts that are relevant to the query.

        There is also a `mixed` key, which tells us the ranking of the search results.

        We may throw some of these back in, in the future. But we're just going to document the behavior here for now.
        """
        
        # Run the search

        search_response = self.api.search(query)

        # Order everything properly

        correct_ordering = search_response.get("mixed", {}).get("main", [])

        # Extract the results

        faq_items = search_response.get("faq", {}).get("results", [])
        news_items = search_response.get("news", {}).get("results", [])
        web_items = search_response.get("web", {}).get("results", [])

        # Get the search results

        search_results: list[WebSearchResult] = []
        async_web_parser_loop = asyncio.get_event_loop()
        web_parsing_tasks = [] # We'll queue up the web parsing tasks here, since they're costly

        for item in correct_ordering:
            item_type = item.get("type")
            if item_type == "web":
                web_item = web_items.pop(0)
                ## We'll add a placeholder search result here, and then replace it with the parsed web result later
                url = web_item.get("url", "")
                placeholder_search_result = WebSearchResult(
                    url=url,
                    content=f"Web Page Title: {web_item.get('title', '')}\nWeb Page URL: {url}\nWeb Page Description: {self.remove_strong(web_item.get('description', ''))}"
                )
                search_results.append(placeholder_search_result)
                ## Queue up the web parsing task
                task = async_web_parser_loop.create_task(self.parse_web(web_item, query))
                web_parsing_tasks.append(task)
            elif item_type == "news":
                parsed_news = self.parse_news(news_items.pop(0))
                if parsed_news is not None:
                    search_results.append(parsed_news)
            elif item_type == "faq":
                parsed_faq = self.parse_faq(faq_items.pop(0))
                search_results.append(parsed_faq)
            if len(search_results) >= n_search_results_to_use:
                break

        ## Replace the placeholder search results with the parsed web results
        web_results = async_web_parser_loop.run_until_complete(asyncio.gather(*web_parsing_tasks))
        web_results_urls = [web_result.url for web_result in web_results]
        for i, search_result in enumerate(search_results):
            url = search_result.url
            if url in web_results_urls:
                search_results[i] = web_results[web_results_urls.index(url)]

        return search_results
    
    def process_raw_search_results(self, results: list[SearchResult]) -> list[str]:
        # We don't need to do any processing here, since we already formatted the results in the `parse` functions.
        processed_search_results = [result.content.strip() for result in results]
        return processed_search_results
claude_retriever/searcher/searchtools/websearch.py
anthropics/anthropic-retrieval-demo
from typing import Optional
from claude_retriever.searcher.types import SearchResult, SearchTool
from anthropic import Anthropic
import wikipedia
from dataclasses import dataclass

import logging
logger = logging.getLogger(__name__)

# Wikipedia Searcher

WIKIPEDIA_DESCRIPTION = """Wikipedia Search Engine Tool: 
The search engine will exclusively search over Wikipedia for pages similar to your query. It returns for each page its title and the full page content. Use this tool to get up-to-date and comprehensive information on a topic. Queries made to this tool should be as atomic as possible. The tool provides broad topic keywords rather than niche search topics. For example, if the query is "Can you tell me about Odysseus's journey in the Odyssey?" the search query you make should be "Odyssey". Here's another example: if the query is "Who created the first neural network?", your first query should be "neural network". As you can see, these queries are quite short. Think generalized keywords, not phrases. 
"""

@dataclass
class WikipediaSearchResult(SearchResult):
    title: str
    
class WikipediaSearchTool(SearchTool):

    def __init__(self,
                 tool_description: str = WIKIPEDIA_DESCRIPTION,
                 truncate_to_n_tokens: Optional[int] = 5000):
        self.tool_description = tool_description
        self.truncate_to_n_tokens = truncate_to_n_tokens
        if truncate_to_n_tokens is not None:
            self.tokenizer = Anthropic().get_tokenizer()

    def raw_search(self, query: str, n_search_results_to_use: int) -> list[WikipediaSearchResult]:
        search_results = self._search(query, n_search_results_to_use)
        return search_results
    
    def process_raw_search_results(self, results: list[WikipediaSearchResult]) -> list[str]:
        processed_search_results = [f'Page Title: {result.title.strip()}\nPage Content:\n{self.truncate_page_content(result.content)}' for result in results]
        return processed_search_results

    def truncate_page_content(self, page_content: str) -> str:
        if self.truncate_to_n_tokens is None:
            return page_content.strip()
        else:
            return self.tokenizer.decode(self.tokenizer.encode(page_content).ids[:self.truncate_to_n_tokens]).strip()
        
    def _search(self, query: str, n_search_results_to_use: int) -> list[WikipediaSearchResult]:
        results: list[str] = wikipedia.search(query)
        search_results: list[WikipediaSearchResult] = []
        for result in results:
            if len(search_results) >= n_search_results_to_use:
                break
            try:
                page = wikipedia.page(result)
            except:
                # The Wikipedia API is a little flaky, so we just skip over pages that fail to load
                continue
            content = page.content
            title = page.title
            search_results.append(WikipediaSearchResult(content=content, title=title))
        return search_results

claude_retriever/searcher/searchtools/wikipedia.py
anthropics/anthropic-retrieval-demo
from dataclasses import dataclass
from abc import ABC, abstractmethod

import sys
import os
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, parent_dir) 

#########################################################
## Embedder: Convert texts to embeddings
#########################################################

@dataclass
class Embedding:
    """
    An embedding of a text, along with the text itself and any metadata associated with it.
    """
    embedding: list[float]
    text: str

@dataclass
class SparseEmbeddingData:
    """
    A sparse embedding index, containing a list of indices and a list of values.
    """
    indices: list[int]
    values: list[float]
    max_index: int

@dataclass
class HybridEmbedding(Embedding):
    """
    A hybrid embedding, containing a dense embedding and a sparse embedding index.
    """
    sparse_embedding: SparseEmbeddingData

class Embedder(ABC):
    """
    An embedder that can embed a single text or a batch of texts.
    """
    dim: int
    
    @abstractmethod
    def embed(self, text: str) -> Embedding:
        """
        Embeds a single text.

        :param text: The text to embed.
        """
        raise NotImplementedError()
    
    @abstractmethod
    def embed_batch(self, texts: list[str]) -> list[Embedding]:
        """
        Embeds a batch of texts.

        :param texts: The texts to embed.
        """
        raise NotImplementedError()


#########################################################
## Search Tool: A wrapper around a searcher with instructions and formatting to help Claude use it
#########################################################

@dataclass
class SearchResult:
    """
    A single search result.
    """
    content: str

class Tool(ABC):
    tool_description: str

class SearchTool(Tool):
    """
    A search tool that can run a query and return a formatted string of search results.
    """

    def __init__(self, tool_description: str):
        self.tool_description = tool_description

    @abstractmethod
    def raw_search(self, query: str, n_search_results_to_use: int) -> list[SearchResult]:
        """
        Runs a query using the searcher, then returns the raw search results without formatting.

        :param query: The query to run.
        :param n_search_results_to_use: The number of results to return.
        """
        raise NotImplementedError()
    
    @abstractmethod
    def process_raw_search_results(
        self, results: list[SearchResult],
    ) -> list[str]:
        """
        Extracts the raw search content from the search results and returns a list of strings that can be passed to Claude.

        :param results: The search results to extract.
        """
        raise NotImplementedError()
    
    def search(self, query: str, n_search_results_to_use: int) -> str:
        from utils import format_results_full # Avoids circular import

        raw_search_results = self.raw_search(query, n_search_results_to_use)
        processed_search_results = self.process_raw_search_results(raw_search_results)
        displayable_search_results = format_results_full(processed_search_results)
        return displayable_search_results 



#########################################################
## VectorStore: An interface to a vector store that can upsert embeddings and run queries
#########################################################

class VectorStore(ABC):
    """
    An interface to a vector store that can upsert embeddings and run queries.
    """
    
    @abstractmethod
    def upsert(self, embeddings: list[Embedding]) -> None:
        """
        Upserts a list of embeddings into the vector store.

        :param embeddings: The embeddings to upsert.
        """
        raise NotImplementedError()

    @abstractmethod
    def query(self, query_embedding: Embedding, n_search_results_to_use: int = 10) -> list[SearchResult]:
        """
        Runs a query using the vector store and returns the results.

        :param query_embedding: The embedding to query with.
        :param n_search_results_to_use: The number of results to return.
        """
        raise NotImplementedError()
claude_retriever/searcher/types.py
anthropics/anthropic-retrieval-demo
import os
from claude_retriever.searcher.types import Embedding, HybridEmbedding, SparseEmbeddingData, SearchResult, VectorStore
from scipy.sparse import csr_matrix
import numpy as np
import json


import logging
logger = logging.getLogger(__name__)

############################################
# Local VectorStore implementations
############################################

class LocalVectorStore(VectorStore):
    '''
    Local vectorstores maintain a single embedding matrix.

    How it works:
    - On init, the datastore is loaded from disk if it exists. If it doesn't exist, then an empty datastore is created.
    - When upserting embeddings, the embeddings are written to disk and the datastore is reloaded into memeory.
    -- On disk, the embeddings are stored in JSONL format, where each line has the embedding and text data for an embedding.
    -- In memory, the embeddings are stored as a numpy array.
    - When querying, the query embedding is compared to all embeddings in the datastore using cosine similarity.

    Note that the vectorstore does not contain any logic for creating embeddings. It is assumed that the embeddings are created elsewhere
    using Embedders and passed to the vectorstore for storage and retrieval. The utils.embed_and_upload() is a wrapper to help do this.
    '''
    def __init__(self, disk_path: str):
        if not disk_path.endswith(".jsonl"):
            raise ValueError("disk_path must be .jsonl")
        self.disk_path = disk_path
        self.embeddings = []
        # if there is no file at disk_path, then create an empty file
        if not os.path.exists(disk_path):
            logger.info(f"Creating empty datastore at {disk_path}")
            open(disk_path, "w").close()
        logger.info(f"Loading datastore from {disk_path}")
        self._load_data(disk_path) # If the file is empty, then will load empty embeddings

    def _create_normed_embedding_matrix(self):
        self.embedding_matrix = np.array([emb.embedding for emb in self.embeddings])
        if len(self.embedding_matrix) > 0:
            self.embedding_norms = np.linalg.norm(self.embedding_matrix, axis=1)
        else:
            self.embedding_norms = np.array([])

    def _load_data(self, disk_path: str):
        file_type = disk_path.split(".")[-1]
        if file_type == "jsonl":
            with open(disk_path, "r") as f:
                for line in f:
                    data = json.loads(line)
                    text = data["text"]
                    embedding = data["embedding"]
                    self.embeddings.append(Embedding(embedding, text))
        else:
            raise ValueError("Invalid file_type. Supported types: 'jsonl'")
        self._create_normed_embedding_matrix()

    def _nearest_neighbors(self, query_embedding: Embedding, n_search_results_to_use: int) -> list[Embedding]:
        query_vector = query_embedding.embedding
        query_vector = np.array(query_vector)
        query_vector_norm = np.linalg.norm(query_vector)
        cosine_similarities = np.dot(self.embedding_matrix, query_vector) / (self.embedding_norms * query_vector_norm)
        sorted_indices = np.argsort(-cosine_similarities)
        nearest_neighbors = [self.embeddings[i] for i in sorted_indices[:min(n_search_results_to_use, len(self.embeddings))]]
        return nearest_neighbors

    def query(self, query_embedding: Embedding, n_search_results_to_use: int) -> list[SearchResult]:
        if len(self.embeddings) == 0:
            logger.warning("No embeddings in datastore. Returning empty list.")
            return []
        best_matches = self._nearest_neighbors(query_embedding, n_search_results_to_use)
        results = [SearchResult(content=match.text) for match in best_matches]
        return results
    
    def upsert(self, embeddings: list[Embedding]) -> None:
        with open(self.disk_path, "a") as file:
            for item in embeddings:
                output_data = {"embedding": item.embedding, "text": item.text}
                json_str = json.dumps(output_data)
                file.write(json_str + "\n")
        self._load_data(self.disk_path) # Once the embeddings have been written to disk, reload the datastore

class LocalHybridVectorStore(VectorStore):

    '''
    Hybrid vectorstores maintain two embedding matrices: one for dense embeddings and one for sparse embeddings.

    How it works:
    - On init, the datastore is loaded from disk if it exists. If it doesn't exist, then an empty datastore is created.
    - When upserting embeddings, the embeddings are written to disk and the datastore is reloaded into memeory.
    -- On disk, the embeddings are stored in JSONL format, where each line has the dense, sparse and text data for an embedding.
    -- In memory, the embeddings are stored as two separate matrices: one for dense embeddings and one for sparse embeddings. The latter is stored as a scipy.sparse.csr_matrix.
    - When querying, the query embedding is compared to all dense embeddings and all sparse embeddings in the datastore. They are combined using either:
    -- split_return: half of the expected number of results are returned from the top dense matches and half from the top sparse matches (after deduplication).
    -- weighted: the top n_search_results_to_use results are returned from the weighted combination of dense and sparse cosine similarities.

    Note that the vectorstore does not contain any logic for creating embeddings. It is assumed that the embeddings are created elsewhere 
    using Embedders and passed to the vectorstore for storage and retrieval. The utils.embed_and_upload() is a wrapper to help do this.
    '''

    def __init__(self, disk_path: str):
        if not disk_path.endswith(".jsonl"):
            raise ValueError("disk_path must be .jsonl")
        self.disk_path = disk_path
        self.embeddings = []
        # if there is no file at disk_path, then create an empty file
        if not os.path.exists(disk_path):
            logger.info(f"Creating empty datastore at {disk_path}")
            open(disk_path, "w").close()
        logger.info(f"Loading datastore from {disk_path}")
        self._load_data(disk_path)

    @staticmethod
    def sparse_embedding_data_to_sparse_matrix(sparse_embedding_index: list[SparseEmbeddingData]) -> csr_matrix:
        """
        Converts a list of SparseEmbeddingData to a sparse matrix.

        The sparse matrix format is a scipy.sparse.csr_matrix, which is a sparse matrix stored in Compressed Sparse Row format.
        """
        if not sparse_embedding_index:
            return csr_matrix((0, 0))

        # Check that all SparseEmbeddingData objects have the same max_index
        max_index = sparse_embedding_index[0].max_index
        for sparse_embedding in sparse_embedding_index:
            if sparse_embedding.max_index != max_index:
                raise ValueError("SparseEmbeddingData objects must have the same max_index")

        # Create sparse matrix
        rows, cols, data = [], [], []
        for idx, sparse_embedding in enumerate(sparse_embedding_index):
            rows.extend([idx] * len(sparse_embedding.indices))
            cols.extend(sparse_embedding.indices)
            data.extend(sparse_embedding.values)

        sparse_matrix = csr_matrix((data, (rows, cols)), shape=(len(sparse_embedding_index), max_index + 1))
        return sparse_matrix

    def _create_normed_embedding_matrices(self):

        '''
        Hybrid vectorstores maintain two embedding matrices: one for dense embeddings and one for sparse embeddings.
        '''

        self.dense_embedding_matrix = np.array([emb.embedding for emb in self.embeddings])
        if len(self.dense_embedding_matrix) > 0:
            self.dense_embedding_norms = np.linalg.norm(self.dense_embedding_matrix, axis=1)
        else:
            self.dense_embedding_norms = np.array([])

        self.sparse_embedding_matrix = self.sparse_embedding_data_to_sparse_matrix([emb.sparse_embedding for emb in self.embeddings])
        if self.sparse_embedding_matrix.shape[0] > 0:
            self.sparse_embedding_norms = np.linalg.norm(self.sparse_embedding_matrix.toarray(), axis=1)
        else:
            self.sparse_embedding_norms = np.array([])


    def _load_data(self, disk_path: str):
        file_type = disk_path.split(".")[-1]
        if file_type == "jsonl":
            with open(disk_path, "r") as f:
                for (i, line) in enumerate(f):
                    data = json.loads(line)
                    if "sparse_embedding" not in data.keys():
                        raise ValueError(f"Invalid jsonl file. 'sparse_embedding' key is missing on line {i}.")
                    text = data["text"]
                    dense_embedding = data["embedding"]
                    sparse_embedding_idx = data["sparse_embedding"]
                    expected_keys = {"indices", "values", "max_index"}
                    if not expected_keys.issubset(set(sparse_embedding_idx.keys())):
                        raise ValueError(f"Invalid jsonl file. Some of expected keys {expected_keys} are missing from sparse_embedding on line {i}.")
                    sparse_embedding = SparseEmbeddingData(sparse_embedding_idx["indices"],
                                                           sparse_embedding_idx["values"],
                                                           sparse_embedding_idx["max_index"])
                    self.embeddings.append(HybridEmbedding(embedding=dense_embedding, sparse_embedding=sparse_embedding, text=text))
        else:
            raise ValueError("Invalid file_type. Supported types: 'jsonl'")
        self._create_normed_embedding_matrices()

    def _get_dense_cosine_similarities(self, query_embedding: HybridEmbedding) -> np.ndarray:

        '''
        This method returns the cosine similarities between the query embedding and all dense embeddings in the datastore.
        '''

        query_dense_vector = query_embedding.embedding
        query_dense_vector = np.array(query_dense_vector)
        query_dense_vector_norm = np.linalg.norm(query_dense_vector)
        dense_inner_prod = np.dot(self.dense_embedding_matrix, query_dense_vector)
        dense_cosine_similarities = dense_inner_prod / (self.dense_embedding_norms * query_dense_vector_norm)
        return dense_cosine_similarities
    
    def _get_sparse_cosine_similarities(self, query_embedding: HybridEmbedding) -> np.ndarray:

        '''
        This method returns the cosine similarities between the query embedding and all sparse embeddings in the datastore.
        '''

        query_sparse_vector = self.sparse_embedding_data_to_sparse_matrix([query_embedding.sparse_embedding])
        query_sparse_vector_norm = np.linalg.norm(query_sparse_vector.toarray())
        sparse_inner_prod = np.dot(self.sparse_embedding_matrix, query_sparse_vector.T).T.toarray().flatten()
        sparse_cosine_similarities = sparse_inner_prod / (self.sparse_embedding_norms * query_sparse_vector_norm)
        return sparse_cosine_similarities

    def _query_split_return(self, query_embedding: HybridEmbedding, n_search_results_to_use: int) -> list[HybridEmbedding]:

        '''
        This method returns half of the expected number of results from the top dense matches and half from the top sparse matches.
        '''

        n_dense_results = int(n_search_results_to_use / 2)
        n_sparse_results = n_search_results_to_use - n_dense_results

        dense_cosine_results = self._get_dense_cosine_similarities(query_embedding)
        sparse_cosine_results = self._get_sparse_cosine_similarities(query_embedding)

        dense_sorted_indices = np.argsort(-dense_cosine_results)
        sparse_sorted_indices = np.argsort(-sparse_cosine_results)
        
        sorted_indices = dense_sorted_indices[:min(n_dense_results, len(self.embeddings))]
        sparse_sorted_indices = [i for i in sparse_sorted_indices if i not in sorted_indices] # remove the already chosen dense results
        sorted_indices = np.concatenate((sorted_indices, sparse_sorted_indices[:min(n_sparse_results, len(sparse_sorted_indices))]))

        nearest_neighbors = [self.embeddings[i] for i in sorted_indices]
        return nearest_neighbors
    
    def _query_weighted(self, query_embedding: HybridEmbedding, n_search_results_to_use: int, sparse_weight: float = 0.5) -> list[HybridEmbedding]:

        '''
        This method returns the top n_search_results_to_use results from the weighted combination of dense and sparse cosine similarities.
        '''

        if sparse_weight < 0.0 or sparse_weight > 1.0:
            raise ValueError("sparse_weight must be between 0.0 and 1.0")
        
        dense_cosine_results = self._get_dense_cosine_similarities(query_embedding)
        sparse_cosine_results = self._get_sparse_cosine_similarities(query_embedding)

        cosine_similarities = (1-sparse_weight) * dense_cosine_results + sparse_weight * sparse_cosine_results

        sorted_indices = np.argsort(-cosine_similarities)
        nearest_neighbors = [self.embeddings[i] for i in sorted_indices[:min(n_search_results_to_use, len(self.embeddings))]]
        return nearest_neighbors


    def query(self, query_embedding: HybridEmbedding, n_search_results_to_use: int, query_strategy: str = "split_return", sparse_weight: float = 0.5) -> list[SearchResult]:
        if len(self.embeddings) == 0:
            logger.warning("No embeddings in datastore. Returning empty list.")
            return []
        if query_strategy == "split_return":
            best_matches = self._query_split_return(query_embedding, n_search_results_to_use)
        elif query_strategy == "weighted":
            best_matches = self._query_weighted(query_embedding, n_search_results_to_use, sparse_weight=sparse_weight)
        else:
            best_matches = self._query_weighted(query_embedding, n_search_results_to_use, sparse_weight=0.0) # default to dense only
        results = [SearchResult(content=match.text) for match in best_matches]
        return results
    
    def upsert(self, embeddings: list[HybridEmbedding]) -> None:
        with open(self.disk_path, "a") as file:
            for item in embeddings:
                sparse_embedding_idx = {"indices": item.sparse_embedding.indices, "values": item.sparse_embedding.values, "max_index": item.sparse_embedding.max_index}
                output_data = {"embedding": item.embedding, "sparse_embedding": sparse_embedding_idx, "text": item.text}
                json_str = json.dumps(output_data)
                file.write(json_str + "\n")
        self._load_data(self.disk_path) # Once the embeddings have been written to disk, reload the datastore

claude_retriever/searcher/vectorstores/local.py
anthropics/anthropic-retrieval-demo
from claude_retriever.searcher.types import Embedding, HybridEmbedding, SearchResult, VectorStore
import pinecone
from more_itertools import chunked

import logging
logger = logging.getLogger(__name__)


############################################
# Pinecone VectorStore implementations
############################################

class PineconeVectorStore(VectorStore):
    '''
    Pinecone vectorstores maintain a single embedding matrix.
    
    How it works:
    - On init, the Pinecone index is loaded (this assumes that the Pinecone index already exists).
    - When upserting embeddings, the embeddings are upserted into the Pinecone index.
    -- The embeddings are stored as a list of ids, vectors, and metadatas. Metadatas are used to store the text data for each embedding; Pinecone indices do not store text data by default.
    -- The ids are the index of the embedding in the Pinecone index.
    - When querying, the query embedding is compared to all embeddings in the Pinecone index using the similarity specified when the index was created.

    Note that the vectorstore does not contain any logic for creating embeddings. It is assumed that the embeddings are created elsewhere
    using Embedders and passed to the vectorstore for storage and retrieval. The utils.embed_and_upload() is a wrapper to help do this.
    '''
    def __init__(self, api_key: str, environment: str, index: str):
        self.api_key = api_key
        self.environment = environment
        self.index = index
        self.pinecone_index = self._init_pinecone_index()
        self.pinecone_index_dimensions = self.pinecone_index.describe_index_stats().dimension

    def _init_pinecone_index(self):
        pinecone.init(
            api_key=self.api_key,
            environment=self.environment,
        )
        if self.index not in pinecone.list_indexes():
            raise ValueError(f"Pinecone index {self.index} does not exist")
        return pinecone.Index(self.index)

    def query(self, query_embedding: Embedding, n_search_results_to_use: int = 10) -> list[SearchResult]:
        if len(query_embedding.embedding) != self.pinecone_index_dimensions:
            raise ValueError(f"Query embedding dimension {len(query_embedding.embedding)} does not match Pinecone index dimension {self.pinecone_index_dimensions}")
        results = self.pinecone_index.query(
            vector=query_embedding.embedding, top_k=n_search_results_to_use, include_metadata=True
        )
        results=[SearchResult(content=match['metadata']['text']) for match in results.matches]
        return results

    def upsert(self, embeddings: list[Embedding], upsert_batch_size: int = 128) -> None:
        '''
        This method upserts embeddings into the Pinecone index in batches of size upsert_batch_size.

        Since Pinecone indices uniquely identify embeddings by their ids,
        we need to keep track of the current index size and update the id counter correspondingly.
        '''
        embedding_chunks = chunked(embeddings, n=upsert_batch_size) # split embeddings into chunks of size upsert_batch_size
        current_index_size = self.pinecone_index.describe_index_stats()['total_vector_count'] # get the current index size from Pinecone
        i = 0 # keep track of the current index in the current batch
        for emb_chunk in embedding_chunks:
            # for each chunk of size upsert_batch_size, create a list of ids, vectors, and metadatas, and upsert them into the Pinecone index
            ids = [str(current_index_size+1+i) for i in range(i,i+len(emb_chunk))]
            vectors = [emb.embedding for emb in emb_chunk]
            metadatas = [{'text': emb.text} for emb in emb_chunk]
            records = list(zip(ids, vectors, metadatas))
            self.pinecone_index.upsert(vectors=records)
            i += len(emb_chunk) 

class PineconeHybridVectorStore(PineconeVectorStore):

    '''
    Pinecone hybrid vectorstores maintain two embedding matrices: one for dense embeddings and one for sparse embeddings.
    
    How it works:
    - On init, the Pinecone index is loaded (this assumes that the Pinecone index already exists).
    - When upserting embeddings, the embeddings are upserted into the Pinecone index.
    -- The embeddings are stored as a list of ids, vectors, and metadatas. Metadatas are used to store the text data for each embedding; Pinecone indices do not store text data by default.
    -- The ids are the index of the embedding in the Pinecone index.
    -- The sparse embeddings are stored as a list of sparse vectors, where each sparse vector is a dict of indices and values.
    - When querying, the query embedding is compared to all dense embeddings and all sparse embeddings in the Pinecone index. They are combined using either:
    -- split_return: half of the expected number of results are returned from the top dense matches and half from the top sparse matches (after deduplication).
    -- weighted: the top n_search_results_to_use results are returned from the weighted combination of dense and sparse cosine similarities.

    Note that the vectorstore does not contain any logic for creating embeddings. It is assumed that the embeddings are created elsewhere
    using Embedders and passed to the vectorstore for storage and retrieval. The utils.embed_and_upload() is a wrapper to help do this.
    '''

    def _query_split_return(self, query_embedding: HybridEmbedding, n_search_results_to_use: int) -> list[SearchResult]:

        '''
        This method returns half of the expected number of results from the top dense matches and half from the top sparse matches.
        '''

        dense_results = self.pinecone_index.query(
            vector=query_embedding.embedding, top_k=n_search_results_to_use, include_metadata=True
        )

        sparse_results = self.pinecone_index.query(
            vector=[0.]*len(query_embedding.embedding),
            sparse_vector={
                'indices': query_embedding.sparse_embedding.indices,
                'values': query_embedding.sparse_embedding.values,
            },
            top_k=n_search_results_to_use,
            include_metadata=True
        )

        n_dense_results = int(n_search_results_to_use / 2)
        n_sparse_results = n_search_results_to_use - n_dense_results

        dense_results = dense_results.matches[:min(n_dense_results, len(dense_results.matches))]
        dense_results_texts = [result['metadata']['text'] for result in dense_results]

        # Remove the already chosen dense results
        sparse_results = [sparse_result for sparse_result in sparse_results.matches if sparse_result['metadata']['text'] not in dense_results_texts]
        sparse_results = sparse_results[:min(n_sparse_results, len(sparse_results))]

        results = dense_results + sparse_results
        results = [SearchResult(content=match['metadata']['text']) for match in results]

        return results
    
    def _query_weighted(self, query_embedding: HybridEmbedding, n_search_results_to_use: int, sparse_weight: float = 0.5) -> list[SearchResult]:

        '''
        This method returns the top n_search_results_to_use results from the weighted combination of dense and sparse cosine similarities.
        '''

        def hybrid_score_norm(dense, sparse, sparse_weight: float):
            """Hybrid score using a convex combination

            (1 - sparse_weight) * dense + sparse_weight * sparse

            Args:
                dense: Array of floats representing
                sparse: a dict of `indices` and `values`
                sparse_weight: scale between 0 and 1
            """
            if sparse_weight < 0 or sparse_weight > 1:
                raise ValueError("Alpha must be between 0 and 1")
            hs = {
                'indices': sparse['indices'],
                'values':  [v * (1 - sparse_weight) for v in sparse['values']]
            }
            return [v * sparse_weight for v in dense], hs

        sparse_query_vector = {
            'indices': query_embedding.sparse_embedding.indices,
            'values': query_embedding.sparse_embedding.values
        }

        dense_query_vector = query_embedding.embedding

        hdense, hsparse = hybrid_score_norm(dense_query_vector, sparse_query_vector, sparse_weight)        

        results = self.pinecone_index.query(
            vector=hdense,
            sparse_vector=hsparse,
            top_k=n_search_results_to_use,
            include_metadata=True
        )

        results = [SearchResult(content=match['metadata']['text']) for match in results.matches]

        return results
        
    def query(self, query_embedding: HybridEmbedding, n_search_results_to_use: int = 10, query_strategy: str = "split_return", sparse_weight: float = 0.5) -> list[SearchResult]:
        if len(query_embedding.embedding) != self.pinecone_index_dimensions:
            raise ValueError(f"Query embedding dimension {len(query_embedding.embedding)} does not match Pinecone index dimension {self.pinecone_index_dimensions}")
        if query_strategy == "split_return":
            results = self._query_split_return(query_embedding, n_search_results_to_use)
        elif query_strategy == "weighted":
            results = self._query_weighted(query_embedding, n_search_results_to_use, sparse_weight=sparse_weight)
        else:
            results = self._query_weighted(query_embedding, n_search_results_to_use, sparse_weight=0.0) # default to dense only
        return results

    def upsert(self, embeddings: list[HybridEmbedding], upsert_batch_size: int = 128) -> None:
        '''
        This method upserts embeddings into the Pinecone index in batches of size upsert_batch_size.

        Since Pinecone indices uniquely identify embeddings by their ids,
        we need to keep track of the current index size and update the id counter correspondingly.
        '''
        embedding_chunks = chunked(embeddings, n=upsert_batch_size) # split embeddings into chunks of size upsert_batch_size
        current_index_size = self.pinecone_index.describe_index_stats()['total_vector_count'] # get the current index size from Pinecone
        i = 0 # keep track of the current index in the current batch
        for emb_chunk in embedding_chunks:
            # for each chunk of size upsert_batch_size, create a list of ids, vectors, and metadatas, and upsert them into the Pinecone index
            ids = [str(current_index_size + 1 + i) for i in range(i,i+len(emb_chunk))]
            vectors = [emb.embedding for emb in emb_chunk]
            metadatas = [{'text': emb.text} for emb in emb_chunk]
            sparse_values = [{'indices': emb.sparse_embedding.indices, 'values': emb.sparse_embedding.values} for emb in emb_chunk]
            records = [{'id': i,
                        'values': v,
                        'metadata': m,
                        'sparse_values': s} for i,v,m,s
                        in zip(ids, vectors, metadatas, sparse_values)]
            self.pinecone_index.upsert(vectors=records)
            i += len(emb_chunk)
claude_retriever/searcher/vectorstores/pinecone.py
anthropics/anthropic-retrieval-demo
import os
import json
import re
import anthropic
from anthropic import Anthropic, AsyncAnthropic
from typing import Optional
from searcher.types import Embedder, VectorStore, Tool, SearchTool, SearchResult
from searcher.embedders.local import LocalEmbedder
from constants import DEFAULT_EMBEDDER
from tqdm import tqdm
from dataclasses import dataclass
import aiohttp

import logging
logger = logging.getLogger(__name__)

# Formatting search results
def format_results(extracted: list[str]) -> str:
        """
        Joins and formats the extracted search results as a string.

        :param extracted: The extracted search results to format.
        """
        result = "\n".join(
            [
                f'<item index="{i+1}">\n<page_content>\n{r}\n</page_content>\n</item>'
                for i, r in enumerate(extracted)
            ]
        )
        return result

def format_results_full(extracted: list[str]) -> str:
    """
    Formats the extracted search results as a string, including the <search_results> tags.

    :param extracted: The extracted search results to format.
    """
    return f"\n<search_results>\n{format_results(extracted)}\n</search_results>"


# Chunking and uploading

@dataclass
class Document:
    """
    A single document.
    """
    text: str
    metadata: Optional[dict] = None

## Embedding and uploading

def embed_and_upload(
        input_file: str,
        vectorstore: VectorStore,
        embedder: Optional[Embedder] = None,
        tokens_per_chunk: int = 384,
        stride: Optional[int] = None,
        batch_size: int = 128) -> None:
    
    if embedder is None:
        logger.info(f"Using default embedder: {DEFAULT_EMBEDDER}")
        embedder = LocalEmbedder(DEFAULT_EMBEDDER)

    # Load the documents
    documents: list[Document] = []
    file_type = input_file.split(".")[-1]
    if file_type == "jsonl":
        with open(input_file, "r") as f:
            for i, line in enumerate(f):
                data = json.loads(line)
                text = data["text"]
                if text is None:
                    raise ValueError(f"Invalid jsonl file. 'text' key is missing on line {i}")
                metadata = data.get("metadata", None)
                doc = Document(text=text, metadata=metadata)
                documents.append(doc)
    else:
        raise ValueError("Invalid file_type. Supported types: 'jsonl'")
    
    # Chunk the documents
    chunked_documents = []
    for document in documents:
        chunks = chunk_document(document, tokens_per_chunk, stride)
        chunked_documents += chunks

    # Embed and upload the documents
    bar = tqdm(total=len(chunked_documents), desc="Embedding and uploading documents", leave=True)
    for i in range(0, len(chunked_documents), batch_size):
        batch = chunked_documents[i:i + batch_size]
        batch_embeddings = embedder.embed_batch([doc.text for doc in batch])
        vectorstore.upsert(batch_embeddings)
        bar.update(len(batch))


def chunk_document(document: Document, tokens_per_chunk: int, stride: Optional[int] = None) -> list[Document]:

    if stride is None:
        stride = tokens_per_chunk

    tok = Anthropic().get_tokenizer()

    raw_text = document.text
    tokenized_text = tok.encode(raw_text).ids

    chunks = []
    for i in range(0, len(tokenized_text), stride):
        chunk = tokenized_text[i:i + tokens_per_chunk]
        chunk_text = tok.decode(chunk)
        chunk_document = Document(text=chunk_text, metadata=document.metadata)
        chunks.append(chunk_document)
    return chunks

## Elasticsearch uploading

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk

def upload_to_elasticsearch(
        input_file: str,
        index_name: str,
        cloud_id: str,
        api_key_id: str,
        api_key: str) -> None:
    
    # Load the documents

    documents: list[Document] = []
    file_type = input_file.split(".")[-1]
    if file_type == "jsonl":
        with open(input_file, "r") as f:
            for i, line in enumerate(f):
                data = json.loads(line)
                text = data["text"]
                if text is None:
                    raise ValueError(f"Invalid jsonl file. 'text' key is missing on line {i}")
                metadata = data.get("metadata", None)
                doc = Document(text=text, metadata=metadata)
                documents.append(doc)
    else:
        raise ValueError("Invalid file_type. Supported types: 'jsonl'")
    
    # Upload the documents

    ## Create the Elasticsearch client

    es = Elasticsearch(
        cloud_id=cloud_id,
        api_key=(api_key_id, api_key),
    )

    ## Upload the documents

    def docs_to_generator():
        for i, document in enumerate(documents):
            yield {
                "_index": index_name,
                "_id": i,
                "text": document.text,
                "metadata": document.metadata
            }
    
    bulk(es, docs_to_generator())
    es.indices.refresh(index=index_name)

# Extract content, potentially using language models

from bs4 import BeautifulSoup

async def scrape_url(url: str, summarize_with_claude: bool = False,
               query: Optional[str] = None,
               anthropic_api_key: Optional[str] = None,
               missing_content_placeholder: str = "CONTENT NOT AVAILABLE") -> str:
    content = await get_url_content(url)
    if content:
        if summarize_with_claude:
            if anthropic_api_key is None:
                try:
                    anthropic_api_key = os.environ["ANTHROPIC_API_KEY"]
                except KeyError:
                    raise ValueError("anthropic_api_key must be provided if llm_extract is True")
            try:
                content = await claude_extract(content, query, anthropic_api_key)
            except Exception as e:
                logger.warning(f"Failed to extract with Claude. Falling back to raw content. Error: {e}")
    else:
        content = missing_content_placeholder
    return content

async def get_url_content(url: str) -> Optional[str]:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            if response.status == 200:
                html = await response.text()
                soup = BeautifulSoup(html, 'html.parser')
                text = soup.get_text(strip=True, separator='\n')
                return text
    return None

async def claude_extract(content: str, query: Optional[str], anthropic_api_key: str, max_tokens_to_read: int = 20_000) -> str:

    # Get first max_tokens_to_read words tokens of content

    client = Anthropic(api_key=anthropic_api_key)
    tokenizer = client.get_tokenizer()
    tokenized_content = tokenizer.encode(content).ids
    if len(tokenized_content) > max_tokens_to_read:
        logger.info(f"Truncating content from {len(tokenized_content)} tokens to {max_tokens_to_read} tokens")
        content = tokenizer.decode(tokenized_content[:max_tokens_to_read]).strip()

    # Generate prompt

    prompt = f"""{anthropic.HUMAN_PROMPT} Here is the content of a web page:
<content>
{content}
</content>"""

    if query:
        prompt += f"""
Here is a search query a user made which resulted in this page:

<query>
{query}
</query>

<instructions>
* Please provide a summary of the web page that is relevant to the query.
* Please made the summary as concise as possible, in bullet points, and include all the information that might be relevant to the query.
"""

    else:

        prompt += f"""

<instructions>
* Please provide a summary of the web page, as concisely as possible, in bullet points."""

    prompt += f"""
* Please do not introduce any additional information beyond what's in the web page.
* Please do not make any guesses about the contents of the web page. Simply summarize based on the given information content of the web page.
* If the content of the web page is not meaningful, relevant, or understandable then don't write anything.
* IMPORTANT: You are going to simulate the output of a web page. Therefore your response should look indistinguishable from what one might read on a web page. This means you should neither refer to yourself in any way in the response nor make it apparent that you are providing a summary. You should not explicitly mention in any way that you are providing a simulated output of a web page.
* IMPORTANT: Please do not ask for feedback on your summary in any way at the end of your response.
* IMPORTANT: Please do not mention that you are providing a summary, or mention "summary" in any other way.
</instructions>{anthropic.AI_PROMPT} Based on the given content{' and query' if query else ''}, the summary of the page would be:
<summary>"""

    logger.info(f"Triggering a Claude extract for a {len(tokenizer.encode(content).ids)} token document")

    client = AsyncAnthropic(api_key=anthropic_api_key)
    response = await client.completions.create(
        prompt=prompt,
        max_tokens_to_sample=512,
        temperature=0.0,
        model="claude-instant-v1",
        stop_sequences=["</summary>", anthropic.HUMAN_PROMPT]
    )

    # Extract summary

    completion = response.completion
    if not completion.endswith("</summary>"):
        completion += "</summary>"
    return '<summary>'+completion

claude_retriever/utils.py
anthropics/anthropic-retrieval-demo
from setuptools import setup, find_packages

VERSION = "0.1.1"
DESCRIPTION = "Claude Search and Retrieval Demo"
LONG_DESCRIPTION = "Lightweight wrapper around the Anthropic Python SDK to experiment with Claude's Search and Retrieval capabilities."

# Setting up
setup(
        name="claude_retriever", 
        version=VERSION,
        author="Anthropic",
        author_email="support@anthropic.com",
        description=DESCRIPTION,
        long_description=LONG_DESCRIPTION,
        packages=find_packages(),
        install_requires=[
    "aiohttp",
    "aiosignal",
    "anthropic",
    "anyio",
    "beautifulsoup4",
    "certifi", 
    "click",
    "elasticsearch",
    "filelock",
    "huggingface-hub",
    "idna",
    "joblib",
    "loguru",
    "mpmath",
    "multidict",
    "nest-asyncio",  
    "nltk",
    "numpy",
    "packaging",
    "pinecone-client",
    "psutil",
    "pydantic",
    "python-dateutil",
    "python-dotenv",
    "pyyaml",
    "requests",
    "scikit-learn",
    "scipy",
    "sentence-transformers",
    "sentencepiece",
    "six",
    "soupsieve",
    "sympy",
    "tenacity",
    "torch",
    "tornado",
    "transformers",
    "urllib3",
    "wikipedia"
],        
        keywords=["python", "anthropic"],
        classifiers= [
            "Development Status :: 3 - Alpha",
            "Intended Audience :: Education",
            "Programming Language :: Python :: 2",
            "Programming Language :: Python :: 3",
            "Operating System :: MacOS :: MacOS X",
            "Operating System :: Microsoft :: Windows",
            "Private :: Do Not Upload"
        ]
)


setup.py
anthropics/anthropic-retrieval-demo
import sys, os
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import unittest

from claude_retriever.searcher.types import Embedding
from claude_retriever.searcher.embedders.local import LocalEmbedder
from claude_retriever.searcher.embedders.huggingface import HuggingFaceEmbedder

import dotenv
dotenv.load_dotenv()

DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"

class TestLocalEmbedder(unittest.TestCase):
    def setUp(self):
        self.embedder = LocalEmbedder(model_name=DEFAULT_EMBEDDING_MODEL)

    def test_embed(self):
        query = "This is a test sentence."
        result = self.embedder.embed(query)
        self.assertIsInstance(result, Embedding)
        self.assertEqual(result.text, query)
        self.assertIsInstance(result.embedding, list)

    def test_embed_batch(self):
        queries = ["This is a test sentence.", "Another test sentence."]
        results = self.embedder.embed_batch(queries)
        self.assertIsInstance(results, list)
        for result, query in zip(results, queries):
            self.assertIsInstance(result, Embedding)
            self.assertEqual(result.text, query)
            self.assertIsInstance(result.embedding, list)

class TestHuggingFaceEmbedder(unittest.TestCase):
    def setUp(self):
        self.api_key = os.getenv("HUGGINGFACE_API_KEY")
        assert self.api_key is not None, "Huggingface API key is not set."
        self.embedder = HuggingFaceEmbedder(api_key=self.api_key, model_name=DEFAULT_EMBEDDING_MODEL)

    def test_embed(self):
        query = "This is a test sentence."
        result = self.embedder.embed(query)
        self.assertIsInstance(result, Embedding)
        self.assertEqual(result.text, query)
        self.assertIsInstance(result.embedding, list)

    def test_embed_batch(self):
        queries = ["This is a test sentence.", "Another test sentence."]
        results = self.embedder.embed_batch(queries)
        self.assertIsInstance(results, list)
        for result, query in zip(results, queries):
            self.assertIsInstance(result, Embedding)
            self.assertEqual(result.text, query)
            self.assertIsInstance(result.embedding, list)

if __name__ == "__main__":
    unittest.main()

tests/test_embedders.py
anthropics/anthropic-retrieval-demo
import sys, os
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import unittest
import numpy as np

from claude_retriever.searcher.types import SearchResult, Embedder, Embedding
from claude_retriever.searcher.vectorstores.local import LocalVectorStore
from claude_retriever.searcher.searchtools.embeddings import EmbeddingSearchTool

class TestLocalSearcher(unittest.TestCase):

    def setUp(self):

        class DummyEmbedder(Embedder):
            def __init__(self, dim: int = 3):
                self.dim = dim
            def embed(self, text: str) -> Embedding:
                emb = self.embed_batch([text])
                return emb[0]
            def embed_batch(self, texts: list[str]) -> list[Embedding]:
                embeddings = np.random.rand(len(texts), self.dim).tolist()
                return [Embedding(embedding=embedding, text=text) for embedding, text in zip(embeddings, texts)]
            def tokenizer_encode(self, text: str) -> list[int]:
                return [1, 2, 3]
            def tokenizer_decode(self, tokens: list[int]) -> str:
                return "This is a test sentence."
    
        self.embedder = DummyEmbedder()
        self.vector_store = LocalVectorStore("tests/data/local_db.jsonl")
        self.searchtool = EmbeddingSearchTool(tool_description="This is a test search tool.",
                                              vector_store=self.vector_store,
                                              embedder=self.embedder)

    def test_search(self):
        query = "This is a test sentence."
        n_search_results_to_use = 2
        search_results = self.searchtool.raw_search(query, n_search_results_to_use)
        self.assertIsInstance(search_results, list)
        self.assertEqual(len(search_results), n_search_results_to_use)
        for result in search_results:
            self.assertIsInstance(result, SearchResult)
            self.assertIsInstance(result.content, str)

if __name__ == "__main__":
    unittest.main()

tests/test_searchers.py
anthropics/anthropic-retrieval-demo
import sys, os
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import unittest
from unittest.mock import MagicMock, patch, mock_open

from claude_retriever.utils import chunk_document, embed_and_upload, Document, Embedder, VectorStore

class TestChunkDocument(unittest.TestCase):

    def test_empty_document(self):
        document = Document(text="")
        result = chunk_document(document, 5)
        self.assertEqual(len(result), 0)

    def test_tokens_per_chunk_larger_than_document(self):
        document = Document(text="This is a test.")
        result = chunk_document(document, 100)
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].text, "This is a test.")
        self.assertEqual(result[0].metadata, document.metadata)

    def test_tokens_per_chunk_equal_to_document(self):
        document = Document(text="This is a test.")
        result = chunk_document(document, 5)
        self.assertEqual(len(result), 1)
        self.assertEqual(result[0].text, "This is a test.")
        self.assertEqual(result[0].metadata, document.metadata)

    def test_tokens_per_chunk_smaller_than_document(self):
        document = Document(text="This is a test.", metadata={"id": 1})
        result = chunk_document(document, 2)
        self.assertEqual(len(result), 3)
        self.assertEqual(result[0].text, "This is")
        self.assertEqual(result[1].text, " a test")
        self.assertEqual(result[2].text, ".")
        self.assertEqual(result[0].metadata, document.metadata)
        self.assertEqual(result[1].metadata, document.metadata)
        self.assertEqual(result[2].metadata, document.metadata)

    def test_tokens_per_chunk_smaller_than_document_with_stride(self):
        document = Document(text="This is a test.", metadata={"id": 1})
        result = chunk_document(document, 3, stride=2)
        self.assertEqual(len(result), 3)
        self.assertEqual(result[0].text, "This is a")
        self.assertEqual(result[1].text, " a test.")
        self.assertEqual(result[2].text, ".")

class TestEmbedAndUpload(unittest.IsolatedAsyncioTestCase):
    def test_embed_and_upload(self):
        input_file = "test.jsonl"
        vectorstore = MagicMock(spec=VectorStore)
        embedder = MagicMock(spec=Embedder)

        with patch("builtins.open", mock_open(read_data='{"text": "Sample text", "metadata": {"id": 1}}\n')) as mock_file:
            with patch("claude_retriever.utils.chunk_document") as mock_chunk_document:
                mock_chunk_document.return_value = [Document(text="Sample text", metadata={"id": 1})]

                embed_and_upload(input_file, vectorstore, embedder, tokens_per_chunk=4, stride=2, batch_size=1)

                mock_file.assert_called_once_with(input_file, "r")
                mock_chunk_document.assert_called_once()
                embedder.embed_batch.assert_called_once_with(["Sample text"])
                vectorstore.upsert.assert_called_once()

if __name__ == "__main__":
    unittest.main()

tests/test_utils.py
anthropics/anthropic-retrieval-demo
import sys, os
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import unittest
import tempfile
from scipy.sparse import csr_matrix
import numpy as np

from claude_retriever.searcher.vectorstores.local import LocalVectorStore, LocalHybridVectorStore
from claude_retriever.searcher.types import Embedding, SparseEmbeddingData, HybridEmbedding

class TestLocalVectorStore(unittest.TestCase):

    def test_load_jsonl(self):
        vector_store = LocalVectorStore("tests/data/local_db.jsonl")
        self.assertEqual(len(vector_store.embeddings), 3)
        self.assertEqual(vector_store.embeddings[0].text, "text1")
        self.assertEqual(vector_store.embeddings[1].text, "text2")
        self.assertEqual(vector_store.embeddings[2].text, "text3")

    def test_query(self):
        vector_store = LocalVectorStore("tests/data/local_db.jsonl")
        query_embedding = Embedding([1, 2, 3], "query")
        n_search_results_to_use_results = vector_store.query(query_embedding, n_search_results_to_use=2)
        self.assertEqual(len(n_search_results_to_use_results), 2)
        self.assertEqual(n_search_results_to_use_results[0].content, "text1")
        self.assertEqual(n_search_results_to_use_results[1].content, "text2")

class TestLocalHybridVectorStore(unittest.TestCase):

    def setUp(self):
        self.temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jsonl")
        self.temp_file.close()
        self.disk_path = self.temp_file.name

        self.dense_embedding1 = [0.5, 0.5, 0.5, 0.5]
        self.dense_embedding2 = [0.0, 1.0, 0.0, 1.0]
        self.dense_embedding3 = [0.25, 0.75, 0.25, 0.75]

        self.sparse_embedding1 = SparseEmbeddingData([0, 1], [1.0, 1.0], 3)
        self.sparse_embedding2 = SparseEmbeddingData([2, 3], [1.0, 1.0], 3)
        self.sparse_embedding3 = SparseEmbeddingData([1], [0.5], 3)

        self.hybrid_embedding1 = HybridEmbedding(embedding=self.dense_embedding1, sparse_embedding=self.sparse_embedding1, text="Hello")
        self.hybrid_embedding2 = HybridEmbedding(embedding=self.dense_embedding2, sparse_embedding=self.sparse_embedding2, text="World")
        self.hybrid_embedding3 = HybridEmbedding(embedding=self.dense_embedding3, sparse_embedding=self.sparse_embedding3, text="Pandas are cute")

        self.hybrid_embedding4 = HybridEmbedding(embedding=self.dense_embedding1, sparse_embedding=self.sparse_embedding2, text="Hello World")

    def tearDown(self):
        os.remove(self.disk_path)

    def test_init(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        self.assertEqual(vector_store.disk_path, self.disk_path)
        self.assertEqual(len(vector_store.embeddings), 0)

    def test_upsert_and_load_data(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        vector_store.upsert([self.hybrid_embedding1, self.hybrid_embedding2, self.hybrid_embedding3])
        self.assertEqual(len(vector_store.embeddings), 3)

    def test_sparse_embedding_data_to_sparse_matrix(self):
        sparse_matrix = LocalHybridVectorStore.sparse_embedding_data_to_sparse_matrix([self.sparse_embedding1, self.sparse_embedding2])
        expected_matrix = csr_matrix(([1.0, 1.0, 1.0, 1.0], ([0, 0, 1, 1], [0, 1, 2, 3])), shape=(2, 4))
        self.assertTrue(np.array_equal(sparse_matrix.toarray(), expected_matrix.toarray()))

    def test_query(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        vector_store.upsert([self.hybrid_embedding1, self.hybrid_embedding2, self.hybrid_embedding3])

        query_embedding = self.hybrid_embedding4

        results = vector_store.query(query_embedding, 2, query_strategy="split_return")
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0].content, "Hello")
        self.assertEqual(results[1].content, "World")

        results = vector_store.query(query_embedding, 2, query_strategy="weighted", sparse_weight=0.9)
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0].content, "World")
        self.assertEqual(results[1].content, "Hello")

    def test_query_empty_store(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        query_embedding = HybridEmbedding(embedding=[0.25, 0.75, 0.25, 0.75], sparse_embedding=SparseEmbeddingData([1, 2], [0.5, 0.5], 3), text="Query")

        results = vector_store.query(query_embedding, 2, query_strategy="split_return")
        self.assertEqual(len(results), 0)

        results = vector_store.query(query_embedding, 2, query_strategy="weighted")
        self.assertEqual(len(results), 0)

    def test_query_with_different_strategies(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        vector_store.upsert([self.hybrid_embedding1, self.hybrid_embedding2])

        query_embedding = HybridEmbedding(embedding=[0.25, 0.75, 0.25, 0.75], sparse_embedding=SparseEmbeddingData([1, 2], [0.5, 0.5], 3), text="Query")

        results_split_return = vector_store.query(query_embedding, 2, query_strategy="split_return")
        results_weighted = vector_store.query(query_embedding, 2, query_strategy="weighted")
        self.assertEqual(len(results_split_return), len(results_weighted))
        self.assertEqual(results_split_return[0].content, results_weighted[0].content)
        self.assertEqual(results_split_return[1].content, results_weighted[1].content)

    def test_query_with_invalid_strategy(self):
        vector_store = LocalHybridVectorStore(self.disk_path)
        vector_store.upsert([self.hybrid_embedding1, self.hybrid_embedding2])

        query_embedding = self.hybrid_embedding4

        results = vector_store.query(query_embedding, 2, query_strategy="invalid_strategy")
        self.assertEqual(len(results), 2)
        self.assertEqual(results[0].content, "Hello")

if __name__ == '__main__':
    unittest.main()
tests/test_vectorstores.py
anthropics/anthropic-retrieval-demo