agentjit 0.1.0__py3-none-any.whl

agentjit/__init__.py

@@ -0,0 +1,43 @@
+"""AgentJIT: Just-In-Time Compiler for AI Agent Trajectories.
+
+Compiles non-deterministic, expensive multi-step LLM workflows into ultra-fast
+deterministic Python code with zero token costs and automatic speculative fallbacks.
+"""
+
+from agentjit.codegen import GuardViolation
+from agentjit.compiler import compile_trajectory, create_compiled_pipeline
+from agentjit.decorators import jit, JITWrapper
+from agentjit.integrations.base import ToolDispatcher
+from agentjit.runtime import CompiledPipeline
+from agentjit.tracer import Tracer, trace_tool
+from agentjit.types import (
+    CompilationResult,
+    ExecutionMetrics,
+    ExecutionNode,
+    Guard,
+    SourceType,
+    TraceStep,
+    Trajectory,
+    ValueSource,
+)
+
+__version__ = "0.1.0"
+__all__ = [
+    "jit",
+    "JITWrapper",
+    "Tracer",
+    "trace_tool",
+    "compile_trajectory",
+    "create_compiled_pipeline",
+    "CompiledPipeline",
+    "GuardViolation",
+    "Trajectory",
+    "TraceStep",
+    "ExecutionNode",
+    "Guard",
+    "ValueSource",
+    "SourceType",
+    "CompilationResult",
+    "ExecutionMetrics",
+    "ToolDispatcher",
+]

agentjit/analyzer.py

@@ -0,0 +1,302 @@
+"""Data-flow dependency analyzer and graph synthesis for agent trajectories."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from agentjit.types import (
+    ExecutionNode,
+    Guard,
+    SourceType,
+    TraceStep,
+    Trajectory,
+    ValueSource,
+)
+
+
+class TrajectoryAnalyzer:
+    """Analyzes a recorded agent trajectory to construct a deterministic execution DAG."""
+
+    def __init__(self, trajectory: Trajectory) -> None:
+        self.trajectory = trajectory
+        self.entry_args = trajectory.entry_args
+        self.steps = trajectory.steps
+
+    def analyze(self) -> Tuple[List[ExecutionNode], List[Guard], ValueSource]:
+        """Perform data-flow analysis and return execution nodes, guards, and final output source."""
+        nodes: List[ExecutionNode] = []
+        step_outputs: Dict[int, Any] = {}
+
+        # 1. Analyze each step and determine argument provenance
+        for step in self.steps:
+            bindings: Dict[str, ValueSource] = {}
+            for param_name, param_val in step.inputs.items():
+                source = self._resolve_value_source(param_val, step.step_id, step_outputs)
+                bindings[param_name] = source
+
+            node = ExecutionNode(
+                node_id=step.step_id,
+                tool_name=step.tool_name,
+                argument_bindings=bindings,
+                output_var_name=f"step_{step.step_id}_out",
+            )
+            nodes.append(node)
+            step_outputs[step.step_id] = step.output
+
+        # 2. Determine where the final result comes from
+        final_source = self._resolve_final_output(step_outputs)
+
+        # 3. Generate speculative execution guards
+        guards = self._generate_guards()
+
+        return nodes, guards, final_source
+
+    def _resolve_value_source(
+        self,
+        target_val: Any,
+        current_step_id: int,
+        step_outputs: Dict[int, Any],
+    ) -> ValueSource:
+        """Trace backwards to locate the origin of target_val."""
+        # 1. Direct match with entry argument
+        for arg_name, arg_val in self.entry_args.items():
+            if self._safe_equals(target_val, arg_val):
+                return ValueSource(source_type=SourceType.ENTRY_ARG, name=arg_name)
+
+        # 2. Direct match with output of an earlier step
+        for prev_id in range(current_step_id - 1, 0, -1):
+            if prev_id in step_outputs and self._safe_equals(target_val, step_outputs[prev_id]):
+                return ValueSource(source_type=SourceType.STEP_OUTPUT, step_id=prev_id)
+
+        # 3. Nested match within entry arguments (e.g. user['id'])
+        for arg_name, arg_val in self.entry_args.items():
+            found, path = self._search_nested(target_val, arg_val)
+            if found:
+                return ValueSource(
+                    source_type=SourceType.NESTED_ACCESS,
+                    name=arg_name,
+                    key_path=path,
+                )
+
+        # 4. Nested match within previous step outputs (e.g. order['customer_id'])
+        for prev_id in range(current_step_id - 1, 0, -1):
+            if prev_id in step_outputs:
+                found, path = self._search_nested(target_val, step_outputs[prev_id])
+                if found:
+                    return ValueSource(
+                        source_type=SourceType.NESTED_ACCESS,
+                        step_id=prev_id,
+                        key_path=path,
+                    )
+
+        # 5. Arithmetic binary operation synthesis (e.g. unit_price * quantity)
+        if isinstance(target_val, (int, float)) and not isinstance(target_val, bool):
+            bin_op = self._try_synthesize_arithmetic(target_val, current_step_id, step_outputs)
+            if bin_op is not None:
+                return bin_op
+
+        # 6. Nested dictionary structure check
+        if isinstance(target_val, dict):
+            dict_sources = {
+                k: self._resolve_value_source(v, current_step_id, step_outputs)
+                for k, v in target_val.items()
+            }
+            return ValueSource(
+                source_type=SourceType.DICT_SYNTHESIS,
+                dict_fields=dict_sources,
+            )
+
+        # 7. Nested list structure check
+        if isinstance(target_val, (list, tuple)):
+            list_sources = [
+                self._resolve_value_source(item, current_step_id, step_outputs)
+                for item in target_val
+            ]
+            return ValueSource(
+                source_type=SourceType.LIST_SYNTHESIS,
+                list_items=list_sources,
+            )
+
+        # 8. Invariant literal
+        return ValueSource(source_type=SourceType.LITERAL, literal_value=target_val)
+
+    def _try_synthesize_arithmetic(
+        self,
+        target_val: Union[int, float],
+        current_step_id: int,
+        step_outputs: Dict[int, Any],
+    ) -> Optional[ValueSource]:
+        """Attempt to synthesize a binary arithmetic operation matching target_val."""
+        candidates: List[Tuple[Union[int, float], ValueSource]] = []
+
+        # Collect numeric candidates from entry args
+        for arg_name, arg_val in self.entry_args.items():
+            if isinstance(arg_val, (int, float)) and not isinstance(arg_val, bool):
+                candidates.append(
+                    (arg_val, ValueSource(source_type=SourceType.ENTRY_ARG, name=arg_name))
+                )
+
+        # Collect numeric candidates from previous step outputs
+        for prev_id in range(1, current_step_id):
+            out = step_outputs.get(prev_id)
+            if isinstance(out, (int, float)) and not isinstance(out, bool):
+                candidates.append(
+                    (out, ValueSource(source_type=SourceType.STEP_OUTPUT, step_id=prev_id))
+                )
+            elif isinstance(out, dict):
+                for k, v in out.items():
+                    if isinstance(v, (int, float)) and not isinstance(v, bool):
+                        candidates.append(
+                            (
+                                v,
+                                ValueSource(
+                                    source_type=SourceType.NESTED_ACCESS,
+                                    step_id=prev_id,
+                                    key_path=[k],
+                                ),
+                            )
+                        )
+
+        # Test pairwise operations (prefer multiplication and addition)
+        ops = [
+            ("*", lambda a, b: a * b),
+            ("+", lambda a, b: a + b),
+            ("-", lambda a, b: a - b),
+            ("/", lambda a, b: (a / b) if b != 0 else None),
+        ]
+
+        for i in range(len(candidates)):
+            for j in range(len(candidates)):
+                if i == j:
+                    continue
+                v1, src1 = candidates[i]
+                v2, src2 = candidates[j]
+
+                for op_symbol, op_fn in ops:
+                    try:
+                        computed = op_fn(v1, v2)
+                        if computed is not None and math.isclose(
+                            computed, target_val, rel_tol=1e-5, abs_tol=1e-5
+                        ):
+                            return ValueSource(
+                                source_type=SourceType.BINARY_OP,
+                                left=src1,
+                                right=src2,
+                                operator=op_symbol,
+                            )
+                    except Exception:
+                        continue
+
+        return None
+
+    def _resolve_final_output(self, step_outputs: Dict[int, Any]) -> ValueSource:
+        """Find the source expression for the trajectory's final result."""
+        final_res = self.trajectory.final_result
+        if not self.steps:
+            return ValueSource(source_type=SourceType.LITERAL, literal_value=final_res)
+
+        # 1. Check if final result directly equals a step's output
+        for step_id in reversed(sorted(step_outputs.keys())):
+            if self._safe_equals(final_res, step_outputs[step_id]):
+                return ValueSource(source_type=SourceType.STEP_OUTPUT, step_id=step_id)
+
+        # 2. Check if final result is a dict composed of step outputs/entry args
+        if isinstance(final_res, dict):
+            dict_sources = {
+                k: self._resolve_value_source(v, len(self.steps) + 1, step_outputs)
+                for k, v in final_res.items()
+            }
+            return ValueSource(
+                source_type=SourceType.DICT_SYNTHESIS,
+                dict_fields=dict_sources,
+            )
+
+        # 3. Check if final result is a list composed of step outputs/entry args
+        if isinstance(final_res, (list, tuple)):
+            list_sources = [
+                self._resolve_value_source(item, len(self.steps) + 1, step_outputs)
+                for item in final_res
+            ]
+            return ValueSource(
+                source_type=SourceType.LIST_SYNTHESIS,
+                list_items=list_sources,
+            )
+
+        # 4. Check if nested in the last step output
+        last_id = self.steps[-1].step_id
+        if last_id in step_outputs:
+            found, path = self._search_nested(final_res, step_outputs[last_id])
+            if found:
+                return ValueSource(
+                    source_type=SourceType.NESTED_ACCESS,
+                    step_id=last_id,
+                    key_path=path,
+                )
+
+        # Default to output of last step
+        return ValueSource(source_type=SourceType.STEP_OUTPUT, step_id=last_id)
+
+    def _generate_guards(self) -> List[Guard]:
+        """Generate speculative validation guards for the entry arguments."""
+        guards: List[Guard] = []
+        for arg_name, arg_val in self.entry_args.items():
+            if arg_val is not None:
+                # Non-null guard
+                guards.append(
+                    Guard(
+                        target_param=arg_name,
+                        guard_type="non_null",
+                        condition_code=f"{arg_name} is not None",
+                        description=f"Argument '{arg_name}' must not be None",
+                    )
+                )
+                # Type guard
+                type_name = type(arg_val).__name__
+                guards.append(
+                    Guard(
+                        target_param=arg_name,
+                        guard_type="type_match",
+                        condition_code=f"isinstance({arg_name}, {type_name})",
+                        description=f"Argument '{arg_name}' must be of type {type_name}",
+                    )
+                )
+        return guards
+
+    def _search_nested(
+        self,
+        target: Any,
+        haystack: Any,
+        max_depth: int = 4,
+    ) -> Tuple[bool, List[Union[str, int]]]:
+        """Recursively search for target inside a nested dictionary/list structure."""
+        if max_depth <= 0 or target is None or haystack is None:
+            return False, []
+
+        if isinstance(haystack, dict):
+            for k, v in haystack.items():
+                if self._safe_equals(target, v):
+                    return True, [k]
+                found, subpath = self._search_nested(target, v, max_depth - 1)
+                if found:
+                    return True, [k] + subpath
+
+        elif isinstance(haystack, (list, tuple)):
+            for idx, item in enumerate(haystack):
+                if self._safe_equals(target, item):
+                    return True, [idx]
+                found, subpath = self._search_nested(target, item, max_depth - 1)
+                if found:
+                    return True, [idx] + subpath
+
+        return False, []
+
+    @staticmethod
+    def _safe_equals(a: Any, b: Any) -> bool:
+        """Safe equality comparison avoiding numpy/pandas ambiguity."""
+        if a is b:
+            return True
+        try:
+            return bool(a == b)
+        except Exception:
+            return False

agentjit/codegen.py

@@ -0,0 +1,137 @@
+"""AST and Python source code generator for compiled agent trajectories."""
+
+from __future__ import annotations
+
+import ast
+from typing import Any, Callable, Dict, List, Optional
+
+from agentjit.types import (
+    CompilationResult,
+    ExecutionNode,
+    Guard,
+    Trajectory,
+    ValueSource,
+)
+
+
+class GuardViolation(Exception):
+    """Raised at runtime when a compiled pipeline's speculative guard fails."""
+    def __init__(self, message: str, param: str = "") -> None:
+        super().__init__(message)
+        self.param = param
+
+
+class CodeGenerator:
+    """Synthesizes human-readable, typed, and executable Python code from an execution DAG."""
+
+    def __init__(
+        self,
+        entry_params: List[str],
+        nodes: List[ExecutionNode],
+        guards: List[Guard],
+        final_source: ValueSource,
+        tool_registry: Dict[str, Callable],
+        func_name: str = "compiled_pipeline",
+    ) -> None:
+        self.entry_params = entry_params
+        self.nodes = nodes
+        self.guards = guards
+        self.final_source = final_source
+        self.tool_registry = tool_registry
+        self.func_name = func_name
+
+    def generate_source(self) -> str:
+        """Generate formatted Python source code for the compiled pipeline."""
+        lines: List[str] = []
+        params_str = ", ".join(self.entry_params)
+        lines.append(f"def {self.func_name}({params_str}):")
+        lines.append('    """JIT-compiled trajectory pipeline generated by AgentJIT.')
+        lines.append("    Executes deterministically in sub-millisecond time with zero token cost.")
+        lines.append('    """')
+
+        # 1. Speculative Guards
+        if self.guards:
+            lines.append("    # --- Speculative Guards ---")
+            for g in self.guards:
+                # Add validation statement
+                lines.append(f"    if not ({g.condition_code}):")
+                escaped_desc = g.description.replace('"', '\\"')
+                lines.append(
+                    f'        raise GuardViolation("{escaped_desc}", param="{g.target_param}")'
+                )
+
+        lines.append("")
+        lines.append("    # --- Execution Steps ---")
+        var_map: Dict[int, str] = {}
+
+        # 2. Tool invocations
+        for node in self.nodes:
+            var_name = node.output_var_name
+            var_map[node.node_id] = var_name
+
+            args_parts: List[str] = []
+            for param_k, val_src in node.argument_bindings.items():
+                arg_expr = val_src.render_expression(var_map)
+                args_parts.append(f"{param_k}={arg_expr}")
+
+            args_str = ", ".join(args_parts)
+            lines.append(
+                f"    {var_name} = _tools[{repr(node.tool_name)}]({args_str})"
+            )
+
+        # 3. Final return
+        lines.append("")
+        lines.append("    # --- Return Final Result ---")
+        return_expr = self.final_source.render_expression(var_map)
+        lines.append(f"    return {return_expr}")
+
+        return "\n".join(lines)
+
+    def compile(self) -> CompilationResult:
+        """Generate source code, validate AST, and compile into an executable callable."""
+        source = self.generate_source()
+
+        # Validate syntax via AST
+        try:
+            tree = ast.parse(source)
+        except SyntaxError as e:
+            return CompilationResult(
+                success=False,
+                source_code=source,
+                error=f"Syntax error in synthesized code: {e}",
+            )
+
+        # Execute compilation in controlled namespace
+        exec_globals: Dict[str, Any] = {
+            "_tools": self.tool_registry,
+            "GuardViolation": GuardViolation,
+            "isinstance": isinstance,
+            "int": int,
+            "float": float,
+            "str": str,
+            "bool": bool,
+            "dict": dict,
+            "list": list,
+        }
+        exec_locals: Dict[str, Any] = {}
+
+        try:
+            code_obj = compile(tree, filename="<agentjit_compiled>", mode="exec")
+            exec(code_obj, exec_globals, exec_locals)
+            compiled_func = exec_locals.get(self.func_name)
+        except Exception as e:
+            return CompilationResult(
+                success=False,
+                source_code=source,
+                error=f"Execution error during module compilation: {e}",
+            )
+
+        return CompilationResult(
+            success=True,
+            source_code=source,
+            compiled_callable=compiled_func,
+            nodes=self.nodes,
+            guards=self.guards,
+            entry_parameters=self.entry_params,
+            output_expression=self.final_source.render_expression({}),
+        )

agentjit/compiler.py

@@ -0,0 +1,60 @@
+"""High-level compiler coordinating analysis, code generation, and pipeline assembly."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, Optional
+
+from agentjit.analyzer import TrajectoryAnalyzer
+from agentjit.codegen import CodeGenerator
+from agentjit.runtime import CompiledPipeline
+from agentjit.types import CompilationResult, Trajectory
+
+
+def compile_trajectory(
+    trajectory: Trajectory,
+    tool_registry: Dict[str, Callable],
+    function_name: str = "compiled_agent_pipeline",
+) -> CompilationResult:
+    """Compile a recorded agent trajectory into an optimized Python pipeline.
+
+    Args:
+        trajectory: The traced sequence of tool calls and entry parameters.
+        tool_registry: Mapping from tool names to their Python callable implementations.
+        function_name: Name of the generated Python function.
+
+    Returns:
+        CompilationResult containing synthesized source code and executable function.
+    """
+    analyzer = TrajectoryAnalyzer(trajectory)
+    nodes, guards, final_source = analyzer.analyze()
+
+    entry_params = list(trajectory.entry_args.keys())
+
+    generator = CodeGenerator(
+        entry_params=entry_params,
+        nodes=nodes,
+        guards=guards,
+        final_source=final_source,
+        tool_registry=tool_registry,
+        func_name=function_name,
+    )
+
+    return generator.compile()
+
+
+def create_compiled_pipeline(
+    trajectory: Trajectory,
+    tool_registry: Dict[str, Callable],
+    fallback_fn: Optional[Callable] = None,
+    function_name: str = "compiled_agent_pipeline",
+) -> CompiledPipeline:
+    """Compile a trajectory and wrap it in an executable CompiledPipeline."""
+    compilation = compile_trajectory(trajectory, tool_registry, function_name)
+    if not compilation.success:
+        raise RuntimeError(f"Failed to compile trajectory: {compilation.error}")
+
+    return CompiledPipeline(
+        compilation=compilation,
+        fallback_fn=fallback_fn,
+        estimated_uncompiled_latency_ms=max(100.0, trajectory.total_duration_ms),
+    )

agentjit/decorators.py

@@ -0,0 +1,129 @@
+"""High-level @jit decorator for automatic tracing, compilation, and execution."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Callable, Dict, Optional
+
+from agentjit.compiler import create_compiled_pipeline
+from agentjit.runtime import CompiledPipeline
+from agentjit.tracer import Tracer
+
+
+class JITWrapper:
+    """Manages the lifecycle of a JIT-compiled agent function."""
+
+    def __init__(
+        self,
+        func: Callable,
+        warmup_runs: int = 1,
+        auto_compile: bool = True,
+        function_name: Optional[str] = None,
+    ) -> None:
+        self.func = func
+        self.warmup_runs = max(1, warmup_runs)
+        self.auto_compile = auto_compile
+        self.function_name = function_name or f"compiled_{func.__name__}"
+
+        self._runs_completed = 0
+        self._pipeline: Optional[CompiledPipeline] = None
+        self._last_trajectory = None
+
+        functools.update_wrapper(self, func)
+
+    @property
+    def is_compiled(self) -> bool:
+        """True if the agent workflow has been successfully compiled."""
+        return self._pipeline is not None
+
+    @property
+    def pipeline(self) -> Optional[CompiledPipeline]:
+        """Access the underlying CompiledPipeline instance."""
+        return self._pipeline
+
+    @property
+    def source_code(self) -> Optional[str]:
+        """View the compiled Python code."""
+        if self._pipeline:
+            return self._pipeline.source_code
+        return None
+
+    @property
+    def stats(self) -> Dict[str, Any]:
+        """Access performance and execution statistics."""
+        if self._pipeline:
+            return self._pipeline.stats
+        return {"status": "uncompiled", "warmup_progress": f"{self._runs_completed}/{self.warmup_runs}"}
+
+    def __call__(self, *args, **kwargs) -> Any:
+        # If already compiled, execute optimized pipeline
+        if self._pipeline is not None:
+            return self._pipeline(*args, **kwargs)
+
+        # Warmup / Tracing phase
+        sig = inspect.signature(self.func)
+        bound = sig.bind(*args, **kwargs)
+        bound.apply_defaults()
+        entry_args = dict(bound.arguments)
+
+        with Tracer(entry_args=entry_args) as tracer:
+            result = self.func(*args, **kwargs)
+            tracer.set_final_result(result)
+            self._last_trajectory = tracer.trajectory
+
+        self._runs_completed += 1
+
+        # Check if warmup threshold reached
+        if self.auto_compile and self._runs_completed >= self.warmup_runs:
+            self._compile_from_tracer(tracer)
+
+        return result
+
+    def _compile_from_tracer(self, tracer: Tracer) -> None:
+        """Trigger compilation using the captured trajectory and tools."""
+        try:
+            self._pipeline = create_compiled_pipeline(
+                trajectory=tracer.trajectory,
+                tool_registry=tracer._registered_tools,
+                fallback_fn=self.func,
+                function_name=self.function_name,
+            )
+        except Exception as err:
+            # If compilation fails, remain in uncompiled fallback mode
+            self._pipeline = None
+
+
+def jit(
+    func: Optional[Callable] = None,
+    *,
+    warmup_runs: int = 1,
+    auto_compile: bool = True,
+    function_name: Optional[str] = None,
+) -> Any:
+    """Decorator to JIT-compile an AI agent workflow.
+
+    Usage:
+        @jit
+        def my_agent(order_id: int, country: str):
+            order = fetch_order(order_id)
+            tax = calculate_vat(order["amount"], country)
+            return update_db(order_id, tax)
+
+        # 1st run: Executes dynamic agent, records trajectory, compiles to pure Python
+        result1 = my_agent(402, "Germany")
+
+        # 2nd run: Executes in <1ms directly with ZERO LLM calls and zero token cost!
+        result2 = my_agent(403, "France")
+    """
+    def decorator(fn: Callable) -> JITWrapper:
+        return JITWrapper(
+            fn,
+            warmup_runs=warmup_runs,
+            auto_compile=auto_compile,
+            function_name=function_name,
+        )
+
+    if func is not None:
+        return decorator(func)
+    return decorator

agentjit/integrations/__init__.py

@@ -0,0 +1,5 @@
+"""Integrations package for AgentJIT."""
+
+from agentjit.integrations.base import ToolDispatcher
+
+__all__ = ["ToolDispatcher"]