""" Inference Script: RAG-Augmented Code Generation with Tool Calling Combines: - Fine-tuned Qwen2.5-Coder-7B (code gen + tool-calling) - RAG pipeline (AST-aware code search) - ReAct-style reasoning loop for multi-step tasks Usage: python inference.py --model your-username/qwen25-coder-7b-code-toolcall --repo /path/to/codebase python inference.py --model your-username/qwen25-coder-7b-code-toolcall --repo /path/to/codebase \ --query "Implement a caching layer for the user service" """ import os, json, re, subprocess from typing import Optional from rag_pipeline import CodebaseIndexer, CodeRetriever, ContextBuilder TOOLS = [ {"type": "function", "function": {"name": "search_codebase", "description": "Search the internal codebase using semantic search.", "parameters": {"type": "object", "properties": { "query": {"type": "string"}, "file_pattern": {"type": "string"}, "max_results": {"type": "integer", "default": 5}}, "required": ["query"]}}}, {"type": "function", "function": {"name": "execute_python", "description": "Execute Python code in a sandboxed environment.", "parameters": {"type": "object", "properties": { "code": {"type": "string"}, "timeout": {"type": "integer", "default": 30}}, "required": ["code"]}}}, {"type": "function", "function": {"name": "read_file", "description": "Read contents of a file from the codebase.", "parameters": {"type": "object", "properties": { "path": {"type": "string"}, "start_line": {"type": "integer"}, "end_line": {"type": "integer"}}, "required": ["path"]}}}, {"type": "function", "function": {"name": "run_tests", "description": "Run unit tests for a module or file.", "parameters": {"type": "object", "properties": { "test_path": {"type": "string"}, "verbose": {"type": "boolean", "default": True}}, "required": ["test_path"]}}} ] SYSTEM_PROMPT = '''You are an expert Python programmer with access to our internal codebase via tools. When helping with code tasks: 1. ALWAYS search the codebase first to understand existing patterns 2. Reference and integrate with existing code 3. Follow the codebase's conventions 4. Write comprehensive docstrings and type hints 5. Consider edge cases and error handling 6. For complex tasks, break them into steps and use tools iteratively Think step by step.''' class ToolExecutor: def __init__(self, retriever: CodeRetriever, repo_path: str): self.retriever = retriever self.repo_path = repo_path def execute(self, tool_name: str, arguments: dict) -> str: handlers = {"search_codebase": self._search, "execute_python": self._execute, "read_file": self._read, "run_tests": self._test} return handlers.get(tool_name, lambda **kw: f"Unknown tool: {tool_name}")(**arguments) def _search(self, query, file_pattern=None, max_results=5): import fnmatch results = self.retriever.search(query, top_k=max_results) output = [{"file": c.file_path, "name": c.name, "type": c.chunk_type, "score": round(s, 3), "signature": c.signature, "content": c.content[:500]} for c, s in results if not file_pattern or fnmatch.fnmatch(c.file_path, file_pattern)] return json.dumps({"results": output}, indent=2) def _execute(self, code, timeout=30): try: r = subprocess.run(["python", "-c", code], capture_output=True, text=True, timeout=timeout, cwd=self.repo_path) return (r.stdout or "") + (r.stderr or "") or "Success (no output)" except subprocess.TimeoutExpired: return f"Timed out after {timeout}s" def _read(self, path, start_line=None, end_line=None): fp = os.path.join(self.repo_path, path) if not os.path.exists(fp): return f"File not found: {path}" with open(fp) as f: lines = f.readlines() if start_line: lines = lines[start_line-1:end_line] return "".join(lines) def _test(self, test_path, verbose=True): cmd = ["python", "-m", "pytest", test_path] + (["-v"] if verbose else []) try: r = subprocess.run(cmd, capture_output=True, text=True, timeout=120, cwd=self.repo_path) return r.stdout + r.stderr except subprocess.TimeoutExpired: return "Tests timed out" class CodeAgent: """ReAct-style agent combining fine-tuned LLM with tool execution.""" def __init__(self, model_id, retriever, repo_path, max_turns=10, device="auto"): self.tool_executor = ToolExecutor(retriever, repo_path) self.context_builder = ContextBuilder(retriever) self.max_turns = max_turns from transformers import AutoModelForCausalLM, AutoTokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) self.model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map=device, trust_remote_code=True) def run(self, user_query, current_file=None): rag_context = self.context_builder.build_context(user_query) user_content = user_query if rag_context: user_content += f"\\n\\n--- Code context ---\\n{rag_context}\\n--- End ---" messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_content}] for turn in range(self.max_turns): response = self._generate(messages) tool_calls = self._extract_tool_calls(response) if not tool_calls: return response messages.append({"role": "assistant", "content": response}) for tc in tool_calls: result = self.tool_executor.execute(tc["name"], tc["arguments"]) messages.append({"role": "tool", "name": tc["name"], "content": result[:3000]}) return response def _generate(self, messages): import torch text = self.tokenizer.apply_chat_template(messages, tools=TOOLS, tokenize=False, add_generation_prompt=True) inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device) with torch.no_grad(): outputs = self.model.generate(**inputs, max_new_tokens=4096, temperature=0.7, top_p=0.9, do_sample=True) return self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True).strip() def _extract_tool_calls(self, response): tool_calls = [] for pattern in [r'\\{"name":\\s*"(\\w+)",\\s*"arguments":\\s*(\\{[^}]+\\})\\}', r'\\n?(.*?)\\n?']: for m in re.finditer(pattern, response, re.DOTALL): try: if m.lastindex == 2: tool_calls.append({"name": m.group(1), "arguments": json.loads(m.group(2))}) else: data = json.loads(m.group(1)) if "name" in data: tool_calls.append(data) except: pass return tool_calls if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--model", required=True) parser.add_argument("--repo", required=True) parser.add_argument("--query", "-q", default=None) parser.add_argument("--index-dir", default=None) parser.add_argument("--embedding-model", default="jinaai/jina-embeddings-v2-base-code") args = parser.parse_args() if args.index_dir and os.path.exists(os.path.join(args.index_dir, "chunks.json")): retriever = CodeRetriever(args.embedding_model) retriever.load_index(args.index_dir) else: retriever = CodebaseIndexer(args.repo, embedding_model=args.embedding_model).index() if args.index_dir: retriever.save_index(args.index_dir) agent = CodeAgent(args.model, retriever, args.repo) if args.query: print(agent.run(args.query)) else: print("Code Assistant (type 'quit' to exit)") while True: q = input("> ").strip() if q.lower() in ("quit", "exit", "q"): break if q: print(agent.run(q))