chore: import upstream snapshot with attribution

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