chore: import upstream snapshot with attribution

autoagent/memory/code_memory.py

@@ -0,0 +1,174 @@
+import os
+from typing import List, Dict
+from autoagent.memory.rag_memory import Memory, Reranker
+from litellm import completion
+import re
+
+class CodeMemory(Memory):
+    def __init__(self, project_path: str, db_name: str = '.sa', platform: str = 'OpenAI', api_key: str = None, embedding_model: str = "text-embedding-ada-002"):
+        super().__init__(project_path, db_name, platform, api_key, embedding_model)
+        self.collection_name = 'code_memory'
+
+    def add_code_files(self, directory: str, exclude_prefix: List[str] = ["workplace_"]):
+        """
+        Add all code files in the specified directory to the memory.
+        
+        Args:
+            directory (str): The directory path containing the code files to add.
+        """
+        code_files = []
+        for root, _, files in os.walk(directory):
+            root_name = str(root)
+            if any(prefix in root_name for prefix in exclude_prefix):
+                continue
+            for file in files:
+                
+                if file.endswith(('.py', '.js', '.java', '.cpp', '.h', '.c', '.html', '.css')):  # add more file types if needed
+                    file_path = os.path.join(root, file)
+                    with open(file_path, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    code_files.append({
+                        "query": f"File: {file_path}\n\nContent:\n{content}",
+                        "response": f"This is the content of file {file_path}"
+                    })
+        self.add_query(code_files, self.collection_name)
+
+    def query_code(self, query_text: str, n_results: int = 5) -> List[Dict]:
+        """
+        Query the code memory.
+        
+        Args:
+            query_text (str): The query text
+            n_results (int): The number of results to return
+        
+        Returns:
+            List[Dict]: The query results list
+        """
+        results = self.query([query_text], self.collection_name, n_results)
+        return [
+            {
+                "file": doc.split('\n')[0].replace("File: ", ""),
+                "content": '\n'.join(doc.split('\n')[3:]),
+                "metadata": metadata
+            }
+            for doc, metadata in zip(results['documents'][0], results['metadatas'][0])
+        ]
+    
+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...
+        ---
+        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/agent_prompts/__init__.py
+        Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
+        ...
+        ---
+        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
+        Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
+        ...
+        ---
+        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
+        Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
+        ...
+        ---
+        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agent_prompts/__init__.py
+        Content: from .BasePrompt import BasePromptGen, ManagerPromptGen, PromptGen
+        ...
+        ---
+        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
+        Content: from typing import List
+
+        from sa.actions import BaseAction, FinishAct, ThinkAct, PlanAct
+        from sa.age...
+        ---
+        File: /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
+        Content: from typing import List
+
+        from sa.actions import BaseAction, FinishAct, ThinkAct, PlanAct
+        from sa.age...
+        ---
+        [Reranked 5 code files]:
+        1. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/BaseAgent.py
+        2. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/__init__.py
+        3. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/ABCAgent.py
+        4. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/ManagerAgent.py
+        5. /Users/tangjiabin/Documents/reasoning/SelfAgent/sa/agents/AgentLogger.py
+        """
+        wrapped_query_results = self.wrap_query_results(query_results)
+        user_prompt = \
+        """
+        [Query]: \n{query_text}
+        [Code files]: \n{query_results}
+        [Reranked 5 code files]:
+        """.format(query_text=query_text, query_results=wrapped_query_results)
+        chat_history = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt}
+        ]
+        create_params = {
+            "model": self.model,
+            "messages": chat_history,
+            "stream": False,
+        }
+        response = completion(**create_params)
+        reranked_results = self.parse_results(response.choices[0].message.content)
+        reranked_results = self.wrap_reranked_results(reranked_results)
+        return reranked_results

autoagent/memory/code_tree/code_parser.py

@@ -0,0 +1,123 @@
+import dataclasses
+from tree_sitter import Language
+import tree_sitter
+import glob
+import uuid
+from loguru import logger
+
+@dataclasses.dataclass
+class Snippet:
+    """Dataclass for storing Embedded Snippets"""
+
+    id: str
+    embedding: list[float] | None
+    snippet: str
+    filename: str
+    language: str
+
+
+class CodeParser:
+    """Code Parser Class."""
+
+    def __init__(self, language: str, node_types: list[str], path_to_object_file: str):
+        self.node_types = node_types
+        self.language = language
+        try:
+            self.parser = tree_sitter.Parser()
+            self.parser.set_language(
+                tree_sitter.Language(f"{path_to_object_file}/my-languages.so", language)
+            )
+        except Exception as e:
+            logger.exception("failed to build %s parser: ", e)
+
+    def parse_file(self, content: str, filename: str):
+        """
+        Parse code snippets from single code file.
+
+        Args:
+            content: The content of the file.
+            filename: The name of the code file.
+
+        Returns:
+        List of Parsed Snippets
+        """
+        try:
+            tree = self.parser.parse(content)
+        except Exception as e:
+            logger.error(f"Failed to parse snippet: {filename} \n Error: {e}")
+            return
+
+        cursor = tree.walk()
+        parsed_snippets = []
+
+        # Walking nodes from abstract syntax tree
+        while cursor.goto_first_child():
+            if cursor.node.type in self.node_types:
+                parsed_snippets.append(
+                    Snippet(
+                        id=str(uuid.uuid4()),
+                        snippet=cursor.node.text,
+                        filename=filename,
+                        language=self.language,
+                        embedding=None,
+                    )
+                )
+
+            while cursor.goto_next_sibling():
+                if cursor.node.type in self.node_types:
+                    parsed_snippets.append(
+                        Snippet(
+                            id=str(uuid.uuid4()),
+                            snippet=cursor.node.text,
+                            filename=filename,
+                            language=self.language,
+                            embedding=None,
+                        )
+                    )
+        return parsed_snippets
+
+    def parse_directory(self, code_directory_path):
+        """
+        Parse code snippets from all files in directory.
+
+        Args:
+            code_directory_path: Directory path containing code files.
+
+        Returns:
+        List of Parsed Snippets
+        """
+        parsed_contents = []
+        for filename in glob.glob(f"{code_directory_path}/**/*.py", recursive=True):
+            # print(filename)
+            with open(filename, "rb") as codefile:
+                code_content = codefile.read()
+
+            parsed_content = self.parse_file(code_content, filename)
+            parsed_contents.extend(parsed_content)
+
+        return parsed_contents
+    
+def to_dataframe_row(embedded_snippets: list[Snippet]):
+    """
+    Helper function to convert Embedded Snippet object to a dataframe row
+    in dictionary format.
+
+    Args:
+        embedded_snippets: List of Snippets to be converted
+
+    Returns:
+        List of Dictionaries
+    """
+    outputs = []
+    for embedded_snippet in embedded_snippets:
+        output = {
+            "ids": embedded_snippet.id,
+            "embeddings": embedded_snippet.embedding,
+            "snippets": embedded_snippet.snippet,
+            "metadatas": {
+                "filenames": embedded_snippet.filename,
+                "languages": embedded_snippet.language,
+            },
+        }
+        outputs.append(output)
+    return outputs

autoagent/memory/codetree_memory.py

@@ -0,0 +1,118 @@
+import os
+from typing import List, Dict
+from autoagent.memory.rag_memory import Memory, Reranker
+import openai
+import re
+from autoagent.memory.code_tree.code_parser import CodeParser, to_dataframe_row
+from tree_sitter import Language
+from loguru import logger
+from openai import OpenAI
+import pandas as pd
+class CodeTreeMemory(Memory):
+    def __init__(self, project_path: str, db_name: str = '.code_tree', platform: str = 'OpenAI', api_key: str = None, embedding_model: str = "text-embedding-ada-002"):
+        super().__init__(project_path, db_name, platform, api_key, embedding_model)
+        self.collection_name = 'code_tree_memory'
+        self.embedder = OpenAI(api_key=api_key)
+        
+
+    def add_code_files(self, directory: str, exclude_prefix: List[str] = ["workplace_"]):
+        """
+        将指定目录下的所有代码文件添加到内存中。
+        
+        Args:
+            directory (str): 要添加的代码文件所在的目录路径
+        """
+        tree_sitter_parent_dir = os.path.dirname(os.getcwd())
+        # Build Tree sitter Parser object
+        Language.build_library(
+            f"{tree_sitter_parent_dir}/my-languages.so",
+            [
+                f"{tree_sitter_parent_dir}/tree-sitter-python",
+            ],
+        )
+        parser = CodeParser(
+            language="python",
+            node_types=["class_definition", "function_definition"],
+            path_to_object_file=tree_sitter_parent_dir,
+        )
+        logger.info("Parsing Code...")
+        parsed_snippets = parser.parse_directory(
+            directory
+        )
+        snippet_texts = list(map(lambda x: x.snippet.decode("ISO-8859-1"), parsed_snippets))
+        embedded_texts = self.embedder.embeddings.create(input=snippet_texts, model="text-embedding-3-small").data
+        embedded_snippets = []
+        for code_text, embedding, snippet in zip(
+            snippet_texts, embedded_texts, parsed_snippets
+        ):
+            snippet.snippet = code_text
+            snippet.embedding = embedding.embedding
+            embedded_snippets.append(snippet)
+
+        # Convert Snippets to DataFrame for ChromaDB Ingestion
+        data = pd.DataFrame(to_dataframe_row(embedded_snippets))
+        collection = self.client.get_or_create_collection(
+            name=self.collection_name, metadata={"hnsw:space": "cosine"}
+        )
+        logger.info(
+            f"Adding {data.shape[0]} Code snippets and embedding to "
+            "local chroma db collection..."
+        )
+        collection.add(
+            documents=data["snippets"].tolist(),
+            embeddings=data["embeddings"].tolist(),
+            metadatas=data["metadatas"].tolist(),
+            ids=data["ids"].tolist(),
+        )
+        
+
+    def query_code(self, query_text: str, n_results: int = 5) -> List[Dict]:
+        """
+        Query the code memory.
+        
+        Args:
+            query_text (str): The query text
+            n_results (int): The number of results to return
+        
+        Returns:
+            List[Dict]: The query results list
+        """
+        query_embedding = self.embedder.embeddings.create(input=[query_text], model="text-embedding-3-small").data[0].embedding
+        results = self.client.get_or_create_collection(self.collection_name).query(query_embeddings=[query_embedding], n_results=n_results)
+        return [
+            {
+                "file": metadata['filenames'],
+                "content": doc
+            }
+            for doc, metadata in zip(results['documents'][0], results['metadatas'][0])
+        ]
+class DummyReranker(Reranker):
+    def __init__(self, model: str = None) -> None:
+        super().__init__(model)
+    def rerank(self, query_results: List[Dict]) -> List[Dict]:
+        wrapped_reranked_results = "[Referenced code files]:"
+        result_path = []
+        for result in query_results:
+            if result['file'] in result_path:
+                continue
+            else:
+                result_path.append(result['file'])
+            wrapped_reranked_results = f"Code path: {result['file']}\n"
+            wrapped_reranked_results += f"Code content:\n{result['content']}...\n"
+            wrapped_reranked_results += "---\n"
+        return wrapped_reranked_results
+
+# 使用示例
+if __name__ == "__main__":
+    code_memory = CodeTreeMemory(project_path = './code_db', db_name='code_tree', platform='OpenAI', api_key='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')
+    
+    # 添加代码文件到内存
+    code_memory.add_code_files("/Users/tangjiabin/Documents/reasoning/SelfAgent/workplace_test/SelfAgent", exclude_prefix=['workplace_', '__pycache__', 'code_db', '.git'])
+    
+    # 查询代码
+    query_results = code_memory.query_code("The definition of BaseAction", n_results=10)
+    
+    for result in query_results:
+        print(f"File: {result['file']}")
+        print(f"Content: {result['content'][:100]}...")  # 只打印前100个字符
+        print("---")

autoagent/memory/paper_memory.py

@@ -0,0 +1,80 @@
+import pandas as pd
+from typing import List, Dict
+from autoagent.memory.rag_memory import Memory, Reranker
+import json
+import math
+import os
+from litellm import completion
+from autoagent.memory.utils import chunking_by_token_size
+
+class TextMemory(Memory):
+    def __init__(
+        self,
+        project_path: str,
+        db_name: str = '.text_table',
+        platform: str = 'OpenAI',
+        api_key: str = None,
+        embedding_model: str = "text-embedding-3-small",
+    ):
+        super().__init__(
+            project_path=project_path,
+            db_name=db_name,
+            platform=platform,
+            api_key=api_key,
+            embedding_model=embedding_model
+        )
+        self.collection_name = 'text_memory'
+
+    def add_text_content(self, paper_content: str, batch_size: int = 100, collection = None):
+        assert collection is not None, "Collection is required. Should be the path of the paper."
+        queries = []
+        content_chunks = chunking_by_token_size(paper_content, max_token_size=4096)
+
+        idx_list = ["chunk_" + str(chunk['chunk_order_index']) for chunk in content_chunks]
+        for chunk in content_chunks:
+            query = {
+                'query': chunk['content'],
+                'response': chunk['content']
+            }
+            queries.append(query)
+        
+        # self.add_query(queries, collection=collection)
+        print(f'Adding {len(queries)} queries to {collection} with batch size {batch_size}')
+        num_batches = math.ceil(len(queries) / batch_size)
+    
+        for i in range(num_batches):
+            start_idx = i * batch_size
+            end_idx = min((i + 1) * batch_size, len(queries))
+            batch_queries = queries[start_idx:end_idx]
+            batch_idx = idx_list[start_idx:end_idx]
+            
+            # Add the current batch of queries
+            self.add_query(batch_queries, collection=collection, idx=batch_idx)
+            
+            print(f"Batch {i+1}/{num_batches} added")
+
+    def query_text_content(
+        self, 
+        query_text: str, 
+        collection: str = None, 
+        n_results: int = 5
+    ) -> List[str]:
+        """
+        Query the table and return the results
+        """
+        assert collection is not None, "Collection is required. Should be the path of the paper."
+        results = self.query([query_text], collection=collection, n_results=n_results)
+        
+        metadata_results = results['metadatas'][0]
+        
+        results = [item['response'] for item in metadata_results]
+        return results
+
+    def peek_table(self, collection: str = None, n_results: int = 20) -> pd.DataFrame:
+        """
+        Peek at the data in the table
+        """
+        assert collection is not None, "Collection is required. Should be the path of the paper."
+        raw_results = self.peek(collection=collection, n_results=n_results)
+        results = [item['response'] for item in raw_results['metadatas']]
+        return results

autoagent/memory/rag_memory.py

@@ -0,0 +1,179 @@
+import uuid
+import os.path
+from datetime import datetime
+from typing import List, Dict
+import chromadb
+from chromadb.utils import embedding_functions
+from abc import ABC, abstractmethod
+from openai import OpenAI
+import numpy as np
+from chromadb.api.types import QueryResult
+chromadb.logger.setLevel(chromadb.logging.ERROR)
+
+
+class Memory:
+    def __init__(
+            self,
+            project_path: str,
+            db_name: str = '.sa',
+            platform: str = 'OpenAI', 
+            api_key: str = None, 
+            embedding_model: str = "text-embedding-3-small"
+    ):
+        """
+        Memory: memory and external knowledge management.
+        Args:
+            project_path: the path to store the data.
+            embedding_model: the embedding model to use, default will use the embedding model from ChromaDB,
+             if the OpenAI has been set in the configuration, it will use the OpenAI embedding model
+             "text-embedding-ada-002".
+        """
+        self.db_name = db_name
+        self.collection_name = 'memory'
+        self.client = chromadb.PersistentClient(path=os.path.join(project_path, self.db_name))
+        self.client.get_or_create_collection(
+                self.collection_name,
+            ) 
+        # use the OpenAI embedding function if the openai section is set in the configuration.
+        if platform == 'OpenAI':
+            openai_client = OpenAI(api_key=api_key or os.environ["OPENAI_API_KEY"])
+            self.embedder = lambda x: [i.embedding for i in openai_client.embeddings.create(input=x, model=embedding_model).data]
+        else:
+            # self.embedder = embedding_functions.DefaultEmbeddingFunction()
+            self.embedder = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
+
+    def add_query(
+            self,
+            queries: List[Dict[str, str]],
+            collection: str = None,
+            idx: List[str] = None
+    ):
+        """
+        add_query: add the queries to the memery.
+        Args:
+            queries: the queries to add to the memery. Should be in the format of
+                {
+                    "query": "the query",
+                    "response": "the response"
+                }
+            collection: the name of the collection to add the queries.
+            idx: the ids of the queries, should be in the same length as the queries.
+            If not provided, the ids will be generated by UUID.
+
+        Return: A list of generated IDs.
+        """
+        if idx:
+            ids = idx
+        else:
+            ids = [str(uuid.uuid4()) for _ in range(len(queries))]
+
+        if not collection:
+            collection = self.collection_name
+
+        query_list = [query['query'] for query in queries]
+        embeddings = self.embedder(query_list)
+        added_time = datetime.now().isoformat()
+        resp_list = [{'response': query['response'], 'created_at': added_time} for query in queries]
+        # insert the record into the database
+        self.client.get_or_create_collection(collection).add(
+            documents=query_list,
+            metadatas=resp_list,
+            ids=ids, 
+            embeddings=embeddings
+        )
+
+        return ids
+
+    def query(self, query_texts: List[str], collection: str = None, n_results: int = 5) -> QueryResult:
+        """
+        query: query the memery.
+        Args:
+            query_texts: the query texts to search in the memery.
+            collection: the name of the collection to search.
+            n_results: the number of results to return.
+
+        Returns: QueryResult
+        class QueryResult(TypedDict):
+            ids: List[IDs]
+            embeddings: Optional[
+                Union[
+                    List[Embeddings],
+                    List[PyEmbeddings],
+                    List[NDArray[Union[np.int32, np.float32]]],
+                ]
+            ]
+            documents: Optional[List[List[Document]]]
+            uris: Optional[List[List[URI]]]
+            data: Optional[List[Loadable]]
+            metadatas: Optional[List[List[Metadata]]]
+            distances: Optional[List[List[float]]]
+            included: Include
+        """
+        if not collection:
+            collection = self.collection_name
+        query_embedding = self.embedder(query_texts)
+        return self.client.get_or_create_collection(collection).query(query_embeddings=query_embedding, n_results=n_results)
+
+    def peek(self, collection: str = None, n_results: int = 20):
+        """
+        peek: peek the memery.
+        Args:
+            collection: the name of the collection to peek.
+            n_results: the number of results to return.
+
+        Returns: the top k results.
+        """
+        if not collection:
+            collection = self.collection_name
+        return self.client.get_or_create_collection(collection).peek(limit=n_results)
+
+    def get(self, collection: str = None, record_id: str = None):
+        """
+        get: get the record by the id.
+        Args:
+            record_id: the id of the record.
+            collection: the name of the collection to get the record.
+
+        Returns: the record.
+        """
+        if not collection:
+            collection = self.collection_name
+        collection = self.client.get_collection(collection)
+        if not record_id:
+            return collection.get()
+
+        return collection.get(record_id)
+
+    def delete(self, collection_name=None):
+        """
+        delete: delete the memery collections.
+        Args:
+            collection_name: the name of the collection to delete.
+        """
+        if not collection_name:
+            collection_name = self.collection_name
+        return self.client.delete_collection(name=collection_name)
+
+    def count(self, collection_name=None):
+        """
+        count: count the number of records in the memery.
+        Args:
+            collection_name: the name of the collection to count.
+        """
+        if not collection_name:
+            collection_name = self.collection_name
+        return self.client.get_or_create_collection(name=collection_name).count()
+
+    def reset(self):
+        """
+        reset: reset the memory.
+        Notice: You may need to set the environment variable `ALLOW_RESET` to `TRUE` to enable this function.
+        """
+        self.client.reset()
+
+class Reranker:
+    def __init__(self, model: str) -> None:
+        self.model = model
+    @abstractmethod
+    def rerank(self, query_text: str, query_results: List[Dict]) -> List[Dict]:
+        raise NotImplementedError("Reranker is not implemented")

autoagent/memory/tool_memory.py

@@ -0,0 +1,164 @@
+import pandas as pd
+from typing import List, Dict
+from autoagent.memory.rag_memory import Memory, Reranker
+import json
+import math
+import os
+from litellm import completion
+from pydantic import BaseModel
+"""
+Category | Tool_Name | Tool_Description | API_Name | API_Description | Method | API_Details | Required_API_Key | Platform
+"""
+class ToolMemory(Memory):
+    def __init__(
+        self,
+        project_path: str,
+        db_name: str = '.tool_table',
+        platform: str = 'OpenAI',
+        api_key: str = None,
+        embedding_model: str = "text-embedding-3-small",
+    ):
+        super().__init__(
+            project_path=project_path,
+            db_name=db_name,
+            platform=platform,
+            api_key=api_key,
+            embedding_model=embedding_model
+        )
+        self.collection_name = 'tool_memory'
+
+    def add_dataframe(self, df: pd.DataFrame, collection: str = None, batch_size: int = 100):
+        if not collection:
+            collection = self.collection_name
+        queries = []
+        for idx, row in df.iterrows():
+            query = {
+                'query': ' '.join(row[['Tool_Name', 'Tool_Description', 'API_Name', 'API_Description']].astype(str)),
+                'response': row.to_json()
+            }
+            queries.append(query)
+        
+        # self.add_query(queries, collection=collection)
+        print(f'Adding {len(queries)} queries to {collection} with batch size {batch_size}')
+        num_batches = math.ceil(len(queries) / batch_size)
+    
+        for i in range(num_batches):
+            start_idx = i * batch_size
+            end_idx = min((i + 1) * batch_size, len(queries))
+            batch_queries = queries[start_idx:end_idx]
+            
+            # Add the current batch of queries
+            self.add_query(batch_queries, collection=collection)
+            
+            print(f"Batch {i+1}/{num_batches} added")
+
+    def query_table(
+        self, 
+        query_text: str, 
+        collection: str = None, 
+        n_results: int = 5
+    ) -> pd.DataFrame:
+        """
+        Query the table and return the results
+        """
+        if not collection:
+            collection = self.collection_name
+        results = self.query([query_text], collection=collection, n_results=n_results)
+        
+        metadata_results = results['metadatas'][0]
+        
+        df_results = pd.DataFrame([json.loads(item['response']) for item in metadata_results])
+        return df_results
+
+    def peek_table(self, collection: str = None, n_results: int = 20) -> pd.DataFrame:
+        """
+        Peek at the data in the table
+        """
+        if not collection:
+            collection = self.collection_name
+        results = self.peek(collection=collection, n_results=n_results)
+        df_results = pd.DataFrame([json.loads(item['response']) for item in results['metadatas']])
+        return df_results
+
+class ToolReranker(Reranker):
+    def rerank(self, query_text: str, query_df: pd.DataFrame) -> str:
+        system_prompt = \
+        """
+        You are a helpful assistant that reranks the given API table based on the query.
+        You should select the top 5 APIs to answer the query in the given format.
+        You can only select APIs I give you.
+        Directly give the answer without any other words.
+        """
+        # Use the DataFrame's to_dict method to convert all rows to a list of dictionaries
+        # print('query_df', query_df)
+        api_data = query_df.to_dict(orient='records')
+        
+        # Use a list comprehension and f-string to format each API's data
+        api_prompts = [f"\n\nAPI {i+1}:\n{api}" for i, api in enumerate(api_data)]
+        
+        # add the query text to the prompt
+        prompt = ''.join(api_prompts)
+        prompt = f"The query is: {query_text}\n\n{prompt}"
+        message = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt}
+        ]
+        
+        class Tools(BaseModel):
+            tool_name: str
+            api_name: str
+            rank: int
+
+        class RerankResult(BaseModel):
+            tools: list[Tools]
+        
+        create_params = {
+            "model": self.model,
+            "messages": message,
+            "stream": False,
+            "response_format": RerankResult
+        }
+        response = completion(**create_params).choices[0].message.content
+        print(response)
+        rerank_result = json.loads(response)["tools"]
+        print(rerank_result)
+        if len(rerank_result) == 0:
+            return "Fail to retrieve the relevant information from the tool documentation."
+        try:
+            return self.wrap_rerank_result(rerank_result, query_df)
+        except Exception as e:
+            raise ValueError(f"Failed to wrap rerank result: {e}")
+        
+    def wrap_rerank_result(self, rerank_result: List[pd.DataFrame], query_df: pd.DataFrame) -> str:
+        res = ""
+        res_tmp = """
+The rank {rank} referenced tool documentation is:
+API Name: {api_name}
+API Description: {api_description}
+API Details: {api_details}
+Required API Key: {required_api_key}
+Platform: {platform}
+"""
+        try:
+            for tool_api in rerank_result:
+                tool_name = tool_api['tool_name']
+                api_name = tool_api['api_name']
+                matched_rows = query_df[(query_df['API_Name'] == api_name) & (query_df['Tool_Name'] == tool_name)]
+                if not matched_rows.empty:
+                    res = res + res_tmp.format(rank=tool_api['rank'], api_name=matched_rows['API_Name'].values[0], api_description=matched_rows['API_Description'].values[0], api_details=matched_rows['API_Details'].values[0], required_api_key=matched_rows['Required_API_Key'].values[0], platform=matched_rows['Platform'].values[0])
+            return res
+        except Exception as e:
+            raise ValueError(f"Failed to wrap rerank result: {e}")
+    def dummy_rerank(self, query_text: str, query_df: pd.DataFrame) -> str:
+        res = ""
+        res_tmp = """
+The rank {rank} referenced tool documentation is:
+API Name: {api_name}
+API Description: {api_description}
+API Details: {api_details}
+Required API Key: {required_api_key}
+Platform: {platform}
+"""
+        for i in range(len(query_df)):
+            res = res + res_tmp.format(rank=i+1, api_name=query_df['API_Name'].values[i], api_description=query_df['API_Description'].values[i], api_details=query_df['API_Details'].values[i], required_api_key=query_df['Required_API_Key'].values[i], platform=query_df['Platform'].values[i])
+        return res

autoagent/memory/utils.py

@@ -0,0 +1,36 @@
+import tiktoken
+ENCODER = None
+
+def encode_string_by_tiktoken(content: str, model_name: str = "gpt-4o"):
+    global ENCODER
+    if ENCODER is None:
+        ENCODER = tiktoken.encoding_for_model(model_name)
+    tokens = ENCODER.encode(content)
+    return tokens
+
+
+def decode_tokens_by_tiktoken(tokens: list[int], model_name: str = "gpt-4o"):
+    global ENCODER
+    if ENCODER is None:
+        ENCODER = tiktoken.encoding_for_model(model_name)
+    content = ENCODER.decode(tokens)
+    return content
+def chunking_by_token_size(
+    content: str, overlap_token_size=128, max_token_size=1024, tiktoken_model="gpt-4o"
+):
+    tokens = encode_string_by_tiktoken(content, model_name=tiktoken_model)
+    results = []
+    for index, start in enumerate(
+        range(0, len(tokens), max_token_size - overlap_token_size)
+    ):
+        chunk_content = decode_tokens_by_tiktoken(
+            tokens[start : start + max_token_size], model_name=tiktoken_model
+        )
+        results.append(
+            {
+                "tokens": min(max_token_size, len(tokens) - start),
+                "content": chunk_content.strip(),
+                "chunk_order_index": index,
+            }
+        )
+    return results