Task memory retriever with pre-configured parameters. Initialize once with all configuration, then retrieve with only query parameter.
| 74 | |
| 75 | |
| 76 | class TaskMemoryRetriever: |
| 77 | """ |
| 78 | Task memory retriever with pre-configured parameters. |
| 79 | |
| 80 | Initialize once with all configuration, then retrieve with only query parameter. |
| 81 | """ |
| 82 | |
| 83 | def __init__( |
| 84 | self, |
| 85 | task_name: str, |
| 86 | memory_dir: str = "./config/mem_store", |
| 87 | top_k: int = 5, |
| 88 | alpha: float = 0.5, |
| 89 | label_filter: Optional[int] = None, |
| 90 | min_score: float = 0.0, |
| 91 | include_details: bool = True, |
| 92 | config: Optional[Dict[str, Any]] = None |
| 93 | ): |
| 94 | """ |
| 95 | Initialize memory retriever with configuration |
| 96 | |
| 97 | Args: |
| 98 | task_name: Name of the task (e.g., "AutoCls2D") |
| 99 | memory_dir: Base memory directory |
| 100 | top_k: Number of results to retrieve (max: 20) |
| 101 | alpha: Weight for BM25 vs vector (0-1) |
| 102 | label_filter: Filter by label (1/0/-1) |
| 103 | min_score: Minimum similarity threshold |
| 104 | include_details: Include detailed experiment info |
| 105 | config: Optional full config dict (for embedding model configuration) |
| 106 | """ |
| 107 | self.task_name = task_name |
| 108 | self.memory_dir = memory_dir |
| 109 | self.top_k = min(max(1, top_k), 20) |
| 110 | self.alpha = max(0.0, min(1.0, alpha)) |
| 111 | self.label_filter = label_filter |
| 112 | self.min_score = min_score |
| 113 | self.include_details = include_details |
| 114 | self.config = config |
| 115 | |
| 116 | logger.info(f"TaskMemoryRetriever initialized: task={task_name}, " |
| 117 | f"dir={memory_dir}, top_k={top_k}, alpha={alpha}") |
| 118 | |
| 119 | async def retrieve(self, query) -> Dict[str, Any]: |
| 120 | """ |
| 121 | Retrieve similar experiment results with only query parameter |
| 122 | |
| 123 | Args: |
| 124 | query: Research idea query (str or dict with goal info) |
| 125 | |
| 126 | Returns: |
| 127 | Dictionary containing retrieval results |
| 128 | """ |
| 129 | # Convert query to string if it's a dict (goal object) |
| 130 | if isinstance(query, dict): |
| 131 | query_parts = [] |
| 132 | if query.get("description"): |
| 133 | query_parts.append(query["description"]) |
no outgoing calls
no test coverage detected