agentforge-ml 0.1.0__py3-none-any.whl

agentforge/eval/report.py

@@ -0,0 +1,55 @@
+"""Pretty-printable evaluation report."""
+
+from __future__ import annotations
+
+import json
+import statistics
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class EvalReport:
+    n: int
+    means: dict[str, float]
+    per_sample: dict[str, list[float]]
+    latencies_ms: list[float] = field(default_factory=list)
+    extras: dict = field(default_factory=dict)
+
+    @property
+    def p50_ms(self) -> float:
+        return statistics.median(self.latencies_ms) if self.latencies_ms else 0.0
+
+    @property
+    def p95_ms(self) -> float:
+        if not self.latencies_ms:
+            return 0.0
+        s = sorted(self.latencies_ms)
+        return s[int(0.95 * (len(s) - 1))]
+
+    def as_table(self) -> str:
+        width = max((len(m) for m in self.means), default=14)
+        bar = "+" + "-" * (width + 4) + "+--------+"
+        lines = [bar, f"| {'metric':<{width + 2}} | mean   |", bar]
+        for name, val in self.means.items():
+            lines.append(f"| {name:<{width + 2}} | {val:.3f} |")
+        lines.append(bar)
+        if self.latencies_ms:
+            lines.append(f"n={self.n}  ·  p50={self.p50_ms:.0f}ms  ·  p95={self.p95_ms:.0f}ms")
+        else:
+            lines.append(f"n={self.n}")
+        return "\n".join(lines)
+
+    def save(self, path: str | Path) -> None:
+        Path(path).write_text(
+            json.dumps(
+                {
+                    "n": self.n,
+                    "means": self.means,
+                    "per_sample": self.per_sample,
+                    "latencies_ms": self.latencies_ms,
+                    "extras": self.extras,
+                },
+                indent=2,
+            )
+        )

agentforge/llm/__init__.py

@@ -0,0 +1,7 @@
+"""LLM backends."""
+
+from agentforge.llm.base import LLM
+from agentforge.llm.hf import HFLLM
+from agentforge.llm.quantized import QuantizedHFLLM
+
+__all__ = ["HFLLM", "LLM", "QuantizedHFLLM"]

agentforge/llm/base.py

@@ -0,0 +1,16 @@
+"""LLM protocol — anything that maps a prompt to text and honors stop strings."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class LLM(Protocol):
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        stop: list[str] | None = None,
+    ) -> str: ...

agentforge/llm/hf.py

@@ -0,0 +1,83 @@
+"""HuggingFace causal LM backend.
+
+ReAct-friendly: applies the model's chat template, accepts stop strings,
+and truncates the output at the first stop marker so the next loop iteration
+gets a clean continuation.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+class HFLLM:
+    def __init__(
+        self,
+        model_id: str = "Qwen/Qwen2.5-3B-Instruct",
+        *,
+        dtype: str = "auto",
+        device_map: str | dict | None = "auto",
+        trust_remote_code: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        self.model_id = model_id
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id, trust_remote_code=trust_remote_code
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=dtype,
+            device_map=device_map,
+            trust_remote_code=trust_remote_code,
+            **kwargs,
+        )
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        stop: list[str] | None = None,
+    ) -> str:
+        text = self._apply_chat_template(prompt)
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
+        gen_kwargs: dict[str, Any] = {
+            "max_new_tokens": max_new_tokens,
+            "pad_token_id": self.tokenizer.pad_token_id,
+        }
+        if temperature > 0:
+            gen_kwargs.update(do_sample=True, temperature=temperature)
+        else:
+            gen_kwargs.update(do_sample=False)
+
+        with torch.no_grad():
+            out = self.model.generate(**inputs, **gen_kwargs)
+        new_ids = out[0, inputs.input_ids.shape[-1] :]
+        answer = self.tokenizer.decode(new_ids, skip_special_tokens=True)
+        return _truncate_at_stop(answer, stop)
+
+    def _apply_chat_template(self, prompt: str) -> str:
+        if hasattr(self.tokenizer, "apply_chat_template") and self.tokenizer.chat_template:
+            messages = [{"role": "user", "content": prompt}]
+            return self.tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
+        return prompt
+
+
+def _truncate_at_stop(text: str, stop: list[str] | None) -> str:
+    if not stop:
+        return text.strip()
+    earliest = len(text)
+    for s in stop:
+        idx = text.find(s)
+        if idx != -1:
+            earliest = min(earliest, idx)
+    return text[:earliest].strip()

agentforge/llm/quantized.py

@@ -0,0 +1,39 @@
+"""Quantized LLM via turboquant-ml — same trick as ragforge-ml."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from agentforge.llm.hf import HFLLM
+
+
+class QuantizedHFLLM(HFLLM):
+    """:class:`HFLLM` that applies a TurboQuant method before serving.
+
+    Example
+    -------
+    >>> llm = QuantizedHFLLM("meta-llama/Llama-3.2-3B-Instruct", method="bnb-nf4")
+    """
+
+    def __init__(
+        self,
+        model_id: str = "meta-llama/Llama-3.2-3B-Instruct",
+        *,
+        method: str = "bnb-nf4",
+        quant_kwargs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(model_id, **kwargs)
+        self.method = method
+        self.model = _quantize(self.model, method=method, **(quant_kwargs or {}))
+
+
+def _quantize(model, *, method: str, **kw):
+    try:
+        from turboquant import quantize
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "turboquant-ml is required for QuantizedHFLLM. "
+            'Install with `pip install "agentforge-ml[quantized]"`.'
+        ) from e
+    return quantize(model, method=method, **kw)

agentforge/memory/__init__.py

@@ -0,0 +1,7 @@
+"""Conversation memory backends."""
+
+from agentforge.memory.base import Memory, Message
+from agentforge.memory.conversation import ConversationMemory
+from agentforge.memory.persistent import PersistentMemory
+
+__all__ = ["ConversationMemory", "Memory", "Message", "PersistentMemory"]

agentforge/memory/base.py

@@ -0,0 +1,23 @@
+"""Memory protocol + Message type."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Protocol
+
+
+@dataclass
+class Message:
+    role: str
+    content: str
+    ts: float = field(default_factory=time.time)
+    extras: dict = field(default_factory=dict)
+
+
+class Memory(Protocol):
+    def add(self, role: str, content: str, **extras: object) -> None: ...
+
+    def get(self, *, limit: int | None = None) -> list[Message]: ...
+
+    def clear(self) -> None: ...

agentforge/memory/conversation.py

@@ -0,0 +1,30 @@
+"""In-memory conversation history with a rolling window."""
+
+from __future__ import annotations
+
+from collections import deque
+
+from agentforge.memory.base import Message
+
+
+class ConversationMemory:
+    """FIFO window of the last ``max_messages`` messages."""
+
+    def __init__(self, max_messages: int = 64) -> None:
+        self.max_messages = max_messages
+        self._buf: deque[Message] = deque(maxlen=max_messages)
+
+    def add(self, role: str, content: str, **extras: object) -> None:
+        self._buf.append(Message(role=role, content=content, extras=dict(extras)))
+
+    def get(self, *, limit: int | None = None) -> list[Message]:
+        msgs = list(self._buf)
+        if limit:
+            msgs = msgs[-limit:]
+        return msgs
+
+    def clear(self) -> None:
+        self._buf.clear()
+
+    def __len__(self) -> int:
+        return len(self._buf)

agentforge/memory/persistent.py

@@ -0,0 +1,80 @@
+"""SQLite-backed persistent memory.
+
+Sessions are keyed by an arbitrary string id — useful for multi-user agents
+where each user has their own running history. No ORM, no migration tooling:
+a single ``messages`` table is created on first use.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import time
+from pathlib import Path
+
+from agentforge.memory.base import Message
+
+
+class PersistentMemory:
+    def __init__(
+        self, db_path: str | Path = "agentforge_memory.db", *, session: str = "default"
+    ) -> None:
+        self.db_path = str(db_path)
+        self.session = session
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+        self._init()
+
+    def _init(self) -> None:
+        with sqlite3.connect(self.db_path) as cx:
+            cx.execute(
+                """
+                CREATE TABLE IF NOT EXISTS messages (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    session TEXT NOT NULL,
+                    role TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    ts REAL NOT NULL,
+                    extras_json TEXT
+                )
+                """
+            )
+            cx.execute(
+                "CREATE INDEX IF NOT EXISTS idx_messages_session_ts ON messages(session, ts)"
+            )
+
+    def add(self, role: str, content: str, **extras: object) -> None:
+        with sqlite3.connect(self.db_path) as cx:
+            cx.execute(
+                "INSERT INTO messages (session, role, content, ts, extras_json) VALUES (?, ?, ?, ?, ?)",
+                (self.session, role, content, time.time(), json.dumps(extras) if extras else None),
+            )
+
+    def get(self, *, limit: int | None = None) -> list[Message]:
+        with sqlite3.connect(self.db_path) as cx:
+            cx.row_factory = sqlite3.Row
+            q = "SELECT role, content, ts, extras_json FROM messages WHERE session = ? ORDER BY ts ASC"
+            rows = cx.execute(q, (self.session,)).fetchall()
+        msgs = [
+            Message(
+                role=r["role"],
+                content=r["content"],
+                ts=r["ts"],
+                extras=json.loads(r["extras_json"]) if r["extras_json"] else {},
+            )
+            for r in rows
+        ]
+        if limit:
+            msgs = msgs[-limit:]
+        return msgs
+
+    def clear(self) -> None:
+        with sqlite3.connect(self.db_path) as cx:
+            cx.execute("DELETE FROM messages WHERE session = ?", (self.session,))
+
+    def __len__(self) -> int:
+        with sqlite3.connect(self.db_path) as cx:
+            return int(
+                cx.execute(
+                    "SELECT COUNT(*) FROM messages WHERE session = ?", (self.session,)
+                ).fetchone()[0]
+            )

agentforge/serve/__init__.py

@@ -0,0 +1,5 @@
+"""FastAPI serve module."""
+
+from agentforge.serve.app import build_app
+
+__all__ = ["build_app"]

agentforge/serve/app.py

@@ -0,0 +1,83 @@
+"""FastAPI app for AgentForge."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from agentforge.core.agent import Agent
+
+
+class AskRequest(BaseModel):
+    question: str = Field(..., min_length=1)
+    max_steps: int | None = Field(None, ge=1, le=20)
+    max_new_tokens: int = Field(256, ge=16, le=2048)
+
+
+class StepModel(BaseModel):
+    thought: str
+    tool: str | None
+    action_input: str | None
+    observation: str | None
+    elapsed_ms: float
+
+
+class AskResponse(BaseModel):
+    question: str
+    final_answer: str
+    n_steps: int
+    success: bool
+    latency_ms: float
+    steps: list[StepModel]
+
+
+class ToolModel(BaseModel):
+    name: str
+    description: str
+
+
+def build_app(agent: Agent) -> Any:
+    try:
+        from fastapi import FastAPI
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            'FastAPI is required. Install with `pip install "agentforge-ml[serve]"`.'
+        ) from e
+
+    app = FastAPI(title="AgentForge", version="0.1.0")
+
+    @app.get("/health")
+    def health() -> dict:
+        return {"status": "ok", "n_tools": len(agent.tools)}
+
+    @app.get("/tools", response_model=list[ToolModel])
+    def list_tools() -> list[ToolModel]:
+        return [ToolModel(name=t.name, description=t.description) for t in agent.tools]
+
+    @app.post("/ask", response_model=AskResponse)
+    def ask(req: AskRequest) -> AskResponse:
+        result = agent.run(
+            req.question,
+            max_steps=req.max_steps,
+            max_new_tokens=req.max_new_tokens,
+        )
+        return AskResponse(
+            question=result.question,
+            final_answer=result.final_answer,
+            n_steps=result.n_steps,
+            success=result.success,
+            latency_ms=result.latency_ms,
+            steps=[
+                StepModel(
+                    thought=s.thought,
+                    tool=s.tool,
+                    action_input=s.action_input,
+                    observation=s.observation,
+                    elapsed_ms=s.elapsed_ms,
+                )
+                for s in result.steps
+            ],
+        )
+
+    return app

agentforge/tools/__init__.py

@@ -0,0 +1,18 @@
+"""Built-in tool registry."""
+
+from agentforge.tools.base import Tool, ToolRegistry
+from agentforge.tools.calculator import Calculator
+from agentforge.tools.python_repl import PythonREPL
+from agentforge.tools.rag import RAGTool
+from agentforge.tools.sql import SQLTool
+from agentforge.tools.web_search import WebSearch
+
+__all__ = [
+    "Calculator",
+    "PythonREPL",
+    "RAGTool",
+    "SQLTool",
+    "Tool",
+    "ToolRegistry",
+    "WebSearch",
+]

agentforge/tools/base.py

@@ -0,0 +1,55 @@
+"""Tool protocol and registry.
+
+A tool is *anything* with three things: a name (so the LLM can request it), a
+description (so the LLM knows what it does), and a ``run(input_str) -> str``
+method (so the orchestrator can call it). We avoid JSON schemas on purpose —
+small open models often emit malformed JSON and the cost/benefit isn't there.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Tool(Protocol):
+    name: str
+    description: str
+
+    def run(self, input_str: str) -> str: ...
+
+
+class ToolRegistry:
+    def __init__(self, tools: list[Tool] | None = None) -> None:
+        self._tools: dict[str, Tool] = {}
+        for t in tools or []:
+            self.register(t)
+
+    def register(self, tool: Tool) -> None:
+        if not getattr(tool, "name", None):
+            raise ValueError(f"Tool has no `.name`: {tool!r}")
+        self._tools[tool.name] = tool
+
+    def get(self, name: str) -> Tool | None:
+        # Lenient lookup so the LLM can write "Calculator", "calculator", or "calc".
+        if name in self._tools:
+            return self._tools[name]
+        lc = name.strip().lower()
+        for k, v in self._tools.items():
+            if k.lower() == lc:
+                return v
+        return None
+
+    def __iter__(self) -> Iterator[Tool]:
+        return iter(self._tools.values())
+
+    def __len__(self) -> int:
+        return len(self._tools)
+
+    def __contains__(self, name: str) -> bool:
+        return self.get(name) is not None
+
+    @property
+    def names(self) -> list[str]:
+        return list(self._tools)

agentforge/tools/calculator.py

@@ -0,0 +1,115 @@
+"""Calculator tool — safe arithmetic via AST whitelist.
+
+We deliberately avoid ``eval()`` and ``sympy.sympify`` directly. The whitelist
+covers the operators a calculator actually needs: + - * / // % **, unary +/-,
+parentheses, int/float literals, and a fixed set of math functions.
+"""
+
+from __future__ import annotations
+
+import ast
+import math
+import operator as op
+
+_BIN_OPS = {
+    ast.Add: op.add,
+    ast.Sub: op.sub,
+    ast.Mult: op.mul,
+    ast.Div: op.truediv,
+    ast.FloorDiv: op.floordiv,
+    ast.Mod: op.mod,
+    ast.Pow: op.pow,
+}
+
+_UNARY_OPS = {
+    ast.UAdd: op.pos,
+    ast.USub: op.neg,
+}
+
+_FUNCS = {
+    "sqrt": math.sqrt,
+    "log": math.log,
+    "ln": math.log,
+    "log2": math.log2,
+    "log10": math.log10,
+    "exp": math.exp,
+    "sin": math.sin,
+    "cos": math.cos,
+    "tan": math.tan,
+    "asin": math.asin,
+    "acos": math.acos,
+    "atan": math.atan,
+    "abs": abs,
+    "round": round,
+    "floor": math.floor,
+    "ceil": math.ceil,
+    "max": max,
+    "min": min,
+}
+
+_CONSTS = {
+    "pi": math.pi,
+    "e": math.e,
+    "tau": math.tau,
+}
+
+
+class Calculator:
+    name = "calculator"
+    description = (
+        "Safely evaluate an arithmetic expression. "
+        "Supports + - * / // % ** parentheses and functions "
+        "sqrt, log, exp, sin/cos/tan, abs, round, floor, ceil, max, min. "
+        "Constants: pi, e. Example: '47 * 1337'."
+    )
+
+    def run(self, input_str: str) -> str:
+        expr = input_str.strip()
+        if not expr:
+            return "Error: empty expression"
+        try:
+            tree = ast.parse(expr, mode="eval")
+            value = _eval(tree.body)
+        except Exception as e:
+            return f"Error: {type(e).__name__}: {e}"
+        return _fmt(value)
+
+
+def _eval(node):
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, (int, float)):
+            return node.value
+        raise ValueError(f"unsupported literal: {node.value!r}")
+    if isinstance(node, ast.BinOp):
+        left = _eval(node.left)
+        right = _eval(node.right)
+        op_fn = _BIN_OPS.get(type(node.op))
+        if op_fn is None:
+            raise ValueError(f"unsupported operator: {type(node.op).__name__}")
+        return op_fn(left, right)
+    if isinstance(node, ast.UnaryOp):
+        op_fn = _UNARY_OPS.get(type(node.op))
+        if op_fn is None:
+            raise ValueError(f"unsupported unary operator: {type(node.op).__name__}")
+        return op_fn(_eval(node.operand))
+    if isinstance(node, ast.Call):
+        if not isinstance(node.func, ast.Name):
+            raise ValueError("only simple function calls allowed")
+        name = node.func.id
+        if name not in _FUNCS:
+            raise ValueError(f"unknown function: {name}")
+        args = [_eval(a) for a in node.args]
+        return _FUNCS[name](*args)
+    if isinstance(node, ast.Name):
+        if node.id not in _CONSTS:
+            raise ValueError(f"unknown identifier: {node.id}")
+        return _CONSTS[node.id]
+    raise ValueError(f"unsupported expression: {ast.dump(node)}")
+
+
+def _fmt(v) -> str:
+    if isinstance(v, float) and v.is_integer():
+        return str(int(v))
+    if isinstance(v, float):
+        return f"{v:.6g}"
+    return str(v)