Class ToolMemory

autoagent/memory/tool_memory.py:12–81

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]