codenib 0.1.0__py3-none-any.whl

codenib/dataset/synthesize/_agent.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: 2025-2026 CodeNib Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared agent runner for the query synthesis pipeline."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+from typing import List, Optional
+
+from claude_agent_sdk import ClaudeAgentOptions, query
+
+from codenib.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class AgentRunner:
+    """Thin wrapper around the Claude Agent SDK for running LLM calls.
+
+    Each synthesis module (ContextLoader, QueryCurator, Verifier) gets its
+    own ``AgentRunner`` instance, potentially with different ``allowed_tools``
+    configurations.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str = "opus",
+        max_turns: int = 10,
+        system_prompt: str = "",
+        allowed_tools: Optional[List[str]] = None,
+        permission_mode: str = "bypassPermissions",
+    ) -> None:
+        self.model = model
+        self.max_turns = max_turns
+        self.system_prompt = system_prompt
+        self.allowed_tools = allowed_tools or []
+        self.permission_mode = permission_mode
+
+    def run(self, prompt: str, *, cwd: Optional[str] = None) -> str:
+        """Synchronous agent call (delegates to async via ``asyncio.run``)."""
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(self.run_async(prompt, cwd=cwd))
+        raise RuntimeError(
+            "Running inside an active event loop. Use run_async instead."
+        )
+
+    async def run_async(self, prompt: str, *, cwd: Optional[str] = None) -> str:
+        """Asynchronous agent call — streams messages from Claude."""
+        options = ClaudeAgentOptions(
+            max_turns=self.max_turns,
+            system_prompt=self.system_prompt,
+            allowed_tools=self.allowed_tools,
+            permission_mode=self.permission_mode,
+            model=self.model,
+            cwd=cwd,
+        )
+        text_chunks: List[str] = []
+        final_text: Optional[str] = None
+        query_gen = query(prompt=prompt, options=options)
+        completed = False
+
+        try:
+            async for message in query_gen:
+                msg_type = message.__class__.__name__
+                logger.debug(f"Agent message type: {msg_type}")
+
+                if msg_type == "ResultMessage":
+                    result = getattr(message, "result", None)
+                    completed = True
+                    if result and isinstance(result, str):
+                        final_text = result.strip()
+                    else:
+                        final_text = "\n".join(text_chunks).strip()
+                    continue
+
+                if msg_type == "ErrorMessage":
+                    error_msg = getattr(message, "error", "Unknown error")
+                    logger.error(f"Agent error message: {error_msg}")
+                    raise RuntimeError(f"Agent returned error: {error_msg}")
+
+                if msg_type == "AssistantMessage":
+                    content = getattr(message, "content", None)
+                    if not content:
+                        continue
+                    block_texts: List[str] = []
+                    for block in content:
+                        text = getattr(block, "text", None)
+                        if isinstance(text, str) and text.strip():
+                            block_texts.append(text.strip())
+                    if block_texts:
+                        text_chunks.append("\n".join(block_texts))
+        finally:
+            if not completed:
+                await query_gen.aclose()
+
+        return final_text if final_text is not None else "\n".join(text_chunks).strip()
+
+    @staticmethod
+    def extract_json(text: str) -> str:
+        """Extract the first valid JSON object from *text*."""
+        text = text.strip()
+        if not text:
+            return text
+
+        fenced = re.search(r"```(?:json)?\s*(\{[\s\S]*\})\s*```", text)
+        if fenced:
+            text = fenced.group(1).strip()
+
+        decoder = json.JSONDecoder()
+        for idx, ch in enumerate(text):
+            if ch != "{":
+                continue
+            try:
+                _obj, end = decoder.raw_decode(text[idx:])
+                return text[idx : idx + end]
+            except json.JSONDecodeError:
+                continue
+        return text

codenib/dataset/synthesize/_types.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: 2025-2026 CodeNib Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared data classes for the query synthesis pipeline."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from pydantic import BaseModel, Field
+
+from codenib.dataset.utils import CodeLocation
+from codenib.types import (
+    NODE_TYPE_CLASS,
+    NODE_TYPE_FUNCTION,
+    NODE_TYPE_METHOD,
+    NodeInfo,
+)
+
+
+class QuerySynthesisResult(BaseModel):
+    """Structured output for synthesized queries."""
+
+    question: str = Field(
+        description="Single natural-language question describing the issue."
+    )
+    focus: Optional[str] = Field(
+        default=None,
+        description="Optional short phrase highlighting the behavior focus.",
+    )
+    hints: Optional[List[str]] = Field(
+        default=None,
+        description="Optional progressive hints that could help locate the code.",
+    )
+
+
+class TargetDiscoveryResult(BaseModel):
+    """Structured output for repository target discovery."""
+
+    target_files: List[str] = Field(default_factory=list)
+    target_symbols: List[str] = Field(default_factory=list)
+    target_symbol_nodes: List[NodeInfo] = Field(default_factory=list)
+    rationale: Optional[str] = Field(default=None)
+
+
+class BehavioralSelectionResult(BaseModel):
+    """Structured output for behavior-first synthesis from sampled code blocks."""
+
+    question: str = Field(description="Natural-language behavioral question.")
+    focus: Optional[str] = Field(default=None)
+    required_block_ids: List[str] = Field(default_factory=list)
+    rationale: Optional[str] = Field(default=None)
+
+
+@dataclass
+class RepoSnapshot:
+    root: Path
+    top_level: List[str]
+    languages: List[Tuple[str, int]]
+    agent_summary: Optional[str] = None
+
+    def format_summary(self) -> str:
+        parts: List[str] = []
+        if self.top_level:
+            parts.append("Top-level entries: " + ", ".join(self.top_level))
+        if self.languages:
+            lang_summary = ", ".join(
+                f"{ext} ({count})" for ext, count in self.languages
+            )
+            parts.append(f"File extensions (top): {lang_summary}")
+        if self.agent_summary:
+            parts.append(f"Agent exploration summary:\n{self.agent_summary}")
+        return "\n\n".join(parts)
+
+
+@dataclass
+class SampledCodeBlock:
+    block_id: str
+    node_id: int
+    node_name: str
+    file_path: str
+    node_type: str
+    start_line: int
+    end_line: int
+    content: str
+    char_count: int
+    line_count: int
+
+    def to_node_info(self, *, include_content: bool = True) -> NodeInfo:
+        return NodeInfo(
+            node_name=self.node_name,
+            type=self.node_type,
+            file=self.file_path,
+            start_line=self.start_line,
+            end_line=self.end_line,
+            content=self.content if include_content else None,
+        )
+
+    def to_code_location(self) -> CodeLocation:
+        symbol_type = None
+        if self.node_type in (NODE_TYPE_FUNCTION, NODE_TYPE_METHOD):
+            symbol_type = "function"
+        elif self.node_type == NODE_TYPE_CLASS:
+            symbol_type = "class"
+        return CodeLocation(
+            file_path=self.file_path,
+            symbol_type=symbol_type,
+            start_line=self.start_line,
+            end_line=self.end_line,
+        )
+
+
+@dataclass
+class BehavioralContext:
+    core_block: SampledCodeBlock
+    candidate_blocks: List[SampledCodeBlock]
+    neighborhood_blocks: List[SampledCodeBlock]
+    # {(file_path, leaf_symbol): (start_line, end_line)} over ALL graph symbol
+    # nodes (0-based, graph attrs) — used to resolve real line spans for the
+    # hint/reasoning path, whose ground truth carries only file:symbol strings.
+    # NOT candidate_blocks: those are randomly downsampled to max_candidate_blocks.
+    symbol_spans: Dict[Tuple[str, str], Tuple[int, int]] = field(default_factory=dict)
+
+
+def format_prompt_block(
+    block: SampledCodeBlock,
+    *,
+    is_core: bool = False,
+    max_block_chars: int = 1800,
+) -> str:
+    """Format a SampledCodeBlock for inclusion in an LLM prompt."""
+    content = block.content
+    if len(content) > max_block_chars:
+        content = (
+            content[:max_block_chars] + "\n# ... truncated for synthesis context ..."
+        )
+    core_tag = " **CORE**" if is_core else ""
+    return (
+        f"[{block.block_id}]{core_tag} type={block.node_type} "
+        f"lines={block.start_line}-{block.end_line} "
+        f"chars={block.char_count}\n"
+        f"{content}"
+    )