codenib 0.1.0__py3-none-any.whl

codenib/agent/extract_agent.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: 2025-2026 CodeNib Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Keyword extraction agent for problem statements.
+This module extracts key terms from problem statements using llama_index.
+"""
+
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from ..llm.litellm_chat import LiteLLMChat, human_message
+from ..log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+# Define Pydantic model for structured output
+class KeywordExtraction(BaseModel):
+    """Model for keyword extraction output."""
+
+    keywords: List[str] = Field(description="List of extracted keywords")
+
+
+class KeywordExtractor:
+    """Agent for extracting keywords from problem statements."""
+
+    def __init__(
+        self,
+        llm: LiteLLMChat,
+    ):
+        """Initialize the keyword extractor.
+
+        Args:
+            llm: A configured LiteLLMChat instance.
+        """
+        self.llm = llm
+        self.structured_llm = self.llm.with_structured_output(KeywordExtraction)
+
+    def extract_keywords(self, problem_statement: str) -> KeywordExtraction:
+        """
+        Extract keywords from a problem statement.
+
+        Args:
+            problem_statement (str): The problem statement to extract keywords from
+
+        Returns:
+            KeywordExtraction: Structured output with extracted keywords
+        """
+        # Create prompt with detailed instructions
+        prompt = (
+            "You are a keyword extraction specialist. "
+            "Your task is to extract important keywords "
+            "from problem statements. Focus on identifying "
+            "technical terms, function names, class names, "
+            "modules, file paths, and concepts that would be "
+            "useful for searching in a codebase. "
+            "\n\n"
+            "Guidelines for extraction:\n"
+            "1. Extract file paths and file names "
+            "(e.g., 'django/db/models/expressions.py'"
+            "-> and 'expressions.py')\n"
+            "2. Extract function and method names "
+            "(e.g., 'separability_matrix', 'run_validators')\n"
+            "3. Extract class names and module names\n"
+            "4. Prefer precise terms over general ones\n"
+            "5. Remove common stopwords and general "
+            "programming terms\n"
+            "\n\n"
+            "Please extract the key technical terms and "
+            "concepts from the following problem statement:"
+            f"\n\n{problem_statement}\n\n"
+            "Return only the essential terms that would be "
+            "most useful for searching in a codebase."
+        )
+
+        # Use structured LLM to get output directly as a KeywordExtraction object
+        input_msg = human_message(prompt)
+        result = self.structured_llm.invoke([input_msg])
+        logger.debug(f"Extracted keywords: {result}")
+        return result
+
+
+def extract_keywords_from_statement(
+    problem_statement: str,
+    llm: LiteLLMChat,
+) -> KeywordExtraction:
+    """
+    Extract keywords from a problem statement.
+
+    Args:
+        problem_statement: The problem statement to extract keywords from.
+        llm: A configured LiteLLMChat instance.
+
+    Returns:
+        KeywordExtraction: Structured output with extracted keywords
+    """
+    extractor = KeywordExtractor(llm=llm)
+    return extractor.extract_keywords(problem_statement)

codenib/agent/harness.py

@@ -0,0 +1,326 @@
+# SPDX-FileCopyrightText: 2025-2026 CodeNib Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Declarative agent harness configuration.
+
+The runner owns execution. This module owns the reusable contract that says
+which tools, prompts, and loop controls a harness exposes. Dataset selection,
+benchmark scoring, model-specific routing, and experiment policy stay outside
+this module.
+"""
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import (
+    Any,
+    Collection,
+    Dict,
+    FrozenSet,
+    Iterable,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Union,
+)
+
+from ..llm.usage import TokenUsage
+from .route_context import normalize_lsp_route_seed_policy
+
+USAGE_TOTAL_KEYS = (
+    "prompt_tokens",
+    "completion_tokens",
+    "total_tokens",
+    "cost_usd",
+    "cache_read_input_tokens",
+    "cache_creation_input_tokens",
+)
+
+
+@contextmanager
+def agent_working_directory(cwd: Optional[Union[str, Path]]) -> Iterator[None]:
+    """Temporarily run agent default tools relative to *cwd*.
+
+    The built-in read/grep/glob tools intentionally accept repo-relative paths.
+    Harnesses that run against many repos should scope the process cwd around
+    each runner invocation and always restore it afterwards.
+    """
+    if cwd is None:
+        yield
+        return
+
+    previous = os.getcwd()
+    os.chdir(cwd)
+    try:
+        yield
+    finally:
+        os.chdir(previous)
+
+
+def run_agent_in_directory(
+    runner: Any,
+    prompt: str,
+    cwd: Optional[Union[str, Path]],
+) -> Any:
+    """Run an ``AgentRunner``-like object with repo-relative default tools."""
+    with agent_working_directory(cwd):
+        return runner.run(prompt)
+
+
+def _normalize_names(
+    values: Optional[Iterable[str]],
+    *,
+    field_name: str,
+    empty_is_none: bool,
+) -> Optional[FrozenSet[str]]:
+    if values is None:
+        return None
+    names = []
+    for value in values:
+        if not isinstance(value, str):
+            raise TypeError(f"{field_name} entries must be strings")
+        name = value.strip()
+        if not name:
+            raise ValueError(f"{field_name} entries must be non-empty strings")
+        names.append(name)
+    if not names and empty_is_none:
+        return None
+    return frozenset(names)
+
+
+@dataclass(frozen=True)
+class AgentHarnessSpec:
+    """Stable runner-facing harness contract.
+
+    This deliberately mirrors generic :class:`AgentRunner` controls only. It
+    does not encode benchmark arms, scorer fields, model names, or dataset
+    identifiers, so experiments can compare policies without leaking those
+    policies into the core runtime API.
+    """
+
+    max_turns: int = 10
+    max_context_tokens: Optional[int] = None
+    allow_skills: Optional[Collection[str]] = None
+    exclude_skills: Optional[Collection[str]] = None
+    include_default_tools: bool = True
+    default_tool_ids: Optional[Collection[str]] = None
+    system_prompt: Optional[str] = None
+    first_turn_tool_choice: Optional[str] = None
+    force_first_turn_only: bool = False
+    force_localization_contract: bool = False
+    compact_after_read: bool = False
+    compact_keep_reads: int = 0
+    enable_lsp_route_context: bool = False
+    lsp_route_seed_limit: int = 8
+    lsp_route_seed_policy: str = "all"
+    lsp_route_query_fallback: bool = False
+    lsp_route_top_k: int = 12
+    lsp_route_include_neighbors: bool = True
+
+    def __post_init__(self) -> None:
+        if self.max_turns <= 0:
+            raise ValueError("max_turns must be positive")
+        if self.max_context_tokens is not None and self.max_context_tokens <= 0:
+            raise ValueError("max_context_tokens must be positive when set")
+        if self.compact_keep_reads < 0:
+            raise ValueError("compact_keep_reads must be non-negative")
+        if self.lsp_route_seed_limit <= 0:
+            raise ValueError("lsp_route_seed_limit must be positive")
+        if self.lsp_route_top_k <= 0:
+            raise ValueError("lsp_route_top_k must be positive")
+        object.__setattr__(
+            self,
+            "lsp_route_seed_policy",
+            normalize_lsp_route_seed_policy(self.lsp_route_seed_policy),
+        )
+
+        object.__setattr__(
+            self,
+            "allow_skills",
+            _normalize_names(
+                self.allow_skills,
+                field_name="allow_skills",
+                empty_is_none=False,
+            ),
+        )
+        object.__setattr__(
+            self,
+            "exclude_skills",
+            _normalize_names(
+                self.exclude_skills,
+                field_name="exclude_skills",
+                empty_is_none=True,
+            ),
+        )
+        object.__setattr__(
+            self,
+            "default_tool_ids",
+            _normalize_names(
+                self.default_tool_ids,
+                field_name="default_tool_ids",
+                empty_is_none=True,
+            ),
+        )
+
+    def with_overrides(self, **overrides: Any) -> "AgentHarnessSpec":
+        """Return a copy with selected fields replaced and revalidated."""
+        data = {field.name: getattr(self, field.name) for field in fields(self)}
+        unknown = set(overrides) - set(data)
+        if unknown:
+            raise TypeError(f"Unknown harness option(s): {sorted(unknown)}")
+        data.update(overrides)
+        return type(self)(**data)
+
+    def to_runner_kwargs(
+        self,
+        *,
+        session_ctx: Optional[Any] = None,
+        manifest: Optional[Any] = None,
+        compile_table: Optional[Any] = None,
+        extra: Optional[Mapping[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Return ``AgentRunner`` keyword arguments for this harness.
+
+        Collection fields are copied into mutable sets because ``AgentRunner``
+        treats them as constructor inputs and may derive internal sets from
+        them. The spec remains immutable and reusable across cells/subagents.
+        """
+        kwargs: Dict[str, Any] = {
+            "max_turns": self.max_turns,
+            "max_context_tokens": self.max_context_tokens,
+            "allow_skills": (
+                set(self.allow_skills) if self.allow_skills is not None else None
+            ),
+            "exclude_skills": (
+                set(self.exclude_skills) if self.exclude_skills is not None else None
+            ),
+            "include_default_tools": self.include_default_tools,
+            "default_tool_ids": (
+                set(self.default_tool_ids)
+                if self.default_tool_ids is not None
+                else None
+            ),
+            "system_prompt": self.system_prompt,
+            "first_turn_tool_choice": self.first_turn_tool_choice,
+            "force_first_turn_only": self.force_first_turn_only,
+            "force_localization_contract": self.force_localization_contract,
+            "compact_after_read": self.compact_after_read,
+            "compact_keep_reads": self.compact_keep_reads,
+            "enable_lsp_route_context": self.enable_lsp_route_context,
+            "lsp_route_seed_limit": self.lsp_route_seed_limit,
+            "lsp_route_seed_policy": self.lsp_route_seed_policy,
+            "lsp_route_query_fallback": self.lsp_route_query_fallback,
+            "lsp_route_top_k": self.lsp_route_top_k,
+            "lsp_route_include_neighbors": self.lsp_route_include_neighbors,
+            "session_ctx": session_ctx,
+            "manifest": manifest,
+            "compile_table": compile_table,
+        }
+        if extra:
+            kwargs.update(extra)
+        return kwargs
+
+    def create_runner(
+        self,
+        *,
+        llm: Optional[Any] = None,
+        model: Optional[str] = None,
+        registry: Optional[Any] = None,
+        session_ctx: Optional[Any] = None,
+        manifest: Optional[Any] = None,
+        compile_table: Optional[Any] = None,
+        **overrides: Any,
+    ) -> Any:
+        """Build an ``AgentRunner`` using this harness.
+
+        ``overrides`` are one-off runner keyword overrides for cases like
+        subagents with a shorter turn budget. They do not mutate the spec.
+        """
+        from .runner import AgentRunner
+
+        kwargs = self.to_runner_kwargs(
+            session_ctx=session_ctx,
+            manifest=manifest,
+            compile_table=compile_table,
+            extra=overrides,
+        )
+        return AgentRunner(llm=llm, model=model, registry=registry, **kwargs)
+
+
+@dataclass
+class AgentRunAccumulator:
+    """Accumulate usage and turn counts across related agent runs.
+
+    A harness may charge one logical cell for several LLM interactions: routing
+    gates, isolated subagents, verify retries, or a final convergence run. This
+    helper keeps that accounting generic so experiment scripts do not each
+    reimplement token/turn summation.
+    """
+
+    _usage_values: Dict[str, List[float]] = field(
+        default_factory=lambda: {key: [] for key in USAGE_TOTAL_KEYS}
+    )
+    _turns: List[int] = field(default_factory=list)
+
+    def add_usage(self, usage: Optional[Any]) -> None:
+        """Add a ``TokenUsage`` or mapping with flat/nested token fields."""
+        if usage is None:
+            return
+        if isinstance(usage, TokenUsage):
+            payload: Mapping[str, Any] = usage.to_dict()
+        elif isinstance(usage, Mapping):
+            nested = usage.get("token_usage")
+            payload = nested if isinstance(nested, Mapping) else usage
+        else:
+            payload = {
+                key: getattr(usage, key, None)
+                for key in USAGE_TOTAL_KEYS
+                if getattr(usage, key, None) is not None
+            }
+
+        for key in USAGE_TOTAL_KEYS:
+            value = payload.get(key)
+            if value is None:
+                continue
+            try:
+                self._usage_values[key].append(float(value))
+            except (TypeError, ValueError):
+                continue
+
+    def add_turns(self, turns: Optional[int]) -> None:
+        """Add one run's turn count when it is known."""
+        if turns is None:
+            return
+        self._turns.append(int(turns))
+
+    def add_result(self, result: Any) -> None:
+        """Add accounting fields from an ``AgentResult``-like object."""
+        self.add_usage(getattr(result, "usage", None))
+        self.add_turns(getattr(result, "total_turns", None))
+
+    def usage_sum(self, key: str) -> Optional[float]:
+        """Return the sum for one usage field, or ``None`` if never recorded."""
+        if key not in self._usage_values:
+            raise KeyError(f"Unknown usage field: {key}")
+        values = self._usage_values[key]
+        if not values:
+            return None
+        total = sum(values)
+        if key != "cost_usd" and total.is_integer():
+            return int(total)
+        return total
+
+    def usage_totals(self) -> Dict[str, Optional[float]]:
+        """Return flat totals with ``None`` for fields never observed."""
+        return {key: self.usage_sum(key) for key in USAGE_TOTAL_KEYS}
+
+    def total_turns(self, *, fallback: Optional[int] = None) -> Optional[int]:
+        """Return summed turns, or ``fallback`` when no turn count was added."""
+        if not self._turns:
+            return fallback
+        return sum(self._turns)

codenib/agent/history.py

@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: 2025-2026 CodeNib Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Token-budgeted chat history for the agent loop (issue #109, phase 4).
+
+The agent loop in :mod:`codenib.agent.runner` appends every assistant
+turn and every tool result to a flat ``messages`` list. On long runs that
+list grows without bound and eventually overflows the model's context
+window.
+
+:class:`TokenBudgetedChatHistory` is a drop-in container that caps the
+*total* context tokens. When adding a message would exceed the budget it
+evicts the **oldest non-system messages first**, always preserving:
+
+* every ``system`` message (the system prompt — pinned, never evicted), and
+* the most recent messages (the live working set the model needs to act).
+
+It is deliberately a plain container (``add_message`` / ``get_messages`` /
+``clear``) so the runner can swap a bare ``list`` for it with no other
+change. Token counting uses ``litellm.token_counter`` when available and
+falls back to a cheap character-based heuristic so the class stays usable
+in unit tests and offline environments.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from ..log_utils import get_logger
+
+logger = get_logger(__name__)
+
+# Rough chars-per-token used by the offline heuristic. ~4 chars/token is the
+# conventional English-text approximation; intentionally conservative so the
+# fallback over-counts slightly rather than under-counts (eviction is safer
+# than overflow).
+_CHARS_PER_TOKEN = 4
+
+
+def _message_text(message: Dict[str, Any]) -> str:
+    """Flatten a chat message dict into the text we count tokens for.
+
+    Covers the three shapes the runner produces: a plain ``content`` string,
+    ``content=None`` with ``tool_calls`` (assistant tool requests), and tool
+    result messages. Best-effort — anything unrecognized is ``str()``-ified.
+    """
+    parts: List[str] = []
+    content = message.get("content")
+    if isinstance(content, str):
+        parts.append(content)
+    elif content is not None:
+        parts.append(str(content))
+
+    for tc in message.get("tool_calls") or []:
+        fn = (tc or {}).get("function") if isinstance(tc, dict) else None
+        if isinstance(fn, dict):
+            parts.append(str(fn.get("name", "")))
+            parts.append(str(fn.get("arguments", "")))
+    return "\n".join(p for p in parts if p)
+
+
+def count_message_tokens(
+    messages: List[Dict[str, Any]],
+    *,
+    model: Optional[str] = None,
+) -> int:
+    """Best-effort token count for a list of chat messages.
+
+    Uses ``litellm.token_counter`` (which understands the model's real
+    tokenizer and per-message role overhead) when a ``model`` is supplied
+    and litellm is importable. Falls back to a ``chars / 4`` heuristic
+    otherwise, so this never raises and stays usable in offline tests.
+    """
+    if not messages:
+        return 0
+    if model:
+        try:
+            import litellm
+
+            return int(litellm.token_counter(model=model, messages=messages))
+        except Exception as exc:  # offline / unknown model / litellm bug
+            logger.debug("token_counter unavailable (%s); using char heuristic", exc)
+
+    chars = sum(len(_message_text(m)) for m in messages)
+    # +4 per message approximates the role / formatting overhead litellm adds.
+    return chars // _CHARS_PER_TOKEN + 4 * len(messages)
+
+
+class TokenBudgetedChatHistory:
+    """Chat-history container that caps total context tokens.
+
+    Pinned ``system`` messages are always retained. When appending a message
+    would push the estimated token total above ``max_tokens``, the oldest
+    *non-system* messages are evicted one at a time (oldest first) until the
+    history fits again, so the system prompt plus the most recent turns are
+    preserved.
+
+    ``keep_last`` guards the tail: the most recent ``keep_last`` non-system
+    messages are never evicted even under budget pressure, so the model
+    always sees the immediate context it needs to act. A single message that
+    on its own exceeds the budget is kept (we never drop the message just
+    added) and a warning is logged.
+
+    The container is intentionally minimal — ``add_message`` /
+    ``get_messages`` / ``clear`` / ``__len__`` / ``__iter__`` — so the runner
+    can use it in place of a bare ``list``.
+    """
+
+    def __init__(
+        self,
+        max_tokens: int,
+        *,
+        model: Optional[str] = None,
+        keep_last: int = 2,
+    ) -> None:
+        if max_tokens <= 0:
+            raise ValueError(f"max_tokens must be positive, got {max_tokens}")
+        self.max_tokens = max_tokens
+        self.model = model
+        self.keep_last = max(0, keep_last)
+        self._messages: List[Dict[str, Any]] = []
+
+    # -- container protocol -------------------------------------------------
+
+    def __len__(self) -> int:
+        return len(self._messages)
+
+    def __iter__(self):
+        return iter(self._messages)
+
+    def get_messages(self) -> List[Dict[str, Any]]:
+        """Return the live message list (the object the LLM is called with)."""
+        return self._messages
+
+    def clear(self) -> None:
+        """Drop all messages."""
+        self._messages.clear()
+
+    # -- mutation -----------------------------------------------------------
+
+    def add_message(self, message: Dict[str, Any]) -> None:
+        """Append *message*, then evict oldest non-system messages if over budget."""
+        self._messages.append(message)
+        self._enforce_budget()
+
+    def extend(self, messages: List[Dict[str, Any]]) -> None:
+        """Append several messages, enforcing the budget after each one."""
+        for message in messages:
+            self.add_message(message)
+
+    def total_tokens(self) -> int:
+        """Estimated token total of the current history."""
+        return count_message_tokens(self._messages, model=self.model)
+
+    # -- internals ----------------------------------------------------------
+
+    def _enforce_budget(self) -> None:
+        """Evict oldest non-system messages until within ``max_tokens``.
+
+        Walks from the front, skipping pinned ``system`` messages and the
+        protected ``keep_last`` tail, dropping the oldest eligible message
+        each pass until the total fits or nothing more is evictable.
+        """
+        while self.total_tokens() > self.max_tokens:
+            idx = self._oldest_evictable_index()
+            if idx is None:
+                # Only system + protected tail remain (or the single message
+                # is itself over budget). Keep what we have; the budget is a
+                # best-effort cap, not a hard guarantee against a giant turn.
+                logger.warning(
+                    "TokenBudgetedChatHistory: %d tokens still over budget %d "
+                    "with only system + last %d message(s) left; not evicting "
+                    "further",
+                    self.total_tokens(),
+                    self.max_tokens,
+                    self.keep_last,
+                )
+                return
+            evicted = self._messages.pop(idx)
+            logger.debug(
+                "TokenBudgetedChatHistory: evicted oldest %s message to stay "
+                "within %d-token budget",
+                evicted.get("role", "?"),
+                self.max_tokens,
+            )
+
+    def _oldest_evictable_index(self) -> Optional[int]:
+        """Index of the oldest message that may be evicted, or ``None``.
+
+        A message is evictable when it is non-system and not inside the
+        protected ``keep_last`` tail of non-system messages.
+        """
+        n = len(self._messages)
+        # The last ``keep_last`` *non-system* messages are pinned (the live
+        # working set the model needs to act). Everything older than that and
+        # non-system is evictable, oldest first.
+        non_system_positions = [
+            i for i, m in enumerate(self._messages) if m.get("role") != "system"
+        ]
+        if self.keep_last:
+            protected = set(non_system_positions[-self.keep_last :])
+        else:
+            protected = set()
+        for i in range(n):
+            m = self._messages[i]
+            if m.get("role") == "system":
+                continue
+            if i in protected:
+                continue
+            return i
+        return None
+
+
+class PlainChatHistory:
+    """Unbounded chat history with the same interface as the budgeted one.
+
+    Lets :class:`~codenib.agent.runner.AgentRunner` use a single code path
+    whether or not a token budget is configured. Behaviour is identical to a
+    bare ``list`` of message dicts — nothing is ever evicted.
+    """
+
+    def __init__(self, messages: Optional[List[Dict[str, Any]]] = None) -> None:
+        self._messages: List[Dict[str, Any]] = list(messages or [])
+
+    def __len__(self) -> int:
+        return len(self._messages)
+
+    def __iter__(self):
+        return iter(self._messages)
+
+    def get_messages(self) -> List[Dict[str, Any]]:
+        return self._messages
+
+    def add_message(self, message: Dict[str, Any]) -> None:
+        self._messages.append(message)
+
+    def extend(self, messages: List[Dict[str, Any]]) -> None:
+        self._messages.extend(messages)
+
+    def clear(self) -> None:
+        self._messages.clear()