File size: 429 Bytes
f62f25f 9fb6f96 f62f25f 9fb6f96 f62f25f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | from langchain.tools import BaseTool
class CodeSearchTool(BaseTool):
name = "Code Search Tool"
description = "Searches relevant code snippets"
def __init__(self, vector_db):
super().__init__()
self.vector_db = vector_db
def _run(self, query: str):
docs = self.vector_db.similarity_search(query, k=5)
return "\n\n".join(
[doc.page_content for doc in docs]
) |