Update scripts/custom_retriever.py

scripts/custom_retriever.py

@@ -4,15 +4,21 @@ import time
 import traceback
 from typing import List, Optional
 
+import tiktoken
 from cohere import AsyncClient
 from dotenv import load_dotenv
-from llama_index.core import QueryBundle
+from llama_index.core import Document, QueryBundle
+from llama_index.core.async_utils import run_async_tasks
 from llama_index.core.retrievers import (
     BaseRetriever,
+    KeywordTableSimpleRetriever,
     VectorIndexRetriever,
 )
 from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle
 from llama_index.core.vector_stores import (
+    FilterCondition,
+    FilterOperator,
+    MetadataFilter,
     MetadataFilters,
 )
 from llama_index.postprocessor.cohere_rerank import CohereRerank
@@ -36,10 +42,13 @@ class AsyncCohereRerank(CohereRerank):
         nodes: List[NodeWithScore],
         query_bundle: Optional[QueryBundle] = None,
     ) -> List[NodeWithScore]:
-        if query_bundle is None or len(nodes) == 0:
+        if query_bundle is None:
+            raise ValueError("Query bundle must be provided.")
+        if len(nodes) == 0:
             return []
 
         async_client = AsyncClient(api_key=self._api_key)
+
         texts = [
             node.node.get_content(metadata_mode=MetadataMode.EMBED)
             for node in nodes
@@ -52,13 +61,19 @@ class AsyncCohereRerank(CohereRerank):
             documents=texts,
         )
 
-        return [
-            NodeWithScore(node=nodes[result.index].node, score=result.relevance_score)
-            for result in results.results
-        ]
+        new_nodes = []
+        for result in results.results:
+            new_node_with_score = NodeWithScore(
+                node=nodes[result.index].node, score=result.relevance_score
+            )
+            new_nodes.append(new_node_with_score)
+
+        return new_nodes
 
 
 class CustomRetriever(BaseRetriever):
+    """Custom retriever that performs both semantic search and hybrid search."""
+
     def __init__(
         self,
         vector_retriever: VectorIndexRetriever,
@@ -66,95 +81,115 @@ class CustomRetriever(BaseRetriever):
         keyword_retriever=None,
         mode: str = "AND",
     ) -> None:
-        super().__init__()
         self._vector_retriever = vector_retriever
         self._document_dict = document_dict
         self._keyword_retriever = keyword_retriever
+        if mode not in ("AND", "OR"):
+            raise ValueError("Invalid mode. Use 'AND' or 'OR'")
         self._mode = mode
-
-    def retrieve(self, query: str, filters: Optional[MetadataFilters] = None) -> List[NodeWithScore]:
-        query_bundle = QueryBundle(query_str=query)
-        if filters:
-            self._vector_retriever.filters = filters
-        return self._retrieve(query_bundle)
-
-    async def aretrieve(self, query: str, filters: Optional[MetadataFilters] = None) -> List[NodeWithScore]:
-        query_bundle = QueryBundle(query_str=query)
-        if filters:
-            self._vector_retriever.filters = filters
-        return await self._aretrieve(query_bundle)
-
-    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
-        return asyncio.run(self._process_retrieval(query_bundle, is_async=False))
-
-    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
-        return await self._process_retrieval(query_bundle, is_async=True)
+        super().__init__()
 
     async def _process_retrieval(
         self, query_bundle: QueryBundle, is_async: bool = True
     ) -> List[NodeWithScore]:
+
+        if not isinstance(query_bundle, QueryBundle):
+            raise TypeError(f"Expected QueryBundle, got {type(query_bundle)}")
+
+        query_bundle.query_str = query_bundle.query_str.replace("\ninput is ", "").rstrip()
+
         start = time.time()
+
         if is_async:
             nodes = await self._vector_retriever.aretrieve(query_bundle)
         else:
             nodes = self._vector_retriever.retrieve(query_bundle)
 
+        keyword_nodes = []
         if self._keyword_retriever:
-            keyword_nodes = (
-                await self._keyword_retriever.aretrieve(query_bundle)
-                if is_async else self._keyword_retriever.retrieve(query_bundle)
-            )
-        else:
-            keyword_nodes = []
+            if is_async:
+                keyword_nodes = await self._keyword_retriever.aretrieve(query_bundle)
+            else:
+                keyword_nodes = self._keyword_retriever.retrieve(query_bundle)
 
+        vector_ids = {n.node.node_id for n in nodes}
+        keyword_ids = {n.node.node_id for n in keyword_nodes}
         combined_dict = {n.node.node_id: n for n in nodes}
         combined_dict.update({n.node.node_id: n for n in keyword_nodes})
 
-        if self._keyword_retriever:
-            if self._mode == "AND":
-                ids = set(combined_dict) & {n.node.node_id for n in keyword_nodes}
-            else:
-                ids = set(combined_dict)
+        if not self._keyword_retriever or not keyword_nodes:
+            retrieve_ids = vector_ids
         else:
-            ids = set(combined_dict)
+            retrieve_ids = (
+                vector_ids.intersection(keyword_ids)
+                if self._mode == "AND"
+                else vector_ids.union(keyword_ids)
+            )
 
-        filtered_nodes = [combined_dict[i] for i in ids]
+        nodes = [combined_dict[rid] for rid in retrieve_ids]
+        nodes = self._filter_nodes_by_unique_doc_id(nodes)
 
-        # Restore full text if `retrieve_doc` is True
-        for node in filtered_nodes:
+        for node in nodes:
             doc_id = node.node.source_node.node_id
-            if node.metadata.get("retrieve_doc"):
-                doc = self._document_dict.get(doc_id)
-                if doc:
-                    node.node.text = doc.text
+            if node.metadata["retrieve_doc"]:
+                doc = self._document_dict[doc_id]
+                node.node.text = doc.text
             node.node.node_id = doc_id
 
-        # Optional: rerank using Cohere
         try:
             reranker = (
-                AsyncCohereRerank(top_n=5)
-                if is_async else CohereRerank(top_n=5)
+                AsyncCohereRerank(top_n=5, model="rerank-english-v3.0")
+                if is_async
+                else CohereRerank(top_n=5, model="rerank-english-v3.0")
             )
-            filtered_nodes = (
-                await reranker.apostprocess_nodes(filtered_nodes, query_bundle)
-                if is_async else reranker.postprocess_nodes(filtered_nodes, query_bundle)
+            nodes = (
+                await reranker.apostprocess_nodes(nodes, query_bundle)
+                if is_async
+                else reranker.postprocess_nodes(nodes, query_bundle)
             )
         except Exception as e:
-            print(f"Reranking failed: {type(e).__name__}: {e}")
+            print(f"Error during reranking: {type(e).__name__}: {str(e)}")
             traceback.print_exc()
 
-        filtered = self._limit_results_by_length(filtered_nodes)
-        print(f"✅ Retrieved in {time.time() - start:.2f}s")
-        return filtered
+        nodes_filtered = self._filter_by_score_and_tokens(nodes)
+
+        duration = time.time() - start
+        print(f"Retrieving nodes took {duration:.2f}s")
+
+        return nodes_filtered[:5]
 
-    def _limit_results_by_length(self, nodes: List[NodeWithScore]) -> List[NodeWithScore]:
-        total_chars = 0
-        limited = []
+    def _filter_nodes_by_unique_doc_id(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        unique_nodes = {}
         for node in nodes:
-            length = len(node.node.text)
-            if total_chars + length > 60_000:  # rough char limit to stay safe with Gemini context
+            doc_id = node.node.source_node.node_id
+            if doc_id is not None and doc_id not in unique_nodes:
+                unique_nodes[doc_id] = node
+        return list(unique_nodes.values())
+
+    def _filter_by_score_and_tokens(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        nodes_filtered = []
+        total_tokens = 0
+        enc = tiktoken.encoding_for_model("gpt-4")  # tokenizer model name is fine for now
+
+        for node in nodes:
+            if node.score < 0.10:
+                continue
+
+            node_tokens = len(enc.encode(node.node.text))
+            if total_tokens + node_tokens > 100_000:
                 break
-            total_chars += length
-            if node.score >= 0.1:
-                limited.append(node)
-        return limited
+
+            total_tokens += node_tokens
+            nodes_filtered.append(node)
+
+        return nodes_filtered
+
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        return await self._process_retrieval(query_bundle, is_async=True)
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        return asyncio.run(self._process_retrieval(query_bundle, is_async=False))