agentensor 0.0.1__tar.gz

agentensor-0.0.1/.bumpversion.cfg

@@ -0,0 +1,12 @@
+[bumpversion]
+current_version = 0.0.1
+commit = True
+tag = True
+
+[bumpversion:file:pyproject.toml]
+search = version = "{current_version}"
+replace = version = "{new_version}"
+
+[bumpversion:file:uv.lock]
+search = version = "{current_version}"
+replace = version = "{new_version}"

agentensor-0.0.1/.github/workflows/after-ci.yml

@@ -0,0 +1,50 @@
+name: After CI
+
+on:
+  workflow_run:
+    workflows: [CI]
+    types: [completed]
+
+permissions:
+  statuses: write
+  pull-requests: write
+
+jobs:
+  smokeshow:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+          python-version: '3.12'
+
+      - uses: dawidd6/action-download-artifact@v6
+        with:
+          workflow: ci.yml
+          name: '(diff-)?coverage-html.*'
+          name_is_regexp: true
+          commit: ${{ github.event.workflow_run.head_sha }}
+          allow_forks: true
+          workflow_conclusion: completed
+          if_no_artifact_found: warn
+
+      - run: uvx smokeshow upload coverage-html
+        if: hashFiles('coverage-html/*.html') != ''
+        env:
+          SMOKESHOW_GITHUB_STATUS_DESCRIPTION: Coverage {coverage-percentage}
+          SMOKESHOW_GITHUB_COVERAGE_THRESHOLD: 50
+          SMOKESHOW_GITHUB_CONTEXT: coverage
+          SMOKESHOW_GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          SMOKESHOW_GITHUB_PR_HEAD_SHA: ${{ github.event.workflow_run.head_sha }}
+          SMOKESHOW_AUTH_KEY: ${{ secrets.SMOKESHOW_AUTH_KEY }}
+
+      - run: uvx smokeshow upload diff-coverage-html
+        if: hashFiles('diff-coverage-html/*.html') != ''
+        env:
+          SMOKESHOW_GITHUB_STATUS_DESCRIPTION: Diff coverage {coverage-percentage}
+          SMOKESHOW_GITHUB_COVERAGE_THRESHOLD: 95
+          SMOKESHOW_GITHUB_CONTEXT: diff-coverage
+          SMOKESHOW_GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          SMOKESHOW_GITHUB_PR_HEAD_SHA: ${{ github.event.workflow_run.head_sha }}
+          SMOKESHOW_AUTH_KEY: ${{ secrets.SMOKESHOW_AUTH_KEY }}

agentensor-0.0.1/.github/workflows/ci.yml

@@ -0,0 +1,103 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - 'v*'
+  pull_request:
+    branches:
+      - main
+    types:
+      - opened
+      - synchronize
+      - ready_for_review
+
+permissions:
+  id-token: write  # Required for PyPI trusted publishing
+  contents: read
+
+jobs:
+  lint:
+    if: github.event.pull_request.draft == false
+    runs-on: ubuntu-latest
+
+    strategy:
+      fail-fast: false
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync --all-groups
+
+      - name: Check style against standards
+        run: uv run make lint
+
+  coverage:
+    if: github.event.pull_request.draft == false
+    runs-on: ubuntu-latest
+
+    strategy:
+      fail-fast: false
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Install dependencies
+        run: uv sync --all-groups
+
+      - name: Run unit tests with coverage
+        run: |
+          uv run coverage run -m pytest
+          uv run coverage html --show-contexts --title "Coverage for ${{ github.sha }}"
+          uv run coverage xml
+          uv run diff-cover coverage.xml --html-report index.html
+
+      - name: Store coverage html
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-html
+          path: htmlcov
+          include-hidden-files: true
+
+      - name: Store diff coverage html
+        uses: actions/upload-artifact@v4
+        with:
+          name: diff-coverage-html
+          path: index.html
+
+      - name: Check coverage
+        run: |
+          uv run coverage report --fail-under 95
+          uv run diff-cover coverage.xml --fail-under 100
+
+  build-and-release:
+    needs: [lint, coverage]
+    if: "success() && startsWith(github.ref, 'refs/tags/')"
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: yezz123/setup-uv@v4
+
+      - name: Build project
+        run: uv build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

agentensor-0.0.1/.gitignore

@@ -0,0 +1,164 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+xunit-result.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# MacOS
+**/.DS_Store

agentensor-0.0.1/.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-json
+      - id: check-case-conflict
+      - id: check-docstring-first
+      - id: check-builtin-literals
+      - id: check-ast
+      - id: check-merge-conflict
+
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: 'v0.4.6'
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
+    rev: v2.13.0
+    hooks:
+      - id: pretty-format-toml
+        args: [--autofix]

agentensor-0.0.1/.python-version

@@ -0,0 +1 @@
+3.12.6

agentensor-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Shaojie Jiang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

agentensor-0.0.1/Makefile

@@ -0,0 +1,15 @@
+lint:
+	ruff check .
+	mypy agentensor examples
+	ruff format . --check
+
+format:
+	ruff format .
+	ruff check . --select I001 --fix
+	ruff check . --select F401 --fix
+
+test:
+	pytest --cov --cov-report term-missing tests/
+
+doc:
+	mkdocs serve --dev-addr=0.0.0.0:8080

agentensor-0.0.1/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: agentensor
+Version: 0.0.1
+Summary: Add your description here
+License: MIT License
+        
+        Copyright (c) 2025 Shaojie Jiang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.12
+Requires-Dist: logfire>=3.14.0
+Requires-Dist: pydantic-ai>=0.0.55
+Description-Content-Type: text/markdown
+
+# AgenTensor
+
+[![CI](https://github.com/ShaojieJiang/agentensor/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/ShaojieJiang/agentensor/actions/workflows/ci.yml?query=branch%3Amain)
+[![Coverage](https://coverage-badge.samuelcolvin.workers.dev/ShaojieJiang/agentensor.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/ShaojieJiang/agentensor)
+[![PyPI](https://img.shields.io/pypi/v/agentensor.svg)](https://pypi.python.org/pypi/agentensor)
+
+## TODO
+
+- [ ] Add parameter saving

agentensor-0.0.1/README.md

@@ -0,0 +1,9 @@
+# AgenTensor
+
+[![CI](https://github.com/ShaojieJiang/agentensor/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/ShaojieJiang/agentensor/actions/workflows/ci.yml?query=branch%3Amain)
+[![Coverage](https://coverage-badge.samuelcolvin.workers.dev/ShaojieJiang/agentensor.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/ShaojieJiang/agentensor)
+[![PyPI](https://img.shields.io/pypi/v/agentensor.svg)](https://pypi.python.org/pypi/agentensor)
+
+## TODO
+
+- [ ] Add parameter saving

agentensor-0.0.1/agentensor/__init__.py

@@ -0,0 +1 @@
+"""Example module."""

agentensor-0.0.1/agentensor/loss.py

@@ -0,0 +1,44 @@
+"""Loss functions."""
+
+from dataclasses import dataclass
+from typing import Any
+from pydantic_ai import models
+from pydantic_evals.evaluators import EvaluationReason, Evaluator, EvaluatorContext
+from pydantic_evals.evaluators.llm_as_a_judge import judge_input_output, judge_output
+from agentensor.tensor import TextTensor
+
+
+@dataclass
+class LLMTensorJudge(Evaluator[TextTensor, TextTensor, Any]):
+    """LLM judge for text tensors.
+
+    Adapted from pydantic_evals.evaluators.common.LLMJudge.
+    """
+
+    rubric: str
+    model: models.Model | models.KnownModelName | None = None
+    include_input: bool = True
+
+    async def evaluate(
+        self,
+        ctx: EvaluatorContext[TextTensor, TextTensor, Any],
+    ) -> EvaluationReason:
+        """Evaluate the text tensor."""
+        if self.include_input:
+            grading_output = await judge_input_output(
+                ctx.inputs.text, ctx.output.text, self.rubric, self.model
+            )
+        else:
+            grading_output = await judge_output(
+                ctx.output.text, self.rubric, self.model
+            )
+        return EvaluationReason(
+            value=grading_output.pass_, reason=grading_output.reason
+        )
+
+    def build_serialization_arguments(self) -> dict[str, Any]:
+        """Build serialization arguments."""
+        result = super().build_serialization_arguments()
+        if (model := result.get("model")) and isinstance(model, models.Model):
+            result["model"] = f"{model.system}:{model.model_name}"
+        return result

agentensor-0.0.1/agentensor/module.py

@@ -0,0 +1,26 @@
+"""Module class."""
+
+from dataclasses import dataclass
+from pydantic_graph.nodes import BaseNode, DepsT, NodeRunEndT, StateT
+from agentensor.tensor import TextTensor
+
+
+@dataclass
+class ModuleState:
+    """State of the graph."""
+
+    input: TextTensor
+
+
+class AgentModule(BaseNode[StateT, DepsT, NodeRunEndT]):
+    """Agent module."""
+
+    @classmethod
+    def get_params(cls) -> list[TextTensor]:
+        """Get the parameters of the module."""
+        params = []
+        for base in cls.__mro__:
+            for _, attr in base.__dict__.items():
+                if isinstance(attr, TextTensor) and attr.requires_grad:
+                    params.append(attr)
+        return params

agentensor-0.0.1/agentensor/optim.py

@@ -0,0 +1,39 @@
+"""Optimizer module."""
+
+from pydantic_ai import Agent
+from pydantic_graph import Graph
+from agentensor.module import AgentModule
+from agentensor.tensor import TextTensor
+
+
+class Optimizer:
+    """Optimizer class."""
+
+    def __init__(self, graph: Graph) -> None:
+        """Initialize the optimizer."""
+        self.params: list[TextTensor] = [
+            param
+            for node in graph.get_nodes()
+            for param in node.get_params()  # type: ignore[attr-defined]
+            if issubclass(node, AgentModule)
+        ]
+        self.agent: Agent = Agent(
+            model="openai:gpt-4o-mini",
+            system_prompt="Rewrite the system prompt given the feedback.",
+        )
+
+    def step(self) -> None:
+        """Step the optimizer."""
+        for param in self.params:
+            if not param.text_grad:
+                continue
+            param.text = self.optimize(param.text, param.text_grad)
+
+    def zero_grad(self) -> None:
+        """Zero the gradients."""
+        for param in self.params:
+            param.text_grad = ""
+
+    def optimize(self, text: str, grad: str) -> str:
+        """Optimize the text."""
+        return self.agent.run_sync(f"Feedback: {grad}\nText: {text}").data

agentensor-0.0.1/agentensor/tensor.py

@@ -0,0 +1,53 @@
+"""Example module."""
+
+from __future__ import annotations
+from pydantic_ai import Agent
+
+
+class TextTensor:
+    """A tensor that represents a text."""
+
+    def __init__(
+        self,
+        text: str,
+        parents: list[TextTensor] | None = None,
+        requires_grad: bool = False,
+    ) -> None:
+        """Initialize a TextTensor."""
+        self.text = text
+        self.requires_grad = requires_grad
+        self.text_grad = ""
+        self.agent = Agent(
+            model="openai:gpt-4o-mini", system_prompt="Answer the user's question."
+        )
+        self.parents: list[TextTensor] = parents or []
+
+    def backward(self, grad: str = "") -> None:
+        """Backward pass for the TextTensor.
+
+        Args:
+            grad (str, optional): The gradient to backpropagate. Defaults to "".
+        """
+        if not grad:
+            return
+
+        if self.requires_grad:
+            self.text_grad = grad
+            for parent in self.parents:
+                if not parent.requires_grad:
+                    continue
+                grad_to_parent = self.calc_grad(parent.text, self.text, grad)
+                parent.backward(grad_to_parent)
+
+    def calc_grad(self, input_text: str, output_text: str, grad: str) -> str:
+        """Calculate the gradient for the TextTensor."""
+        return self.agent.run_sync(
+            f"Here is the input: \n\n>{input_text}\n\nI got this "
+            f"output: \n\n>{output_text}\n\nHere is the feedback: \n\n"
+            f">{grad}\n\nHow should I improve the input to get a "
+            f"better output?"
+        ).data
+
+    def __str__(self) -> str:
+        """Return the text as a string."""
+        return self.text

agentensor-0.0.1/agentensor/train.py

@@ -0,0 +1,66 @@
+"""Trainer."""
+
+from typing import Any
+from pydantic_evals import Dataset
+from pydantic_graph import Graph
+from agentensor.module import AgentModule, ModuleState
+from agentensor.optim import Optimizer
+from agentensor.tensor import TextTensor
+
+
+class Trainer:
+    """Trainer."""
+
+    def __init__(
+        self,
+        graph: Graph[ModuleState, None, TextTensor],
+        start_node: type[AgentModule],
+        dataset: Dataset[TextTensor, TextTensor, Any],
+        optimizer: Optimizer,
+        epochs: int,
+        stop_threshold: float = 0.95,
+    ):
+        """Initialize the trainer."""
+        self.graph = graph
+        self.start_node = start_node
+        self.dataset = dataset
+        self.optimizer = optimizer
+        self.epochs = epochs
+        self.stop_threshold = stop_threshold
+
+    async def step(self, x: TextTensor) -> TextTensor:
+        """Step the optimizer."""
+        state = ModuleState(input=x)
+        result = await self.graph.run(self.start_node(), state=state)  # type: ignore[arg-type]
+        return result.output
+
+    def train(self) -> None:
+        """Train the model."""
+        for i in range(self.epochs):
+            report = self.dataset.evaluate_sync(self.step)
+            report.print(
+                include_input=True, include_output=True, include_durations=True
+            )
+
+            # Backward those failed cases
+            for case in report.cases:
+                losses = []
+                for evaluator in case.assertions.values():
+                    if not evaluator.value:
+                        assert evaluator.reason
+                        losses.append(evaluator.reason)
+                if losses:
+                    case.output.backward(" ".join(losses))
+
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+            print(f"Epoch {i + 1}")
+            for param in self.optimizer.params:
+                print(param.text)  # pragma: no cover
+            print()
+            performance = report.averages().assertions
+            assert performance is not None
+            if performance >= self.stop_threshold:
+                print("Optimization complete.")
+                break