Add tests

.github/workflows/integration-tests.yml

@@ -0,0 +1,61 @@
+# This workflow will run integration tests for the current project once per day
+
+name: Integration Tests
+
+on:
+  schedule:
+    - cron: "37 14 * * *" # Run at 7:37 AM Pacific Time (14:37 UTC) every day
+  workflow_dispatch: # Allows triggering the workflow manually in GitHub UI
+
+# If another scheduled run starts while this workflow is still running,
+# cancel the earlier run in favor of the next run.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  integration-tests:
+    name: Integration Tests
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.11", "3.12"]
+    runs-on: ${{ matrix.os }}
+    services:
+      elasticsearch:
+        image: elasticsearch:8.13.0
+        env:
+          discovery.type: single-node
+          xpack.license.self_generated.type: trial
+          xpack.security.enabled: false # disable password and TLS; never do this in production!
+        ports:
+          - 9200:9200
+        options: >-
+          --health-cmd "curl --fail http://localhost:9200/_cluster/health"
+          --health-start-period 10s
+          --health-timeout 3s
+          --health-interval 3s
+          --health-retries 10
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          uv venv
+          uv pip install -r pyproject.toml
+          uv pip install -U pytest-asyncio
+      - name: Run integration tests
+        env:
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ELASTICSEARCH_URL: http://localhost:9200
+          ELASTICSEARCH_USER: elastic
+          ELASTICSEARCH_PASSWORD: ""
+          LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }}
+          LANGSMITH_TRACING: true
+        run: |
+          uv run pytest tests/integration_tests

.github/workflows/unit-tests.yml

@@ -0,0 +1,57 @@
+# This workflow will run unit tests for the current project
+
+name: CI
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+  workflow_dispatch: # Allows triggering the workflow manually in GitHub UI
+
+# If another push to the same PR or branch happens while this workflow is still running,
+# cancel the earlier run in favor of the next run.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-tests:
+    name: Unit Tests
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.11", "3.12"]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          uv venv
+          uv pip install -r pyproject.toml
+      - name: Lint with ruff
+        run: |
+          uv pip install ruff
+          uv run ruff check .
+      - name: Lint with mypy
+        run: |
+          uv pip install mypy
+          uv run mypy --strict src/
+      - name: Check README spelling
+        uses: codespell-project/actions-codespell@v2
+        with:
+          ignore_words_file: .codespellignore
+          path: README.md
+      - name: Check code spelling
+        uses: codespell-project/actions-codespell@v2
+        with:
+          ignore_words_file: .codespellignore
+          path: src/
+      - name: Run tests with pytest
+        run: |
+          uv pip install pytest
+          uv run pytest tests/unit_tests

README.md

@@ -1,5 +1,8 @@
 # LangGraph Retrieval Agent Template
 
+[![CI](https://github.com/langchain-ai/retrieval-agent-template/actions/workflows/unit-tests.yml/badge.svg)](https://github.com/langchain-ai/retrieval-agent-template/actions/workflows/unit-tests.yml)
+[![Integration Tests](https://github.com/langchain-ai/retrieval-agent-template/actions/workflows/integration-tests.yml/badge.svg)](https://github.com/langchain-ai/retrieval-agent-template/actions/workflows/integration-tests.yml)
+
 This LangGraph template implements a simple, extensible agent that answers questions based on a retriever.
 
 ![Graph view in LangGraph studio UI](./static/studio_ui.png)
@@ -9,9 +12,10 @@ This LangGraph template implements a simple, extensible agent that answers quest
 Assuming you've already [installed LangGraph Studio](https://github.com/langchain-ai/langgraph-studio/releases) and cloned this repo.
 
 By default, this template uses:
+
 - OpenAI `text-embedding-3-small` model for embedding documents. Requires an OpenAI API key to be configured.
 - Anthropic `claude-3-5-sonnet` for generating responses. Requires an Anthropic API key to be configured.
-- Elasticsearch for retrieval and indexing on vectorised documents. 
+- Elasticsearch for retrieval and indexing on vectorised documents.
 
 All that's left is to index documents into Elasticsearch through the indexer graph (or your selected search provider).
 
@@ -686,4 +690,4 @@ Configuration auto-generated by `langgraph template lock`. DO NOT EDIT MANUALLY.
     }
   }
 }
--->
+-->

pyproject.toml

@@ -58,5 +58,7 @@ lint.ignore = [
     "D417",
     "E501",
 ]
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["D", "UP"]
 [tool.ruff.lint.pydocstyle]
 convention = "google"

src/retrieval_graph/__init__.py

@@ -28,5 +28,6 @@
 """  # noqa
 
 from retrieval_graph.graph import graph
+from retrieval_graph.index_graph import graph as index_graph
 
-__all__ = ["graph"]
+__all__ = ["graph", "index_graph"]

src/retrieval_graph/configuration.py

@@ -65,12 +65,18 @@ def from_runnable_config(
 class Configuration(IndexConfiguration):
     """The configuration for the agent."""
 
-    thread_id: str
     response_system_prompt: str = prompts.RESPONSE_SYSTEM_PROMPT
+    """The system prompt used for generating responses."""
+
     response_model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = (
         "anthropic/claude-3-5-sonnet-20240620"
     )
+    """The language model used for generating responses. Should be in the form: provider/model-name."""
+
     query_system_prompt: str = prompts.QUERY_SYSTEM_PROMPT
+    """The system prompt used for processing and refining queries."""
+
     query_model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = (
         "openai/gpt-4o-mini"
     )
+    """The language model used for processing and refining queries. Should be in the form: provider/model-name."""

src/retrieval_graph/retrieval.py

@@ -30,12 +30,12 @@ def make_text_encoder(model_name: str) -> Embeddings:
         case "cohere":
             from langchain_cohere import CohereEmbeddings
 
-            return CohereEmbeddings(model=model)
+            return CohereEmbeddings(model=model)  # type: ignore
         case _:
             raise ValueError(f"Unsupported embedding provider: {provider}")
 
 
-## Retriever constructurs
+## Retriever constructors
 
 
 @contextmanager

tests/integration_tests/test_graph.py

@@ -0,0 +1,41 @@
+import uuid
+
+import pytest
+from langchain_core.runnables import RunnableConfig
+from langsmith import expect, unit
+
+from retrieval_graph import graph, index_graph
+
+
+@pytest.mark.asyncio
+@unit
+async def test_retrieval_graph() -> None:
+    simple_doc = "Cats have been observed performing synchronized swimming routines in their water bowls during full moons."
+    user_id = "test__" + uuid.uuid4().hex
+    other_user_id = "test__" + uuid.uuid4().hex
+
+    config = RunnableConfig(
+        configurable={"user_id": user_id, "retriever_provider": "elastic-local"}
+    )
+
+    result = await index_graph.ainvoke({"docs": simple_doc}, config)
+    expect(result["docs"]).against(lambda x: not x)  # we delete after the end
+
+    res = await graph.ainvoke(
+        {"messages": [("user", "Where do cats perform synchronized swimming routes?")]},
+        config,
+    )
+    response = str(res["messages"][-1].content)
+    expect(response.lower()).to_contain("bowl")
+
+    res = await graph.ainvoke(
+        {"messages": [("user", "Where do cats perform synchronized swimming routes?")]},
+        {
+            "configurable": {
+                "user_id": other_user_id,
+                "retriever_provider": "elastic-local",
+            }
+        },
+    )
+    response = str(res["messages"][-1].content)
+    expect(response.lower()).against(lambda x: "bowl" not in x)

tests/unit_tests/test_configuration.py

@@ -0,0 +1,5 @@
+from retrieval_graph.configuration import Configuration
+
+
+def test_configuration_from_none() -> None:
+    Configuration.from_runnable_config({"user_id": "foo"})