File size: 434 Bytes
9fb6f96
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
from crewai.tools import BaseTool


class CodeSearchTool(BaseTool):
    name: str = "Code Search Tool"
    description: str = "Searches relevant code snippets from vector database"

    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])