""" Agent tools. Two callable tools the LLM can invoke to gather context before writing tests: - search_code(query) : semantic+keyword search over the repo - get_definition(name) : fetch the full source of a function/class by name Both are thin wrappers over the same retrieval/index the Q&A side uses. The CodeTools object holds the live index so the tool functions can reach it. """ # OpenAI tool schemas (function-calling format). TOOL_SCHEMAS = [ { "type": "function", "function": { "name": "search_code", "description": "Search the repository for code relevant to a query. " "Returns matching definitions with their file and line range.", "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "what to look for"} }, "required": ["query"], }, }, }, { "type": "function", "function": { "name": "get_definition", "description": "Return the full source code of a function or class by its name " "(e.g. 'login_user' or 'UserService.create_token').", "parameters": { "type": "object", "properties": { "name": {"type": "string", "description": "function/class name"} }, "required": ["name"], }, }, }, ] class CodeTools: """Holds the live index and exposes the tool implementations.""" def __init__(self, chunks, embedder, hybrid, reranker): self.chunks = chunks self.embedder = embedder self.hybrid = hybrid self.reranker = reranker def search_code(self, query, k=5): emb = self.embedder.create_embeddings([query])[0] results = self.reranker.rerank(query, self.hybrid.search(query, emb, k=10)) out = [] for r in results[:k]: d = r["document"] out.append({ "name": d["name"], "file": d["file"], "lines": f"{d['start_line']}-{d['end_line']}", "code": d["code"], }) return out def get_definition(self, name): name_l = name.lower() # exact match first, then suffix match (so "create_token" finds "UserService.create_token") matches = [c for c in self.chunks if c["name"].lower() == name_l] if not matches: matches = [c for c in self.chunks if c["name"].lower().endswith("." + name_l) or name_l in c["name"].lower()] if not matches: return {"error": f"No definition found named '{name}'."} d = matches[0] return { "name": d["name"], "file": d["file"], "lines": f"{d['start_line']}-{d['end_line']}", "code": d["code"], } def dispatch(self, tool_name, args): """Route a tool call from the LLM to the right implementation.""" if tool_name == "search_code": return self.search_code(**args) if tool_name == "get_definition": return self.get_definition(**args) return {"error": f"Unknown tool '{tool_name}'."}