(self, df: pd.DataFrame, collection: str = None, batch_size: int = 100)
| 28 | self.collection_name = 'tool_memory' |
| 29 | |
| 30 | def add_dataframe(self, df: pd.DataFrame, collection: str = None, batch_size: int = 100): |
| 31 | if not collection: |
| 32 | collection = self.collection_name |
| 33 | queries = [] |
| 34 | for idx, row in df.iterrows(): |
| 35 | query = { |
| 36 | 'query': ' '.join(row[['Tool_Name', 'Tool_Description', 'API_Name', 'API_Description']].astype(str)), |
| 37 | 'response': row.to_json() |
| 38 | } |
| 39 | queries.append(query) |
| 40 | |
| 41 | # self.add_query(queries, collection=collection) |
| 42 | print(f'Adding {len(queries)} queries to {collection} with batch size {batch_size}') |
| 43 | num_batches = math.ceil(len(queries) / batch_size) |
| 44 | |
| 45 | for i in range(num_batches): |
| 46 | start_idx = i * batch_size |
| 47 | end_idx = min((i + 1) * batch_size, len(queries)) |
| 48 | batch_queries = queries[start_idx:end_idx] |
| 49 | |
| 50 | # Add the current batch of queries |
| 51 | self.add_query(batch_queries, collection=collection) |
| 52 | |
| 53 | print(f"Batch {i+1}/{num_batches} added") |
| 54 | |
| 55 | def query_table( |
| 56 | self, |
no test coverage detected