Class CodeReranker

autoagent/memory/code_memory.py:57–174

        ]
    
class CodeReranker(Reranker):
    def __init__(self, model: str) -> None:
        super().__init__(model)
    def wrap_query_results(self, query_results: List[Dict]) -> str:
        wrapped_query_results = ""
        for result in query_results:
            wrapped_query_results += f"File: {result['file']}\n"
            wrapped_query_results += f"Content: {result['content'][:300]}...\n"
            wrapped_query_results += "---"
        return wrapped_query_results
    def wrap_reranked_results(self, reranked_paths: List[str]) -> str:
        wrapped_reranked_results = "[Referenced code files]:"
        for path in reranked_paths:
            wrapped_reranked_results += f"Code path: {path}\n"
            try:
                with open(path, 'r', encoding='utf-8') as file:
                    content = file.read()
                wrapped_reranked_results += f"Code content:\n{content}\n"
            except Exception as e:
                wrapped_reranked_results += f"Error reading file: {str(e)}\n"
            wrapped_reranked_results += "---\n"
        return wrapped_reranked_results
    def parse_results(self, reranked_results: str) -> List[str]:
        lines = reranked_results.strip().split('\n')
    
        # get the last 5 lines
        last_lines = lines[-5:]
        
        # remove the number and dot at the beginning of each line
        cleaned_lines = [re.sub(r'^\d+\.\s*', '', line.strip()) for line in last_lines]
        unique_lines = list(dict.fromkeys(cleaned_lines))
        
        return unique_lines
    def rerank(self, query_text: str, query_results: List[Dict]) -> List[Dict]:
        system_prompt = \
        """
        You are a helpful assistant that reranks the given code files (containing the path of files and Overview of the content of files) based on the query.
        You should rerank the code files based on the query, and the most relevant code files should be ranked on the top. 
        You should select the top 5 code files to answer the query, by giving the file path of the code files.

        Example: 
        [Query]: "The definition of 'BaseAgent'"
        [Code files]: 
        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
        Content: from .ABCAgent import ABCAgent
        from .BaseAgent import BaseAgent
        from .ManagerAgent import ManagerAge...
        ---
        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
        Content: from .ABCAgent import ABCAgent
        from .BaseAgent import BaseAgent
        from .ManagerAgent import ManagerAge...
        ---
        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
        Content: from .ABCAgent import ABCAgent
        from .BaseAgent import BaseAgent
        from .ManagerAgent import ManagerAge...
        ---