prefactor-langchain 0.2.2__tar.gz → 0.2.4__tar.gz

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/.gitignore

@@ -69,3 +69,4 @@ htmlcov/
 # Env
 .env
 mise.local.toml
+prefactor.json

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.4
 Name: prefactor-langchain
-Version: 0.2.2
+Version: 0.2.4
 Summary: LangChain integration for Prefactor observability
 Author-email: Prefactor Pty Ltd <josh@prefactor.tech>
 License: MIT
 Requires-Python: <4.0.0,>=3.11.0
 Requires-Dist: langchain-core>=1.0.0
-Requires-Dist: prefactor-core>=0.2.1
+Requires-Dist: prefactor-core>=0.2.2
 Description-Content-Type: text/markdown
 
 # prefactor-langchain
@@ -24,13 +24,26 @@ pip install prefactor-langchain
 ### Factory pattern (quickest setup)
 
 ```python
-from prefactor_langchain import PrefactorMiddleware
+from prefactor_langchain import LangChainToolSchemaConfig, PrefactorMiddleware
 
 middleware = PrefactorMiddleware.from_config(
     api_url="https://api.prefactor.ai",
     api_token="your-api-token",
     agent_id="my-agent",
     agent_name="My Agent",  # optional
+    tool_schemas={
+        "send_email": LangChainToolSchemaConfig(
+            span_type="send-email",
+            input_schema={
+                "type": "object",
+                "properties": {
+                    "to": {"type": "string", "format": "email"},
+                    "subject": {"type": "string"},
+                },
+                "required": ["to", "subject"],
+            },
+        )
+    },
 )
 
 # Use with LangChain's create_agent()
@@ -38,6 +51,7 @@ middleware = PrefactorMiddleware.from_config(
 # - Agent execution (langchain:agent)
 # - LLM calls (langchain:llm)
 # - Tool executions (langchain:tool)
+# - Tool-specific executions (for example langchain:tool:send-email)
 
 result = agent.invoke({"messages": [...]})
 
@@ -73,6 +87,47 @@ await middleware.close()  # closes the agent instance only
 await client.close()
 ```
 
+### SchemaRegistry composition
+
+Use a shared ``SchemaRegistry`` when you want custom workflow span types and
+LangChain tool schemas to be published together.
+
+```python
+from prefactor_core import SchemaRegistry
+from prefactor_langchain import (
+    LangChainToolSchemaConfig,
+    PrefactorMiddleware,
+    register_langchain_schemas,
+)
+
+registry = SchemaRegistry()
+registry.register_type(
+    name="workflow:run",
+    params_schema={"type": "object"},
+    result_schema={"type": "object"},
+)
+register_langchain_schemas(
+    registry,
+    tool_schemas={
+        "send_email": LangChainToolSchemaConfig(
+            span_type="send-email",
+            input_schema={
+                "type": "object",
+                "properties": {"to": {"type": "string"}},
+                "required": ["to"],
+            },
+        )
+    },
+)
+
+middleware = PrefactorMiddleware.from_config(
+    api_url="https://api.prefactor.ai",
+    api_token="your-api-token",
+    agent_id="my-agent",
+    schema_registry=registry,
+)
+```
+
 ### Pre-configured instance (spans outside the agent)
 
 Pass an `AgentInstanceHandle` you created yourself when you also need to
@@ -111,6 +166,12 @@ await instance.finish()
 await client.close()
 ```
 
+If you pass ``tool_schemas=...`` with a pre-created instance, the middleware
+uses those mappings to emit the right tool span types at runtime. The instance's
+already-registered schema version is not mutated, so you must register matching
+tool schemas before creating the instance if you want those per-tool span types
+to appear in the backend schema version.
+
 ## Span Types
 
 This package creates LangChain-specific spans with the `langchain:*` namespace:

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/README.md

@@ -13,13 +13,26 @@ pip install prefactor-langchain
 ### Factory pattern (quickest setup)
 
 ```python
-from prefactor_langchain import PrefactorMiddleware
+from prefactor_langchain import LangChainToolSchemaConfig, PrefactorMiddleware
 
 middleware = PrefactorMiddleware.from_config(
     api_url="https://api.prefactor.ai",
     api_token="your-api-token",
     agent_id="my-agent",
     agent_name="My Agent",  # optional
+    tool_schemas={
+        "send_email": LangChainToolSchemaConfig(
+            span_type="send-email",
+            input_schema={
+                "type": "object",
+                "properties": {
+                    "to": {"type": "string", "format": "email"},
+                    "subject": {"type": "string"},
+                },
+                "required": ["to", "subject"],
+            },
+        )
+    },
 )
 
 # Use with LangChain's create_agent()
@@ -27,6 +40,7 @@ middleware = PrefactorMiddleware.from_config(
 # - Agent execution (langchain:agent)
 # - LLM calls (langchain:llm)
 # - Tool executions (langchain:tool)
+# - Tool-specific executions (for example langchain:tool:send-email)
 
 result = agent.invoke({"messages": [...]})
 
@@ -62,6 +76,47 @@ await middleware.close()  # closes the agent instance only
 await client.close()
 ```
 
+### SchemaRegistry composition
+
+Use a shared ``SchemaRegistry`` when you want custom workflow span types and
+LangChain tool schemas to be published together.
+
+```python
+from prefactor_core import SchemaRegistry
+from prefactor_langchain import (
+    LangChainToolSchemaConfig,
+    PrefactorMiddleware,
+    register_langchain_schemas,
+)
+
+registry = SchemaRegistry()
+registry.register_type(
+    name="workflow:run",
+    params_schema={"type": "object"},
+    result_schema={"type": "object"},
+)
+register_langchain_schemas(
+    registry,
+    tool_schemas={
+        "send_email": LangChainToolSchemaConfig(
+            span_type="send-email",
+            input_schema={
+                "type": "object",
+                "properties": {"to": {"type": "string"}},
+                "required": ["to"],
+            },
+        )
+    },
+)
+
+middleware = PrefactorMiddleware.from_config(
+    api_url="https://api.prefactor.ai",
+    api_token="your-api-token",
+    agent_id="my-agent",
+    schema_registry=registry,
+)
+```
+
 ### Pre-configured instance (spans outside the agent)
 
 Pass an `AgentInstanceHandle` you created yourself when you also need to
@@ -100,6 +155,12 @@ await instance.finish()
 await client.close()
 ```
 
+If you pass ``tool_schemas=...`` with a pre-created instance, the middleware
+uses those mappings to emit the right tool span types at runtime. The instance's
+already-registered schema version is not mutated, so you must register matching
+tool schemas before creating the instance if you want those per-tool span types
+to appear in the backend schema version.
+
 ## Span Types
 
 This package creates LangChain-specific spans with the `langchain:*` namespace:

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/examples/simple_agent.py

@@ -50,7 +50,7 @@ def _safe_eval(node: ast.expr) -> float:
     if isinstance(node, ast.Constant):
         return node.n  # type: ignore[return-value]
     if isinstance(node, ast.BinOp):
-        return _OPS[type(node.op)](_safe_eval(node.left), _safe_eval(node.right))
+        return _OPS[type(node.op)](_safe_eval(node.left), _safe_eval(node.right))  # type: ignore[invalid-argument-type]
     if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
         return -_safe_eval(node.operand)
     raise ValueError(f"Unsupported expression: {ast.dump(node)}")

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "prefactor-langchain"
-version = "0.2.2"
+version = "0.2.4"
 description = "LangChain integration for Prefactor observability"
 readme = "README.md"
 license = { text = "MIT" }
@@ -9,7 +9,7 @@ authors = [
 ]
 requires-python = ">=3.11.0, <4.0.0"
 dependencies = [
-    "prefactor-core>=0.2.1",
+    "prefactor-core>=0.2.2",
     "langchain-core>=1.0.0",
 ]
 

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/src/prefactor_langchain/__init__.py

@@ -1,14 +1,19 @@
 """Prefactor LangChain - LangChain integration for Prefactor observability."""
 
+from __future__ import annotations
+
 from prefactor_langchain.metadata_extractor import (
     extract_error_info,
     extract_token_usage,
 )
 from prefactor_langchain.middleware import PrefactorMiddleware
 from prefactor_langchain.schemas import (
+    DEFAULT_LANGCHAIN_AGENT_SCHEMA,
     LANGCHAIN_AGENT_SCHEMA,
     LANGCHAIN_LLM_SCHEMA,
     LANGCHAIN_TOOL_SCHEMA,
+    LangChainToolSchemaConfig,
+    compile_langchain_agent_schema,
     register_langchain_schemas,
 )
 from prefactor_langchain.spans import (
@@ -20,7 +25,7 @@ from prefactor_langchain.spans import (
     ToolSpan,
 )
 
-__version__ = "0.2.2"
+__version__ = "0.2.4"
 
 __all__ = [
     # Middleware
@@ -30,6 +35,9 @@ __all__ = [
     "extract_error_info",
     # Schemas
     "register_langchain_schemas",
+    "compile_langchain_agent_schema",
+    "LangChainToolSchemaConfig",
+    "DEFAULT_LANGCHAIN_AGENT_SCHEMA",
     "LANGCHAIN_AGENT_SCHEMA",
     "LANGCHAIN_LLM_SCHEMA",
     "LANGCHAIN_TOOL_SCHEMA",

{prefactor_langchain-0.2.2 → prefactor_langchain-0.2.4}/src/prefactor_langchain/middleware.py

@@ -1,9 +1,12 @@
 """LangChain middleware for automatic tracing via prefactor-core."""
 
+from __future__ import annotations
+
 import asyncio
 import hashlib
 import json
 import logging
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, AsyncContextManager, Callable
 
 from langchain.agents.middleware import AgentMiddleware
@@ -20,7 +23,7 @@ if TYPE_CHECKING:
     pass
 
 from .metadata_extractor import extract_token_usage
-from .schemas import register_langchain_schemas
+from .schemas import LangChainToolSchemaConfig, register_langchain_schemas
 from .spans import AgentSpan, LLMSpan, ToolSpan
 
 logger = logging.getLogger("prefactor_langchain.middleware")
@@ -111,6 +114,8 @@ class PrefactorMiddleware(AgentMiddleware):
         agent_id: str = "langchain-agent",
         agent_name: str | None = None,
         instance: AgentInstanceHandle | None = None,
+        tool_schemas: Mapping[str, LangChainToolSchemaConfig | Mapping[str, Any]]
+        | None = None,
     ):
         """Initialize the middleware with a pre-configured client or instance.
 
@@ -148,6 +153,9 @@ class PrefactorMiddleware(AgentMiddleware):
             self._current_parent_span_id: str | None = None
             self._loop: asyncio.AbstractEventLoop | None = None
             self._pending_emit_futures: list[asyncio.Task[None]] = []
+            self._tool_span_types = (
+                self._register_tool_schemas(None, tool_schemas) if tool_schemas else {}
+            )
             return
 
         if client is None:
@@ -179,6 +187,10 @@ class PrefactorMiddleware(AgentMiddleware):
         self._current_parent_span_id: str | None = None
         self._loop: asyncio.AbstractEventLoop | None = None
         self._pending_emit_futures: list[asyncio.Task[None]] = []
+        self._tool_span_types = {}
+
+        if tool_schemas:
+            self._tool_span_types = self._register_tool_schemas(client, tool_schemas)
 
     @classmethod
     def from_config(
@@ -189,6 +201,8 @@ class PrefactorMiddleware(AgentMiddleware):
         agent_name: str | None = None,
         schema_registry: SchemaRegistry | None = None,
         include_langchain_schemas: bool = True,
+        tool_schemas: Mapping[str, LangChainToolSchemaConfig | Mapping[str, Any]]
+        | None = None,
     ) -> "PrefactorMiddleware":
         """Factory method to create middleware from configuration.
 
@@ -203,6 +217,8 @@ class PrefactorMiddleware(AgentMiddleware):
             schema_registry: Optional SchemaRegistry for registering span schemas.
             include_langchain_schemas: If True and schema_registry is provided,
                 automatically register LangChain-specific schemas.
+            tool_schemas: Optional per-tool schema configuration for tool-specific
+                span types.
 
         Returns:
             A configured PrefactorMiddleware instance with lazy initialization.
@@ -223,8 +239,12 @@ class PrefactorMiddleware(AgentMiddleware):
 
         # Create or augment schema registry
         registry = schema_registry or SchemaRegistry()
-        if include_langchain_schemas and not registry.has_schema("langchain:llm"):
-            register_langchain_schemas(registry)
+        tool_span_types: dict[str, str] = {}
+        if include_langchain_schemas or tool_schemas:
+            tool_span_types = register_langchain_schemas(
+                registry,
+                tool_schemas=tool_schemas,
+            )
 
         config = PrefactorCoreConfig(
             http_config=http_config,
@@ -245,9 +265,32 @@ class PrefactorMiddleware(AgentMiddleware):
         middleware._current_parent_span_id = None
         middleware._loop = None
         middleware._pending_emit_futures = []
+        middleware._tool_span_types = tool_span_types
         logger.debug("PrefactorMiddleware created via from_config()")
         return middleware
 
+    def _register_tool_schemas(
+        self,
+        client: PrefactorCoreClient | None,
+        tool_schemas: Mapping[str, LangChainToolSchemaConfig | Mapping[str, Any]],
+    ) -> dict[str, str]:
+        """Register tool schemas into a client registry when available."""
+        if client is None:
+            return register_langchain_schemas(
+                SchemaRegistry(),
+                tool_schemas=tool_schemas,
+            )
+
+        registry = client._config.schema_registry
+        if registry is None:
+            registry = SchemaRegistry()
+            client._config.schema_registry = registry
+
+        return register_langchain_schemas(
+            registry,
+            tool_schemas=tool_schemas,
+        )
+
     async def _ensure_initialized(self) -> AgentInstanceHandle:
         """Ensure the agent instance is initialized (lazily created if needed).
 
@@ -419,6 +462,18 @@ class PrefactorMiddleware(AgentMiddleware):
                     "tool_name": tool_call.get("name", "unknown"),
                     "arguments": tool_call.get("args", {}),
                 }
+            if hasattr(request, "tool") and hasattr(request.tool, "name"):
+                return {
+                    "tool_name": getattr(request.tool, "name", "unknown"),
+                    "arguments": getattr(request, "args", {}),
+                }
+            if hasattr(request, "name"):
+                return {
+                    "tool_name": request.name,
+                    "arguments": getattr(
+                        request, "input", getattr(request, "args", {})
+                    ),
+                }
             return {}
         except Exception as e:
             logger.error("Error extracting tool inputs: %s", e, exc_info=True)
@@ -460,12 +515,14 @@ class PrefactorMiddleware(AgentMiddleware):
             result["error"] = span_data.error.to_dict()
         return result
 
-    def _build_tool_params(self, span_data: Any) -> dict[str, Any]:
+    def _build_tool_params(self, span_data: Any, schema_name: str) -> dict[str, Any]:
         """Build the payload (params) dict for a tool span."""
         params: dict[str, Any] = {}
         if span_data.tool_name:
             params["tool_name"] = span_data.tool_name
-        if span_data.inputs:
+        if schema_name != "langchain:tool":
+            params["inputs"] = span_data.arguments
+        elif span_data.inputs:
             params["inputs"] = span_data.inputs
         return params
 
@@ -526,7 +583,15 @@ class PrefactorMiddleware(AgentMiddleware):
         if loop is None or loop.is_closed():
             return ""
 
-        import asyncio
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        if running_loop is not None and running_loop is loop:
+            raise RuntimeError(
+                "_create_agent_span_sync must be called from a non-event-loop thread"
+            )
 
         # Use whatever parent was set by set_parent_span() from the async side.
         # We cannot read the SpanContextStack here because run_in_executor does
@@ -546,7 +611,7 @@ class PrefactorMiddleware(AgentMiddleware):
         self._current_parent_span_id = span_id
         return span_id
 
-    def _emit_child_span(self, span_data: Any) -> None:
+    def _emit_child_span(self, span_data: Any, schema_name: str | None = None) -> None:
         """Schedule a child span (LLM/tool) emission with parent tracking.
 
         Uses the current parent span ID from _current_parent_span_id.
@@ -558,15 +623,15 @@ class PrefactorMiddleware(AgentMiddleware):
             return
 
         instance = self._instance
-        schema_name: str = span_data.type
+        schema_name = schema_name or span_data.type
         parent_span_id = self._current_parent_span_id
         pending = self._pending_emit_futures
 
         if schema_name == "langchain:llm":
             params = self._build_llm_params(span_data)
             result = self._build_llm_result(span_data)
-        elif schema_name == "langchain:tool":
-            params = self._build_tool_params(span_data)
+        elif schema_name.startswith("langchain:tool"):
+            params = self._build_tool_params(span_data, schema_name)
             result = self._build_tool_result(span_data)
         else:
             return
@@ -664,6 +729,17 @@ class PrefactorMiddleware(AgentMiddleware):
             async def _finish() -> None:
                 await instance.finish_span(span_id, result_payload=result)
 
+            try:
+                running_loop = asyncio.get_running_loop()
+            except RuntimeError:
+                running_loop = None
+
+            if running_loop is not None and running_loop is loop:
+                raise RuntimeError(
+                    "after_agent must be called from a non-event-loop thread; "
+                    "use aafter_agent for async execution"
+                )
+
             future = asyncio.run_coroutine_threadsafe(_finish(), loop)
             future.result(timeout=5.0)
 
@@ -718,8 +794,16 @@ class PrefactorMiddleware(AgentMiddleware):
             # SpanContextStack.peek() here returns the outer workflow span
             # because we are executing in the async event loop context.
             self._agent_span_cm = instance.span("langchain:agent")
-            self._agent_span_context = await self._agent_span_cm.__aenter__()
-            await self._agent_span_context.start(params)
+            try:
+                self._agent_span_context = await self._agent_span_cm.__aenter__()
+                await self._agent_span_context.start(params)
+            except BaseException:
+                import sys
+
+                await self._agent_span_cm.__aexit__(*sys.exc_info())
+                self._agent_span_cm = None
+                self._agent_span_context = None
+                raise
 
             self._agent_span_id = self._agent_span_context.id
             logger.debug("Created agent span %s (async)", self._agent_span_id)
@@ -872,7 +956,9 @@ class PrefactorMiddleware(AgentMiddleware):
             type="langchain:tool",
             inputs=inputs,
             tool_name=tool_name,
+            arguments=inputs.get("arguments", {}),
         )
+        schema_name = self._resolve_tool_schema_name(tool_name)
 
         try:
             response = handler(request)
@@ -880,13 +966,13 @@ class PrefactorMiddleware(AgentMiddleware):
             outputs = self._extract_tool_output(response)
             span_data.complete(outputs=outputs)
 
-            self._emit_child_span(span_data)
+            self._emit_child_span(span_data, schema_name=schema_name)
             logger.debug("Tool call completed (%s)", tool_name)
             return response
 
         except Exception as e:
             span_data.fail(e)
-            self._emit_child_span(span_data)
+            self._emit_child_span(span_data, schema_name=schema_name)
             logger.error("Tool call failed: %s", e, exc_info=True)
             raise
 
@@ -914,10 +1000,12 @@ class PrefactorMiddleware(AgentMiddleware):
             type="langchain:tool",
             inputs=inputs,
             tool_name=tool_name,
+            arguments=inputs.get("arguments", {}),
         )
+        schema_name = self._resolve_tool_schema_name(tool_name)
 
-        async with instance.span("langchain:tool") as ctx:
-            await ctx.start(self._build_tool_params(span_data))
+        async with instance.span(schema_name) as ctx:
+            await ctx.start(self._build_tool_params(span_data, schema_name))
             try:
                 response = await handler(request)
 
@@ -932,3 +1020,7 @@ class PrefactorMiddleware(AgentMiddleware):
                 await ctx.fail(self._build_tool_result(span_data))
                 logger.error("Tool call failed: %s", e, exc_info=True)
                 raise
+
+    def _resolve_tool_schema_name(self, tool_name: str) -> str:
+        """Resolve a tool's configured span type, if one exists."""
+        return self._tool_span_types.get(tool_name, "langchain:tool")