agentensor 0.0.2__tar.gz → 0.0.4__tar.gz

{agentensor-0.0.2 → agentensor-0.0.4}/.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.2
+current_version = 0.0.4
 commit = True
 tag = True
 

{agentensor-0.0.2 → agentensor-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentensor
-Version: 0.0.2
+Version: 0.0.4
 Summary: Add your description here
 License: MIT License
         
@@ -26,8 +26,12 @@ License: MIT License
 License-File: LICENSE
 Requires-Python: >=3.12
 Requires-Dist: datasets>=3.5.0
+Requires-Dist: langchain-ollama>=0.3.3
+Requires-Dist: langchain-openai>=0.3.18
+Requires-Dist: langchain>=0.3.25
+Requires-Dist: langgraph>=0.4.5
 Requires-Dist: logfire>=3.14.0
-Requires-Dist: pydantic-ai>=0.1.3
+Requires-Dist: pydantic-ai>=0.2.4
 Description-Content-Type: text/markdown
 
 # AgenTensor

{agentensor-0.0.2 → agentensor-0.0.4}/agentensor/loss.py

@@ -16,7 +16,7 @@ class LLMTensorJudge(Evaluator[TextTensor, TextTensor, Any]):
     """
 
     rubric: str
-    model: models.KnownModelName | None = None
+    model: models.Model | models.KnownModelName | None = None
     include_input: bool = True
 
     async def evaluate(

agentensor-0.0.4/agentensor/module.py

@@ -0,0 +1,62 @@
+"""Module class."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+from langchain.chat_models import init_chat_model
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+from langgraph.graph.graph import CompiledGraph
+from pydantic import BaseModel, ConfigDict
+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
+    llm: str | BaseChatModel = "gpt-4o-mini"
+
+    def model_post_init(self, __context: Any) -> None:
+        """Post initialization hook."""
+        if isinstance(self.llm, str):  # pragma: no cover
+            self.llm = init_chat_model(self.llm)
+
+    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"]
+        try:
+            result = await self.agent.ainvoke(
+                {"messages": [HumanMessage(content=state["output"].text)]}
+            )
+            output = str(
+                result.get("structured_response", result["messages"][-1].content)
+            )  # prioritize structured response over raw response
+        except UnexpectedModelBehavior:  # pragma: no cover
+            output = "Error"
+
+        output_tensor = TextTensor(
+            output,
+            parents=[state["output"], self.system_prompt],
+            requires_grad=True,
+            model=self.llm,
+        )
+
+        return {"output": output_tensor}
+
+    @property
+    @abstractmethod
+    def agent(self) -> CompiledGraph:
+        """Get agent instance."""
+        pass  # pragma: no cover

agentensor-0.0.4/agentensor/optim.py

@@ -0,0 +1,57 @@
+"""Optimizer module."""
+
+from langchain.chat_models import init_chat_model
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+from langgraph.graph import StateGraph
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import create_react_agent
+from agentensor.module import AgentModule
+from agentensor.tensor import TextTensor
+
+
+class Optimizer:
+    """Optimizer class."""
+
+    def __init__(
+        self,
+        graph: StateGraph,
+        model: str | BaseChatModel = "gpt-4o-mini",
+    ) -> None:
+        """Initialize the optimizer."""
+        self.params: list[TextTensor] = [
+            param
+            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]
+        ]
+        if isinstance(model, str):
+            self.model = init_chat_model(model)
+        else:  # pragma: no cover
+            self.model = model
+
+    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.zero_grad()
+
+    def optimize(self, text: str, grad: str) -> str:
+        """Optimize the text."""
+        result = self.agent.invoke(
+            {"messages": [HumanMessage(content=f"Feedback: {grad}\nText: {text}")]}
+        )
+        return result["messages"][-1].content
+
+    @property
+    def agent(self) -> CompiledGraph:
+        """Get the agent."""
+        return create_react_agent(
+            self.model, tools=[], prompt="Rewrite the system prompt given the feedback."
+        )

{agentensor-0.0.2 → agentensor-0.0.4}/agentensor/tensor.py

@@ -1,7 +1,11 @@
 """Example module."""
 
 from __future__ import annotations
-from pydantic_ai import Agent, models
+from langchain.chat_models import init_chat_model
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import create_react_agent
 
 
 class TextTensor:
@@ -12,17 +16,17 @@ class TextTensor:
         text: str,
         parents: list[TextTensor] | None = None,
         requires_grad: bool = False,
-        model: models.Model | models.KnownModelName | str | None = None,
+        model: str | BaseChatModel = "gpt-4o-mini",
     ) -> None:
         """Initialize a TextTensor."""
         self.text = text
         self.requires_grad = requires_grad
         self.gradients: list[str] = []
-        self.agent = Agent(
-            model=model or "openai:gpt-4o-mini",
-            system_prompt="Answer the user's question.",
-        )
         self.parents: list[TextTensor] = parents or []
+        if isinstance(model, str):
+            self.model = init_chat_model(model)
+        else:
+            self.model = model
 
     def backward(self, grad: str = "") -> None:
         """Backward pass for the TextTensor.
@@ -43,12 +47,21 @@ class TextTensor:
 
     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
+        result = self.agent.invoke(
+            {
+                "messages": [
+                    HumanMessage(
+                        content=(
+                            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?"
+                        )
+                    )
+                ]
+            }
+        )
+        return result["messages"][-1].content
 
     @property
     def text_grad(self) -> str:
@@ -62,3 +75,10 @@ class TextTensor:
     def __str__(self) -> str:
         """Return the text as a string."""
         return self.text
+
+    @property
+    def agent(self) -> CompiledGraph:
+        """Get the agent."""
+        return create_react_agent(
+            self.model, tools=[], prompt="Answer the user's question."
+        )

{agentensor-0.0.2 → agentensor-0.0.4}/agentensor/train.py

@@ -1,11 +1,9 @@
 """Trainer."""
 
 from typing import Any, Literal
+from langgraph.graph.graph import CompiledGraph
 from pydantic_evals import Dataset
 from pydantic_evals.reporting import EvaluationReport
-from pydantic_graph import Graph
-from pydantic_graph.nodes import DepsT, StateT
-from agentensor.module import AgentModule, ModuleState
 from agentensor.optim import Optimizer
 from agentensor.tensor import TextTensor
 
@@ -15,9 +13,8 @@ class Trainer:
 
     def __init__(
         self,
-        graph: Graph[StateT, DepsT, TextTensor],
-        graph_state: ModuleState,
-        start_node: type[AgentModule],
+        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,
@@ -27,8 +24,7 @@ class Trainer:
     ):
         """Initialize the trainer."""
         self.graph = graph
-        self.graph_state = graph_state
-        self.start_node = start_node
+        self.graph_recursion_limit = graph_recursion_limit
         self.optimizer = optimizer
         self.epochs = epochs
         self.stop_threshold = stop_threshold
@@ -38,9 +34,10 @@ class Trainer:
 
     async def forward(self, x: TextTensor) -> TextTensor:
         """Forward the graph."""
-        self.graph_state.input = x
-        result = await self.graph.run(self.start_node(), state=self.graph_state)  # type: ignore[arg-type]
-        return result.output
+        result = await self.graph.ainvoke(
+            {"output": x}, {"recursion_limit": self.graph_recursion_limit}
+        )
+        return result["output"]
 
     def train(self) -> None:
         """Train the graph."""

{agentensor-0.0.2 → agentensor-0.0.4}/examples/evaluate.py

@@ -3,16 +3,17 @@
 from __future__ import annotations
 import json
 from dataclasses import dataclass
+from typing import TypedDict
 from datasets import load_dataset
+from langchain.chat_models import init_chat_model
+from langchain_core.language_models import BaseChatModel
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import create_react_agent
 from pydantic import BaseModel
-from pydantic_ai import Agent, models
-from pydantic_ai.exceptions import UnexpectedModelBehavior
-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 pydantic_graph import End, Graph, GraphRunContext
-from agentensor.module import AgentModule, ModuleState
+from agentensor.module import AgentModule
 from agentensor.tensor import TextTensor
 from agentensor.train import Trainer
 
@@ -48,11 +49,10 @@ class MultiLabelClassificationAccuracy(Evaluator):
         return set(output) == set(expected)  # type: ignore[arg-type]
 
 
-@dataclass
-class EvaluateState(ModuleState):
+class EvaluateState(TypedDict):
     """State of the graph."""
 
-    agent_prompt: TextTensor = TextTensor(text="")
+    output: TextTensor
 
 
 class ClassificationResults(BaseModel, use_attribute_docstrings=True):
@@ -73,12 +73,15 @@ class HFMultiClassClassificationTask:
         self,
         task_repo: str,
         evaluators: list[Evaluator],
-        model: models.Model | models.KnownModelName | str | None = None,
+        model: BaseChatModel | str = "gpt-4o-mini",
     ) -> None:
         """Initialize the multi-class classification task."""
         self.task_repo = task_repo
         self.evaluators = evaluators
-        self.model = model
+        if isinstance(model, str):
+            self.model = init_chat_model(model)
+        else:
+            self.model = model
         self.dataset = self._prepare_dataset()
 
     def _prepare_dataset(self) -> dict[str, Dataset]:
@@ -101,60 +104,50 @@ class HFMultiClassClassificationTask:
         return dataset
 
 
-class AgentNode(AgentModule[EvaluateState, None, TextTensor]):
+class AgentNode(AgentModule):
     """Agent node."""
 
-    async def run(self, ctx: GraphRunContext[EvaluateState, None]) -> End[TextTensor]:  # type: ignore[override]
-        """Run the agent node."""
-        agent = Agent(
-            model=model,
-            system_prompt=ctx.state.agent_prompt.text,
-            output_type=ClassificationResults,
-        )
-        assert ctx.state.input
-        try:
-            result = await agent.run(ctx.state.input.text)
-            output = result.output
-        except UnexpectedModelBehavior:
-            output = "Error"  # type: ignore[assignment]
-
-        output_tensor = TextTensor(
-            str(output),
-            parents=[ctx.state.input, ctx.state.agent_prompt],
-            requires_grad=True,
+    @property
+    def agent(self) -> CompiledGraph:
+        """Get agent instance."""
+        return create_react_agent(
+            self.llm,
+            tools=[],
+            prompt=self.system_prompt.text,
+            response_format=ClassificationResults,
         )
 
-        return End(output_tensor)
-
 
 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"
+    model = init_chat_model("llama3.2:1b", model_provider="ollama")
+    # model = "gpt-4o-mini"
 
     task = HFMultiClassClassificationTask(
         task_repo="knowledgator/events_classification_biotech",
         evaluators=[GenerationTimeout(), MultiLabelClassificationAccuracy()],
         model=model,
     )
-    state = EvaluateState(
-        agent_prompt=TextTensor(
-            (
-                "Classify the following text into one of the following "
-                "categories: [expanding industry, new initiatives or programs, "
-                "article publication, other]"
+    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,
             ),
-            requires_grad=True,
-            model=model,
-        )
+            llm=model,
+        ),
     )
-    graph = Graph(nodes=[AgentNode])
+    graph.add_edge(START, "agent")
+    graph.add_edge("agent", END)
+    compiled_graph = graph.compile()
     trainer = Trainer(
-        graph,
-        state,
-        AgentNode,  # type: ignore[arg-type]
+        compiled_graph,
         train_dataset=task.dataset["train"],
         test_dataset=task.dataset["test"],
     )

agentensor-0.0.4/examples/train.py

@@ -0,0 +1,118 @@
+"""Example usage of agentensor."""
+
+from __future__ import annotations
+import os
+from dataclasses import dataclass
+from typing import Any, TypedDict
+from langchain.chat_models import init_chat_model
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.graph import CompiledGraph
+from langgraph.prebuilt import create_react_agent
+from pydantic_ai import 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."""
+
+    @property
+    def agent(self) -> CompiledGraph:
+        """Get agent instance."""
+        return create_react_agent(
+            self.llm,
+            tools=[],
+            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",
+        )
+    eval_model = OpenAIModel(
+        model_name="llama3.2:1b",
+        provider=OpenAIProvider(base_url="http://localhost:11434/v1", api_key="ollama"),
+    )
+    model = init_chat_model("llama3.2:1b", model_provider="ollama")
+    # eval_model = "gpt-4o-mini"
+    # model = "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=eval_model),
+            FormatJudge(model=eval_model),
+        ],
+    )
+
+    graph = StateGraph(TrainState)
+    graph.add_node(
+        "agent",
+        AgentNode(
+            system_prompt=TextTensor(
+                "You are a helpful assistant.", requires_grad=True, model=model
+            ),
+            llm=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.2 → agentensor-0.0.4}/pyproject.toml

@@ -27,8 +27,12 @@ docs = [
 [project]
 dependencies = [
   "datasets>=3.5.0",
+  "langchain>=0.3.25",
+  "langchain-ollama>=0.3.3",
+  "langchain-openai>=0.3.18",
+  "langgraph>=0.4.5",
   "logfire>=3.14.0",
-  "pydantic-ai>=0.1.3"
+  "pydantic-ai>=0.2.4"
 ]
 description = "Add your description here"
 license = {file = "LICENSE"}
@@ -36,7 +40,7 @@ name = "agentensor"
 readme = "README.md"
 requires-python = ">=3.12"
 url = "https://github.com/ShaojieJiang/agentensor"
-version = "0.0.2"
+version = "0.0.4"
 
 [tool.coverage.report]
 exclude_lines = [

{agentensor-0.0.2 → agentensor-0.0.4}/tests/test_loss.py

@@ -33,8 +33,12 @@ def mock_judge_output():
 
 
 @pytest.fixture
-def evaluator_context(mock_openai):
+@patch("agentensor.tensor.init_chat_model")
+def evaluator_context(mock_init_chat_model, mock_openai):
     """Create a test evaluator context."""
+    # Mock the model initialization
+    mock_init_chat_model.return_value = MagicMock()
+
     return EvaluatorContext(
         name="test_evaluator",
         inputs=TextTensor(text="test input"),