agentensor 0.0.1__tar.gz → 0.0.3__tar.gz

{agentensor-0.0.1 → agentensor-0.0.3}/.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.1
+current_version = 0.0.3
 commit = True
 tag = True
 

{agentensor-0.0.1 → agentensor-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentensor
-Version: 0.0.1
+Version: 0.0.3
 Summary: Add your description here
 License: MIT License
         
@@ -25,8 +25,10 @@ License: MIT License
         SOFTWARE.
 License-File: LICENSE
 Requires-Python: >=3.12
+Requires-Dist: datasets>=3.5.0
+Requires-Dist: langgraph>=0.4.5
 Requires-Dist: logfire>=3.14.0
-Requires-Dist: pydantic-ai>=0.0.55
+Requires-Dist: pydantic-ai>=0.2.4
 Description-Content-Type: text/markdown
 
 # AgenTensor

agentensor-0.0.3/agentensor/module.py

@@ -0,0 +1,49 @@
+"""Module class."""
+
+from abc import ABC, abstractmethod
+from pydantic import BaseModel, ConfigDict
+from pydantic_ai import Agent, models
+from pydantic_ai.exceptions import UnexpectedModelBehavior
+from agentensor.tensor import TextTensor
+
+
+class AgentModule(BaseModel, ABC):
+    """Agent module."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    system_prompt: TextTensor
+    model: models.Model | models.KnownModelName | str = "openai:gpt-4o"
+
+    def get_params(self) -> list[TextTensor]:
+        """Get the parameters of the module."""
+        params = []
+        for field_name in self.__class__.model_fields.keys():
+            field = getattr(self, field_name)
+            if isinstance(field, TextTensor) and field.requires_grad:
+                params.append(field)
+        return params
+
+    async def __call__(self, state: dict) -> dict:
+        """Run the agent node."""
+        assert state["output"]
+        agent = self.get_agent()
+        try:
+            result = await agent.run(state["output"].text)
+            output = str(result.output)
+        except UnexpectedModelBehavior:  # pragma: no cover
+            output = "Error"
+
+        output_tensor = TextTensor(
+            output,
+            parents=[state["output"], self.system_prompt],
+            requires_grad=True,
+            model=self.model,
+        )
+
+        return {"output": output_tensor}
+
+    @abstractmethod
+    def get_agent(self) -> Agent:
+        """Get agent instance."""
+        pass  # pragma: no cover

{agentensor-0.0.1 → agentensor-0.0.3}/agentensor/optim.py

@@ -1,7 +1,7 @@
 """Optimizer module."""
 
-from pydantic_ai import Agent
-from pydantic_graph import Graph
+from langgraph.graph import StateGraph
+from pydantic_ai import Agent, models
 from agentensor.module import AgentModule
 from agentensor.tensor import TextTensor
 
@@ -9,16 +9,20 @@ from agentensor.tensor import TextTensor
 class Optimizer:
     """Optimizer class."""
 
-    def __init__(self, graph: Graph) -> None:
+    def __init__(
+        self,
+        graph: StateGraph,
+        model: models.Model | models.KnownModelName | str | None = None,
+    ) -> 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)
+            for node in graph.nodes.values()
+            if isinstance(node.runnable.afunc, AgentModule)  # type: ignore[attr-defined]
+            for param in node.runnable.afunc.get_params()  # type: ignore[attr-defined]
         ]
         self.agent: Agent = Agent(
-            model="openai:gpt-4o-mini",
+            model=model or "openai:gpt-4o-mini",
             system_prompt="Rewrite the system prompt given the feedback.",
         )
 
@@ -32,7 +36,7 @@ class Optimizer:
     def zero_grad(self) -> None:
         """Zero the gradients."""
         for param in self.params:
-            param.text_grad = ""
+            param.zero_grad()
 
     def optimize(self, text: str, grad: str) -> str:
         """Optimize the text."""

{agentensor-0.0.1 → agentensor-0.0.3}/agentensor/tensor.py

@@ -1,7 +1,7 @@
 """Example module."""
 
 from __future__ import annotations
-from pydantic_ai import Agent
+from pydantic_ai import Agent, models
 
 
 class TextTensor:
@@ -12,13 +12,15 @@ class TextTensor:
         text: str,
         parents: list[TextTensor] | None = None,
         requires_grad: bool = False,
+        model: models.Model | models.KnownModelName | str | None = None,
     ) -> None:
         """Initialize a TextTensor."""
         self.text = text
         self.requires_grad = requires_grad
-        self.text_grad = ""
+        self.gradients: list[str] = []
         self.agent = Agent(
-            model="openai:gpt-4o-mini", system_prompt="Answer the user's question."
+            model=model or "openai:gpt-4o-mini",
+            system_prompt="Answer the user's question.",
         )
         self.parents: list[TextTensor] = parents or []
 
@@ -32,7 +34,7 @@ class TextTensor:
             return
 
         if self.requires_grad:
-            self.text_grad = grad
+            self.gradients.append(grad)
             for parent in self.parents:
                 if not parent.requires_grad:
                     continue
@@ -46,7 +48,16 @@ class TextTensor:
             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
+        ).output
+
+    @property
+    def text_grad(self) -> str:
+        """String representation of the gradients."""
+        return " ".join(self.gradients)
+
+    def zero_grad(self) -> None:
+        """Zero the gradients."""
+        self.gradients = []
 
     def __str__(self) -> str:
         """Return the text as a string."""

agentensor-0.0.3/agentensor/train.py

@@ -0,0 +1,93 @@
+"""Trainer."""
+
+from typing import Any, Literal
+from langgraph.graph.graph import CompiledGraph
+from pydantic_evals import Dataset
+from pydantic_evals.reporting import EvaluationReport
+from agentensor.optim import Optimizer
+from agentensor.tensor import TextTensor
+
+
+class Trainer:
+    """Trainer."""
+
+    def __init__(
+        self,
+        graph: CompiledGraph,
+        graph_recursion_limit: int = 25,
+        train_dataset: Dataset[TextTensor, TextTensor, Any] | None = None,
+        eval_dataset: Dataset[TextTensor, TextTensor, Any] | None = None,
+        test_dataset: Dataset[TextTensor, TextTensor, Any] | None = None,
+        optimizer: Optimizer | None = None,
+        epochs: int = 10,
+        stop_threshold: float = 0.95,
+    ):
+        """Initialize the trainer."""
+        self.graph = graph
+        self.graph_recursion_limit = graph_recursion_limit
+        self.optimizer = optimizer
+        self.epochs = epochs
+        self.stop_threshold = stop_threshold
+        self.train_dataset = train_dataset
+        self.eval_dataset = eval_dataset
+        self.test_dataset = test_dataset
+
+    async def forward(self, x: TextTensor) -> TextTensor:
+        """Forward the graph."""
+        result = await self.graph.ainvoke(
+            {"output": x}, {"recursion_limit": self.graph_recursion_limit}
+        )
+        return result["output"]
+
+    def train(self) -> None:
+        """Train the graph."""
+        assert self.train_dataset, "Train dataset is required"
+        assert self.optimizer, "Optimizer is required"
+        for i in range(self.epochs):
+            report = self.evaluate("train")
+            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
+
+    def evaluate(
+        self,
+        data_split: Literal["train", "eval", "test"] = "eval",
+        limit_cases: int | None = None,
+    ) -> EvaluationReport:
+        """Evaluate the graph."""
+        dataset = getattr(self, f"{data_split}_dataset")
+        assert dataset, f"{data_split} dataset is required"
+        if limit_cases:  # pragma: no cover
+            limited_cases = dataset.cases[:limit_cases]
+            dataset = Dataset(cases=limited_cases, evaluators=dataset.evaluators)
+        report = dataset.evaluate_sync(self.forward)
+
+        return report
+
+    def test(self, limit_cases: int | None = None) -> None:
+        """Test the graph."""
+        report = self.evaluate("test", limit_cases=limit_cases)
+        report.print(include_input=True, include_output=True, include_durations=True)

agentensor-0.0.3/examples/evaluate.py

@@ -0,0 +1,151 @@
+"""Tasks."""
+
+from __future__ import annotations
+import json
+from dataclasses import dataclass
+from typing import TypedDict
+from datasets import load_dataset
+from langgraph.graph import END, START, StateGraph
+from pydantic import BaseModel
+from pydantic_ai import Agent, models
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_evals import Case, Dataset
+from pydantic_evals.evaluators import EvaluationReason, Evaluator, EvaluatorContext
+from agentensor.module import AgentModule
+from agentensor.tensor import TextTensor
+from agentensor.train import Trainer
+
+
+@dataclass
+class GenerationTimeout(Evaluator[str, bool]):
+    """The generation took too long."""
+
+    threshold: float = 10.0
+
+    async def evaluate(self, ctx: EvaluatorContext[str, bool]) -> EvaluationReason:
+        """Evaluate the time taken to generate the output."""
+        return EvaluationReason(
+            value=ctx.duration <= self.threshold,
+            reason=(
+                f"The generation took {ctx.duration} seconds, which is longer "
+                f"than the threshold of {self.threshold} seconds."
+            ),
+        )
+
+
+@dataclass
+class MultiLabelClassificationAccuracy(Evaluator):
+    """Classification accuracy evaluator."""
+
+    async def evaluate(self, ctx: EvaluatorContext) -> bool:
+        """Evaluate the accuracy of the classification."""
+        try:
+            output = json.loads(ctx.output.text)
+        except json.JSONDecodeError:
+            return False
+        expected = ctx.expected_output
+        return set(output) == set(expected)  # type: ignore[arg-type]
+
+
+class EvaluateState(TypedDict):
+    """State of the graph."""
+
+    output: TextTensor
+
+
+class ClassificationResults(BaseModel, use_attribute_docstrings=True):
+    """Classification result for a data."""
+
+    labels: list[str]
+    """labels for this data point."""
+
+    def __str__(self) -> str:
+        """Return the string representation of the classification results."""
+        return json.dumps(self.labels)
+
+
+class HFMultiClassClassificationTask:
+    """Multi-class classification task from Hugging Face."""
+
+    def __init__(
+        self,
+        task_repo: str,
+        evaluators: list[Evaluator],
+        model: models.Model | models.KnownModelName | str | None = None,
+    ) -> None:
+        """Initialize the multi-class classification task."""
+        self.task_repo = task_repo
+        self.evaluators = evaluators
+        self.model = model
+        self.dataset = self._prepare_dataset()
+
+    def _prepare_dataset(self) -> dict[str, Dataset]:
+        """Return the Pydantic Evals dataset."""
+        hf_dataset = load_dataset(self.task_repo, trust_remote_code=True)
+        dataset = {}
+        for split in hf_dataset.keys():
+            cases = []
+            for example in hf_dataset[split]:
+                cases.append(
+                    Case(
+                        inputs=TextTensor(
+                            f"Title: {example['title']}\nContent: {example['content']}",
+                            model=self.model,
+                        ),
+                        expected_output=example["all_labels"],
+                    )
+                )
+            dataset[split] = Dataset(cases=cases, evaluators=self.evaluators)
+        return dataset
+
+
+class AgentNode(AgentModule):
+    """Agent node."""
+
+    def get_agent(self) -> Agent:
+        """Get agent instance."""
+        return Agent(
+            model=self.model or "openai:gpt-4o-mini",
+            system_prompt=self.system_prompt.text,
+            output_type=ClassificationResults,  # type: ignore[arg-type]
+        )
+
+
+if __name__ == "__main__":
+    model = OpenAIModel(
+        model_name="llama3.2:1b",
+        provider=OpenAIProvider(base_url="http://localhost:11434/v1", api_key="ollama"),
+    )
+    # model = "openai:gpt-4o-mini"
+
+    task = HFMultiClassClassificationTask(
+        task_repo="knowledgator/events_classification_biotech",
+        evaluators=[GenerationTimeout(), MultiLabelClassificationAccuracy()],
+        model=model,
+    )
+    graph = StateGraph(EvaluateState)
+    graph.add_node(
+        "agent",
+        AgentNode(
+            system_prompt=TextTensor(
+                (
+                    "Classify the following text into one of the following "
+                    "categories: [expanding industry, new initiatives or programs, "
+                    "article publication, other]"
+                ),
+                requires_grad=True,
+                model=model,
+            ),
+            model=model,
+        ),
+    )
+    graph.add_edge(START, "agent")
+    graph.add_edge("agent", END)
+    compiled_graph = graph.compile()
+    trainer = Trainer(
+        compiled_graph,
+        train_dataset=task.dataset["train"],
+        test_dataset=task.dataset["test"],
+    )
+    trainer.test(limit_cases=10)

agentensor-0.0.3/examples/train.py

@@ -0,0 +1,111 @@
+"""Example usage of agentensor."""
+
+from __future__ import annotations
+import os
+from dataclasses import dataclass
+from typing import Any, TypedDict
+from langgraph.graph import END, START, StateGraph
+from pydantic_ai import Agent, models
+from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_evals import Case, Dataset
+from agentensor.loss import LLMTensorJudge
+from agentensor.module import AgentModule
+from agentensor.optim import Optimizer
+from agentensor.tensor import TextTensor
+from agentensor.train import Trainer
+
+
+@dataclass
+class ChineseLanguageJudge(LLMTensorJudge):
+    """Chinese language judge."""
+
+    rubric: str = "The output should be in Chinese."
+    model: models.Model | models.KnownModelName = "openai:gpt-4o-mini"
+    include_input = True
+
+
+@dataclass
+class FormatJudge(LLMTensorJudge):
+    """Format judge."""
+
+    rubric: str = "The output should start by introducing itself."
+    model: models.Model | models.KnownModelName = "openai:gpt-4o-mini"
+    include_input = True
+
+
+class TrainState(TypedDict):
+    """State of the graph."""
+
+    output: TextTensor
+
+
+class AgentNode(AgentModule):
+    """Agent node."""
+
+    def get_agent(self) -> Agent:
+        """Get agent instance."""
+        return Agent(
+            model=self.model or "openai:gpt-4o-mini",
+            system_prompt=self.system_prompt.text,
+        )
+
+
+def main() -> None:
+    """Main function."""
+    if os.environ.get("LOGFIRE_TOKEN", None):
+        import logfire
+
+        logfire.configure(
+            send_to_logfire="if-token-present",
+            environment="development",
+            service_name="evals",
+        )
+    model = OpenAIModel(
+        model_name="llama3.2:1b",
+        provider=OpenAIProvider(base_url="http://localhost:11434/v1", api_key="ollama"),
+    )
+    # model="openai:gpt-4o-mini"
+
+    dataset = Dataset[TextTensor, TextTensor, Any](
+        cases=[
+            Case(
+                inputs=TextTensor("Hello, how are you?", model=model),
+                metadata={"language": "English"},
+            ),
+            Case(
+                inputs=TextTensor("こんにちは、元気ですか？", model=model),
+                metadata={"language": "Japanese"},
+            ),
+        ],
+        evaluators=[
+            ChineseLanguageJudge(model=model),
+            FormatJudge(model=model),
+        ],
+    )
+
+    graph = StateGraph(TrainState)
+    graph.add_node(
+        "agent",
+        AgentNode(
+            system_prompt=TextTensor(
+                "You are a helpful assistant.", requires_grad=True, model=model
+            ),
+            model=model,
+        ),
+    )
+    graph.add_edge(START, "agent")
+    graph.add_edge("agent", END)
+    compiled_graph = graph.compile()
+    optimizer = Optimizer(graph, model=model)
+    trainer = Trainer(
+        compiled_graph,
+        train_dataset=dataset,
+        optimizer=optimizer,
+        epochs=15,
+    )
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()

{agentensor-0.0.1 → agentensor-0.0.3}/pyproject.toml

@@ -14,7 +14,7 @@ dev = [
   "pytest-cov>=4.1.0",
   "ruff>=0.11.3",
   "smokeshow>=0.5.0",
-  "types-requests",
+  "types-requests"
 ]
 docs = [
   "mkdocs",
@@ -26,17 +26,18 @@ docs = [
 
 [project]
 dependencies = [
-    "logfire>=3.14.0",
-    "pydantic-ai>=0.0.55",
+  "datasets>=3.5.0",
+  "langgraph>=0.4.5",
+  "logfire>=3.14.0",
+  "pydantic-ai>=0.2.4"
 ]
 description = "Add your description here"
 license = {file = "LICENSE"}
 name = "agentensor"
 readme = "README.md"
 requires-python = ">=3.12"
-version = "0.0.1"
 url = "https://github.com/ShaojieJiang/agentensor"
-
+version = "0.0.3"
 
 [tool.coverage.report]
 exclude_lines = [

{agentensor-0.0.1 → agentensor-0.0.3}/tests/test_module.py

@@ -1,8 +1,8 @@
 """Test module for the Module class."""
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
-from agentensor.module import AgentModule, ModuleState
+from agentensor.module import AgentModule
 from agentensor.tensor import TextTensor
 
 
@@ -15,28 +15,17 @@ def mock_agent():
         yield mock_agent
 
 
-def test_module_state_initialization(mock_agent):
-    """Test ModuleState initialization."""
-    input_tensor = TextTensor("test input")
-    state = ModuleState(input=input_tensor)
-
-    assert isinstance(state.input, TextTensor)
-    assert state.input.text == "test input"
-
-
-def test_module_get_params(mock_agent):
+def test_module_get_params():
     """Test AgentModule.get_params() method."""
 
     class TestModule(AgentModule):
-        param1 = TextTensor("param1", requires_grad=True)
-        param2 = TextTensor("param2", requires_grad=False)
-        param3 = TextTensor("param3", requires_grad=True)
-        non_param = "not a tensor"
-
-        def __init__(self):
-            pass
+        system_prompt: TextTensor = TextTensor("param1", requires_grad=True)
+        param2: TextTensor = TextTensor("param2", requires_grad=False)
+        param3: TextTensor = TextTensor("param3", requires_grad=True)
+        model: str = "openai:gpt-4o"
+        non_param: str = "not a tensor"
 
-        def run(self, state: ModuleState) -> None:
+        def get_agent(self):
             """Dummy run method for testing."""
             pass
 
@@ -54,13 +43,10 @@ def test_module_get_params_empty(mock_agent):
     """Test AgentModule.get_params() with no parameters."""
 
     class EmptyModule(AgentModule):
-        non_param = "not a tensor"
-        param = TextTensor("param", requires_grad=False)
+        system_prompt: TextTensor = TextTensor("param", requires_grad=False)
+        non_param: str = "not a tensor"
 
-        def __init__(self):
-            pass
-
-        def run(self, state: ModuleState) -> None:
+        def get_agent(self):
             """Dummy run method for testing."""
             pass
 
@@ -70,26 +56,20 @@ def test_module_get_params_empty(mock_agent):
     assert len(params) == 0
 
 
-def test_module_get_params_inheritance(mock_agent):
+def test_module_get_params_inheritance():
     """Test AgentModule.get_params() with inheritance."""
 
     class ParentModule(AgentModule):
-        parent_param = TextTensor("parent", requires_grad=True)
-
-        def __init__(self):
-            pass
+        system_prompt: TextTensor = TextTensor("parent", requires_grad=True)
 
-        def run(self, state: ModuleState) -> None:
+        def get_agent(self):
             """Dummy run method for testing."""
             pass
 
     class ChildModule(ParentModule):
-        child_param = TextTensor("child", requires_grad=True)
+        child_param: TextTensor = TextTensor("child", requires_grad=True)
 
-        def __init__(self):
-            super().__init__()
-
-        def run(self, state: ModuleState) -> None:
+        def get_agent(self):
             """Dummy run method for testing."""
             pass
 
@@ -100,3 +80,23 @@ def test_module_get_params_inheritance(mock_agent):
     assert all(isinstance(p, TextTensor) for p in params)
     assert all(p.requires_grad for p in params)
     assert {p.text for p in params} == {"parent", "child"}
+
+
+@pytest.mark.asyncio
+async def test_module_call():
+    class TestModule(AgentModule):
+        system_prompt: TextTensor = TextTensor("system prompt", requires_grad=True)
+
+        def get_agent(self):
+            """Dummy run method for testing."""
+            mock_agent = AsyncMock()
+            run_output = MagicMock()
+            run_output.output = "Output text"
+            mock_agent.run.return_value = run_output
+            return mock_agent
+
+    module = TestModule()
+
+    result = await module({"output": TextTensor("Input text")})
+    assert isinstance(result["output"], TextTensor)
+    assert result["output"].text == "Output text"