Upload inference.py

inference.py

@@ -0,0 +1,178 @@
+"""
+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'<tool_call>\\n?(.*?)\\n?</tool_call>']:
+            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))