diff --git a/graphify/extract.py b/graphify/extract.py index aeaca1a5..21c1508c 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -1960,6 +1960,7 @@ def extract_go(path: Path) -> dict: edges: list[dict] = [] seen_ids: set[str] = set() function_bodies: list[tuple[str, object]] = [] + go_imported_pkgs: set[str] = set() # local names of imported packages def add_node(nid: str, label: str, line: int) -> None: if nid not in seen_ids: @@ -2057,12 +2058,21 @@ def extract_go(path: Path) -> dict: # don't collide with local files of the same basename. tgt_nid = _make_id("go", "pkg", raw) add_edge(file_nid, tgt_nid, "imports_from", spec.start_point[0] + 1) + # Track local name (alias or last path segment) + alias = spec.child_by_field_name("name") + local_name = _read_text(alias, source) if alias else raw.split("/")[-1] + if local_name and local_name != "_" and local_name != ".": + go_imported_pkgs.add(local_name) elif child.type == "import_spec": path_node = child.child_by_field_name("path") if path_node: raw = _read_text(path_node, source).strip('"') tgt_nid = _make_id("go", "pkg", raw) add_edge(file_nid, tgt_nid, "imports_from", child.start_point[0] + 1) + alias = child.child_by_field_name("name") + local_name = _read_text(alias, source) if alias else raw.split("/")[-1] + if local_name and local_name != "_" and local_name != ".": + go_imported_pkgs.add(local_name) return for child in node.children: @@ -2090,8 +2100,12 @@ def extract_go(path: Path) -> dict: if func_node.type == "identifier": callee_name = _read_text(func_node, source) elif func_node.type == "selector_expression": - is_member_call = True field = func_node.child_by_field_name("field") + operand = func_node.child_by_field_name("operand") + receiver_name = _read_text(operand, source) if operand else "" + # Package-qualified call (e.g. fmt.Println) → allow cross-file resolution. + # Receiver method call (e.g. s.logger.Log) → skip, no import evidence. + is_member_call = receiver_name not in go_imported_pkgs if field: callee_name = _read_text(field, source) if callee_name: diff --git a/graphify/llm.py b/graphify/llm.py index a9df0e79..ec697734 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -6,7 +6,9 @@ from __future__ import annotations import json import os +import sys import time +from collections.abc import Callable from pathlib import Path BACKENDS: dict[str, dict] = { @@ -57,6 +59,20 @@ def _read_files(paths: list[Path], root: Path) -> str: return "\n\n".join(parts) +def _parse_llm_json(raw: str) -> dict: + """Strip optional markdown fences and parse JSON. Returns empty fragment on failure.""" + if raw.startswith("```"): + raw = raw.split("```", 2)[1] + if raw.startswith("json"): + raw = raw[4:] + raw = raw.rsplit("```", 1)[0] + try: + return json.loads(raw.strip()) + except json.JSONDecodeError as exc: + print(f"[graphify] LLM returned invalid JSON, skipping chunk: {exc}", file=sys.stderr) + return {"nodes": [], "edges": [], "hyperedges": []} + + def _call_openai_compat( base_url: str, api_key: str, @@ -82,14 +98,7 @@ def _call_openai_compat( max_completion_tokens=8192, temperature=0, ) - raw = resp.choices[0].message.content or "{}" - # Strip markdown fences if model adds them despite instructions - if raw.startswith("```"): - raw = raw.split("```", 2)[1] - if raw.startswith("json"): - raw = raw[4:] - raw = raw.rsplit("```", 1)[0] - result = json.loads(raw.strip()) + result = _parse_llm_json(resp.choices[0].message.content or "{}") result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0 result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0 result["model"] = model @@ -113,13 +122,7 @@ def _call_claude(api_key: str, model: str, user_message: str) -> dict: system=_EXTRACTION_SYSTEM, messages=[{"role": "user", "content": user_message}], ) - raw = resp.content[0].text if resp.content else "{}" - if raw.startswith("```"): - raw = raw.split("```", 2)[1] - if raw.startswith("json"): - raw = raw[4:] - raw = raw.rsplit("```", 1)[0] - result = json.loads(raw.strip()) + result = _parse_llm_json(resp.content[0].text if resp.content else "{}") result["input_tokens"] = resp.usage.input_tokens if resp.usage else 0 result["output_tokens"] = resp.usage.output_tokens if resp.usage else 0 result["model"] = model @@ -164,7 +167,7 @@ def extract_corpus_parallel( model: str | None = None, root: Path = Path("."), chunk_size: int = 20, - on_chunk_done: object = None, + on_chunk_done: Callable | None = None, ) -> dict: """Extract a corpus in chunks, merging results.