Merge pull request #9 from lambda-feedback/shimmy

Shimmy

Author: m-messer

CLAUDE.md

@@ -4,13 +4,13 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 
 ## Project Overview
 
-This is a chat function connecting students to an AI educational chatbot that is integrated with the **Lambda-Feedback** educational platform. It deploys as an AWS Lambda function (containerized via Docker) that receives student chat messages with educational context and returns LLM-powered chatbot responses. Incoming requests follow the [muEd API](https://mued.org/) schema (`context`, `user`, `messages`).
+This is a chat function connecting students to an AI educational chatbot that is integrated with the **Lambda-Feedback** educational platform. It's containerized via Docker and deployed behind [shimmy](https://github.com/lambda-feedback/shimmy), a shim that spawns this function as a persistent JSON-RPC worker process and exposes it as the muEd `/chat` / `/chat/health` HTTP API (both locally and as an AWS Lambda container). Incoming requests follow the [muEd API](https://mued.org/) schema (`context`, `user`, `messages`).
 
 ## Commands
 
 **Testing:**
 ```bash
-pytest                              # Run all unit tests
+PYTHONPATH=. pytest                 # Run all unit tests (CI sets PYTHONPATH=. too)
 python tests/manual_agent_run.py   # Test agent locally with example inputs
 python tests/manual_agent_requests.py  # Test running Docker container
 ```
@@ -23,39 +23,45 @@ docker run --env-file .env -p 8080:8080 llm_chat
 
 **Manual API test (while Docker is running):**
 ```bash
-curl -X POST http://localhost:8080/2015-03-31/functions/function/invocations \
+curl -X POST http://localhost:8080/chat \
   -H 'Content-Type: application/json' \
-  -d '{"body":"{\"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}]}"}'
+  -H 'X-Api-Version: 0.1.0' \
+  -d '{"messages": [{"role": "USER", "content": "hi"}]}'
+
+curl http://localhost:8080/chat/health -H 'X-Api-Version: 0.1.0'
 ```
 
 **Run a single test:**
 ```bash
 pytest tests/test_module.py        # Run specific test file
-pytest tests/test_index.py::test_function_name  # Run specific test
+pytest tests/test_module.py::TestChatModuleFunction::test_response_format  # Run specific test
 ```
 
 ## Architecture
 
 ### Request Flow
 
 ```
-Lambda event → index.py (handler)
-  → validates via lf_toolkit ChatRequest schema
-  → src/module.py (chat_module)
-    → extracts muEd API context (messages, conversationId, question context, user type)
-    → parses educational context to prompt text via src/agent/context.py
-    → src/agent/agent.py (BaseAgent / LangGraph)
-      → routes to call_llm or summarize_conversation node
-      → calls LLM provider (OpenAI / Google / Azure / Ollama)
-  → returns ChatResponse (output, summary, conversationalStyle, processingTime)
+shimmy (shim, container entrypoint)
+  → spawns index.py as a persistent worker subprocess (lf_toolkit RPC server)
+  → forwards POST /chat / GET /chat/health as JSON-RPC "chat" / "chat/health" calls
+  → index.py registers src/module.py's chat_module / chat_health_module as handlers
+    → lf_toolkit validates the request body against the muEd ChatRequest schema
+    → src/module.py (chat_module)
+      → extracts muEd API context (messages, conversationId, question context, user type)
+      → parses educational context to prompt text via src/agent/context.py
+      → src/agent/agent.py (BaseAgent / LangGraph)
+        → routes to call_llm or summarize_conversation node
+        → calls LLM provider (OpenAI / Google / Azure / Ollama)
+    → returns ChatResponse (output, summary, conversationalStyle, processingTime)
 ```
 
 ### Key Files
 
 | File | Role |
 |------|------|
-| `index.py` | AWS Lambda entry point; parses event body, validates schema |
-| `src/module.py` | Transforms muEd API request → invokes agent → builds ChatResponse |
+| `index.py` | Worker entrypoint; registers `chat_module`/`chat_health_module` with `lf_toolkit`'s RPC server (`create_server()` + `run()`) |
+| `src/module.py` | Transforms muEd API request → invokes agent → builds ChatResponse; also exposes `chat_health_module()` |
 | `src/agent/agent.py` | LangGraph stateful graph; manages message history and summarization |
 | `src/agent/prompts.py` | System prompts for tutor behavior, summarization, style detection |
 | `src/agent/llm_factory.py` | Factory classes for each LLM provider (OpenAI, Google, Azure, Ollama) |

Dockerfile

@@ -1,18 +1,14 @@
-ARG PYTHON_VERSION=3.13
+ARG BASE_VERSION=python:3.12
 
-FROM public.ecr.aws/lambda/python:${PYTHON_VERSION}
+# evaluation-function-base's python image bundles the shimmy binary,
+# the Lambda RIE, and the entrypoint.sh that picks between them.
+FROM ghcr.io/lambda-feedback/evaluation-function-base/${BASE_VERSION}
 
-# Set working directory
-WORKDIR ${LAMBDA_TASK_ROOT}
+RUN apt-get update && apt-get install -y \
+      build-essential \
+    && rm -rf /var/lib/apt/lists/*
 
-RUN pip install --upgrade pip 
-RUN dnf install -y git \
-    && dnf install -y \
-      gcc \
-      gcc-c++ \
-      make \
-      python3-devel \
-    && dnf clean all
+RUN pip install --upgrade pip
 
 COPY requirements.txt .
 RUN pip install -r requirements.txt
@@ -27,5 +23,15 @@ COPY index.py .
 
 COPY tests ./tests
 
-# Set the Lambda function handler
-CMD ["index.handler"]
+# Command shimmy uses to start the chat function worker
+ENV FUNCTION_COMMAND="python"
+
+# Args to start the chat function worker with
+ENV FUNCTION_ARGS="index.py"
+
+# The transport to use for the RPC server
+ENV FUNCTION_RPC_TRANSPORT="ipc"
+
+ENV FUNCTION_WORKER_SEND_TIMEOUT="170s"
+
+ENV LOG_LEVEL="debug"

README.md

@@ -119,7 +119,6 @@ The agent uses **two separate LLM instances** — `self.llm` for chat responses
     ├── manual_agent_run.py               # allows testing of any LLM agent on a couple of example inputs
     ├── utils.py                          # shared test helpers
     ├── test_example_inputs.py            # pytests for the example input files
-    ├── test_index.py                     # pytests
     └── test_module.py                    # pytests
 ```
 
@@ -130,18 +129,18 @@ To test your function, you can run the unit tests, call the code directly throug
 
 ### Run Unit Tests
 
-You can run the unit tests using `pytest`.
+You can run the unit tests using `pytest`. Run it from the repository root with `PYTHONPATH=.` set (as CI does) so the `tests` and `src` packages resolve correctly:
 
 ```bash
-pytest
+PYTHONPATH=. pytest
 ```
 
 ### Run the Chat Script
 
-You can run the Python function itself. Make sure to have a main function in either `src/module.py` or `index.py`.
+You can run the Python function itself directly — `index.py` wires `chat_module`/`chat_health_module` into `lf_toolkit`'s RPC server, the same way shimmy invokes it inside the container. This requires the `EVAL_IO`/`EVAL_RPC_TRANSPORT` environment variables shimmy would normally set (see `lf_toolkit`'s docs), so prefer the Docker or `manual_agent_run.py` routes below for everyday testing.
 
 ```bash
-python src/module.py
+python index.py
 ```
 
 You can also use the `manual_agent_run.py` script to test the agents with example inputs from Lambda Feedback questions and synthetic conversations.
@@ -173,33 +172,41 @@ docker run -e OPENAI_API_KEY={your key} -e OPENAI_MODEL={your LLM model name} -p
 docker run --env-file .env -it --name my-lambda-container -p 8080:8080 llm_chat
 ```
 
-This will start the chat function and expose it on port `8080` and it will be open to be curl:
+This starts shimmy (the [Lambda Feedback shim](https://github.com/lambda-feedback/shimmy)) as the container's entrypoint, which spawns this function as a worker subprocess and exposes it on port `8080` as the muEd chat API:
 
 ```bash
-curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations' \
+curl --location 'http://localhost:8080/chat' \
 --header 'Content-Type: application/json' \
---data '{"body":"{\"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}]}"}'
+--header 'X-Api-Version: 0.1.0' \
+--data '{"messages": [{"role": "USER", "content": "hi"}]}'
+```
+
+Health check:
+
+```bash
+curl --location 'http://localhost:8080/chat/health' \
+--header 'X-Api-Version: 0.1.0'
 ```
 
 #### Call Docker Container
 ##### A. Call Docker with Python Requests
 
-In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
+In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the `/chat` and `/chat/health` routes of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
 
 ##### B. Call Docker Container through API request
 
 POST URL:
 
 ```bash
-http://localhost:8080/2015-03-31/functions/function/invocations
+http://localhost:8080/chat
 ```
 
-Per the [muEd `ChatRequest` schema](https://mued.org/), only `messages` is required; `conversationId`, `user`, `context`, and `configuration` are all optional.
+Per the [muEd `ChatRequest` schema](https://mued.org/), only `messages` is required; `conversationId`, `user`, `context`, and `configuration` are all optional. Requests may include an `X-Api-Version: 0.1.0` header.
 
-**Minimal request — only required components** (stringified within `body` for the AWS Lambda Runtime Interface Emulator):
+**Minimal request — only required components:**
 
 ```JSON
-{"body":"{\"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}]}"}
+{"messages": [{"role": "USER", "content": "hi"}]}
 ```
 
 **Full request as Lambda Feedback sends it** — all optional fields populated:

docs/dev.md

@@ -28,10 +28,10 @@ To test your function, you can run the unit tests, call the code directly throug
 
 ### Run Unit Tests
 
-You can run the unit tests using `pytest`.
+You can run the unit tests using `pytest`. Run it from the repository root with `PYTHONPATH=.` set (as CI does) so the `tests` and `src` packages resolve correctly:
 
 ```bash
-pytest
+PYTHONPATH=. pytest
 ```
 
 ### Run the Chat Script
@@ -65,31 +65,39 @@ docker run -e OPENAI_API_KEY={your key} -e OPENAI_MODEL={your LLM chosen model n
 docker run --env-file .env -it --name my-lambda-container -p 8080:8080 llm_chat
 ```
 
-This will start the chat function and expose it on port `8080` and it will be open to be curl:
+This starts shimmy (the [Lambda Feedback shim](https://github.com/lambda-feedback/shimmy)) as the container's entrypoint, which spawns this function as a worker subprocess and exposes it on port `8080` as the muEd chat API:
 
 ```bash
-curl --location 'http://localhost:8080/2015-03-31/functions/function/invocations' \
+curl --location 'http://localhost:8080/chat' \
 --header 'Content-Type: application/json' \
---data '{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}'
+--header 'X-Api-Version: 0.1.0' \
+--data '{"conversationId": "12345Test", "messages": [{"role": "USER", "content": "hi"}], "user": {"type": "LEARNER"}}'
+```
+
+Health check:
+
+```bash
+curl --location 'http://localhost:8080/chat/health' \
+--header 'X-Api-Version: 0.1.0'
 ```
 
 #### Call Docker Container
 ##### A. Call Docker with Python Requests
 
-In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the POST URL of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
+In the `tests/` folder you can find the `manual_agent_requests.py` script that calls the `/chat` and `/chat/health` routes of the running docker container. It reads any kind of input files with the expected schema. You can use this to test your curl calls of the chatbot.
 
 ##### B. Call Docker Container through API request
 
 POST URL:
 
 ```bash
-http://localhost:8080/2015-03-31/functions/function/invocations
+http://localhost:8080/chat
 ```
 
-Input body (stringified within body for API request):
+Input body (requests must include an `X-Api-Version: 0.1.0` header):
 
 ```JSON
-{"body":"{\"conversationId\": \"12345Test\", \"messages\": [{\"role\": \"USER\", \"content\": \"hi\"}], \"user\": {\"type\": \"LEARNER\"}}"}
+{"conversationId": "12345Test", "messages": [{"role": "USER", "content": "hi"}], "user": {"type": "LEARNER"}}
 ```
 
 Body with optional fields:

index.py

@@ -1,37 +1,14 @@
-import json
-from pydantic import ValidationError
+from lf_toolkit import create_server, run
 
-from lf_toolkit.chat import ChatRequest
-from src.module import chat_module
+from src.module import chat_health_module, chat_module
 
 
-def handler(event, context):
-    """
-    Lambda handler function
-    """
-    print("Received event:", json.dumps(event))
+def main():
+    server = create_server()
+    server.chat(chat_module)
+    server.chat_health(chat_health_module)
+    run(server)
 
-    if "body" in event:
-        try:
-            event = json.loads(event["body"])
-        except json.JSONDecodeError:
-            return {
-                "statusCode": 400,
-                "body": "Invalid JSON format in the body. Please check the input.",
-            }
 
-    try:
-        request = ChatRequest.model_validate(event)
-    except ValidationError as e:
-        return {"statusCode": 400, "body": e.json()}
-
-    try:
-        result = chat_module(request)
-    except Exception as e:
-        return {
-            "statusCode": 500,
-            "body": f"An error occurred within the chat_module(): {str(e)}",
-        }
-
-    response = {"statusCode": 200, "body": result.model_dump_json()}
-    return response
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -10,6 +10,6 @@ langdetect
 langgraph
 langsmith
 
-lf_toolkit[ipc] @ git+https://github.com/lambda-feedback/toolkit-python.git@main
+lf_toolkit[ipc] @ git+https://github.com/lambda-feedback/toolkit-python.git@fix/ipc
 pytest
 flake8

src/module.py

@@ -1,8 +1,8 @@
 import time
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 
-from lf_toolkit.chat import ChatRequest, ChatResponse, Message
-from lf_toolkit.shared.mued_api_v0_1_0 import Role
+from lf_toolkit.chat import ChatCapabilities, ChatHealthResponse, ChatRequest, ChatResponse, Message
+from lf_toolkit.shared.mued_api_v0_1_0 import DataPolicySupport, HealthStatus, Role
 
 from src.agent.context import parse_json_to_prompt
 from src.agent.agent import invoke_base_agent
@@ -61,6 +61,22 @@ def chat_module(request: ChatRequest) -> ChatResponse:
     )
 
 
+def chat_health_module() -> ChatHealthResponse:
+    """
+    Health-check entry point — reports whether this chat function is up and
+    what it supports, for the shim's GET /chat/health.
+    """
+    return ChatHealthResponse(
+        status=HealthStatus.OK,
+        capabilities=ChatCapabilities(
+            supportsChat=True,
+            supportsUserPreferences=True,
+            supportsStreaming=False,
+            supportsDataPolicy=DataPolicySupport.NOT_SUPPORTED,
+        ),
+    )
+
+
 def _to_langchain_messages(messages):
     result = []
     for m in messages:

tests/manual_agent_requests.py

@@ -2,26 +2,34 @@
 import json
 
 """
-Script that sends a request to the local endpoint of the docker container to test the chatbot agent.
+Script that sends requests straight to shimmy's muEd chat routes on the
+locally running docker container (`docker build` and `docker run`) to test
+the chatbot agent end-to-end, behind the shim.
 """
 
-# URL for the local endpoint to docker (`docker build` and `docker run`)
-url = "http://localhost:8080/2015-03-31/functions/function/invocations"
+base_url = "http://localhost:8080"
+
+headers = {
+    'Content-Type': 'application/json',
+    'X-Api-Version': '0.1.0',
+}
+
+# Health check
+health_response = requests.get(f"{base_url}/chat/health", headers=headers)
+print("GET /chat/health ->", health_response.status_code)
+print(health_response.text)
 
 # File path for the input text
 path = "tests/example_inputs/"
 input_file = path + "example_input_1.json"
 
 # Step 1: Read the input file
 with open(input_file, "r") as file:
-    data = file.read()
+    payload = file.read()
 
-payload = json.dumps({"body": data})
 print(payload)
-headers = {
-  'Content-Type': 'application/json'
-}
 
-response = requests.request("POST", url, headers=headers, data=payload)
+response = requests.post(f"{base_url}/chat", headers=headers, data=payload)
 
+print("POST /chat ->", response.status_code)
 print(response.text)