AutoRAG 0.1.3__tar.gz → 0.1.4__tar.gz

{autorag-0.1.3 → autorag-0.1.4}/.github/workflows/test.yml

@@ -7,6 +7,9 @@ on:
   pull_request:
     branches:
       - main
+  pull_request_target:
+    branches:
+      - main
 
 env:
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
@@ -29,7 +32,7 @@ jobs:
           pip install -e .
       - name: Install dependencies
         run: |
-          pip install pytest pytest-xdist
+          pip install pytest pytest-xdist pytest-asyncio
       - name: delete tests package
         run: python3 tests/delete_tests.py
       - name: Run tests

{autorag-0.1.3 → autorag-0.1.4}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.1.3
+Version: 0.1.4
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -500,8 +500,6 @@ node_lines:
 You can check our all supporting Nodes & modules
 at [here](https://edai.notion.site/Supporting-Nodes-modules-0ebc7810649f4e41aead472a92976be4?pvs=4)
 
-
-
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline

{autorag-0.1.3 → autorag-0.1.4}/AutoRAG.egg-info/SOURCES.txt

@@ -1,4 +1,6 @@
 .gitignore
+CODE_OF_CONDUCT.md
+CONTRIBUTING.md
 LICENSE
 README.md
 dev_requirements.txt
@@ -57,12 +59,14 @@ autorag/nodes/generator/vllm.py
 autorag/nodes/passagecompressor/__init__.py
 autorag/nodes/passagecompressor/base.py
 autorag/nodes/passagecompressor/pass_compressor.py
+autorag/nodes/passagecompressor/refine.py
 autorag/nodes/passagecompressor/run.py
 autorag/nodes/passagecompressor/tree_summarize.py
 autorag/nodes/passagefilter/__init__.py
 autorag/nodes/passagefilter/base.py
 autorag/nodes/passagefilter/pass_passage_filter.py
 autorag/nodes/passagefilter/percentile_cutoff.py
+autorag/nodes/passagefilter/recency.py
 autorag/nodes/passagefilter/run.py
 autorag/nodes/passagefilter/threshold_cutoff.py
 autorag/nodes/passagereranker/__init__.py
@@ -170,8 +174,10 @@ docs/source/nodes/generator/generator.md
 docs/source/nodes/generator/llama_index_llm.md
 docs/source/nodes/generator/vllm.md
 docs/source/nodes/passage_compressor/passage_compressor.md
+docs/source/nodes/passage_compressor/refine.md
 docs/source/nodes/passage_compressor/tree_summarize.md
 docs/source/nodes/passage_filter/passage_filter.md
+docs/source/nodes/passage_filter/recency_filter.md
 docs/source/nodes/passage_filter/similarity_percentile_cutoff.md
 docs/source/nodes/passage_filter/similarity_threshold_cutoff.md
 docs/source/nodes/passage_reranker/cohere.md
@@ -243,13 +249,16 @@ tests/autorag/nodes/generator/test_generator_base.py
 tests/autorag/nodes/generator/test_llama_index_llm.py
 tests/autorag/nodes/generator/test_run_generator_node.py
 tests/autorag/nodes/generator/test_vllm.py
+tests/autorag/nodes/passagecompressor/test_base_passage_compressor.py
 tests/autorag/nodes/passagecompressor/test_pass_compressor.py
+tests/autorag/nodes/passagecompressor/test_refine.py
 tests/autorag/nodes/passagecompressor/test_run_passage_compressor_node.py
 tests/autorag/nodes/passagecompressor/test_tree_summarize.py
 tests/autorag/nodes/passagefilter/test_pass_passage_filter.py
 tests/autorag/nodes/passagefilter/test_passage_filter_base.py
 tests/autorag/nodes/passagefilter/test_passage_filter_run.py
 tests/autorag/nodes/passagefilter/test_percentile_cutoff.py
+tests/autorag/nodes/passagefilter/test_recency_filter.py
 tests/autorag/nodes/passagefilter/test_threshold_cutoff.py
 tests/autorag/nodes/passagereranker/test_cohere_reranker.py
 tests/autorag/nodes/passagereranker/test_colbert_reranker.py

autorag-0.1.4/CODE_OF_CONDUCT.md

@@ -0,0 +1,41 @@
+# Code of Conduct
+
+## Request from out side
+
+We, the contributors and maintainers, commit to ensuring that everyone's participation in our project and community is free from harassment, regardless of age, body size, disability, ethnicity, gender identity and expression, experience level, education, socioeconomic status, nationality, personal appearance, race, religion, or sexual orientation. This is done in the spirit of creating a friendly and open environment.
+
+## Our Standards
+
+The following are some instances of actions that support the development of a pleasant environment:
+
+* Speaking in an open and accepting manner 
+* Respecting the opinions and experiences of others 
+* Taking constructive criticism in stride
+* Putting the good of the community first 
+* Demonstrating empathy for other community members
+
+Example of participant behavior that is undesirable include:
+
+* Public or private harassment 
+*The publication of another person's private information, such as a physical or electronic address, without that person's express consent 
+* The use of sexualized language or imagery and unwanted sexual attention or advances 
+* Trolling, offensive or derogatory remarks, and personal or political attacks 
+
+## Our Responsibilities
+
+Project maintainers are expected to take appropriate and equitable remedial action in response to any instances of undesirable behavior, as well as to clarify the standards of acceptable behavior.
+
+The right and obligation of project maintainers is to delete, modify, or reject comments, commits, code, wiki edits, issues, and other contributions that do not follow this code of conduct. They also have the authority to temporarily or permanently ban any contributor for any other actions they believe to be improper, threatening, offensive, or harmful.
+
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces when an individual is representing the project or its community. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event. Representation of a project may be further defined and clarified by project maintainers.
+
+## Enforcement
+
+Reports of abusive, harassing, or otherwise inappropriate behavior can be sent to jeffrey@markr.ai' or 'vkehfdl1@gmail.com, the project team's email address. After each complaint is examined and looked into, a response that is judged essential and fitting for the situation will be given. The project team has a duty to keep the identity of the incident reported discreet. Specific enforcement policies may have additional information posted separately.
+
+Project maintainers may be subject to temporary or permanent consequences, as decided by other project leadership members, for failing to abide by and enforce the Code of Conduct in good faith.
+
+

autorag-0.1.4/CONTRIBUTING.md

@@ -0,0 +1,134 @@
+# Getting Started
+
+Thank you so much for your consideration of contributing to AutoRAG open source project, and a big welcome! Users in our community are the ones who make it a reality—people just like you.
+
+By reading and adhering to these principles, we can ensure that the contribution process is simple and efficient for all parties. Additionally, it conveys your agreement to honour the developers' time as they oversee and work on these open-source projects. We will respect you in return by taking care of your problem, evaluating your changes, and assisting you in completing your pull requests.
+
+
+- View the [README](https://github.com/Marker-Inc-Korea/AutoRAG/blob/main/README.md) or [watch this video](https://youtu.be/2ojK8xjyXAU?si=nJz-IgrXFaMiyyW5) to get your development environment up and running. 
+- Learn how to [format pull requests](#submitting-a-pull-request).
+- Read how to [rebase/merge upstream branches](#configuring-remotes).
+- Follow our [code of conduct](CODE_OF_CONDUCT.md).
+- [Find an issue to work on](https://github.com/Marker-Inc-Korea/AutoRAG/issues) and start smashing!
+  
+# Contributing Guidelines [![contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/Marker-Inc-Korea/AutoRAG/issues))
+
+When contributing to this repository, please first discuss the change you wish to make via an issue.
+
+Remember that this is an inclusive community, committed to creating a safe, positive environment. See the whole [Code of Conduct](CODE_OF_CONDUCT.md) and please follow it in all your interactions with the project.
+
+
+## Submitting or Requesting an Issue/Enhancement
+
+### The best ways to report issues or make requests for improvements are as follows:
+- Please explore the issue tracker before submitting an issue. There may already be an issue for your issue, and the conversation may have made remedies readily apparent.
+- When creating the problem, include the screenshots also.
+
+### Best Practices for getting assigned to work on an Issue/Enhancement:
+- If you would like to work on an issue, inform in the issue ticket by commenting on it.
+- Please be sure that you are able to reproduce the issue, before working on it. If not, please ask for clarification by commenting or asking the issue creator.
+
+**Note:** Please do not work on an issue which is already being worked on by another contributor. We don't encourage creating multiple pull requests for the same issue. Also, please allow the assigned person at least 2 days to work on the issue (The time might vary depending on the difficulty). If there is no progress after the deadline, please comment on the issue asking the contributor whether he/she is still working on it. If there is no reply, then feel free to work on the issue.
+
+
+## Submitting a Pull Request
+
+### Best Practices to send Pull Requests:
+  - Fork the [project](https://github.com/Marker-Inc-Korea/AutoRAG) on GitHub
+  - Clone the project locally into your system.
+```
+git clone https://github.com/Marker-Inc-Korea/AutoRAG.git
+```
+  - Make sure you are in the `main` branch.
+```
+git checkout main
+```
+  - Create a new branch with a meaningful name before adding and committing your changes.
+```
+git checkout -b branch-name
+```
+  - Add the files you changed. (avoid using `git add .`)
+```
+git add file-name
+```
+  - Commit the added files
+```
+git commit
+```
+  - If you forgot to add some changes, you can edit your previous commit message.
+```
+git commit --amend
+```
+  - Squash multiple commits to a single commit. (example: squash last two commits done on this branch into one)
+```
+git rebase --interactive HEAD~2 
+```
+  - Push this branch to your remote repository on GitHub.
+```
+git push origin branch-name
+```
+  - If any of the squashed commits have already been pushed to your remote repository, you need to do a force push.
+```
+git push origin remote-branch-name --force
+```
+  - Follow the Pull request template and submit a pull request with a motive for your change and the method you used to achieve it to be merged with the `main` branch.
+  - If you can, please submit the pull request with the fix or improvements including tests.
+  - During review, if you are requested to make changes, rebase your branch and squash the multiple commits into one. Once you push these changes the pull request will edit automatically.
+
+
+## Configuring remotes
+When a repository is cloned, it has a default remote called `origin` that points to your fork on GitHub, not the original repository it was forked from. To keep track of the original repository, you should add another remote called `upstream`.
+
+1. Set the `upstream`.
+```
+git remote add upstream https://github.com/Marker-Inc-Korea/AutoRAG.git
+```
+2. Use `git remote -v` to check the status. The output must be something like this:
+```
+  > origin    https://github.com/your-username/AutoRAG.git (fetch)
+  > origin    https://github.com/your-username/AutoRAG.git (push)
+  > upstream  https://github.com/Marker-Inc-Korea/AutoRAG.git (fetch)
+  > upstream  https://github.com/Marker-Inc-Korea/AutoRAG.git (push)
+```
+3. To update your local copy with remote changes, run the following: (This will give you an exact copy of the current remote. You should not have any local changes on your main branch, if you do, use rebase instead).
+```
+git fetch upstream
+git checkout main
+git merge upstream/main
+```
+4. Push these merged changes to the main branch on your fork. Ensure to pull in upstream changes regularly to keep your forked repository up to date.
+```
+git push origin main
+```
+5. Switch to the branch you are using for some piece of work.
+```
+git checkout branch-name
+```
+6. Rebase your branch, which means, take in all latest changes and replay your work in the branch on top of this - this produces cleaner versions/history.
+```
+git rebase main
+```
+7. Push the final changes when you're ready.
+```
+git push origin branch-name
+```
+
+## After your Pull Request is merged
+After your pull request is merged, you can safely delete your branch and pull the changes from the main (upstream) repository.
+
+1. Delete the remote branch on GitHub.
+```
+git push origin --delete branch-name
+```
+2. Checkout the main branch.
+```
+git checkout main
+```
+3. Delete the local branch.
+```
+git branch -D branch-name
+```
+4. Update your main branch with the latest upstream version.
+```
+git pull upstream main
+```

{autorag-0.1.3 → autorag-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.1.3
+Version: 0.1.4
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -500,8 +500,6 @@ node_lines:
 You can check our all supporting Nodes & modules
 at [here](https://edai.notion.site/Supporting-Nodes-modules-0ebc7810649f4e41aead472a92976be4?pvs=4)
 
-
-
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline

{autorag-0.1.3 → autorag-0.1.4}/README.md

@@ -240,8 +240,6 @@ node_lines:
 You can check our all supporting Nodes & modules
 at [here](https://edai.notion.site/Supporting-Nodes-modules-0ebc7810649f4e41aead472a92976be4?pvs=4)
 
-
-
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline

autorag-0.1.4/autorag/VERSION

@@ -0,0 +1 @@
+0.1.4

{autorag-0.1.3 → autorag-0.1.4}/autorag/data/qacreation/base.py

@@ -1,3 +1,4 @@
+import logging
 import uuid
 from typing import Callable, Optional
 
@@ -5,6 +6,8 @@ import pandas as pd
 
 from autorag.utils.util import save_parquet_safe
 
+logger = logging.getLogger("AutoRAG")
+
 
 def make_single_content_qa(corpus_df: pd.DataFrame,
                            content_size: int,
@@ -33,6 +36,11 @@ def make_single_content_qa(corpus_df: pd.DataFrame,
     :return: QA dataset dataframe.
         You can save this as parquet file to use at AutoRAG.
     """
+    assert content_size > 0, "content_size must be greater than 0."
+    if content_size > len(corpus_df):
+        logger.warning(f"content_size {content_size} is larger than the corpus size {len(corpus_df)}. "
+                       "Setting content_size to the corpus size.")
+        content_size = len(corpus_df)
     sampled_corpus = corpus_df.sample(n=content_size, random_state=random_state)
     sampled_corpus = sampled_corpus.reset_index(drop=True)
 

{autorag-0.1.3 → autorag-0.1.4}/autorag/data/qacreation/llama_index.py

@@ -3,6 +3,7 @@ import os.path
 import random
 from typing import Optional, List, Dict, Any
 
+import pandas as pd
 from llama_index.core.service_context_elements.llm_predictor import LLMPredictorType
 
 from autorag.utils.util import process_batch
@@ -89,14 +90,20 @@ def generate_qa_llama_index_by_ratio(
     prompts = list(map(lambda path: open(path, 'r').read(), prompts_ratio.keys()))
     assert all([validate_llama_index_prompt(prompt) for prompt in prompts])
 
+    content_indices = list(range(len(contents)))
     random.seed(random_state)
-    random.shuffle(contents)
+    random.shuffle(content_indices)
+
+    slice_content_indices: List[List[str]] = distribute_list_by_ratio(content_indices, list(prompts_ratio.values()))
+    temp_df = pd.DataFrame({'idx': slice_content_indices, 'prompt': prompts})
+    temp_df = temp_df.explode('idx', ignore_index=True)
+    temp_df = temp_df.sort_values(by='idx', ascending=True)
+
+    final_df = pd.DataFrame({'content': contents, 'prompt': temp_df['prompt'].tolist()})
 
-    slice_contents: List[List[str]] = distribute_list_by_ratio(contents, list(prompts_ratio.values()))
     tasks = [
         async_qa_gen_llama_index(content, llm, prompt, question_num_per_content, max_retries)
-        for prompt, type_contents in zip(prompts, slice_contents)
-        for content in type_contents
+        for content, prompt in zip(final_df['content'].tolist(), final_df['prompt'].tolist())
     ]
 
     loops = asyncio.get_event_loop()

{autorag-0.1.3 → autorag-0.1.4}/autorag/evaluate/metric/generation.py

@@ -10,7 +10,7 @@ import sacrebleu
 import torch
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
-from openai import OpenAI
+from openai import AsyncOpenAI
 from rouge_score import tokenizers
 from rouge_score.rouge_scorer import RougeScorer
 
@@ -194,25 +194,39 @@ def sem_score(generation_gt: List[List[str]], generations: List[str],
     return result
 
 
-@generation_metric
-def g_eval(generation_gt: List[str], pred: str,
+def g_eval(generation_gt: List[List[str]], generations: List[str],
            metrics: Optional[List[str]] = None,
            model: str = 'gpt-4-0125-preview',
-           ) -> float:
+           batch_size: int = 8) -> List[float]:
     """
-    Calculate G-Eval score.
-    G-eval is a metric that uses high-performance LLM model to evaluate generation performance.
-    It evaluates the generation result by coherence, consistency, fluency, and relevance.
-    It uses only 'openai' model, and we recommend to use gpt-4 for evaluation accuracy.
+        Calculate G-Eval score.
+        G-eval is a metric that uses high-performance LLM model to evaluate generation performance.
+        It evaluates the generation result by coherence, consistency, fluency, and relevance.
+        It uses only 'openai' model, and we recommend to use gpt-4 for evaluation accuracy.
 
-    :param generation_gt: A list of ground truth.
-    :param pred: Model generation.
-    :param metrics: A list of metrics to use for evaluation.
-        Default is all metrics, which is ['coherence', 'consistency', 'fluency', 'relevance'].
-    :param model: OpenAI model name.
-        Default is 'gpt-4-0125-preview'.
-    :return: G-Eval score.
+        :param generation_gt: A list of ground truth.
+            Must be 2-d list of string.
+            Because it can be a multiple ground truth.
+            It will get the max of g_eval score.
+        :param generations: A list of generations that LLM generated.
+        :param metrics: A list of metrics to use for evaluation.
+            Default is all metrics, which is ['coherence', 'consistency', 'fluency', 'relevance'].
+        :param model: OpenAI model name.
+            Default is 'gpt-4-0125-preview'.
+        :param batch_size: The batch size for processing.
+            Default is 8.
+        :return: G-Eval score.
     """
+    loop = asyncio.get_event_loop()
+    tasks = [async_g_eval(gt, pred, metrics, model) for gt, pred in zip(generation_gt, generations)]
+    result = loop.run_until_complete(process_batch(tasks, batch_size=batch_size))
+    return result
+
+
+async def async_g_eval(generation_gt: List[str], pred: str,
+                       metrics: Optional[List[str]] = None,
+                       model: str = 'gpt-4-0125-preview',
+                       ) -> float:
     available_metrics = ['coherence', 'consistency', 'fluency', 'relevance']
     if metrics is None:
         metrics = available_metrics
@@ -229,13 +243,13 @@ def g_eval(generation_gt: List[str], pred: str,
         "relevance": open(os.path.join(prompt_path, "rel_detailed.txt")).read(),
     }
 
-    client = OpenAI()
+    client = AsyncOpenAI()
 
-    def g_eval_score(prompt: str, gen_gt: List[str], pred: str):
+    async def g_eval_score(prompt: str, gen_gt: List[str], pred: str):
         scores = []
         for gt in gen_gt:
             input_prompt = prompt.replace('{{Document}}', gt).replace('{{Summary}}', pred)
-            response = client.chat.completions.create(
+            response = await client.chat.completions.create(
                 model=model,
                 messages=[
                     {"role": "system", "content": input_prompt}
@@ -265,7 +279,7 @@ def g_eval(generation_gt: List[str], pred: str,
 
         return int(max(target_tokens, key=target_tokens.get))
 
-    g_eval_scores = list(map(lambda x: g_eval_score(g_eval_prompts[x], generation_gt, pred), metrics))
+    g_eval_scores = await asyncio.gather(*(g_eval_score(g_eval_prompts[x], generation_gt, pred) for x in metrics))
     return sum(g_eval_scores) / len(g_eval_scores)
 
 

{autorag-0.1.3 → autorag-0.1.4}/autorag/nodes/passagecompressor/__init__.py

@@ -1,2 +1,3 @@
 from .pass_compressor import pass_compressor
+from .refine import refine
 from .tree_summarize import tree_summarize

{autorag-0.1.3 → autorag-0.1.4}/autorag/nodes/passagecompressor/base.py

@@ -26,7 +26,7 @@ def passage_compressor_node(func):
         retrieved_ids = previous_result['retrieved_ids'].tolist()
         retrieve_scores = previous_result['retrieve_scores'].tolist()
 
-        if func.__name__ == 'tree_summarize':
+        if func.__name__ in ['tree_summarize', 'refine']:
             param_list = ['prompt', 'chat_prompt', 'context_window', 'num_output', 'batch']
             param_dict = dict(filter(lambda x: x[0] in param_list, kwargs.items()))
             kwargs_dict = dict(filter(lambda x: x[0] not in param_list, kwargs.items()))

autorag-0.1.4/autorag/nodes/passagecompressor/refine.py

@@ -0,0 +1,61 @@
+import asyncio
+from typing import List, Optional
+
+from llama_index.core import PromptTemplate
+from llama_index.core.prompts import PromptType
+from llama_index.core.prompts.utils import is_chat_model
+from llama_index.core.response_synthesizers import Refine
+from llama_index.core.service_context_elements.llm_predictor import LLMPredictorType
+
+from autorag.nodes.passagecompressor.base import passage_compressor_node
+from autorag.utils.util import process_batch
+
+
+@passage_compressor_node
+def refine(queries: List[str],
+           contents: List[List[str]],
+           scores,
+           ids,
+           llm: LLMPredictorType,
+           prompt: Optional[str] = None,
+           chat_prompt: Optional[str] = None,
+           batch: int = 16,
+           ) -> List[str]:
+    """
+    Refine a response to a query across text chunks.
+    This function is a wrapper for llama_index.response_synthesizers.Refine.
+    For more information, visit https://docs.llamaindex.ai/en/stable/examples/response_synthesizers/refine/.
+
+    :param queries: The queries for retrieved passages.
+    :param contents: The contents of retrieved passages.
+    :param scores: The scores of retrieved passages.
+        Do not use in this function, so you can pass an empty list.
+    :param ids: The ids of retrieved passages.
+        Do not use in this function, so you can pass an empty list.
+    :param llm: The llm instance that will be used to summarize.
+    :param prompt: The prompt template for refine.
+        If you want to use chat prompt, you should pass chat_prompt instead.
+        At prompt, you must specify where to put 'context_msg' and 'query_str'.
+        Default is None. When it is None, it will use llama index default prompt.
+    :param chat_prompt: The chat prompt template for refine.
+        If you want to use normal prompt, you should pass prompt instead.
+        At prompt, you must specify where to put 'context_msg' and 'query_str'.
+        Default is None. When it is None, it will use llama index default chat prompt.
+    :param batch: The batch size for llm.
+        Set low if you face some errors.
+        Default is 16.
+    :return: The list of compressed texts.
+    """
+    if prompt is not None and not is_chat_model(llm):
+        refine_template = PromptTemplate(prompt, prompt_type=PromptType.REFINE)
+    elif chat_prompt is not None and is_chat_model(llm):
+        refine_template = PromptTemplate(chat_prompt, prompt_type=PromptType.REFINE)
+    else:
+        refine_template = None
+    summarizer = Refine(llm=llm,
+                        refine_template=refine_template,
+                        verbose=True)
+    tasks = [summarizer.aget_response(query, content) for query, content in zip(queries, contents)]
+    loop = asyncio.get_event_loop()
+    results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+    return results

{autorag-0.1.3 → autorag-0.1.4}/autorag/nodes/passagefilter/__init__.py

@@ -1,3 +1,4 @@
 from .percentile_cutoff import similarity_percentile_cutoff
 from .pass_passage_filter import pass_passage_filter
+from .recency import recency_filter
 from .threshold_cutoff import similarity_threshold_cutoff

{autorag-0.1.3 → autorag-0.1.4}/autorag/nodes/passagefilter/base.py

@@ -1,10 +1,11 @@
 import functools
+import os
 from pathlib import Path
 from typing import Union, Tuple, List
 
 import pandas as pd
 
-from autorag.utils import result_to_dataframe, validate_qa_dataset
+from autorag.utils import result_to_dataframe, validate_qa_dataset, fetch_contents
 
 
 # same with passage filter from now
@@ -33,8 +34,15 @@ def passage_filter_node(func):
         assert "retrieved_ids" in previous_result.columns, "previous_result must have retrieved_ids column."
         ids = previous_result["retrieved_ids"].tolist()
 
-        filtered_contents, filtered_ids, filtered_scores = func(queries=queries, contents_list=contents,
-                                                                scores_list=scores, ids_list=ids, *args, **kwargs)
+        if func.__name__ == 'recency_filter':
+            corpus_df = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"))
+            metadatas = fetch_contents(corpus_df, ids, column_name='metadata')
+            times = [[time['last_modified_datetime'] for time in time_list] for time_list in metadatas]
+            filtered_contents, filtered_ids, filtered_scores \
+                = func(contents_list=contents, scores_list=scores, ids_list=ids, time_list=times, *args, **kwargs)
+        else:
+            filtered_contents, filtered_ids, filtered_scores = func(queries=queries, contents_list=contents,
+                                                                    scores_list=scores, ids_list=ids, *args, **kwargs)
 
         return filtered_contents, filtered_ids, filtered_scores
 

autorag-0.1.4/autorag/nodes/passagefilter/recency.py

@@ -0,0 +1,58 @@
+import logging
+from datetime import datetime
+from typing import List, Tuple
+
+from autorag.nodes.passagefilter.base import passage_filter_node
+
+logger = logging.getLogger("AutoRAG")
+
+
+@passage_filter_node
+def recency_filter(contents_list: List[List[str]],
+                   scores_list: List[List[float]], ids_list: List[List[str]],
+                   time_list: List[List[datetime]],
+                   threshold: str,
+                   ) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+    """
+    Filter out the contents that are below the threshold datetime.
+    If all contents are filtered, keep the only one recency content.
+    If the threshold date format is incorrect, return the original contents.
+
+    :param contents_list: The list of lists of contents to filter
+    :param scores_list: The list of lists of scores retrieved
+    :param ids_list: The list of lists of ids retrieved
+    :param time_list: The list of lists of datetime retrieved
+    :param threshold: The threshold to cut off
+    :return: Tuple of lists containing the filtered contents, ids, and scores
+    """
+
+    def sort_row(contents, scores, ids, time, _datetime_threshold):
+        combined = list(zip(contents, scores, ids, time))
+        combined_filtered = [item for item in combined if item[3] >= _datetime_threshold]
+
+        if combined_filtered:
+            remain_contents, remain_scores, remain_ids, _ = zip(*combined_filtered)
+        else:
+            combined.sort(key=lambda x: x[3], reverse=True)
+            remain_contents, remain_scores, remain_ids, _ = zip(*combined[:1])
+
+        return list(remain_contents), list(remain_ids), list(remain_scores)
+
+    def parse_threshold(threshold_str):
+        for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M", "%Y-%m-%d"):
+            try:
+                return datetime.strptime(threshold_str, fmt)
+            except ValueError:
+                continue
+        logger.info("threshold date format is incorrect, "
+                    "should be YYYY-MM-DD or YYYY-MM-DD HH:MM:SS or YYYY-MM-DD HH:MM")
+        return None
+
+    datetime_threshold = parse_threshold(threshold)
+    if datetime_threshold is None:
+        return contents_list, ids_list, scores_list
+
+    remain_contents_list, remain_ids_list, remain_scores_list = zip(
+        *map(sort_row, contents_list, scores_list, ids_list, time_list, [datetime_threshold] * len(contents_list)))
+
+    return remain_contents_list, remain_ids_list, remain_scores_list