ai-codebase-intelligence / tools /code_search_tool.py
pkraman06's picture
Update tools/code_search_tool.py
f62f25f verified
raw
history blame contribute delete
429 Bytes
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]
)