feat: use helicone in example test suite (#940)

Co-authored-by: Apoorv Taneja <purutaneja.com@gmail.com>

.github/workflows/examples.yml

@@ -78,21 +78,23 @@ jobs:
         env:
           DOCKER_USER: ${{secrets.DOCKER_USER}}
           OPENAI_API_KEY: ${{secrets.OPENAI_API_KEY}}
+          HELICONE_API_KEY: ${{secrets.HELICONE_API_KEY}}
           JULEP_API_KEY: ${{secrets.JULEP_API_KEY}}
           JULEP_API_URL: ${{secrets.JULEP_API_URL}}
           LISTENNOTES_API_KEY: ${{secrets.LISTENNOTES_API_KEY}}
           COMPOSIO_API_KEY: ${{ inputs.api_key || secrets.COMPOSIO_API_KEY_RELEASE }}
           COMPOSIO_BASE_URL: ${{ inputs.base_url || secrets.COMPOSIO_BASE_URL || 'https://backend.composio.dev/api' }}
           # TODO(@kaavee): Add Anthropic API key
         run: |
-          python -m pip install --upgrade pip pipenv pytest
-          python -m pip install plugins/${{matrix.package}}
-          python -m pip install '.[all]'
-          python -m pip install 'numpy<2' python-dotenv
-          python -m pip install unstructured
+          python -m pip install uv
+          python -m uv pip install --upgrade pip pytest
+          python -m uv pip install plugins/${{ matrix.package }}
+          python -m uv pip install '.[all]'
+          python -m uv pip install 'numpy<2' python-dotenv
+          python -m uv pip install unstructured
 
-          COMPOSIO_BASE_URL=${{ env.COMPOSIO_BASE_URL }} COMPOSIO_API_KEY=${{ env.COMPOSIO_API_KEY }} COMPOSIO_LOGGING_LEVEL='debug' PLUGIN_TO_TEST=${{matrix.package}} CI=false pytest -vv tests/test_example.py
-          COMPOSIO_BASE_URL=${{ env.COMPOSIO_BASE_URL }} COMPOSIO_API_KEY=${{ env.COMPOSIO_API_KEY }} COMPOSIO_LOGGING_LEVEL='debug' PLUGIN_TO_TEST=${{matrix.package}} CI=false pytest -vv tests/test_load_tools.py
+          COMPOSIO_BASE_URL=${{ env.COMPOSIO_BASE_URL }} COMPOSIO_API_KEY=${{ env.COMPOSIO_API_KEY }} COMPOSIO_LOGGING_LEVEL='debug' CI=false PLUGIN_TO_TEST=${{ matrix.package }} python -m uv run pytest -vv tests/test_example.py
+          COMPOSIO_BASE_URL=${{ env.COMPOSIO_BASE_URL }} COMPOSIO_API_KEY=${{ env.COMPOSIO_API_KEY }} COMPOSIO_LOGGING_LEVEL='debug' CI=false PLUGIN_TO_TEST=${{ matrix.package }} python -m uv run pytest -vv tests/test_load_tools.py
       - name: Slack Notification on Failure
         if: ${{ failure() && github.ref == 'refs/heads/master' && !contains(github.event.head_commit.message, 'release') && !contains(github.event.head_commit.message, 'Release') && !inputs.dont_notify }}
         uses: rtCamp/action-slack-notify@v2

.github/workflows/run_examples.yml

@@ -1,9 +1,7 @@
 name: Run Examples Tests
 on:
   workflow_dispatch:
-  push:
-    branches:
-      - master
+  pull_request:
     paths:
       - 'python/**'
 

python/plugins/langchain/composio_langchain/toolset.py

@@ -118,9 +118,7 @@ def _wrap_tool(
             schema_params=schema_params,
             entity_id=entity_id,
         )
-        parameters = json_schema_to_model(
-            json_schema=schema_params,
-        )
+        parameters = json_schema_to_model(json_schema=schema_params)
         tool = StructuredTool.from_function(
             name=action,
             description=description,

python/tests/test_example.py

@@ -2,6 +2,7 @@
 E2E Tests for plugin demos and examples.
 """
 
+import ast
 import os
 import subprocess
 import sys
@@ -211,9 +212,16 @@ def test_example(
             val is not None
         ), f"Please provide value for `{key}` for testing `{example['file']}`"
 
+    filepath = Path(example["file"])
+    code = filepath.read_text(encoding="utf-8")
+
+    if plugin_to_test != "lyzr":
+        code = add_helicone_headers(code)
+        filepath.write_text(code, encoding="utf-8")
+
     cwd = example.get("cwd", None)
     proc = subprocess.Popen(  # pylint: disable=consider-using-with
-        args=[sys.executable, example["file"]],
+        args=[sys.executable, str(filepath)],
         # TODO(@angryblade): Sanitize the env before running the process.
         env={**os.environ, **example["env"]},
         stdout=subprocess.PIPE,
@@ -240,3 +248,85 @@ def test_example(
     )
     for match in example["match"]["values"]:
         assert match in output
+
+
+# The following will get unparsed to get the nodes to put in the
+# OpenAI function calls
+OS_IMPORT = "import os"
+BASE_URL = "'https://oai.helicone.ai/v1'"
+DEFAULT_HEADERS_DICT = """
+{
+    "Helicone-Auth": f"Bearer {os.environ['HELICONE_API_KEY']}",
+    "Helicone-Cache-Enabled": "true",
+    "Helicone-User-Id": "GitHub-CI-Example-Tests",
+}
+"""
+OS_IMPORT_STMT = ast.parse(OS_IMPORT).body[0]
+BASE_URL_EXPR = ast.parse(BASE_URL, mode="eval").body
+DEFAULT_HEADERS_EXPR = ast.parse(DEFAULT_HEADERS_DICT, mode="eval").body
+
+
+class HeliconeAdder(ast.NodeTransformer):
+    def __init__(self) -> None:
+        self.has_os_import = False
+        self.openai_patched_successfully = False
+
+    def visit_Import(self, node: ast.Import) -> ast.Import:
+        self.generic_visit(node)
+
+        # If the module already imports `os`, then set the flag
+        # to indicate that it has been imported.
+        for alias in node.names:
+            if alias.name == "os":
+                self.has_os_import = True
+
+        return node
+
+    def visit_Call(self, node: ast.Call) -> ast.Call:
+        self.generic_visit(node)
+
+        # If it is a call to the function `OpenAI` or `ChatOpenAI`,
+        # then preserve the original arguments, and add two new
+        # keyword arguments, `base_url` and `default_headers`.
+        if isinstance(node.func, ast.Name) and node.func.id in ("OpenAI", "ChatOpenAI"):
+            new_keywords = [
+                ast.keyword(arg="base_url", value=BASE_URL_EXPR),
+                ast.keyword(arg="default_headers", value=DEFAULT_HEADERS_EXPR),
+            ]
+            node.keywords.extend(new_keywords)
+            self.openai_patched_successfully = True
+
+        return node
+
+    def visit_Dict(self, node: ast.Dict) -> ast.Dict:
+        self.generic_visit(node)
+
+        # If the dictionary has a 'config_list' string key, and its
+        # value is a list, then add the `base_url` and `default_headers`
+        # keys to the dictionary.
+        if any(
+            isinstance(key, ast.Constant)
+            and key.value == "config_list"
+            and isinstance(value, ast.List)
+            for key, value in zip(node.keys, node.values, strict=True)
+        ):
+            node.keys.extend(
+                [ast.Constant("base_url"), ast.Constant("default_headers")]
+            )
+            node.values.extend([BASE_URL_EXPR, DEFAULT_HEADERS_EXPR])
+            self.openai_patched_successfully = True
+
+        return node
+
+
+def add_helicone_headers(source: str) -> str:
+    tree = ast.parse(source)
+    helicone_adder = HeliconeAdder()
+    tree = helicone_adder.visit(tree)
+    assert helicone_adder.openai_patched_successfully
+
+    # If the module does not import `os`, then add the import statement
+    if not helicone_adder.has_os_import:
+        tree.body.insert(0, OS_IMPORT_STMT)
+
+    return ast.unparse(tree)