agentensor 0.0.1__py3-none-any.whl

agentensor/__init__.py

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

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/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/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/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/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

agentensor-0.0.1.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,10 @@
+agentensor/__init__.py,sha256=q_xtykXlbfOMT4e0rlS-LqNVewgcSDKKJN8EAN4ekXA,22
+agentensor/loss.py,sha256=Gtit_dgCRaT4-KzwaahoB9Prik64ZD60CuSaXa6HFLg,1531
+agentensor/module.py,sha256=4sq14LdGUwIZSpxsAQzuSVEZXq0xsGYpuYrfdIx_5Cw,682
+agentensor/optim.py,sha256=_YQDi5AX1vqAlRXEw3TYaty3_BoFVCHrWZIrSv8wXgU,1217
+agentensor/tensor.py,sha256=9yMxugXJy8aLniPPe3TSv52Ui5ll6QHD4IwPzcqRKwM,1686
+agentensor/train.py,sha256=-pdhyJgoqj-wprL9gz1g4-jI0ZOoPwJUBbf_gGVBxpU,2179
+agentensor-0.0.1.dist-info/METADATA,sha256=TXVk6ma6mRb9AtD8wLjMtDavll29bI4ob7pxqncAle8,1989
+agentensor-0.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+agentensor-0.0.1.dist-info/licenses/LICENSE,sha256=UMhWnNnpcD5pmbMMRSZGXPtDuafJDTSZtzuvH245gAo,1070
+agentensor-0.0.1.dist-info/RECORD,,

agentensor-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

agentensor-0.0.1.dist-info/licenses/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.