diff --git a/.github/workflows/production-deploy.yml b/.github/workflows/production-deploy.yml index 9aa6de7..1e97204 100644 --- a/.github/workflows/production-deploy.yml +++ b/.github/workflows/production-deploy.yml @@ -31,7 +31,7 @@ jobs: environment: "production" version-bump: ${{ inputs.version-bump }} branch: ${{ inputs.branch }} - deployed-environment-variables: '[\"OPENAI_API_KEY\",\"OPENAI_MODEL\",\"GOOGLE_AI_API_KEY\",\"GOOGLE_AI_MODEL\"]' + deployed-environment-variables: '[\"OPENAI_API_KEY\",\"OPENAI_MODEL\",\"GOOGLE_AI_API_KEY\",\"GOOGLE_AI_MODEL\",\"OPENROUTER_API_KEY\",\"OPENROUTER_MODEL\",\"OPENROUTER_BASE_URL\"]' secrets: aws-key-id: ${{ secrets.LAMBDA_CONTAINER_PIPELINE_AWS_ID }} aws-secret-key: ${{ secrets.LAMBDA_CONTAINER_PIPELINE_AWS_SECRET }} diff --git a/.github/workflows/staging-deploy.yml b/.github/workflows/staging-deploy.yml index 7810a08..8bc4df4 100644 --- a/.github/workflows/staging-deploy.yml +++ b/.github/workflows/staging-deploy.yml @@ -23,6 +23,9 @@ jobs: OPENAI_MODEL: ${{ vars.OPENAI_MODEL }} GOOGLE_AI_API_KEY: ${{ secrets.GOOGLE_AI_API_KEY }} GOOGLE_AI_MODEL: ${{ vars.GOOGLE_AI_MODEL }} + OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} + OPENROUTER_MODEL: ${{ vars.OPENROUTER_MODEL }} + OPENROUTER_BASE_URL: ${{ vars.OPENROUTER_BASE_URL }} steps: - name: Checkout Code uses: actions/checkout@v4 @@ -67,7 +70,7 @@ jobs: with: template-repository-name: "lambda-feedback/chat-function-boilerplate" environment: "staging" - deployed-environment-variables: '[\"OPENAI_API_KEY\",\"OPENAI_MODEL\",\"GOOGLE_AI_API_KEY\",\"GOOGLE_AI_MODEL\"]' + deployed-environment-variables: '[\"OPENAI_API_KEY\",\"OPENAI_MODEL\",\"GOOGLE_AI_API_KEY\",\"GOOGLE_AI_MODEL\",\"OPENROUTER_API_KEY\",\"OPENROUTER_MODEL\",\"OPENROUTER_BASE_URL\"]' secrets: aws-key-id: ${{ secrets.LAMBDA_CONTAINER_PIPELINE_AWS_ID }} aws-secret-key: ${{ secrets.LAMBDA_CONTAINER_PIPELINE_AWS_SECRET }} diff --git a/.github/workflows/test-lint.yml b/.github/workflows/test-lint.yml index 59115f8..92a5959 100644 --- a/.github/workflows/test-lint.yml +++ b/.github/workflows/test-lint.yml @@ -21,6 +21,9 @@ jobs: OPENAI_MODEL: ${{ vars.OPENAI_MODEL }} GOOGLE_AI_API_KEY: ${{ secrets.GOOGLE_AI_API_KEY }} GOOGLE_AI_MODEL: ${{ vars.GOOGLE_AI_MODEL }} + OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }} + OPENROUTER_MODEL: ${{ vars.OPENROUTER_MODEL }} + OPENROUTER_BASE_URL: ${{ vars.OPENROUTER_BASE_URL }} steps: - name: Checkout uses: actions/checkout@v4 diff --git a/AGENTS.md b/AGENTS.md index d926dc0..117f13e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,13 +4,13 @@ This file provides guidance to AI agents when working with code in this reposito ## 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). It 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`). ## 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,15 +23,18 @@ 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 @@ -39,23 +42,26 @@ pytest tests/test_index.py::test_function_name # Run specific test ### 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) | diff --git a/CLAUDE.md b/CLAUDE.md index e691a20..400387b 100644 --- a/CLAUDE.md +++ b/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). It 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`). ## 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,15 +23,18 @@ 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 @@ -39,23 +42,26 @@ pytest tests/test_index.py::test_function_name # Run specific test ### 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) | diff --git a/Dockerfile b/Dockerfile index 38276cc..72c9d5c 100755 --- a/Dockerfile +++ b/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"] \ No newline at end of file +# 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" diff --git a/README.md b/README.md index 2a99e92..afcba1c 100755 --- a/README.md +++ b/README.md @@ -22,19 +22,15 @@ OPENAI_MODEL GOOGLE_AI_API_KEY GOOGLE_AI_MODEL ``` - -> [!Note] -> If you decide to use another endpoint such as Azure or Ollama or any other, please update the github workflow files to use the right secrets and variables for testing. +> If you use OpenRouter: ```bash -> If you use Azure-OpenAI: -AZURE_OPENAI_API_KEY -AZURE_OPENAI_ENDPOINT -AZURE_OPENAI_API_VERSION -AZURE_OPENAI_CHAT_DEPLOYMENT_NAME -AZURE_OPENAI_EMBEDDING_3072_DEPLOYMENT -AZURE_OPENAI_EMBEDDING_1536_DEPLOYMENT -AZURE_OPENAI_EMBEDDING_3072_MODEL -AZURE_OPENAI_EMBEDDING_1536_MODEL +OPENROUTER_API_KEY +OPENROUTER_MODEL +OPENROUTER_BASE_URL +``` + +> [!NOTE] +> If you decide to use other providers like Azure OpenAI or Ollama, you will need to update the workflow files and the `llm_factory.py` file to include the necessary environment variables for those providers. > For monitoring of the LLM calls (follow instructions on how to set up on langsmith online): LANGCHAIN_TRACING_V2 @@ -117,7 +113,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 ``` @@ -128,18 +123,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. @@ -171,33 +166,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: diff --git a/docs/dev.md b/docs/dev.md index bc420d3..e7bf056 100644 --- a/docs/dev.md +++ b/docs/dev.md @@ -27,10 +27,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 @@ -64,31 +64,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 ``` -Body (stringified within body for API request): +Body (requests may 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: diff --git a/index.py b/index.py index 103969e..fecf53d 100644 --- a/index.py +++ b/index.py @@ -1,35 +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 - """ - 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.", - } +def main(): + server = create_server() + server.chat(chat_module) + server.chat_health(chat_health_module) + run(server) - 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() diff --git a/requirements.txt b/requirements.txt index cf6c893..88d33ce 100644 --- a/requirements.txt +++ b/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@v1.1.1 pytest flake8 \ No newline at end of file diff --git a/src/agent/llm_factory.py b/src/agent/llm_factory.py index 07e887b..956f504 100644 --- a/src/agent/llm_factory.py +++ b/src/agent/llm_factory.py @@ -1,35 +1,27 @@ import os +from typing import Optional -from langchain_openai import AzureChatOpenAI -from langchain_openai import AzureOpenAIEmbeddings -from langchain_community.llms import Ollama -from langchain_community.embeddings import OllamaEmbeddings -from langchain_openai import ChatOpenAI -from langchain_openai import OpenAIEmbeddings -from langchain_google_genai import ChatGoogleGenerativeAI from dotenv import load_dotenv load_dotenv() class AzureLLMs: def __init__(self, temperature: int = 0): + from langchain_openai import AzureChatOpenAI + self._azure_llm = AzureChatOpenAI( openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"], azure_deployment=os.environ["AZURE_OPENAI_CHAT_DEPLOYMENT_NAME"], temperature=temperature, max_tokens=None, ) - self._azure_embedding = AzureOpenAIEmbeddings(azure_deployment=os.environ['AZURE_OPENAI_EMBEDDING_1536_DEPLOYMENT'], - openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"], - model=os.environ["AZURE_OPENAI_EMBEDDING_1536_MODEL"]) - + def get_llm(self): return self._azure_llm - def get_embedding(self): - return self._azure_embedding - class OllamaLLMs: def __init__(self): + from langchain_community.llms import Ollama + self._ollama_llm = Ollama( model=os.environ['OLLAMA_MODEL'], base_url=os.environ['OLLAMA_BASE_URL'], @@ -38,42 +30,26 @@ def __init__(self): }, ) - self._ollama_embedding = OllamaEmbeddings( - model='nomic-embed-text:137m-v1.5-fp16', - base_url=os.environ['OLLAMA_BASE_URL'], - headers={ - 'X-API-Key': os.environ['OLLAMA_API_KEY'], - }, - show_progress=True - ) - def get_llm(self): return self._ollama_llm - def get_embedding(self): - return self._ollama_embedding - class OpenAILLMs: def __init__(self, temperature: int = 0): + from langchain_openai import ChatOpenAI + self._openai_llm = ChatOpenAI( model=os.environ['OPENAI_MODEL'], temperature=temperature, api_key=os.environ["OPENAI_API_KEY"], ) - self._openai_embedding = OpenAIEmbeddings( - model='text-embedding-ada-002', - api_key=os.environ['OPENAI_API_KEY'], - ) - def get_llm(self): return self._openai_llm - def get_embedding(self): - return self._openai_embedding - class GoogleAILLMs: def __init__(self, temperature: int = 0): + from langchain_google_genai import ChatGoogleGenerativeAI + self._google_llm = ChatGoogleGenerativeAI( model=os.environ['GOOGLE_AI_MODEL'], temperature=temperature, @@ -82,3 +58,21 @@ def __init__(self, temperature: int = 0): def get_llm(self): return self._google_llm + +class OpenRouterLLMs: + def __init__(self, temperature: int = 0, model: Optional[str] = None): + from langchain_openai import ChatOpenAI + + model_name = model or os.environ['OPENROUTER_MODEL'] + key = os.environ['OPENROUTER_API_KEY'] + base_url = os.environ['OPENROUTER_BASE_URL'] + + self._openrouter_llm = ChatOpenAI( + model=model_name, + temperature=temperature, + api_key=key, + base_url=base_url, + ) + + def get_llm(self): + return self._openrouter_llm \ No newline at end of file diff --git a/src/module.py b/src/module.py index 133fdc2..121a967 100755 --- a/src/module.py +++ b/src/module.py @@ -1,12 +1,10 @@ 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 - def chat_module(request: ChatRequest) -> ChatResponse: """ @@ -21,6 +19,7 @@ def chat_module(request: ChatRequest) -> ChatResponse: Edit src/agent/prompts.py to change the chatbot's behaviour. Edit src/agent/agent.py to change the agent logic (summarisation threshold, LLM provider, etc.). """ + from src.agent.agent import invoke_base_agent conversation_id = request.conversationId @@ -61,6 +60,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=False, + supportsStreaming=False, + supportsDataPolicy=DataPolicySupport.NOT_SUPPORTED, + ), + ) + + def _to_langchain_messages(messages): result = [] for m in messages: diff --git a/tests/manual_agent_requests.py b/tests/manual_agent_requests.py index 57023b5..03ba27f 100644 --- a/tests/manual_agent_requests.py +++ b/tests/manual_agent_requests.py @@ -1,12 +1,22 @@ import requests -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/" @@ -14,14 +24,11 @@ # 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) diff --git a/tests/test_example_inputs.py b/tests/test_example_inputs.py index 3373fd1..3af5d3d 100644 --- a/tests/test_example_inputs.py +++ b/tests/test_example_inputs.py @@ -1,7 +1,8 @@ import unittest import json import os -from index import handler +from lf_toolkit.chat import ChatRequest +from src.module import chat_module from tests.utils import assert_valid_chat_request, assert_valid_chat_response EXAMPLE_INPUTS_DIR = "tests/example_inputs" @@ -17,9 +18,10 @@ def _test(self, filename: str): with open(os.path.join(EXAMPLE_INPUTS_DIR, filename)) as f: payload = json.load(f) assert_valid_chat_request(self, payload) - result = handler({"body": json.dumps(payload)}, None) + request = ChatRequest.model_validate(payload) + result = chat_module(request) assert_valid_chat_response(self, result) - return payload, json.loads(result["body"]) + return payload, json.loads(result.model_dump_json()) def test_example_input_0_simple(self): self._test("example_input_0.json") diff --git a/tests/test_index.py b/tests/test_index.py deleted file mode 100644 index b6047ec..0000000 --- a/tests/test_index.py +++ /dev/null @@ -1,30 +0,0 @@ -import unittest -import json -from index import handler -from tests.utils import assert_valid_chat_request, assert_valid_chat_response - - -def make_event(body: dict) -> dict: - return {"body": json.dumps(body)} - - -BASE_BODY = { - "messages": [{"role": "USER", "content": "Hello, World"}], - "conversationId": "1234Test", -} - - -class TestChatIndexFunction(unittest.TestCase): - - def test_missing_messages(self): - body = {k: v for k, v in BASE_BODY.items() if k != "messages"} - result = handler(make_event(body), None) - self.assertEqual(result.get("statusCode"), 400) - - def test_invalid_json_body(self): - result = handler({"body": "not valid json"}, None) - self.assertEqual(result.get("statusCode"), 400) - - def test_response_format(self): - assert_valid_chat_request(self, BASE_BODY) - assert_valid_chat_response(self, handler(make_event(BASE_BODY), None)) diff --git a/tests/test_module.py b/tests/test_module.py index 43be647..b7ca462 100755 --- a/tests/test_module.py +++ b/tests/test_module.py @@ -1,6 +1,6 @@ import unittest from lf_toolkit.chat import ChatRequest -from src.module import chat_module +from src.module import chat_health_module, chat_module from tests.utils import assert_valid_chat_response @@ -17,3 +17,11 @@ class TestChatModuleFunction(unittest.TestCase): def test_response_format(self): assert_valid_chat_response(self, chat_module(make_request())) + + +class TestChatHealthModuleFunction(unittest.TestCase): + + def test_reports_healthy_with_chat_capability(self): + result = chat_health_module() + self.assertEqual(result.status, "OK") + self.assertTrue(result.capabilities.supportsChat) diff --git a/tests/utils.py b/tests/utils.py index 113c36e..ba4b291 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -12,16 +12,9 @@ def assert_valid_chat_request(test: unittest.TestCase, payload: dict): test.assertGreater(len(request.messages), 0, "messages must not be empty") -def assert_valid_chat_response(test: unittest.TestCase, result): - """ - Assert a result matches the expected muEd ChatResponse format. - Accepts either a ChatResponse object or a Lambda handler result dict. - """ - if isinstance(result, ChatResponse): - body = json.loads(result.model_dump_json()) - else: - test.assertEqual(result.get("statusCode"), 200) - body = json.loads(result["body"]) +def assert_valid_chat_response(test: unittest.TestCase, result: ChatResponse): + """Assert a ChatResponse matches the expected muEd ChatResponse format.""" + body = json.loads(result.model_dump_json()) output = body.get("output", {}) test.assertEqual(output.get("role"), "ASSISTANT")