File size: 30,814 Bytes
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0200006
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bfa6ae5
 
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
503bd66
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bfa6ae5
 
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fb0178d
 
1b35d41
 
0200006
 
 
 
1b35d41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
"""
Research Agent β€” Autonomous Migration Context Gatherer
======================================================

A separate LLM-driven agent that runs BEFORE the fix agent. It gets up to
12 steps to autonomously investigate what broke and why, using real tools
that fetch live documentation from PyPI, GitHub, Python docs, and the web.

The agent is NOT limited to a hardcoded database. It can:
  - Parse test errors to identify the exact exception and involved packages
  - Look up package info on PyPI (versions, homepage, changelog URL)
  - Fetch actual documentation pages, changelogs, and migration guides
  - Read Python "What's New" pages for specific versions
  - Search GitHub release notes for breaking change details
  - Compare old vs new dependency versions from the task metadata
  - Synthesize all findings into a focused context for the fix agent

The key insight: the model drives the research strategy. We give it good
tools and let it figure out what to look up based on the actual error.
"""

from __future__ import annotations

import json
import logging
import re
import time
from typing import Any, Dict, List, Optional
from urllib.request import urlopen, Request
from urllib.error import URLError

log = logging.getLogger("research_agent")


# ---------------------------------------------------------------------------
# Model family detection
# ---------------------------------------------------------------------------
def _detect_model_family(model_name: str) -> str:
    name_lower = model_name.lower()
    if "gemma" in name_lower:
        return "gemma"
    if "qwen3" in name_lower:
        return "qwen3"
    if "qwen" in name_lower:
        return "qwen2"
    return "unknown"


def _strip_model_artifacts(raw_text: str, family: str) -> str:
    clean = raw_text
    if family == "gemma":
        clean = re.sub(r"<\|channel>thought\n.*?<channel\|>", "", clean, flags=re.DOTALL)
        for tok in ["<turn|>", "<|turn>", "<eos>", "</s>"]:
            clean = clean.replace(tok, "")
    elif family == "qwen3":
        clean = re.sub(r"<think>.*?</think>", "", clean, flags=re.DOTALL)
        # Truncate at first <|im_end|> β€” everything after is hallucinated
        im_end = clean.find("<|im_end|>")
        if im_end != -1:
            clean = clean[:im_end]
        for tok in ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]:
            clean = clean.replace(tok, "")
    elif family == "qwen2":
        im_end = clean.find("<|im_end|>")
        if im_end != -1:
            clean = clean[:im_end]
        for tok in ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]:
            clean = clean.replace(tok, "")
    else:
        for tok in ["<eos>", "</s>", "<|im_end|>", "<|endoftext|>", "<turn|>", "<|turn>"]:
            clean = clean.replace(tok, "")
    return clean.strip()


# ───────────────────────────────────────────────────────────────
# Research Tools β€” real, live, no hardcoded DB
# ───────────────────────────────────────────────────────────────

def tool_parse_error(test_output: str) -> str:
    """Parse test output to extract every error, traceback file, and package."""
    errors = []
    packages = set()

    for m in re.finditer(
        r"(ImportError|ModuleNotFoundError|AttributeError|TypeError|NameError|"
        r"SyntaxError|ValueError|DeprecationWarning|RuntimeError|KeyError): (.+)",
        test_output,
    ):
        errors.append(f"{m.group(1)}: {m.group(2).strip()}")
        for pkg in re.finditer(r"'(\w+(?:\.\w+)*)'", m.group(2)):
            top = pkg.group(1).split(".")[0]
            if top not in ("builtins", "importlib", "usr", "work"):
                packages.add(top)

    files = re.findall(r'File "(/work/\S+\.py)", line (\d+)', test_output)
    if not files:
        files = re.findall(r"(/work/\S+\.py):(\d+)", test_output)

    failed = re.search(r"(\d+) (?:failed|error)", test_output)
    passed = re.search(r"(\d+) passed", test_output)

    out = "=== ERROR ANALYSIS ===\n"
    if errors:
        out += f"Errors found ({len(errors)}):\n"
        for e in errors[:15]:
            out += f"  - {e}\n"
    else:
        out += "Could not parse specific errors. Check the raw test output.\n"

    if files:
        unique = list(dict.fromkeys(f"{f}:{ln}" for f, ln in files))
        out += f"\nSource files referenced ({len(unique)}):\n"
        for f in unique[:15]:
            out += f"  - {f}\n"

    if packages:
        out += f"\nPackages involved: {', '.join(sorted(packages))}\n"

    if failed:
        out += f"\nTest summary: {failed.group(0)}"
        if passed:
            out += f", {passed.group(0)}"
        out += "\n"

    return out


def tool_search_pypi(package_name: str) -> str:
    """Fetch package metadata from PyPI: version, summary, homepage, changelog URL."""
    try:
        url = f"https://pypi.org/pypi/{package_name}/json"
        req = Request(url, headers={"User-Agent": "code-migration-research/1.0"})
        with urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read())

        info = data.get("info", {})
        version = info.get("version", "unknown")
        summary = info.get("summary", "")
        home_page = info.get("home_page", "") or ""
        project_urls = info.get("project_urls") or {}

        # Try to find changelog/release notes URL
        changelog_url = ""
        for key in ["Changelog", "Changes", "Release Notes", "History",
                     "Release notes", "CHANGELOG", "What's New"]:
            if key in project_urls:
                changelog_url = project_urls[key]
                break

        # Also grab the GitHub URL if available
        github_url = ""
        for key in ["Source", "Source Code", "Repository", "Homepage", "GitHub"]:
            if key in project_urls:
                val = project_urls[key]
                if "github.com" in val:
                    github_url = val
                    break
        if not github_url and home_page and "github.com" in home_page:
            github_url = home_page

        result = f"Package: {package_name}\n"
        result += f"Latest version: {version}\n"
        result += f"Summary: {summary}\n"
        if home_page:
            result += f"Homepage: {home_page}\n"
        if changelog_url:
            result += f"Changelog URL: {changelog_url}\n"
        if github_url:
            result += f"GitHub: {github_url}\n"

        # List all project URLs for the agent to explore
        if project_urls:
            result += "All project URLs:\n"
            for k, v in project_urls.items():
                result += f"  {k}: {v}\n"

        return result
    except Exception as e:
        return f"PyPI lookup failed for '{package_name}': {e}"


def tool_fetch_url(url: str) -> str:
    """Fetch and return clean text content from any URL. All HTML is stripped.

    Returns up to 5000 chars of readable text content.
    """
    try:
        req = Request(url, headers={
            "User-Agent": "code-migration-research/1.0",
            "Accept": "text/html,application/xhtml+xml,text/plain,*/*",
        })
        with urlopen(req, timeout=15) as resp:
            raw = resp.read().decode("utf-8", errors="replace")

        # Remove script and style blocks entirely
        text = re.sub(r"<script[^>]*>.*?</script>", "", raw, flags=re.DOTALL | re.IGNORECASE)
        text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags=re.DOTALL | re.IGNORECASE)
        # Remove HTML comments
        text = re.sub(r"<!--.*?-->", "", text, flags=re.DOTALL)
        # Replace block-level tags with newlines for readability
        text = re.sub(r"<(?:br|p|div|li|tr|h\d|dt|dd|section|article)[^>]*>", "\n", text, flags=re.IGNORECASE)
        # Strip all remaining HTML tags
        text = re.sub(r"<[^>]+>", " ", text)
        # Decode common HTML entities
        text = text.replace("&amp;", "&").replace("&lt;", "<").replace("&gt;", ">")
        text = text.replace("&quot;", '"').replace("&#39;", "'").replace("&nbsp;", " ")
        text = re.sub(r"&#\d+;", " ", text)
        text = re.sub(r"&\w+;", " ", text)
        # Collapse inline whitespace but keep newlines
        text = re.sub(r"[^\S\n]+", " ", text)
        # Collapse multiple blank lines into one
        text = re.sub(r"\n\s*\n+", "\n\n", text)
        text = text.strip()

        if not text:
            return f"Page at {url} returned empty content after HTML stripping."

        return text[:5000]
    except Exception as e:
        return f"Failed to fetch {url}: {e}"


def tool_fetch_python_docs(version: str) -> str:
    """Fetch the official Python 'What's New' page for a specific version.

    Example: version="3.12" fetches https://docs.python.org/3.12/whatsnew/3.12.html
    Returns the text content focused on breaking changes and removals.
    """
    minor = version.split(".")
    if len(minor) >= 2:
        ver = f"{minor[0]}.{minor[1]}"
    else:
        ver = version

    url = f"https://docs.python.org/{ver}/whatsnew/{ver}.html"
    try:
        req = Request(url, headers={"User-Agent": "code-migration-research/1.0"})
        with urlopen(req, timeout=15) as resp:
            raw = resp.read().decode("utf-8", errors="replace")

        # Strip HTML
        text = re.sub(r"<script[^>]*>.*?</script>", " ", raw, flags=re.DOTALL)
        text = re.sub(r"<style[^>]*>.*?</style>", " ", text, flags=re.DOTALL)
        text = re.sub(r"<[^>]+>", " ", text)
        text = re.sub(r"&nbsp;", " ", text)
        text = re.sub(r"&[a-z]+;", " ", text)
        text = re.sub(r"\s+", " ", text).strip()

        # Try to extract just the "Removed" and "Deprecated" sections
        # by looking for those keywords and grabbing surrounding context
        sections = []
        for keyword in ["Removed", "removed", "Deprecated", "deprecated",
                        "Breaking", "breaking", "Porting"]:
            for m in re.finditer(keyword, text):
                start = max(0, m.start() - 100)
                end = min(len(text), m.end() + 800)
                chunk = text[start:end].strip()
                if chunk and chunk not in sections:
                    sections.append(chunk)

        if sections:
            result = f"Python {ver} What's New β€” breaking changes:\n\n"
            result += "\n---\n".join(sections[:8])
            return result[:5000]

        # Fallback: return first 5000 chars
        return f"Python {ver} What's New (full page excerpt):\n{text[:5000]}"
    except Exception as e:
        return f"Failed to fetch Python {ver} What's New: {e}"


def tool_search_github_releases(owner: str, repo: str) -> str:
    """Fetch the latest 5 release notes from a GitHub repository.

    Use this to find breaking changes, migration notes, and changelog entries.
    Pass owner and repo separately, e.g. owner="numpy", repo="numpy".
    """
    try:
        url = f"https://api.github.com/repos/{owner}/{repo}/releases?per_page=5"
        req = Request(url, headers={
            "User-Agent": "code-migration-research/1.0",
            "Accept": "application/vnd.github.v3+json",
        })
        with urlopen(req, timeout=10) as resp:
            releases = json.loads(resp.read())

        if not releases:
            return f"No releases found for {owner}/{repo}."

        result = []
        for rel in releases[:5]:
            tag = rel.get("tag_name", "unknown")
            name = rel.get("name", tag)
            body = rel.get("body", "")[:1200]
            result.append(f"=== {name} ({tag}) ===\n{body}")

        return "\n\n".join(result)[:5000]
    except Exception as e:
        return f"GitHub releases lookup failed for {owner}/{repo}: {e}"


def tool_compare_versions(dependency_versions: str, package_name: str) -> str:
    """Extract the installed version of a specific package from the dependency list.

    This helps you know exactly which version is installed so you can look up
    the right migration guide (e.g., numpy 1.x β†’ 2.x has different breaks than 1.24 β†’ 1.26).
    """
    if not dependency_versions:
        return f"No dependency version info available for {package_name}."

    for line in dependency_versions.strip().splitlines():
        line = line.strip()
        if not line:
            continue
        # Parse "package==version" or "package>=version"
        parts = re.split(r"[=><~!]+", line, maxsplit=1)
        if len(parts) >= 2:
            pkg = parts[0].strip().lower().replace("-", "_")
            ver = parts[1].strip()
            if pkg == package_name.lower().replace("-", "_"):
                return f"{package_name} installed version: {ver}"

    return f"Package '{package_name}' not found in dependency list."


def tool_find_changelog(github_url: str) -> str:
    """Given a GitHub repo URL, try to find and fetch the CHANGELOG/CHANGES/HISTORY file.

    Automatically tries common changelog filenames on both main and master branches.
    This is the most reliable way to find breaking changes for a package.

    Example: github_url="https://github.com/numpy/numpy"
    """
    # Extract owner/repo from URL
    m = re.search(r"github\.com/([^/]+)/([^/]+)", github_url)
    if not m:
        return f"Could not parse GitHub owner/repo from: {github_url}"

    owner = m.group(1)
    repo = m.group(2).rstrip("/").replace(".git", "")

    # Try common changelog filenames on common branches
    filenames = [
        "CHANGELOG.md", "CHANGELOG.rst", "CHANGELOG",
        "CHANGES.md", "CHANGES.rst", "CHANGES",
        "HISTORY.md", "HISTORY.rst",
        "NEWS.md", "NEWS.rst",
        "RELEASE_NOTES.md",
    ]
    branches = ["main", "master"]

    for branch in branches:
        for fname in filenames:
            raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/{branch}/{fname}"
            try:
                req = Request(raw_url, headers={"User-Agent": "code-migration-research/1.0"})
                with urlopen(req, timeout=5) as resp:
                    content = resp.read().decode("utf-8", errors="replace")

                if len(content) < 20:
                    continue

                # Found it β€” return a useful chunk
                # Try to find breaking/removed/deprecated sections
                sections = []
                for keyword in ["breaking", "removed", "deprecated", "incompatible",
                                "migration", "upgrade", "BREAKING", "Removed", "Deprecated"]:
                    for km in re.finditer(keyword, content, re.IGNORECASE):
                        start = max(0, km.start() - 200)
                        end = min(len(content), km.end() + 1500)
                        chunk = content[start:end].strip()
                        if chunk not in sections:
                            sections.append(chunk)

                if sections:
                    result = f"Found {fname} in {owner}/{repo} ({branch}):\n\n"
                    result += "\n---\n".join(sections[:6])
                    return result[:5000]

                # No keyword matches β€” return the first 5000 chars
                return f"Found {fname} in {owner}/{repo} ({branch}):\n\n{content[:5000]}"

            except Exception:
                continue

    return f"No changelog file found in {owner}/{repo}. Try search_github_releases or fetch_url with a specific docs URL."


# ───────────────────────────────────────────────────────────────
# Tool registry
# ───────────────────────────────────────────────────────────────

RESEARCH_TOOLS = {
    "parse_error": tool_parse_error,
    "search_pypi": tool_search_pypi,
    "fetch_url": tool_fetch_url,
    "fetch_python_docs": tool_fetch_python_docs,
    "search_github_releases": tool_search_github_releases,
    "compare_versions": tool_compare_versions,
    "find_changelog": tool_find_changelog,
}


# ───────────────────────────────────────────────────────────────
# System prompt β€” the agent drives the research strategy
# ───────────────────────────────────────────────────────────────

RESEARCH_SYSTEM_PROMPT = """You are a research agent. Your job is to investigate WHY tests are failing
after a Python project's dependencies were upgraded, and find the EXACT fix patterns.

You have these tools (output ONE JSON tool call per turn):

1. parse_error(test_output)
   - Extracts all errors, tracebacks, file paths, and package names from test output
   - ALWAYS call this first to understand what broke

2. compare_versions(dependency_versions, package_name)
   - Shows the exact installed version of a package
   - Use this to know if it's a major version jump (e.g., numpy 1.x β†’ 2.x)

3. search_pypi(package_name)
   - Gets package info from PyPI: latest version, homepage, changelog URL, GitHub URL
   - Use this to find where the docs/changelog live

4. find_changelog(github_url)
   - Given a GitHub URL, finds and fetches the CHANGELOG/CHANGES/HISTORY file
   - Automatically searches common filenames and branches
   - Returns sections mentioning breaking changes, removals, deprecations
   - THIS IS YOUR BEST TOOL for finding what changed between versions

5. fetch_url(url)
   - Fetches any URL and returns text content (HTML stripped)
   - Use this to read specific docs pages, migration guides, or changelog URLs from PyPI

6. fetch_python_docs(version)
   - Fetches the official Python "What's New" page for a version (e.g., "3.12")
   - Shows removed modules, deprecated APIs, breaking changes

7. search_github_releases(owner, repo)
   - Gets recent release notes from GitHub
   - Good for finding breaking changes when there's no changelog file

Output format: {"name": "tool_name", "arguments": {"arg": "value"}}

STRATEGY:
1. parse_error β€” understand exactly what errors occurred and which packages are involved
2. compare_versions β€” check what version of the failing package is installed (is it a major bump?)
3. search_pypi β€” find the package's GitHub URL
4. find_changelog β€” fetch the changelog from GitHub and look for breaking changes
5. If changelog doesn't have what you need, try search_github_releases or fetch_url on the docs
6. For Python stdlib issues, use fetch_python_docs
7. IMPORTANT: Also check the dependency versions list for ANY major version bumps
   (e.g., numpy 1.x→2.x, pandas 1.x→2.x, Django 3.x→5.x). These often have
   breaking changes that won't show up in the CURRENT test errors but will appear
   after you fix the first error. Research ALL major-bumped packages, not just
   the one currently failing.
8. When you have enough info: {"name": "done", "arguments": {"summary": "YOUR FINDINGS"}}

Your summary MUST include for EACH error:
- The exact error message
- The root cause (what was removed/renamed/moved in which version)
- The EXACT fix: old code β†’ new code replacement
- Which file(s) need to be changed

Also include any POTENTIAL breaking changes from major version bumps in the
dependency list, even if they don't show up in the current test errors yet.
For example, if pandas went from 1.x to 2.x, mention that DataFrame.append()
was removed and the fix is pd.concat().

RULES:
- NEVER call the same tool with the same arguments twice
- Don't just look things up β€” READ the actual changelog. Use find_changelog.
- Be specific: "replace np.math.factorial with math.factorial" not "check numpy docs"
- You have limited steps β€” be efficient. parse_error β†’ search_pypi β†’ find_changelog β†’ done.
- If the dependency list shows multiple major version bumps, research ALL of them.
"""


# ───────────────────────────────────────────────────────────────
# Research Agent class
# ───────────────────────────────────────────────────────────────

class ResearchAgent:
    """LLM-driven agent that autonomously researches migration breaking changes."""

    def __init__(self, model, tokenizer, max_steps: int = 12, model_name: str = ""):
        self.model = model
        self.tokenizer = tokenizer
        self.max_steps = max_steps
        self.model_family = _detect_model_family(model_name) if model_name else "unknown"
        self.last_research_steps: List[Dict] = []

    def research(
        self,
        repo_name: str,
        old_python: str,
        new_python: str,
        related_modules: str,
        test_output: str,
        dependency_versions: str = "",
        on_step: Any = None,
    ) -> str:
        """Run the research loop and return migration context for the fix agent."""
        import torch

        log.info("  [RESEARCH] Starting for %s (%s β†’ %s, modules=%s)",
                 repo_name, old_python, new_python, related_modules)

        # Build the initial context with everything the agent needs
        dep_summary = ""
        if dependency_versions:
            dep_summary = f"\nInstalled dependency versions:\n{dependency_versions}\n"

        # Highlight which modules are involved so the agent researches all of them
        modules_note = ""
        if related_modules and "," in related_modules:
            mods = [m.strip() for m in related_modules.split(",")]
            modules_note = (
                f"\nIMPORTANT: This project uses multiple upgraded packages: {', '.join(mods)}. "
                f"Research breaking changes for ALL of them, not just the one in the current error. "
                f"After fixing one error, another package's breaking change may surface.\n"
            )

        initial_context = (
            f"Repository: {repo_name}\n"
            f"Python version migration: {old_python} β†’ {new_python}\n"
            f"Related modules: {related_modules}\n"
            f"{modules_note}"
            f"{dep_summary}\n"
            f"Failing test output (last 3000 chars):\n"
            f"{test_output[-3000:]}\n"
        )

        messages = [
            {"role": "system", "content": RESEARCH_SYSTEM_PROMPT},
            {"role": "user", "content": initial_context},
        ]

        gathered_info: List[str] = []
        research_steps: List[Dict] = []
        seen_calls: set = set()  # track ALL previous calls, not just last

        for step in range(1, self.max_steps + 1):
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            # Generate
            try:
                text = self.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True,
                    **({"enable_thinking": False} if self.model_family == "qwen3" else {}),
                )
                if self.model_family == "gemma":
                    text = text.replace("<|think|>", "")

                inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
                input_len = inputs["input_ids"].shape[-1]

                # Build gen kwargs with stop tokens for Qwen
                gen_kwargs = dict(
                    max_new_tokens=400,
                    temperature=0.3,
                    top_p=0.95,
                    do_sample=True,
                    pad_token_id=self.tokenizer.pad_token_id,
                )
                if self.model_family in ("qwen3", "qwen2"):
                    stop_ids = []
                    for tok_str in ["<|im_end|>", "<|endoftext|>"]:
                        tid = self.tokenizer.convert_tokens_to_ids(tok_str)
                        if tid is not None and tid != self.tokenizer.unk_token_id:
                            stop_ids.append(tid)
                    if stop_ids:
                        eos = self.tokenizer.eos_token_id
                        if isinstance(eos, int):
                            stop_ids.append(eos)
                        gen_kwargs["eos_token_id"] = list(set(stop_ids))

                with torch.no_grad():
                    outputs = self.model.generate(**inputs, **gen_kwargs)

                raw = self.tokenizer.decode(outputs[0][input_len:], skip_special_tokens=False)
                del inputs, outputs
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

                clean = _strip_model_artifacts(raw, self.model_family)

            except Exception as e:
                log.warning("  [RESEARCH] Step %d generation failed: %s", step, e)
                break

            # Parse tool call
            parsed = _parse_research_call(clean)
            tool_name = parsed["tool_name"]
            tool_args = parsed["tool_args"]

            # Done signal
            if tool_name == "done":
                summary = tool_args.get("summary", "")
                if summary:
                    gathered_info.append(f"=== RESEARCH SUMMARY ===\n{summary}")
                log.info("  [RESEARCH] Step %d: done (summary=%d chars)", step, len(summary))
                break

            # Anti-repeat: block any call we've already made
            curr_key = f"{tool_name}:{json.dumps(tool_args, sort_keys=True)}"
            if curr_key in seen_calls:
                log.info("  [RESEARCH] Step %d: already called %s with same args β†’ forcing done", step, tool_name)
                break
            seen_calls.add(curr_key)

            # Execute tool
            result = self._execute_tool(
                tool_name, tool_args,
                test_output=test_output,
                dependency_versions=dependency_versions,
                new_python=new_python,
                related_modules=related_modules,
            )

            gathered_info.append(result)
            args_short = json.dumps(tool_args, default=str)[:150]
            log.info("  [RESEARCH] Step %d: %s(%s) β†’ %d chars",
                     step, tool_name, args_short, len(result))

            research_steps.append({
                "step": step,
                "tool": tool_name,
                "args": tool_args,
                "result": result,
                "model_output": clean,
            })

            # Live callback for streaming steps to the frontend
            if on_step:
                on_step(research_steps[-1])

            # Feed result back to the agent
            messages.append({"role": "assistant", "content": clean})
            messages.append({"role": "user", "content": f"Tool result:\n{result}"})

        self.last_research_steps = research_steps

        if not gathered_info:
            log.warning("  [RESEARCH] No info gathered, returning empty context")
            return "Research agent could not gather migration context. Use general debugging."

        return "\n\n".join(gathered_info)

    def _execute_tool(
        self,
        tool_name: str,
        tool_args: Dict[str, Any],
        *,
        test_output: str,
        dependency_versions: str,
        new_python: str,
        related_modules: str,
    ) -> str:
        """Dispatch a tool call with proper argument handling."""
        try:
            if tool_name == "parse_error":
                return tool_parse_error(test_output)

            elif tool_name == "compare_versions":
                return tool_compare_versions(
                    dependency_versions,
                    tool_args.get("package_name", related_modules),
                )

            elif tool_name == "search_pypi":
                return tool_search_pypi(
                    tool_args.get("package_name", related_modules),
                )

            elif tool_name == "fetch_url":
                url = tool_args.get("url", "")
                if not url:
                    return "Error: fetch_url requires a 'url' argument."
                return tool_fetch_url(url)

            elif tool_name == "fetch_python_docs":
                return tool_fetch_python_docs(
                    tool_args.get("version", new_python),
                )

            elif tool_name == "search_github_releases":
                owner = tool_args.get("owner", "")
                repo = tool_args.get("repo", "")
                if not owner or not repo:
                    return "Error: search_github_releases requires 'owner' and 'repo'."
                return tool_search_github_releases(owner, repo)

            elif tool_name == "find_changelog":
                github_url = tool_args.get("github_url", "")
                if not github_url:
                    return "Error: find_changelog requires a 'github_url' argument."
                return tool_find_changelog(github_url)

            else:
                return f"Unknown tool: {tool_name}. Available: {', '.join(RESEARCH_TOOLS.keys())}"

        except Exception as e:
            return f"Tool '{tool_name}' error: {e}"


def _parse_research_call(text: str) -> Dict[str, Any]:
    """Parse a JSON tool call from model output.

    Finds the FIRST complete JSON object to avoid hallucinated multi-turn content.
    """
    text = text.strip()

    # Strip markdown fences
    if text.startswith("```"):
        lines = text.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        text = "\n".join(lines).strip()

    # Find first { and match its closing }
    start = text.find("{")
    if start != -1:
        depth = 0
        for i in range(start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(text[start:i + 1])
                        name = data.get("name", data.get("tool_name", ""))
                        args = data.get("arguments", data.get("tool_args", data.get("parameters", {})))
                        if name:
                            return {"tool_name": name, "tool_args": args}
                    except json.JSONDecodeError:
                        pass
                    break

    # Fallback: start with parse_error
    return {"tool_name": "parse_error", "tool_args": {}}