konash 0.1.0__tar.gz

konash-0.1.0/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.4
+Name: konash
+Version: 0.1.0
+Summary: Train knowledge agents that search, retrieve, compress, and reason — on a single GPU.
+Requires-Python: >=3.11
+Requires-Dist: numpy>=1.24
+Provides-Extra: dev
+Requires-Dist: pytest>=7.4; extra == "dev"
+Provides-Extra: train
+Requires-Dist: torch>=2.1; extra == "train"
+Requires-Dist: transformers>=4.36; extra == "train"
+Requires-Dist: peft>=0.7; extra == "train"
+Requires-Dist: accelerate>=0.25; extra == "train"
+Provides-Extra: quant
+Requires-Dist: bitsandbytes>=0.41; extra == "quant"
+Provides-Extra: all
+Requires-Dist: torch>=2.1; extra == "all"
+Requires-Dist: transformers>=4.36; extra == "all"
+Requires-Dist: peft>=0.7; extra == "all"
+Requires-Dist: accelerate>=0.25; extra == "all"
+Requires-Dist: bitsandbytes>=0.41; extra == "all"

konash-0.1.0/README.md

@@ -0,0 +1,153 @@
+<div align="center">
+
+# KONASH
+
+**Knowledge-grounded Off-policy Networks for Agentic System Harnesses**
+
+<p>
+Train knowledge agents that search, retrieve, compress, and reason — on a single GPU.
+</p>
+
+[![PRs-Welcome](https://img.shields.io/badge/PRs-welcome-blue.svg)](CONTRIBUTING.md)
+[![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)](LICENSE)
+
+[![Join Discord](https://img.shields.io/badge/Join%20Discord-5865F2?style=plastic&logo=discord&logoColor=white)](#)
+[![Documentation](https://img.shields.io/badge/Documentation-orange?style=plastic&logo=gitbook&logoColor=white)](#)
+
+</div>
+
+KONASH trains knowledge agents via reinforcement learning that match or exceed frontier models on grounded reasoning tasks — at a fraction of the cost. **Single-GPU training, open-source models, 1/100th the compute.**
+
+---
+
+## Key Benefits
+
+- **100x cheaper training** — Single GPU replaces multi-node clusters. ~$100–500 per iteration instead of ~$10K–50K.
+- **Higher quality** — RL-trained agents search more efficiently, retrieve more diversely, and reason more accurately than their base models. The gains are algorithmic, not scale-dependent.
+- **Consistent results** — Parallel thinking (N=10–20 rollouts + aggregation) turns probabilistic search into near-deterministic accuracy. Cheap rollouts on a small model mean you can afford this on every query.
+- **Zero lock-in** — Your model, your weights, your infrastructure. Deploy anywhere with vLLM and LoRA hot-swapping.
+
+```python
+# Before: Static RAG — one query, hope for the best
+docs = retriever.search(query, top_k=10)
+answer = llm.generate(f"Answer based on: {docs}\n\n{query}")
+
+# After: KONASH — RL-trained agent that searches iteratively
+agent = konash.Agent("./checkpoints/iter2", corpus="./my_docs")
+answer = agent.solve(query, parallel_rollouts=10)
+```
+
+[Learn more about KONASH](https://kona.sh)
+
+---
+
+## KONASH Overview
+
+KONASH is an open-source RL framework that improves agent reliability by training knowledge agents to search, retrieve, compress, and reason over evidence — all on a single GPU using off-policy RL. For a quick hands-on introduction, run one of the notebooks below. When you're ready to learn more, check out the [docs](https://kona.sh).
+
+### Notebooks
+
+| Agent Task | Example Notebook | Description | Comparative Performance |
+|---|---|---|---|
+| **Trivia Night** | [Train agent](notebooks/trivia_night.ipynb) | Qwen 3.5 7B learns to answer multi-constraint trivia by searching Wikipedia | [Link coming soon] |
+| **20 Questions** | [Train agent](#) | Qwen 3.5 7B learns to identify a mystery entity in 20 yes/no searches | [Link coming soon] |
+| **GeoGuessr** | [Train agent](#) | Qwen 3.5 7B learns to pinpoint locations from landmark and terrain descriptions | [Link coming soon] |
+
+## KONASH News
+
+Explore our latest research and updates on building SOTA knowledge agents.
+
+- **[OAPL: Off-Policy RL That Actually Works on One GPU](#)** — Train knowledge agents without multi-node clusters using large-batch iterative off-policy reinforcement learning.
+- **[Agentic Data Synthesis: Let Your Model Write Its Own Curriculum](#)** — Generate diverse, grounded training data from any corpus — no manual annotation required.
+- **[Parallel Thinking: How a 7B Model Beats Frontier Single-Shot](#)** — Scale quality at inference time with N parallel rollouts and generative aggregation.
+- **[Compression as an RL Skill: Teaching Models What to Remember](#)** — Train context compression end-to-end with task reward, not as a separate summarization step.
+
+[See all blog posts](https://kona.sh/blog)
+
+## Why KONASH?
+
+- KONASH provides a complete pipeline for training knowledge agents on **existing corpora**. We abstract the training, synthesis, and serving into a modular system that your code doesn't need to interface with.
+- **Train from anywhere.** Run the KONASH client on your laptop and let the server kick off training on a single GPU — local or cloud. No multi-node clusters required.
+- Integrations with hosted platforms like W&B and Langfuse provide flexible observability and **simplify debugging** across the full synthesis-train-eval loop.
+- KONASH is customizable with **intelligent defaults**. You can configure OAPL hyperparameters, compression thresholds, and inference engine settings to meet specific needs, or take advantage of defaults optimized for single-GPU training efficiency and stability.
+
+## Installation
+
+KONASH agents can be trained from any client machine that runs Python. To add to an existing project, run this command:
+
+```
+pip install konash
+```
+
+---
+
+## Training Loop Overview
+
+KONASH uses **large-batch iterative off-policy RL** — unlike online RL frameworks, all data is generated upfront and training happens in a single offline pass. Each iteration improves the model, which then generates better data for the next iteration.
+
+1. **Data Synthesis**
+
+   1. KONASH generates training questions from your corpus using an agentic synthesis pipeline — the model explores documents via vector search and proposes grounded QA pairs.
+   2. A deduplication step ensures no overlap with your evaluation set.
+   3. On later iterations, the improved model synthesizes its own curriculum — harder, more diverse questions.
+
+2. **Rollout Generation**
+
+   1. The model (or latest checkpoint) generates multiple rollouts per question, interacting with vector search and compression tools.
+   2. Each rollout is a full multi-step agent trajectory: search queries, retrieved documents, context compression, and a final answer.
+   3. Rewards are computed automatically from answer correctness against ground truth.
+   4. Pass-rate filtering keeps questions at the learning frontier — not too easy, not too hard.
+
+3. **Training**
+
+   1. The full set of trajectories becomes a large offline dataset. Training runs in a single batch — no interleaving with inference.
+   2. The server trains your model using OAPL with QLoRA. Long trajectories are segmented at compression boundaries, and tool outputs are masked from log-prob computation.
+   3. The newly trained LoRA is saved and becomes the starting point for the next iteration.
+
+4. **Iterate**
+
+   1. The trained checkpoint becomes the new reference policy.
+   2. All rollouts are **regenerated from scratch** with the improved model — this is what makes each iteration progressively better.
+   3. Training runs again on the fresh data. 2–3 iterations yields the best results.
+
+## Supported Models
+
+KONASH should work with most vLLM/HuggingFace-transformers compatible causal language models, or at least the ones supported by [Unsloth](https://docs.unsloth.ai/get-started/all-our-models). If any model isn't working for you, please let us know on [Discord](#) or open an issue on [GitHub](https://github.com/konaequity/openkona/issues)!
+
+---
+
+## Contributing
+
+KONASH is in active development and contributions are welcome! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for more information.
+
+---
+
+## Citation
+
+```bibtex
+@misc{konaequity2026konash,
+  author = {Kona Equity},
+  title = {KONASH: Knowledge-grounded Off-policy Networks for Agentic System Harnesses},
+  year = {2026},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/konaequity/openkona}}
+}
+```
+
+---
+
+## License
+
+This repository's source code is available under the [Apache-2.0 License](LICENSE).
+
+---
+
+## Credits
+
+KONASH builds directly on the research and open-source work of:
+
+- [OAPL](https://arxiv.org/abs/2503.01735) — Ritter et al., 2026 (the RL algorithm)
+- [Unsloth](https://github.com/unslothai/unsloth) — Parameter-efficient training backend
+- [vLLM](https://github.com/vllm-project/vllm) — High-throughput inference engine
+- [FAISS](https://github.com/facebookresearch/faiss) — Vector search

konash-0.1.0/konash/__init__.py

@@ -0,0 +1,6 @@
+"""KONASH: Knowledge-grounded Off-policy Networks for Agentic System Harnesses."""
+
+from konash.api import Agent
+from konash.corpus import Corpus
+
+__all__ = ["Agent", "Corpus"]

konash-0.1.0/konash/agent.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import copy
+from typing import Any, Callable, Dict, List, Optional
+
+
+class Agent:
+    """Core agent that wraps an injected LLM client to drive multi-step rollouts.
+
+    The LLM client is any object that exposes a ``generate(messages, **kwargs)``
+    method returning an assistant message dict (at minimum ``{"role": "assistant",
+    "content": "..."}``).  Tool calls, adapters, and history compression are
+    handled at this layer so that the harness environment stays LLM-agnostic.
+    """
+
+    llm_client = None
+
+    def __init__(
+        self,
+        llm_client: Any = None,
+        *,
+        system_prompt: str | None = None,
+        max_steps: int = 20,
+        stop_sequences: List[str] | None = None,
+    ) -> None:
+        self.llm_client = llm_client
+        self.system_prompt = system_prompt
+        self.max_steps = max_steps
+        self.stop_sequences = stop_sequences or []
+        self._active_adapters: List[Dict[str, Any]] = []
+
+    # ------------------------------------------------------------------
+    # Core generation
+    # ------------------------------------------------------------------
+
+    def generate(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+        """Send *messages* to the LLM client and return the raw response.
+
+        The caller is responsible for building the message list (including any
+        system prompt).  Extra ``kwargs`` are forwarded to the client.
+        """
+        if self.llm_client is None:
+            raise RuntimeError("No llm_client configured on this Agent.")
+        return self.llm_client.generate(messages, **kwargs)
+
+    # ------------------------------------------------------------------
+    # Step-level generation
+    # ------------------------------------------------------------------
+
+    def generate_step(
+        self,
+        conversation_history: List[Dict[str, str]],
+        available_tools: List[Dict[str, Any]] | None = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Produce a single agent step (one LLM turn).
+
+        Returns a dict with at least ``{"role": "assistant", "content": ...}``.
+        If the model emits a tool-call the dict may also contain a
+        ``"tool_calls"`` key.
+        """
+        messages = self._build_messages(conversation_history)
+        gen_kwargs: Dict[str, Any] = dict(kwargs)
+        if available_tools:
+            gen_kwargs["tools"] = available_tools
+        if self.stop_sequences:
+            gen_kwargs.setdefault("stop", self.stop_sequences)
+
+        response = self.generate(messages, **gen_kwargs)
+        return response
+
+    # ------------------------------------------------------------------
+    # Full rollout
+    # ------------------------------------------------------------------
+
+    def generate_rollout(
+        self,
+        prompt: str,
+        environment: Any = None,
+        *,
+        max_steps: int | None = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Run a complete episode loop and return the trajectory.
+
+        If an *environment* is provided its ``step`` / ``run_episode`` contract
+        is used; otherwise we fall back to a simple generate-until-done loop.
+        """
+        steps = max_steps or self.max_steps
+
+        if environment is not None:
+            environment.reset(prompt=prompt)
+            result = environment.run_episode(agent=self, max_steps=steps, **kwargs)
+            return result
+
+        # Standalone loop (no environment)
+        history: List[Dict[str, str]] = [{"role": "user", "content": prompt}]
+        trajectory: List[Dict[str, Any]] = []
+
+        for _step_idx in range(steps):
+            response = self.generate_step(history)
+            trajectory.append(response)
+            history.append(response)
+
+            # Check for natural termination
+            if self._is_terminal(response):
+                break
+
+        return {
+            "prompt": prompt,
+            "history": history,
+            "trajectory": trajectory,
+            "final_answer": self.extract_final_answer(history),
+        }
+
+    # ------------------------------------------------------------------
+    # History compression
+    # ------------------------------------------------------------------
+
+    def compress_history(
+        self,
+        conversation_history: List[Dict[str, str]],
+        *,
+        target_tokens: int | None = None,
+        **kwargs,
+    ) -> List[Dict[str, str]]:
+        """Ask the LLM to produce a shorter version of *conversation_history*.
+
+        Returns a new message list whose semantic content is preserved but
+        whose token footprint is reduced.
+        """
+        if not conversation_history:
+            return []
+
+        # Build a compression prompt
+        serialized = "\n".join(
+            f"[{m.get('role', 'unknown')}]: {m.get('content', '')}"
+            for m in conversation_history
+        )
+        compression_prompt = (
+            "Compress the following conversation into a concise summary that "
+            "preserves all critical facts, tool results, and reasoning steps. "
+            "Return ONLY the summary.\n\n" + serialized
+        )
+        if target_tokens is not None:
+            compression_prompt += f"\n\nTarget length: roughly {target_tokens} tokens."
+
+        messages = [{"role": "user", "content": compression_prompt}]
+        response = self.generate(messages, **kwargs)
+        summary_content = response.get("content", "") if isinstance(response, dict) else str(response)
+
+        compressed: List[Dict[str, str]] = [
+            {"role": "system", "content": f"[Compressed history] {summary_content}"},
+        ]
+        return compressed
+
+    # ------------------------------------------------------------------
+    # Answer extraction
+    # ------------------------------------------------------------------
+
+    def extract_final_answer(
+        self,
+        conversation_history: List[Dict[str, str]],
+        **kwargs,
+    ) -> str | None:
+        """Extract the final answer from the conversation history.
+
+        Walks the history backwards looking for the last assistant message
+        that does not appear to be a tool call.
+        """
+        for message in reversed(conversation_history):
+            if message.get("role") != "assistant":
+                continue
+            # Skip messages that are purely tool calls with no textual content
+            if message.get("tool_calls") and not message.get("content"):
+                continue
+            content = message.get("content", "")
+            if content:
+                return content
+        return None
+
+    # ------------------------------------------------------------------
+    # LoRA adapter management
+    # ------------------------------------------------------------------
+
+    def load_adapter(
+        self,
+        adapter_path: str,
+        *,
+        adapter_name: str | None = None,
+        weight: float = 1.0,
+        **kwargs,
+    ) -> None:
+        """Load a LoRA adapter and merge it into the active set."""
+        name = adapter_name or adapter_path.rstrip("/").split("/")[-1]
+        adapter_record = {
+            "path": adapter_path,
+            "name": name,
+            "weight": weight,
+            **kwargs,
+        }
+        self._active_adapters.append(adapter_record)
+
+        # Delegate to the LLM client if it supports adapter loading
+        if hasattr(self.llm_client, "load_adapter"):
+            self.llm_client.load_adapter(adapter_path, name=name, weight=weight, **kwargs)
+
+    def unload_adapter(
+        self,
+        adapter_name: str | None = None,
+        **kwargs,
+    ) -> None:
+        """Unload a LoRA adapter (or all adapters if *adapter_name* is ``None``)."""
+        if adapter_name is None:
+            removed = list(self._active_adapters)
+            self._active_adapters.clear()
+        else:
+            removed = [a for a in self._active_adapters if a["name"] == adapter_name]
+            self._active_adapters = [
+                a for a in self._active_adapters if a["name"] != adapter_name
+            ]
+
+        # Delegate to the LLM client if it supports adapter unloading
+        if hasattr(self.llm_client, "unload_adapter"):
+            for adapter in removed:
+                self.llm_client.unload_adapter(adapter["name"], **kwargs)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _build_messages(
+        self, conversation_history: List[Dict[str, str]]
+    ) -> List[Dict[str, str]]:
+        """Prepend the system prompt (if any) to the conversation history."""
+        messages: List[Dict[str, str]] = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.extend(conversation_history)
+        return messages
+
+    @staticmethod
+    def _is_terminal(response: Dict[str, Any]) -> bool:
+        """Heuristic: a response is terminal if it has content and no tool calls."""
+        has_tool_calls = bool(response.get("tool_calls"))
+        has_content = bool(response.get("content"))
+        return has_content and not has_tool_calls
+
+
+class ValueGuidedAgent:
+    """Agent variant that generates multiple candidate continuations per step,
+    scores each with a value model, and selects the best one.
+
+    Drop-in replacement for ``Agent.generate_step`` when plugged into an
+    environment or strategy that expects the same interface.
+    """
+
+    candidate_width = 2
+    value_model = None
+
+    def __init__(
+        self,
+        llm_client: Any = None,
+        *,
+        candidate_width: int = 2,
+        value_model: Any = None,
+        system_prompt: str | None = None,
+        max_steps: int = 20,
+        stop_sequences: List[str] | None = None,
+    ) -> None:
+        # Compose an inner Agent for actual generation
+        self._agent = Agent(
+            llm_client=llm_client,
+            system_prompt=system_prompt,
+            max_steps=max_steps,
+            stop_sequences=stop_sequences,
+        )
+        self.candidate_width = candidate_width
+        self.value_model = value_model
+        self.llm_client = llm_client
+
+    def generate_step(
+        self,
+        conversation_history: List[Dict[str, str]],
+        available_tools: List[Dict[str, Any]] | None = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """Generate *candidate_width* candidates, score them, pick the best."""
+        candidates: List[Dict[str, Any]] = []
+        for _ in range(self.candidate_width):
+            candidate = self._agent.generate_step(
+                conversation_history,
+                available_tools=available_tools,
+                **kwargs,
+            )
+            candidates.append(candidate)
+
+        if not candidates:
+            raise RuntimeError("No candidates generated.")
+
+        # If no value model is available, fall back to first candidate
+        if self.value_model is None:
+            return candidates[0]
+
+        # Score each candidate
+        best_candidate = candidates[0]
+        best_score = float("-inf")
+        for candidate in candidates:
+            augmented_history = list(conversation_history) + [candidate]
+            score = self.value_model.score_partial_rollout(augmented_history)
+            if score > best_score:
+                best_score = score
+                best_candidate = candidate
+
+        return best_candidate