agentensor 0.0.4__tar.gz → 0.1.0__tar.gz

agentensor-0.1.0/.bumpversion.cfg

@@ -0,0 +1,18 @@
+[bumpversion]
+current_version = 0.1.0
+commit = True
+tag = True
+tag_name = agentensor-v{new_version}
+message = Bump agentensor version to {new_version}
+
+[bumpversion:file:pyproject.toml]
+search = version = "{current_version}"
+replace = version = "{new_version}"
+
+[bumpversion:file:../../uv.lock]
+search = [[package]]
+	name = "agentensor"
+	version = "{current_version}"
+replace = [[package]]
+	name = "agentensor"
+	version = "{new_version}"

{agentensor-0.0.4 → agentensor-0.1.0}/.gitignore

@@ -10,11 +10,10 @@ __pycache__/
 .Python
 build/
 develop-eggs/
-dist/
 downloads/
 eggs/
 .eggs/
-lib/
+/lib/
 lib64/
 parts/
 sdist/
@@ -57,7 +56,6 @@ cover/
 *.pot
 
 # Django stuff:
-*.log
 local_settings.py
 db.sqlite3
 db.sqlite3-journal
@@ -153,12 +151,44 @@ dmypy.json
 # 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/
+# Dependencies
+node_modules/
+
+# Logs
+logs
+*.log
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+pnpm-debug.log*
+lerna-debug.log*
 
-# MacOS
+# Build
+dist/
+dist-ssr/
+*.local
+
+# Editor directories and files
+.vscode/*
+!.vscode/extensions.json
+.idea/
+*.suo
+*.ntvs*
+*.njsproj
+*.sln
+*.sw?
+
+# OS generated files
+.DS_Store
 **/.DS_Store
+
+# LangGraph
+checkpoints.sqlite
+
+# Orcheo
+.orcheo/
+**/workflow_config.json
+
+# Miscellaneous
+homepage/
+orcheo-paper/

agentensor-0.1.0/PKG-INFO

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: agentensor
+Version: 0.1.0
+Summary: Agent prompt tensors, modules, and optimizers for Orcheo workflows
+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>=1.1.3
+Requires-Dist: langgraph>=1.0.5
+Requires-Dist: logfire>=3.14.0
+Requires-Dist: pydantic-ai>=0.2.4
+Provides-Extra: dev
+Requires-Dist: bump2version; extra == 'dev'
+Requires-Dist: diff-cover>=9.2.4; extra == 'dev'
+Requires-Dist: mypy>=1.11.2; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23.8; extra == 'dev'
+Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
+Requires-Dist: ruff>=0.11.3; extra == 'dev'
+Requires-Dist: smokeshow>=0.5.0; extra == 'dev'
+Requires-Dist: types-requests; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == 'docs'
+Requires-Dist: mkdocs-gen-files; extra == 'docs'
+Requires-Dist: mkdocs-jupyter; extra == 'docs'
+Requires-Dist: mkdocs-material; extra == 'docs'
+Requires-Dist: mkdocstrings[python]>=0.28.1; extra == 'docs'
+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.4 → agentensor-0.1.0}/pyproject.toml

@@ -2,11 +2,27 @@
 build-backend = "hatchling.build"
 requires = ["hatchling"]
 
-[dependency-groups]
+[project]
+dependencies = [
+  "datasets>=3.5.0",
+  "langchain>=1.1.3",
+  "langchain-ollama>=0.3.3",
+  "langchain-openai>=0.3.18",
+  "langgraph>=1.0.5",
+  "logfire>=3.14.0",
+  "pydantic-ai>=0.2.4"
+]
+description = "Agent prompt tensors, modules, and optimizers for Orcheo workflows"
+name = "agentensor"
+readme = "README.md"
+requires-python = ">=3.12"
+url = "https://github.com/ShaojieJiang/agentensor"
+version = "0.1.0"
+
+[project.optional-dependencies]
 dev = [
   "bump2version",
   "diff-cover>=9.2.4",
-  "isort",
   "mypy>=1.11.2",
   "pre-commit",
   "pytest",
@@ -24,24 +40,6 @@ docs = [
   "mkdocstrings[python]>=0.28.1"
 ]
 
-[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.2.4"
-]
-description = "Add your description here"
-license = {file = "LICENSE"}
-name = "agentensor"
-readme = "README.md"
-requires-python = ">=3.12"
-url = "https://github.com/ShaojieJiang/agentensor"
-version = "0.0.4"
-
 [tool.coverage.report]
 exclude_lines = [
   "pragma: no cover",
@@ -55,6 +53,9 @@ branch = true
 command_line = "-m pytest"
 source = ["agentensor"]
 
+[tool.hatch.build.targets.wheel]
+packages = ["src/agentensor"]
+
 [tool.mypy]
 disallow_untyped_defs = true
 ignore_missing_imports = true

agentensor-0.1.0/src/agentensor/__init__.py

@@ -0,0 +1,8 @@
+"""Agentensor core primitives."""
+
+from agentensor.loss import LLMTensorJudge
+from agentensor.optim import Optimizer
+from agentensor.train import GraphTrainer, Trainer
+
+
+__all__ = ["GraphTrainer", "LLMTensorJudge", "Optimizer", "Trainer"]

{agentensor-0.0.4 → agentensor-0.1.0/src}/agentensor/module.py

@@ -5,12 +5,15 @@ 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 langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel, ConfigDict
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from agentensor.tensor import TextTensor
 
 
+CompiledGraph = CompiledStateGraph[Any, Any, Any, Any]
+
+
 class AgentModule(BaseModel, ABC):
     """Agent module."""
 

agentensor-0.1.0/src/agentensor/optim.py

@@ -0,0 +1,109 @@
+"""Optimizer module."""
+
+from langchain.agents import create_agent
+from langchain.chat_models import init_chat_model
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import Runnable
+from langgraph.graph import StateGraph
+from agentensor.module import AgentModule
+from agentensor.tensor import TextTensor
+
+
+class Optimizer:
+    """Optimizer class."""
+
+    def __init__(
+        self,
+        graph: StateGraph | None = None,
+        model: str | BaseChatModel = "gpt-4o-mini",
+        params: list[TextTensor] | None = None,
+    ) -> None:
+        """Initialize the optimizer."""
+        self.params: list[TextTensor] = (
+            self._coerce_params(params, source="params") if params is not None else []
+        )
+        if graph is not None and not params:
+            for node in graph.nodes.values():
+                runnable = getattr(node, "runnable", None)
+                module: AgentModule | None = None
+
+                if isinstance(runnable, AgentModule):
+                    module = runnable
+                else:
+                    function = getattr(runnable, "afunc", None)
+                    if isinstance(function, AgentModule):
+                        module = function
+                    else:
+                        bound_self = getattr(function, "__self__", None)
+                        if isinstance(bound_self, AgentModule):
+                            module = bound_self
+
+                if module is not None:
+                    self.params.extend(
+                        self._coerce_params(
+                            module.get_params(),
+                            source=f"{module.__class__.__name__}.get_params()",
+                        )
+                    )
+                    continue
+
+                param_provider = getattr(runnable, "get_params", None)
+                if callable(param_provider):
+                    self.params.extend(
+                        self._coerce_params(
+                            param_provider(),
+                            source=f"{runnable.__class__.__name__}.get_params()",
+                        )
+                    )
+        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) -> Runnable:
+        """Get the agent."""
+        return create_agent(
+            self.model,
+            tools=[],
+            system_prompt="Rewrite the system prompt given the feedback.",
+        )
+
+    @staticmethod
+    def _coerce_params(params: object, *, source: str) -> list[TextTensor]:
+        """Normalize parameters and raise if they are incompatible."""
+        if isinstance(params, TextTensor):
+            return [params]
+        if isinstance(params, list | tuple):
+            if not params:
+                return []
+            invalid = [param for param in params if not isinstance(param, TextTensor)]
+            if invalid:
+                raise TypeError(
+                    f"{source} must contain only TextTensor instances, got "
+                    f"{type(invalid[0]).__name__}."
+                )
+            return list(params)
+        raise TypeError(
+            f"{source} must return a TextTensor or list of TextTensor objects."
+        )

{agentensor-0.0.4 → agentensor-0.1.0/src}/agentensor/tensor.py

@@ -1,11 +1,12 @@
-"""Example module."""
+"""Text tensor primitives."""
 
 from __future__ import annotations
+from typing import Any
+from langchain.agents import create_agent
 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
+from langchain_core.runnables import Runnable
 
 
 class TextTensor:
@@ -16,15 +17,18 @@ class TextTensor:
         text: str,
         parents: list[TextTensor] | None = None,
         requires_grad: bool = False,
-        model: str | BaseChatModel = "gpt-4o-mini",
+        metadata: dict[str, Any] | None = None,
+        model: str | BaseChatModel = "openai:gpt-4o-mini",
+        model_kwargs: dict[str, Any] | None = None,
     ) -> None:
         """Initialize a TextTensor."""
         self.text = text
         self.requires_grad = requires_grad
         self.gradients: list[str] = []
+        self.metadata = dict(metadata) if metadata is not None else {}
         self.parents: list[TextTensor] = parents or []
         if isinstance(model, str):
-            self.model = init_chat_model(model)
+            self.model = init_chat_model(model, **(model_kwargs or {}))
         else:
             self.model = model
 
@@ -77,8 +81,10 @@ class TextTensor:
         return self.text
 
     @property
-    def agent(self) -> CompiledGraph:
+    def agent(self) -> Runnable:
         """Get the agent."""
-        return create_react_agent(
-            self.model, tools=[], prompt="Answer the user's question."
+        return create_agent(
+            self.model,
+            tools=[],
+            system_prompt="Answer the user's question.",
         )

agentensor-0.1.0/src/agentensor/train.py

@@ -0,0 +1,280 @@
+"""Trainer."""
+
+from __future__ import annotations
+import asyncio
+import json
+from collections.abc import Mapping
+from typing import Any, Literal
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph.state import CompiledStateGraph
+from pydantic_evals import Dataset
+from pydantic_evals.reporting import EvaluationReport
+from agentensor.optim import Optimizer
+from agentensor.tensor import TextTensor
+
+
+CompiledGraph = CompiledStateGraph[Any, Any, Any, Any]
+
+
+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."""
+        self._require_dataset("train")
+        optimizer = self._require_optimizer()
+        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: list[str] = []
+                for evaluator in case.assertions.values():
+                    if not evaluator.value:
+                        reason = getattr(evaluator, "reason", None)
+                        if reason is None or (
+                            isinstance(reason, str) and not reason.strip()
+                        ):
+                            losses.append("Evaluation failed without a reason.")
+                        else:
+                            losses.append(str(reason))
+                if losses:
+                    case.output.backward(" ".join(losses))
+
+            optimizer.step()
+            optimizer.zero_grad()
+
+            performance = report.averages()
+            assertions = None if performance is None else performance.assertions
+            self.after_epoch(i, report)
+            if assertions is not None and assertions >= self.stop_threshold:
+                print("Optimization complete.")
+                break
+
+    def evaluate(
+        self,
+        data_split: Literal["train", "eval", "test"] = "eval",
+        limit_cases: int | None = None,
+        max_concurrency: int | None = None,
+        progress: bool = True,
+    ) -> EvaluationReport:
+        """Evaluate the graph."""
+        dataset = self._require_dataset(data_split)
+        if limit_cases:
+            limited_cases = dataset.cases[:limit_cases]
+            dataset = Dataset(cases=limited_cases, evaluators=dataset.evaluators)
+        report = dataset.evaluate_sync(
+            self.forward,
+            max_concurrency=max_concurrency,
+            progress=progress,
+        )
+
+        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)
+
+    def after_epoch(self, epoch_index: int, report: EvaluationReport) -> None:
+        """Optional hook for subclasses to record state."""
+        return None
+
+    def _require_dataset(self, data_split: str) -> Dataset[Any, Any, Any]:
+        """Return the dataset for a split or raise a descriptive error."""
+        dataset = getattr(self, f"{data_split}_dataset", None)
+        if dataset is None:
+            raise ValueError(f"{data_split} dataset is required")
+        return dataset
+
+    def _require_optimizer(self) -> Optimizer:
+        """Return the optimizer or raise a descriptive error."""
+        if self.optimizer is None:
+            raise ValueError("Optimizer is required")
+        return self.optimizer
+
+
+class GraphTrainer(Trainer):
+    """Trainer that runs a compiled graph against mapping inputs."""
+
+    def __init__(
+        self,
+        *,
+        graph: CompiledGraph,
+        dataset: Dataset[Any, Any, Any],
+        optimizer: Optimizer | None = None,
+        epochs: int,
+        runtime_prompts: Mapping[str, TextTensor],
+        base_state: Mapping[str, Any] | None = None,
+        graph_config: RunnableConfig | None = None,
+        max_concurrency: int = 1,
+        case_timeout: int = 30,
+        stop_threshold: float = 2.0,
+        script_format: bool = False,
+    ) -> None:
+        """Initialize the graph trainer."""
+        super().__init__(
+            graph=graph,
+            train_dataset=dataset,
+            eval_dataset=dataset,
+            optimizer=optimizer,
+            epochs=epochs,
+            stop_threshold=stop_threshold,
+        )
+        self.runtime_prompts = runtime_prompts
+        self.base_state = base_state or {}
+        self.graph_config = graph_config
+        self.max_concurrency = max_concurrency
+        self.case_timeout = case_timeout
+        self.reports: list[EvaluationReport] = []
+        self.script_format = script_format
+        self.prompt_history: list[dict[str, str]] = []
+
+    async def forward(self, case_inputs: Mapping[str, Any]) -> TextTensor:  # type: ignore[override]
+        """Execute the compiled graph and return a tensor with the raw payload."""
+        merged_inputs = self._merge_inputs(case_inputs)
+        case_state = self._build_case_state(merged_inputs)
+        output_state = await asyncio.wait_for(
+            self.graph.ainvoke(case_state, config=self.graph_config),
+            timeout=self.case_timeout,
+        )
+        output_payload = self._extract_output(output_state)
+        parents = [
+            prompt for prompt in self.runtime_prompts.values() if prompt.requires_grad
+        ]
+        tensor = TextTensor(
+            text=self._stringify_output(output_payload),
+            requires_grad=True,
+            parents=parents,
+            metadata={"payload": output_payload},
+        )
+        return tensor
+
+    def evaluate(
+        self,
+        data_split: Literal["train", "eval", "test"] = "eval",
+        limit_cases: int | None = None,
+        max_concurrency: int | None = None,
+        progress: bool = False,
+    ) -> EvaluationReport:
+        """Run evaluation without rendering progress to stdout."""
+        dataset = self._require_dataset(data_split)
+        cases = dataset.cases
+        if limit_cases is not None:
+            cases = cases[:limit_cases]
+            dataset = Dataset(cases=cases, evaluators=dataset.evaluators)
+        report = dataset.evaluate_sync(
+            self.forward,
+            max_concurrency=max_concurrency or self.max_concurrency,
+            progress=progress,
+        )
+        self.reports.append(report)
+        return report
+
+    def _merge_inputs(self, inputs: Mapping[str, Any]) -> dict[str, Any]:
+        merged: dict[str, Any] = {}
+        if isinstance(self.base_state, Mapping):
+            base_inputs = self.base_state.get("inputs", self.base_state)
+            if isinstance(base_inputs, Mapping):  # pragma: no branch
+                merged.update(base_inputs)
+        merged.update(inputs)
+        return merged
+
+    def _build_case_state(self, inputs: Mapping[str, Any]) -> dict[str, Any]:
+        runtime_config = (
+            dict(self.base_state.get("config", {}))
+            if isinstance(self.base_state, Mapping)
+            else {}
+        )
+        if self.script_format:
+            state = dict(inputs)
+            state["config"] = runtime_config | {"prompts": self.runtime_prompts}
+            return state
+        return {
+            "messages": [],
+            "results": {},
+            "inputs": dict(inputs),
+            "structured_response": None,
+            "config": runtime_config | {"prompts": self.runtime_prompts},
+        }
+
+    @staticmethod
+    def _stringify_output(output: Any) -> str:
+        if isinstance(output, str):
+            return output
+        try:
+            return json.dumps(output)
+        except TypeError:
+            return str(output)
+
+    @staticmethod
+    def _extract_output(output_state: Any) -> Any:
+        if isinstance(output_state, Mapping):  # pragma: no branch
+            results = output_state.get("results")
+            if isinstance(results, Mapping) and results:
+                return results
+            if "output" in output_state:
+                return output_state["output"]
+            message_output = GraphTrainer._extract_message_output(
+                output_state.get("messages")
+            )
+            if message_output is not None:
+                return message_output
+        return output_state
+
+    @staticmethod
+    def _extract_message_output(messages: Any) -> Any | None:
+        if not isinstance(messages, list):
+            return None
+        fallback: Any | None = None
+        for message in reversed(messages):
+            if isinstance(message, Mapping):
+                content = message.get("content")
+                role = message.get("role") or message.get("type")
+            else:
+                content = getattr(message, "content", None)
+                role = getattr(message, "role", None) or getattr(message, "type", None)
+            if role in {"assistant", "ai"}:
+                return content
+            if fallback is None and content is not None:  # pragma: no branch
+                if not (isinstance(content, str) and not content.strip()):
+                    fallback = content
+        return fallback
+
+    def after_epoch(self, epoch_index: int, report: EvaluationReport) -> None:
+        """Record prompt snapshots after each optimizer step."""
+        snapshot: dict[str, str] = {
+            name: tensor.text for name, tensor in self.runtime_prompts.items()
+        }
+        self.prompt_history.append(snapshot)

agentensor-0.0.4/.bumpversion.cfg

@@ -1,12 +0,0 @@
-[bumpversion]
-current_version = 0.0.4
-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}"