atlanhq
diff --git a/‎.github/workflows/checks.yaml
Lines changed: 1 addition & 1 deletion b/‎.github/workflows/checks.yaml
Lines changed: 1 addition & 1 deletion
diff --git a/‎.github/workflows/docs.yaml
Lines changed: 3 additions & 3 deletions b/‎.github/workflows/docs.yaml
Lines changed: 3 additions & 3 deletions
diff --git a/‎.github/workflows/e2e-integration-tests.yaml
Lines changed: 7 additions & 1 deletion b/‎.github/workflows/e2e-integration-tests.yaml
Lines changed: 7 additions & 1 deletion
diff --git a/‎.github/workflows/tests.yaml
Lines changed: 2 additions & 2 deletions b/‎.github/workflows/tests.yaml
Lines changed: 2 additions & 2 deletions
diff --git a/‎.gitignore
Lines changed: 1 addition & 1 deletion b/‎.gitignore
Lines changed: 1 addition & 1 deletion
diff --git a/‎Makefile
Lines changed: 4 additions & 6 deletions b/‎Makefile
Lines changed: 4 additions & 6 deletions
diff --git a/‎tests/unit/workflows/__init__.py renamed to ‎application_sdk/activities/agents/__init__.py b/‎tests/unit/workflows/__init__.py renamed to ‎application_sdk/activities/agents/__init__.py
diff --git a/‎application_sdk/activities/agents/langgraph.py
Lines changed: 152 additions & 0 deletions b/‎application_sdk/activities/agents/langgraph.py
Lines changed: 152 additions & 0 deletions
diff --git a/‎application_sdk/activities/common/utils.py
Lines changed: 5 additions & 3 deletions b/‎application_sdk/activities/common/utils.py
Lines changed: 5 additions & 3 deletions
@@ -33,5 +33,5 @@ jobs:
     - name: Install dependencies
       run: |
         poetry config virtualenvs.in-project true
-        poetry install --no-root --without dev,test --extras "workflows"
+        poetry install --no-root --without dev,test --extras "workflows langgraph_agent"
     - uses: pre-commit/[email protected]
@@ -50,7 +50,7 @@ jobs:
           git config --global url."https://${{ secrets.ORG_PAT_GITHUB }}".insteadOf "ssh://git";
           poetry install --with "dev, test" --all-extras
           source .venv/bin/activate
-          make sphinx_docs
+          make apidocs
           pyreverse -o html application_sdk
 
       - name: Setup AWS Credentials
@@ -62,14 +62,14 @@ jobs:
       - name: Upload Docs to Kryptonite Bucket
         if: ${{ github.event_name == 'push' }}
         run: |
-          aws s3 sync ./docs/_build/html s3://kryptonite-store/application-sdk/${{ steps.get_branch.outputs.branch }} --delete
+          aws s3 sync ./docs/site s3://kryptonite-store/application-sdk/${{ steps.get_branch.outputs.branch }} --delete
           aws s3 cp packages.html s3://kryptonite-store/application-sdk/${{ steps.get_branch.outputs.branch }}/packages.html
           aws s3 cp classes.html s3://kryptonite-store/application-sdk/${{ steps.get_branch.outputs.branch }}/classes.html
 
       - name: Upload Docs to Kryptonite Bucket
         if: ${{ github.event_name == 'pull_request' }}
         run: |
-          aws s3 sync ./docs/_build/html s3://kryptonite-store/application-sdk/${{ github.head_ref }} --delete
+          aws s3 sync ./docs/site s3://kryptonite-store/application-sdk/${{ github.head_ref }} --delete
           aws s3 cp packages.html s3://kryptonite-store/application-sdk/${{ github.head_ref }}/packages.html
           aws s3 cp classes.html s3://kryptonite-store/application-sdk/${{ github.head_ref }}/classes.html
 
 
@@ -34,12 +34,18 @@ jobs:
             echo "distinct_id=${{ github.sha }}" >> $GITHUB_ENV
           fi
 
+          if [ "${{ github.base_ref }}" ]; then
+            echo "base_branch=refs/heads/${{ github.base_ref }}" >> $GITHUB_ENV
+          else
+            echo "base_branch=${{ github.ref }}" >> $GITHUB_ENV
+          fi
+
       - name: Trigger E2E Integrations Tests
         uses: codex-/return-dispatch@v2
         id: return_dispatch
         with:
           token: "${{ secrets.ORG_PAT_GITHUB }}"
-          ref: "refs/heads/develop"
+          ref: ${{ env.base_branch }}
           repo: ${{ matrix.repo_name }}
           owner: "atlanhq"
           workflow: "e2e-integration-test.yaml"
 
@@ -75,14 +75,14 @@ jobs:
         uses: actions/cache@v4
         with:
           path: .venv
-          key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+          key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock', '**/pyproject.toml') }}
 
       #----------------------------------------------
       # install dependencies if cache does not exist
       #----------------------------------------------
       - name: Install dependencies
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction --no-root --extras "iceberg workflows daft"
+        run: poetry install --no-interaction --no-root --extras "iceberg workflows daft langgraph_agent"
 
       - name: Run Unit and Integration Tests with Coverage
         run: |
 
@@ -137,7 +137,7 @@ venv.bak/
 .ropeproject
 
 # mkdocs documentation
-/site
+docs/site
 
 # mypy
 .mypy_cache/
 
@@ -28,9 +28,7 @@ stop-all:
 	@echo "All detached processes stopped."
 
 
-# Generate Sphinx documentation
-sphinx_docs:
-	mkdir -p docs/reference
-	sphinx-apidoc -o ./docs/reference ./application_sdk --tocfile index --module-first --separate --force
-	cd docs && make html
-
+# Generate documentation
+apidocs:
+	cd docs && mkdocs build
+	cd docs && pydoctor --html-output=site/api ../application_sdk --docformat=google
@@ -0,0 +1,152 @@
+"""
+LangGraph activities interface.
+
+This module provides the activities interface for LangGraph agent operations,
+including graph compilation and task execution.
+"""
+
+import importlib
+from typing import Any, Callable, Dict, Optional, Union
+
+from temporalio import activity
+
+from application_sdk.activities import ActivitiesInterface
+from application_sdk.common.logger_adaptors import get_logger
+
+logger = get_logger(__name__)
+
+try:
+    from langgraph.graph import StateGraph
+
+    from application_sdk.agents.langgraph_agent import LangGraphAgent
+
+    LANGGRAPH_AVAILABLE = True
+except ImportError:
+    logger.warning("LangGraph is not installed, agent functionality will be disabled")
+    LANGGRAPH_AVAILABLE = False
+    StateGraph = Any  # type: ignore
+    LangGraphAgent = Any  # type: ignore
+
+# Registry to store graph builder functions
+_graph_builders: Dict[str, Callable[..., Union[StateGraph, Any]]] = {}
+
+
+def register_graph_builder(
+    name: str, builder_func: Callable[..., Union[StateGraph, Any]]
+) -> None:
+    """Register a graph builder function.
+
+    Args:
+        name: The name to register the builder under
+        builder_func: A function that returns a StateGraph
+    """
+    if not LANGGRAPH_AVAILABLE:
+        logger.warning(
+            "LangGraph is not installed, graph builder registration will be ignored"
+        )
+        return
+
+    _graph_builders[name] = builder_func
+
+
+def get_graph_builder(name: str) -> Optional[Callable[..., Union[StateGraph, Any]]]:
+    """Get a graph builder by name.
+
+    Args:
+        name: The name of the builder to get
+
+    Returns:
+        The builder function or None if not found
+    """
+    if not LANGGRAPH_AVAILABLE:
+        logger.warning("LangGraph is not installed, cannot get graph builder")
+        return None
+
+    # Check if builder is directly in registry
+    if name in _graph_builders:
+        return _graph_builders[name]
+
+    # Try to import from module path (e.g. "my_module.my_builder")
+    try:
+        if "." in name:
+            module_path, func_name = name.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            if hasattr(module, func_name):
+                builder = getattr(module, func_name)
+                # Cache for future use
+                _graph_builders[name] = builder
+                return builder
+    except ImportError:
+        logger.error(f"Could not import graph builder module: {name}")
+    except Exception as e:
+        logger.error(f"Error loading graph builder: {str(e)}")
+
+    return None
+
+
+class LangGraphActivities(ActivitiesInterface):
+    """Activities for LangGraph agent operations.
+
+    This class defines the activities that can be executed as part of
+    LangGraph workflows, including graph compilation and task execution.
+    """
+
+    @activity.defn
+    async def run_agent(self, activity_input: Dict[str, Any]) -> Dict[str, Any]:
+        """Runs the LangGraph agent with the given task.
+
+        Args:
+            activity_input (Dict[str, Any]): Input for the activity,
+                including user query, state, and graph_builder_name.
+
+        Returns:
+            Dict[str, Any]: Result of the agent task execution.
+        """
+        if not LANGGRAPH_AVAILABLE:
+            return {
+                "error": "LangGraph is not installed. Please install the package with 'pip install langgraph' or use the langgraph_agent extra."
+            }
+
+        try:
+            user_query = activity_input.get("user_query")
+            if not user_query:
+                return {"error": "Error: No user query provided."}
+
+            graph_builder_name = activity_input.get("graph_builder_name")
+            if not graph_builder_name:
+                return {"error": "Error: No graph builder name provided."}
+
+            # Get the graph builder function
+            graph_builder = get_graph_builder(graph_builder_name)
+            if not graph_builder:
+                return {
+                    "error": f"Error: Graph builder '{graph_builder_name}' not found."
+                }
+
+            # Build the StateGraph
+            try:
+                state_graph = graph_builder()
+            except Exception as e:
+                logger.error(f"Error building graph: {str(e)}")
+                return {"error": f"Error building graph: {str(e)}"}
+
+            # Initialize the agent
+            agent = LangGraphAgent(
+                state_graph=state_graph,
+                state=activity_input.get("state", {"messages": []}),
+            )
+
+            # Compile and run the graph
+            agent.compile_graph()
+            response = agent.run(user_query)
+            state = agent.state
+
+            return {
+                "result": "Agent execution completed successfully",
+                "response": response,
+                "state": state,
+            }
+        except Exception as e:
+            error_msg = f"Error running agent task: {str(e)}"
+            logger.error(error_msg)
+            return {"error": error_msg}
@@ -51,9 +51,11 @@ def auto_heartbeater(fn: F) -> F:
         F: The decorated activity function.
 
     Note:
-        * We have activities that can run for a long time, in case of a failure (say: worker crash)
+        We have activities that can run for a long time, in case of a failure (say: worker crash).
+
         Temporal will not retry the activity until the configured timeout is reached.
-        * We add auto_heartbeater to activities to ensure an failure is detected earlier
+
+        We add auto_heartbeater to activities to ensure an failure is detected earlier
         and the activity is retried.
 
         - https://temporal.io/blog/activity-timeouts
@@ -63,7 +65,7 @@ def auto_heartbeater(fn: F) -> F:
         >>> @activity.defn
         >>> @auto_heartbeater
         >>> async def my_activity():
-        >>>     pass
+        ...     pass
     """
 
     @wraps(fn)