(self, query: str, topk: int = 3)
| 23 | } |
| 24 | |
| 25 | def query(self, query: str, topk: int = 3) -> List[str]: |
| 26 | query_embedding = self.embed.embed_query(query) |
| 27 | |
| 28 | queue = PriorityQueue() |
| 29 | for tool_name, tool_info in self.documents.items(): |
| 30 | tool_embedding = tool_info["embedding"] |
| 31 | tool_sim = self.similarity(query_embedding, tool_embedding) |
| 32 | queue.put([-tool_sim, tool_name]) |
| 33 | |
| 34 | result = [] |
| 35 | for i in range(min(topk, len(queue.queue))): |
| 36 | tool = queue.get() |
| 37 | result.append(tool[1]) |
| 38 | |
| 39 | return result |
| 40 | |
| 41 | def similarity(self, query: List[float], document: List[float]) -> float: |
| 42 | return sum([i * j for i, j in zip(query, document)]) |
no test coverage detected