| 7 | SLUG = "memory" |
| 8 | |
| 9 | class Memory: |
| 10 | def __init__(self, max_size: int = 100): |
| 11 | self.max_size = max_size |
| 12 | self.items: List[str] = [] |
| 13 | self.vectorizer = TfidfVectorizer() |
| 14 | self.vectors = None |
| 15 | self.completion_tokens = 0 |
| 16 | |
| 17 | def add(self, item: str): |
| 18 | if len(self.items) >= self.max_size: |
| 19 | self.items.pop(0) |
| 20 | self.items.append(item) |
| 21 | self.vectors = None # Reset vectors to force recalculation |
| 22 | |
| 23 | def get_relevant(self, query: str, n: int = 10) -> List[str]: |
| 24 | if not self.items: |
| 25 | return [] |
| 26 | |
| 27 | if self.vectors is None: |
| 28 | self.vectors = self.vectorizer.fit_transform(self.items) |
| 29 | |
| 30 | query_vector = self.vectorizer.transform([query]) |
| 31 | similarities = cosine_similarity(query_vector, self.vectors).flatten() |
| 32 | top_indices = similarities.argsort()[-n:][::-1] |
| 33 | |
| 34 | return [self.items[i] for i in top_indices] |
| 35 | |
| 36 | def extract_query(text: str) -> Tuple[str, str]: |
| 37 | query_index = text.rfind("Query:") |