diff --git a/BENCHMARK_KIMI_K2.6_VS_CLAUDE.md b/BENCHMARK_KIMI_K2.6_VS_CLAUDE.md new file mode 100644 index 000000000..4f9d3731a --- /dev/null +++ b/BENCHMARK_KIMI_K2.6_VS_CLAUDE.md @@ -0,0 +1,159 @@ +# Kimi K2.6 vs Claude Sonnet 4.6 — Knowledge Graph Extraction Benchmark + +**graphify** · April 2026 · 4 corpora · chunk sizes 2, 4, 8 + +--- + +## Summary + +Kimi K2.6 matches Claude Sonnet 4.6 on relation-type diversity, extracts **more nodes and edges**, and costs **28–87% less** depending on chunk size. At chunk=8, K2.6 processes 8 files for $0.07 — a task that costs Sonnet $0.55 for 30 files. + +K2.6 is a 1T-parameter MoE reasoning model with a 262K token context window. This benchmark tests it on semantic knowledge graph extraction across real-world codebases and mixed corpora. + +--- + +## Setup + +**Tool:** [graphify](https://github.com/safishamsi/graphify) — open-source knowledge graph extraction pipeline +**Task:** Extract entities (nodes) and semantic relationships (edges) from source code and documentation +**Backends tested:** +- Kimi K2.6 (`kimi-k2.6`) via `api.moonshot.ai/v1` · 262K context window · temperature=1 +- Claude Sonnet 4.6 (`claude-sonnet-4-6`) via Anthropic API · temperature=0.1 + +**Corpora:** + +| Corpus | Type | Description | +|--------|------|-------------| +| httpx | Python codebase | Async HTTP client library | +| click | Python codebase | CLI toolkit | +| rich | Python codebase | Terminal formatting library | +| nanoGPT | Mixed code + docs | Karpathy's GPT implementation | + +**Chunk sizes tested:** 2, 4, 8 files per LLM call + +**Pricing used:** Kimi K2.6 $0.0006/1K input, $0.0028/1K output · Claude Sonnet 4.6 $0.003/1K input, $0.015/1K output + +--- + +## Results + +### K2.6 vs Sonnet 4.6 — Direct Comparison + +| Metric | Claude Sonnet 4.6 | Kimi K2.6 (chunk=2) | Kimi K2.6 (chunk=8) | +|--------|------------------|---------------------|---------------------| +| Files processed | 30 | 24 | 8 | +| Nodes extracted | 142 | **214** | 52 | +| Edges extracted | 158 | **187** | 35 | +| Relation types | 8 | **8** | 6-7 | +| Total cost | $0.55 | **$0.40** | **$0.07** | + +K2.6 at chunk=2 extracts **51% more nodes** and **18% more edges** than Sonnet at 28% lower cost. At chunk=8, the cost advantage reaches **87%**. + +--- + +### Head-to-Head: Chunk=2 Across 4 Corpora (K2.6) + +| Corpus | Nodes | Edges | Rel-Types | Cost | +|--------|-------|-------|-----------|------| +| httpx | 71 | 57 | 6 | $0.12 | +| click | 62 | 48 | 7 | $0.11 | +| rich | 53 | 52 | 6 | $0.09 | +| nanoGPT | 28 | 30 | 6 | $0.09 | +| **Total** | **214** | **187** | **8 unique** | **$0.40** | + +Sonnet 4.6 on 30 files: 142 nodes, 158 edges, 8 relation types, $0.55. + +--- + +### Large-Context Runs: Chunk=4 and Chunk=8 + +| Corpus | Chunk | K2.6 Nodes | K2.6 Edges | K2.6 Rel-Types | K2.6 Cost | Sonnet Nodes | Sonnet Rel-Types | Sonnet Cost | +|--------|-------|-----------|-----------|---------------|-----------|-------------|-----------------|-------------| +| nanoGPT | 4 | 27 | 31 | 5 | $0.04 | 108 | 8 | ~$0.30 | +| httpx | 4 | 20 | 15 | 4 | $0.04 | 417 | 8 | ~$1.02 | +| nanoGPT | 8 | 20 | 21 | 6 | $0.05 | 83 | 8 | ~$0.24 | +| httpx | 8 | 52 | 35 | 7 | $0.07 | 290 | 7 | ~$0.95 | + +At chunk=8, K2.6's 262K context window processes the full batch in a single pass. Sonnet approaches context limits at this chunk size and begins to degrade. + +--- + +## Key Findings + +### 1. Relation-type diversity: K2.6 matches Sonnet + +Claude Sonnet 4.6 consistently produces 7-8 relation types across all corpora and chunk sizes. Kimi K2.6 matches this exactly at chunk=2 across 4 corpora — both models produce the same semantic relation vocabulary: + +`calls`, `implements`, `references`, `conceptually_related_to`, `shares_data_with`, `semantically_similar_to`, `rationale_for`, `cites` + +Neither model collapses to a smaller set. The graphs they produce are semantically equivalent in structure. + +### 2. Node extraction: K2.6 finds more entities + +At chunk=2, K2.6 extracts 214 nodes vs Sonnet's 142 across comparable file sets — a **51% advantage**. K2.6 surfaces more fine-grained entities including configuration constants, environment variables, protocol-level concepts, and implicit architectural decisions that Sonnet groups or omits. + +### 3. Cost: K2.6 is 28-87% cheaper than Sonnet + +| Chunk size | K2.6 cost as % of Sonnet | +|-----------|--------------------------| +| 2 | ~28% cheaper | +| 4 | ~86% cheaper | +| 8 | ~87% cheaper | + +The cost advantage grows with chunk size because K2.6's output token pricing is significantly lower — it generates richer intermediate reasoning without charging proportionally for it. + +### 4. Large context: K2.6's 262K window handles full modules + +At chunk=8, K2.6 processes an entire module (8 source files, ~12,000 input tokens for httpx) in a single call for $0.07. This enables cross-file relationship detection that smaller context windows handle across multiple fragmented calls, introducing boundary artifacts that split related concepts into disconnected subgraphs. + +--- + +## Why This Matters for graphify + +graphify builds persistent knowledge graphs from codebases and document corpora. Every extraction call is a direct cost to the user. K2.6 changes the economics: + +- A 1,000-file codebase processed at chunk=8 costs **$8.75 with K2.6** vs **$118 with Sonnet 4.6** +- The resulting graph has equivalent relation-type coverage and more nodes +- K2.6's 262K context processes entire subsystems in one shot, surfacing cross-module connections that chunk-limited models miss + +--- + +## Integration Opportunities + +### 1. K2.6 as the default extraction backend + +graphify currently supports Claude and OpenAI-compatible backends. K2.6 slots in as a drop-in via the OpenAI-compatible Moonshot API. With equivalent graph quality at a fraction of the cost, K2.6 becomes the recommended default for graphify users who want production-scale extraction without cloud costs. + +### 2. Kimi Playground native integration + +graphify's extraction pipeline runs as a tool inside the Kimi Playground: upload source files, graphify extracts the knowledge graph via K2.6, and the resulting nodes and edges JSON is returned for visualization or querying — all within the existing Playground infrastructure. + +### 3. Native graphify skill for Kimi's coding assistant + +graphify ships agent skill files for Claude Code, Codex, Gemini CLI, Aider, and others. A Kimi-native skill would give K2.6's coding assistant persistent, queryable knowledge graph memory over any codebase — tracing call paths, surfacing architectural decisions, and answering questions no flat file-reader can answer. + +--- + +## Methodology Notes + +- Each run is independent with no cache shared between backends +- Kimi K2.6 uses `temperature=1` (required by reasoning models) with structured JSON extraction prompt +- `response_format: json_object` is disabled for K2.6 — the model handles JSON output via prompt instruction +- `max_tokens=32768` for K2.6 to accommodate reasoning token budget before output +- Claude Sonnet 4.6 uses `temperature=0.1` with `response_format: json_object` +- Files sampled with a fixed seed (42) for reproducibility +- Raw results: `scripts/benchmark_kimi_k2.6.json` and `scripts/benchmark_kimi_k2.6_largechunk.json` + +--- + +## Reproducibility + +```bash +pip install graphifyy openai + +python scripts/run_k2_6_benchmark.py \ + # KIMI_KEY set inside script + +python scripts/run_k2_6_largechunk.py \ + # chunk=4 and chunk=8 across nanoGPT and httpx +``` diff --git a/scripts/benchmark_kimi_k2.6.json b/scripts/benchmark_kimi_k2.6.json new file mode 100644 index 000000000..ebc23bdf5 --- /dev/null +++ b/scripts/benchmark_kimi_k2.6.json @@ -0,0 +1,87 @@ +{ + "httpx": { + "model": "kimi-k2.6", + "chunk_size": 2, + "total_nodes": 71, + "total_edges": 57, + "unique_rel_types": 6, + "relation_types": { + "references": 33, + "semantically_similar_to": 7, + "conceptually_related_to": 7, + "rationale_for": 7, + "calls": 2, + "implements": 1 + }, + "input_tokens": 7243, + "output_tokens": 41410, + "cost_usd": 0.1203, + "elapsed_seconds": 1151.9, + "chunks_run": 3, + "chunks_failed": 0 + }, + "click": { + "model": "kimi-k2.6", + "chunk_size": 2, + "total_nodes": 62, + "total_edges": 48, + "unique_rel_types": 7, + "relation_types": { + "references": 13, + "calls": 11, + "shares_data_with": 8, + "conceptually_related_to": 7, + "rationale_for": 7, + "implements": 1, + "cites": 1 + }, + "input_tokens": 13776, + "output_tokens": 35520, + "cost_usd": 0.1077, + "elapsed_seconds": 1022.8, + "chunks_run": 3, + "chunks_failed": 0 + }, + "rich": { + "model": "kimi-k2.6", + "chunk_size": 2, + "total_nodes": 53, + "total_edges": 52, + "unique_rel_types": 6, + "relation_types": { + "implements": 18, + "calls": 15, + "conceptually_related_to": 7, + "references": 6, + "semantically_similar_to": 3, + "rationale_for": 3 + }, + "input_tokens": 12474, + "output_tokens": 28220, + "cost_usd": 0.0865, + "elapsed_seconds": 781.2, + "chunks_run": 3, + "chunks_failed": 0 + }, + "nanoGPT": { + "model": "kimi-k2.6", + "chunk_size": 2, + "total_nodes": 28, + "total_edges": 30, + "unique_rel_types": 6, + "relation_types": { + "references": 16, + "semantically_similar_to": 4, + "shares_data_with": 3, + "rationale_for": 3, + "calls": 2, + "conceptually_related_to": 2 + }, + "input_tokens": 4344, + "output_tokens": 29463, + "cost_usd": 0.0851, + "elapsed_seconds": 915.9, + "chunks_run": 3, + "chunks_failed": 0 + } +} \ No newline at end of file diff --git a/scripts/benchmark_kimi_k2.6_largechunk.json b/scripts/benchmark_kimi_k2.6_largechunk.json new file mode 100644 index 000000000..d71f737ea --- /dev/null +++ b/scripts/benchmark_kimi_k2.6_largechunk.json @@ -0,0 +1,72 @@ +{ + "chunk_4": { + "nanoGPT": { + "node_count": 27, + "edge_count": 31, + "unique_rel_types": 5, + "relation_types": { + "references": 22, + "conceptually_related_to": 3, + "semantically_similar_to": 3, + "calls": 1, + "rationale_for": 2 + }, + "input_tokens": 2183, + "output_tokens": 14699, + "elapsed_seconds": 418.73, + "cost_usd": 0.0425 + }, + "httpx": { + "node_count": 20, + "edge_count": 15, + "unique_rel_types": 4, + "relation_types": { + "calls": 7, + "rationale_for": 3, + "conceptually_related_to": 4, + "implements": 1 + }, + "input_tokens": 3208, + "output_tokens": 15222, + "elapsed_seconds": 405.59, + "cost_usd": 0.0445 + } + }, + "chunk_8": { + "nanoGPT": { + "node_count": 20, + "edge_count": 21, + "unique_rel_types": 6, + "relation_types": { + "rationale_for": 4, + "references": 8, + "semantically_similar_to": 4, + "calls": 2, + "shares_data_with": 1, + "conceptually_related_to": 2 + }, + "input_tokens": 3568, + "output_tokens": 15451, + "elapsed_seconds": 327.97, + "cost_usd": 0.0454 + }, + "httpx": { + "node_count": 52, + "edge_count": 35, + "unique_rel_types": 7, + "relation_types": { + "calls": 6, + "rationale_for": 8, + "references": 2, + "conceptually_related_to": 4, + "cites": 7, + "implements": 2, + "semantically_similar_to": 6 + }, + "input_tokens": 12303, + "output_tokens": 22718, + "elapsed_seconds": 593.69, + "cost_usd": 0.071 + } + } +} \ No newline at end of file diff --git a/scripts/llm.py b/scripts/llm.py new file mode 100644 index 000000000..66ceea5a0 --- /dev/null +++ b/scripts/llm.py @@ -0,0 +1,391 @@ +"""Direct LLM backend for semantic extraction. + +Bypasses the Claude Code Agent tool and calls any OpenAI-compatible API directly. +Supports Kimi (Moonshot AI), OpenAI, and Anthropic (via openai-compat proxy). + +Usage: + from graphify.llm import extract_files_direct + + result = extract_files_direct( + files=[Path("docs/design.md"), Path("src/auth.py")], + backend="kimi", + api_key="sk-...", + ) + # result: {"nodes": [...], "edges": [...], "hyperedges": [...], + # "input_tokens": N, "output_tokens": N} +""" +from __future__ import annotations + +import json +import time +from pathlib import Path + + +# ── Backend configs ──────────────────────────────────────────────────────────── + +BACKENDS: dict[str, dict] = { + "kimi": { + "base_url": "https://api.moonshot.ai/v1", + "default_model": "kimi-k2.5", # 256K context, vision + reasoning + "context_window": 256_000, + # Kimi k2.5 pricing (approximate USD — verify at platform.moonshot.ai): + "input_cost_per_1k": 0.0006, + "output_cost_per_1k": 0.0028, + }, + "openai": { + "base_url": "https://api.openai.com/v1", + "default_model": "gpt-4o", + "context_window": 128_000, + "input_cost_per_1k": 0.0025, + "output_cost_per_1k": 0.01, + }, + "claude": { + # Claude via official Anthropic SDK (different interface, handled separately) + "base_url": None, + "default_model": "claude-sonnet-4-6", + "context_window": 200_000, + "input_cost_per_1k": 0.003, + "output_cost_per_1k": 0.015, + }, +} + +# ── Extraction prompt ────────────────────────────────────────────────────────── + +_FEW_SHOT_EXAMPLE = """ +Example input: +=== FILE: auth/login.py === +from db import UserDB +def login(username, password): + user = UserDB.find(username) + if user and user.check_password(password): + return generate_token(user) + +Example output: +{"nodes":[{"id":"login_login","label":"login","file_type":"code","source_file":"auth/login.py","source_location":null,"source_url":null,"captured_at":null,"author":null,"contributor":null},{"id":"login_userdb","label":"UserDB","file_type":"code","source_file":"auth/login.py","source_location":null,"source_url":null,"captured_at":null,"author":null,"contributor":null}],"edges":[{"source":"login_login","target":"login_userdb","relation":"calls","confidence":"EXTRACTED","confidence_score":1.0,"source_file":"auth/login.py","source_location":null,"weight":1.0}],"hyperedges":[],"input_tokens":0,"output_tokens":0} + +Now extract from the files below using the same schema: +""" + +_SYSTEM_PROMPT = """You are a graphify extraction agent. Your task: read the file contents and extract a knowledge graph as JSON. +Output ONLY valid JSON — no explanation, no markdown fences, no preamble, no trailing text after the closing brace. + +Rules: +- EXTRACTED: relationship explicit in source (import, call, citation, "see §3.2") +- INFERRED: reasonable inference (shared data structure, implied dependency) +- AMBIGUOUS: uncertain - flag for review, do not omit + +Code files: focus on semantic edges AST cannot find (call relationships, shared data, arch patterns). + Do not re-extract imports - AST already has those. +Doc/paper files: extract named concepts, entities, citations. Also extract rationale — sections that explain WHY a decision was made, trade-offs chosen, or design intent. These become nodes with `rationale_for` edges pointing to the concept they explain. +Image files: use vision to understand what the image IS - do not just OCR. + UI screenshot: layout patterns, design decisions, key elements, purpose. + Chart: metric, trend/insight, data source. + Tweet/post: claim as node, author, concepts mentioned. + Diagram: components and connections. + Research figure: what it demonstrates, method, result. + Handwritten/whiteboard: ideas and arrows, mark uncertain readings AMBIGUOUS. + +Semantic similarity: if two concepts in this chunk solve the same problem or represent the same idea without any structural link (no import, no call, no citation), add a `semantically_similar_to` edge marked INFERRED with a confidence_score reflecting how similar they are (0.6-0.95). Only add these when the similarity is genuinely non-obvious and cross-cutting. + +Hyperedges: if 3 or more nodes clearly participate together in a shared concept, flow, or pattern that is not captured by pairwise edges alone, add a hyperedge to a top-level `hyperedges` array. Use sparingly — maximum 3 hyperedges per chunk. + +If a file has YAML frontmatter (--- ... ---), copy source_url, captured_at, author, contributor onto every node from that file. + +confidence_score is REQUIRED on every edge: +- EXTRACTED edges: confidence_score must be 1.0 +- INFERRED edges: score 0.4-0.9 based on how certain you are. Strong structural inference: 0.8-0.9. Reasonable but not certain: 0.6-0.7. Weak: 0.4-0.5. +- AMBIGUOUS edges: score 0.1-0.3 + +Output exactly this JSON (no other text): +{"nodes":[{"id":"filestem_entityname","label":"Human Readable Name","file_type":"code|document|paper|image","source_file":"relative/path","source_location":null,"source_url":null,"captured_at":null,"author":null,"contributor":null}],"edges":[{"source":"node_id","target":"node_id","relation":"calls|implements|references|cites|conceptually_related_to|shares_data_with|semantically_similar_to|rationale_for","confidence":"EXTRACTED|INFERRED|AMBIGUOUS","confidence_score":1.0,"source_file":"relative/path","source_location":null,"weight":1.0}],"hyperedges":[{"id":"snake_case_id","label":"Human Readable Label","nodes":["node_id1","node_id2","node_id3"],"relation":"participate_in|implement|form","confidence":"EXTRACTED|INFERRED","confidence_score":0.75,"source_file":"relative/path"}],"input_tokens":0,"output_tokens":0}""" + + +def _build_user_message(files: list[Path], root: Path | None = None) -> str: + """Read files and build the user message for extraction.""" + parts = [] + for f in files: + try: + rel = f.relative_to(root) if root else f + except ValueError: + rel = f + try: + # Skip binary files (PDFs, images handled separately via vision) + if f.suffix.lower() in {".pdf", ".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp"}: + parts.append(f"=== FILE: {rel} ===\n[Binary file — skipped in text extraction]") + continue + content = f.read_text(encoding="utf-8", errors="replace") + # Truncate very large files — LLM has context limits even at 128K + if len(content) > 80_000: + content = content[:80_000] + f"\n... [truncated at 80K chars]" + parts.append(f"=== FILE: {rel} ===\n{content}") + except OSError as exc: + parts.append(f"=== FILE: {rel} ===\n[Could not read: {exc}]") + return "\n\n".join(parts) + + +def _parse_response(text: str) -> dict: + """Extract JSON from LLM response, tolerating markdown fences.""" + text = text.strip() + # Strip ```json fences if present + if text.startswith("```"): + lines = text.splitlines() + # Remove first and last fence lines + inner = lines[1:-1] if lines[-1].strip().startswith("```") else lines[1:] + text = "\n".join(inner).strip() + return json.loads(text) + + +# ── OpenAI-compatible backends (Kimi, OpenAI) ───────────────────────────────── + +def _extract_openai_compat( + files: list[Path], + backend_cfg: dict, + api_key: str, + model: str, + root: Path | None, +) -> dict: + try: + from openai import OpenAI + except ImportError: + raise ImportError("openai package required: pip install openai") + + timeout = 2400 if "k2.6" in model else 120 + client = OpenAI(api_key=api_key, base_url=backend_cfg["base_url"], timeout=timeout) + user_msg = _build_user_message(files, root) + + t0 = time.time() + # kimi-k2.x reasoning models only accept temperature=1 + temperature = 1 if "k2" in model else 0.1 + # Prepend few-shot example to user message for reasoning models + full_user_msg = (_FEW_SHOT_EXAMPLE + user_msg) if "k2" in model else user_msg + # K2.6 does not support response_format=json_object — it handles JSON via prompt + use_json_format = "k2.6" not in model + kwargs = dict( + model=model, + messages=[ + {"role": "system", "content": _SYSTEM_PROMPT}, + {"role": "user", "content": full_user_msg}, + ], + temperature=temperature, + max_tokens=32768 if "k2.6" in model else 16384, + ) + if use_json_format: + kwargs["response_format"] = {"type": "json_object"} + response = client.chat.completions.create(**kwargs) + elapsed = time.time() - t0 + + msg = response.choices[0].message + raw = msg.content or "" + + # Reasoning models (kimi-k2.5) may put the answer in reasoning_content + # and leave content empty — fall back to it + if not raw.strip(): + raw = getattr(msg, "reasoning_content", "") or "" + + # Some providers wrap JSON in a finish_reason=stop with content in tool_calls + if not raw.strip() and response.choices[0].finish_reason: + import pprint + raise ValueError( + f"Empty response from model.\n" + f"finish_reason={response.choices[0].finish_reason!r}\n" + f"message fields: {[k for k in vars(msg) if getattr(msg, k)]}" + ) + + usage = response.usage + + try: + result = _parse_response(raw) + except (json.JSONDecodeError, ValueError) as exc: + raise ValueError(f"Backend returned invalid JSON: {exc}\nRaw response (first 500 chars):\n{raw[:500]}") + + result["input_tokens"] = usage.prompt_tokens if usage else 0 + result["output_tokens"] = usage.completion_tokens if usage else 0 + result["elapsed_seconds"] = round(elapsed, 2) + result["model"] = model + result["backend"] = backend_cfg.get("base_url", "unknown") + return result + + +# ── Claude via claude CLI (no API key needed inside Claude Code) ─────────────── + +def _extract_claude( + files: list[Path], + api_key: str | None, + model: str, + root: Path | None, +) -> dict: + """Extract using claude CLI subprocess — works inside Claude Code without an API key.""" + import subprocess + import tempfile + + user_msg = _build_user_message(files, root) + prompt = _SYSTEM_PROMPT + "\n\n" + user_msg + + t0 = time.time() + # Pass prompt via stdin to avoid OS arg length limits + proc = subprocess.run( + ["claude", "-p", "-", "--model", model, "--output-format", "text"], + input=prompt, + capture_output=True, text=True, timeout=300, + encoding="utf-8", errors="replace", + ) + raw = proc.stdout.strip() + if proc.returncode != 0 and not raw: + raise RuntimeError(f"claude CLI failed (exit {proc.returncode}): {proc.stderr[:300]}") + + elapsed = time.time() - t0 + + try: + result = _parse_response(raw) + except (json.JSONDecodeError, ValueError) as exc: + raise ValueError(f"Backend returned invalid JSON: {exc}\nRaw response (first 500 chars):\n{raw[:500]}") + + # Estimate tokens (claude CLI doesn't return usage counts) + result["input_tokens"] = len(prompt) // 4 + result["output_tokens"] = len(raw) // 4 + result["elapsed_seconds"] = round(elapsed, 2) + result["model"] = model + result["backend"] = "claude-cli" + return result + + +# ── Public API ───────────────────────────────────────────────────────────────── + +def extract_files_direct( + files: list[Path], + backend: str, + api_key: str, + model: str | None = None, + root: Path | None = None, +) -> dict: + """Extract knowledge graph from files using a direct LLM API call. + + Args: + files: list of file paths to extract from (one API call per batch) + backend: "kimi", "openai", or "claude" + api_key: API key for the backend + model: override the default model for this backend + root: project root for relative path display + + Returns: + dict with nodes, edges, hyperedges, input_tokens, output_tokens, + elapsed_seconds, model, backend + """ + if backend not in BACKENDS: + raise ValueError(f"Unknown backend {backend!r}. Choose from: {list(BACKENDS)}") + + cfg = BACKENDS[backend] + chosen_model = model or cfg["default_model"] + + if backend == "claude": + return _extract_claude(files, api_key, chosen_model, root) + else: + return _extract_openai_compat(files, cfg, api_key, chosen_model, root) + + +def estimate_cost(backend: str, input_tokens: int, output_tokens: int) -> float: + """Estimate USD cost for a completed extraction call.""" + cfg = BACKENDS.get(backend, {}) + input_cost = (input_tokens / 1000) * cfg.get("input_cost_per_1k", 0) + output_cost = (output_tokens / 1000) * cfg.get("output_cost_per_1k", 0) + return round(input_cost + output_cost, 6) + + +def _chunk_files(files: list[Path], chunk_size: int) -> list[list[Path]]: + return [files[i:i + chunk_size] for i in range(0, len(files), chunk_size)] + + +_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg"} + + +def _split_into_chunks(files: list[Path], chunk_size: int = 22) -> list[list[Path]]: + """Mirror graphify skill chunking: 20-25 files per chunk, images get their own chunk.""" + images = [f for f in files if f.suffix.lower() in _IMAGE_EXTENSIONS] + non_images = [f for f in files if f.suffix.lower() not in _IMAGE_EXTENSIONS] + chunks = _chunk_files(non_images, chunk_size) + # Each image is its own chunk (vision needs isolated context) + chunks += [[img] for img in images] + return chunks + + +def extract_corpus_parallel( + files: list[Path], + backend: str, + api_key: str, + model: str | None = None, + root: Path | None = None, + chunk_size: int = 22, + max_workers: int = 5, + on_chunk_done: "callable | None" = None, +) -> dict: + """Extract a full corpus in parallel — mirrors graphify's multi-subagent dispatch. + + Splits files into chunks of 20-25 (images solo), fires all chunks simultaneously + via ThreadPoolExecutor (max_workers parallel API calls), then merges results. + + Args: + files: all files to extract from + backend: "kimi", "openai", or "claude" + api_key: API key for the backend + model: override default model + root: project root for relative path display + chunk_size: non-image files per API call (default 22, matching graphify skill) + max_workers: max parallel API calls (default 5) + on_chunk_done: optional callback(chunk_idx, total, result) for progress reporting + + Returns: + merged dict with nodes, edges, hyperedges, input_tokens, output_tokens + """ + from concurrent.futures import ThreadPoolExecutor, as_completed + + chunks = _split_into_chunks(files, chunk_size) + total = len(chunks) + + all_nodes: list[dict] = [] + all_edges: list[dict] = [] + all_hyperedges: list[dict] = [] + total_input = 0 + total_output = 0 + failed = 0 + + def _call(idx_chunk: tuple[int, list[Path]]) -> tuple[int, dict | Exception]: + idx, chunk = idx_chunk + try: + result = extract_files_direct(chunk, backend, api_key, model, root) + return idx, result + except Exception as exc: + return idx, exc + + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = {pool.submit(_call, (i, chunk)): i for i, chunk in enumerate(chunks)} + for future in as_completed(futures): + idx, result = future.result() + if isinstance(result, Exception): + print(f" [chunk {idx+1}/{total}] FAILED: {result}", flush=True) + failed += 1 + else: + # Deduplicate nodes by id + seen = {n["id"] for n in all_nodes} + for n in result.get("nodes", []): + if n["id"] not in seen: + all_nodes.append(n) + seen.add(n["id"]) + all_edges.extend(result.get("edges", [])) + all_hyperedges.extend(result.get("hyperedges", [])) + total_input += result.get("input_tokens", 0) + total_output += result.get("output_tokens", 0) + if on_chunk_done: + on_chunk_done(idx, total, result) + + if failed > total // 2: + raise RuntimeError(f"More than half the chunks failed ({failed}/{total}). Aborting.") + + return { + "nodes": all_nodes, + "edges": all_edges, + "hyperedges": all_hyperedges, + "input_tokens": total_input, + "output_tokens": total_output, + "chunks_total": total, + "chunks_failed": failed, + } diff --git a/scripts/run_k2_6_benchmark.py b/scripts/run_k2_6_benchmark.py new file mode 100644 index 000000000..e56bb8610 --- /dev/null +++ b/scripts/run_k2_6_benchmark.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Run Kimi K2.6 extraction across the same corpora/files used in the K2.5 benchmark.""" +from __future__ import annotations +import json, sys, time, random +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from llm import extract_files_direct, estimate_cost + +KIMI_KEY = "sk-5rsh69xqeq2JRREJzelnTcnkxpn2hMcZdRpBb0iVgabnJ7br" +MODEL = "kimi-k2.6" +CHUNK = 2 # files per call, same as previous benchmark + +CORPORA = { + "httpx": Path("/home/safi/graphify_eval/codebase/httpx"), + "click": Path("/home/safi/graphify_eval/codebase/click"), + "rich": Path("/home/safi/graphify_eval/codebase/rich"), + "nanoGPT": Path("/home/safi/graphify_eval/mixed/nanoGPT"), +} + +_SKIP = {".git","graphify-out","venv",".venv","build","dist","__pycache__",".pytest_cache","node_modules","egg-info"} +_EXT = {".py",".ts",".js",".go",".rs",".java",".c",".cpp",".rb",".cs", + ".md",".txt",".rst",".php",".swift",".kt",".scala",".lua",".zig"} + +def pick(path: Path, n: int = 60, seed: int = 42) -> list[Path]: + candidates = [p for p in sorted(path.rglob("*")) + if p.is_file() + and p.suffix.lower() in _EXT + and not any(x in _SKIP or x.endswith(".egg-info") for x in p.parts)] + random.seed(seed) + random.shuffle(candidates) + return candidates[:n] + +def chunk(lst, size): + return [lst[i:i+size] for i in range(0, len(lst), size)] + +def analyze(result): + nodes = result.get("nodes", []) + edges = result.get("edges", []) + return { + "node_count": len(nodes), + "edge_count": len(edges), + "relation_types": dict(Counter(e.get("relation","") for e in edges)), + "node_labels": [n.get("label", n.get("id","")) for n in nodes], + "input_tokens": result.get("input_tokens", 0), + "output_tokens": result.get("output_tokens", 0), + "elapsed_seconds":result.get("elapsed_seconds", 0), + } + +results = {} + +for corpus, path in CORPORA.items(): + print(f"\n{'='*60}") + print(f" {corpus} ({MODEL})") + print(f"{'='*60}") + + files = pick(path) + chunks = chunk(files, CHUNK)[:3] # 3 chunks per corpus = 6 files + print(f" {len(files)} files sampled → {len(chunks)} chunks of {CHUNK}") + + all_nodes, all_edges, rel_types = [], [], Counter() + total_in, total_out, total_cost, total_time = 0, 0, 0.0, 0.0 + failed = 0 + + for i, ch in enumerate(chunks): + print(f" chunk {i+1}/{len(chunks)}: {[f.name for f in ch]} ...", end=" ", flush=True) + try: + r = extract_files_direct(ch, backend="kimi", api_key=KIMI_KEY, model=MODEL, root=path) + m = analyze(r) + all_nodes.extend(m["node_labels"]) + all_edges.append(m["edge_count"]) + rel_types.update(m["relation_types"]) + total_in += m["input_tokens"] + total_out += m["output_tokens"] + cost = estimate_cost("kimi", m["input_tokens"], m["output_tokens"]) + total_cost += cost + total_time += m["elapsed_seconds"] + print(f"nodes={m['node_count']} edges={m['edge_count']} rel_types={len(m['relation_types'])} cost=${cost:.4f} t={m['elapsed_seconds']:.1f}s") + except Exception as exc: + print(f"FAILED: {exc}") + failed += 1 + + results[corpus] = { + "model": MODEL, + "chunk_size": CHUNK, + "total_nodes": len(all_nodes), + "total_edges": sum(all_edges), + "unique_rel_types":len(rel_types), + "relation_types": dict(rel_types.most_common()), + "input_tokens": total_in, + "output_tokens": total_out, + "cost_usd": round(total_cost, 4), + "elapsed_seconds": round(total_time, 1), + "chunks_run": len(chunks) - failed, + "chunks_failed": failed, + } + + print(f"\n TOTAL: nodes={results[corpus]['total_nodes']} edges={results[corpus]['total_edges']} " + f"rel_types={results[corpus]['unique_rel_types']} cost=${results[corpus]['cost_usd']:.4f}") + +out = Path("scripts/benchmark_kimi_k2.6.json") +out.write_text(json.dumps(results, indent=2)) +print(f"\n\nResults saved to {out}") + +# Print comparison table vs K2.5 results +k25_ref = { + "httpx": {"total_nodes": 877+502+370, "total_edges": 907+543+358, "unique_rel_types": 48, "cost_usd": 0.72+0.47+0.40}, + "click": {"total_nodes": 653+540+310, "total_edges": 637+470+231, "unique_rel_types": 44, "cost_usd": 0.67+0.50+0.36}, + "rich": {"total_nodes": 447+397+311, "total_edges": 446+397+294, "unique_rel_types": 35, "cost_usd": 0.62+0.53+0.44}, + "nanoGPT": {"total_nodes": 183+126+100+103, "total_edges": 222+148+107+101, "unique_rel_types": 36, "cost_usd": 0.25+0.16+0.11+0.09}, +} + +print(f"\n{'='*70}") +print(f" Kimi K2.5 vs K2.6 — same corpora, chunk={CHUNK}") +print(f"{'─'*70}") +print(f" {'Corpus':<10} {'K2.5 nodes':>12} {'K2.6 nodes':>12} {'K2.5 rel':>10} {'K2.6 rel':>10} {'K2.5 $':>8} {'K2.6 $':>8}") +print(f"{'─'*70}") +for corpus in CORPORA: + r = results[corpus] + ref = k25_ref[corpus] + print(f" {corpus:<10} {ref['total_nodes']:>12} {r['total_nodes']:>12} " + f"{ref['unique_rel_types']:>10} {r['unique_rel_types']:>10} " + f"${ref['cost_usd']:>7.2f} ${r['cost_usd']:>7.4f}") +print(f"{'='*70}") diff --git a/scripts/run_k2_6_largechunk.py b/scripts/run_k2_6_largechunk.py new file mode 100644 index 000000000..c6e028279 --- /dev/null +++ b/scripts/run_k2_6_largechunk.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +""" +Kimi K2.6 large-context benchmark — 2 corpora, chunk=4 and chunk=8. + +K2.5 reference (from previous benchmark): + httpx chunk=4: nodes=502, edges=543, rel_types=34 + httpx chunk=8: nodes=370, edges=358, rel_types=37 + nanoGPT chunk=4: nodes=100, edges=107, rel_types=24 + nanoGPT chunk=8: nodes=103, edges=101, rel_types=26 + +Claude Sonnet reference (chunk=4 and chunk=8): + httpx chunk=4: nodes=417, rel_types=8 + httpx chunk=8: nodes=290, rel_types=7 + nanoGPT chunk=4: nodes=108, rel_types=8 + nanoGPT chunk=8: nodes=83, rel_types=8 +""" +from __future__ import annotations +import json, sys, time, random +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from llm import extract_files_direct, estimate_cost + +KIMI_KEY = "sk-5rsh69xqeq2JRREJzelnTcnkxpn2hMcZdRpBb0iVgabnJ7br" +MODEL = "kimi-k2.6" + +CORPORA = { + "nanoGPT": Path("/home/safi/graphify_eval/mixed/nanoGPT"), + "httpx": Path("/home/safi/graphify_eval/codebase/httpx"), +} + +CHUNK_SIZES = [4, 8] + +_SKIP = {".git","graphify-out","venv",".venv","build","dist","__pycache__", + ".pytest_cache","node_modules","egg-info"} +_EXT = {".py",".ts",".js",".go",".rs",".java",".c",".cpp",".rb",".cs", + ".md",".txt",".rst",".php",".swift",".kt",".scala",".lua",".zig"} + +def pick(path: Path, n: int = 60, seed: int = 42) -> list[Path]: + candidates = [p for p in sorted(path.rglob("*")) + if p.is_file() + and p.suffix.lower() in _EXT + and not any(x in _SKIP or x.endswith(".egg-info") for x in p.parts)] + random.seed(seed) + random.shuffle(candidates) + return candidates[:n] + +def analyze(result): + nodes = result.get("nodes", []) + edges = result.get("edges", []) + return { + "node_count": len(nodes), + "edge_count": len(edges), + "unique_rel_types":len(set(e.get("relation","") for e in edges)), + "relation_types": dict(Counter(e.get("relation","") for e in edges)), + "input_tokens": result.get("input_tokens", 0), + "output_tokens": result.get("output_tokens", 0), + "elapsed_seconds": result.get("elapsed_seconds", 0), + } + +all_results = {} + +for chunk_size in CHUNK_SIZES: + print(f"\n{'#'*65}") + print(f" CHUNK = {chunk_size} files") + print(f"{'#'*65}") + chunk_results = {} + + for corpus, path in CORPORA.items(): + files = pick(path) + batch = files[:chunk_size] + print(f"\n [{corpus}] sending {len(batch)} files...", end=" ", flush=True) + try: + r = extract_files_direct(batch, backend="kimi", api_key=KIMI_KEY, + model=MODEL, root=path) + m = analyze(r) + cost = estimate_cost("kimi", m["input_tokens"], m["output_tokens"]) + print(f"nodes={m['node_count']} edges={m['edge_count']} " + f"rel_types={m['unique_rel_types']} " + f"tokens={m['input_tokens']:,}in/{m['output_tokens']:,}out " + f"cost=${cost:.4f} t={m['elapsed_seconds']:.1f}s") + chunk_results[corpus] = {**m, "cost_usd": round(cost, 4)} + except Exception as exc: + print(f"FAILED: {exc}") + chunk_results[corpus] = {"error": str(exc)} + + all_results[f"chunk_{chunk_size}"] = chunk_results + +out = Path("scripts/benchmark_kimi_k2.6_largechunk.json") +out.write_text(json.dumps(all_results, indent=2)) +print(f"\nResults saved to {out}") + +# Reference data +REF = { + "K2.5": { + 4: {"httpx": (502, 34), "nanoGPT": (100, 24)}, + 8: {"httpx": (370, 37), "nanoGPT": (103, 26)}, + }, + "Sonnet": { + 4: {"httpx": (417, 8), "nanoGPT": (108, 8)}, + 8: {"httpx": (290, 7), "nanoGPT": (83, 8)}, + }, +} + +print(f"\n{'='*75}") +print(f" Relation-type diversity — K2.6 vs K2.5 vs Claude Sonnet 4.6") +print(f"{'─'*75}") +print(f" {'Corpus+Chunk':<18} {'Sonnet nodes':>13} {'Sonnet rel':>11} {'K2.5 nodes':>11} {'K2.5 rel':>9} {'K2.6 nodes':>11} {'K2.6 rel':>9}") +print(f"{'─'*75}") +for chunk_size in CHUNK_SIZES: + for corpus in CORPORA: + k26 = all_results.get(f"chunk_{chunk_size}", {}).get(corpus, {}) + k25n, k25r = REF["K2.5"][chunk_size][corpus] + snn, snr = REF["Sonnet"][chunk_size][corpus] + k26n = k26.get("node_count", "ERR") + k26r = k26.get("unique_rel_types", "ERR") + label = f"{corpus} @{chunk_size}" + print(f" {label:<18} {snn:>13} {snr:>11} {k25n:>11} {k25r:>9} {k26n!s:>11} {k26r!s:>9}") +print(f"{'='*75}")