From 541d4be2d4540a5837704f87cfcbb81e5c9ad47b Mon Sep 17 00:00:00 2001 From: Ayush Noori Date: Wed, 11 Mar 2026 15:27:34 -0700 Subject: [PATCH] Draft OpenAI support, not yet working with HMS Azure workspace --- .env.example | 12 ++++++++ .gitignore | 1 + README.md | 26 +++++++++++++++-- package.json | 1 + pnpm-lock.yaml | 49 ++++++++++++++++++++++++++++++++ src/index.tsx | 77 +++++++++++++++++++++++++++++++++++++++++++++----- 6 files changed, 157 insertions(+), 9 deletions(-) diff --git a/.env.example b/.env.example index 8b3e125..9ac673e 100755 --- a/.env.example +++ b/.env.example @@ -1 +1,13 @@ +# LLM provider: "anthropic" (default) or "azure" +# LLM_PROVIDER=anthropic + +# Required when LLM_PROVIDER = "anthropic" or unset ANTHROPIC_API_KEY= + +# Azure OpenAI (when LLM_PROVIDER = "azure") +# Example: HMS endpoint https://azure-ai.hms.edu +# If your server expects paths under /openai, use https://azure-ai.hms.edu/openai +AZURE_OPENAI_ENDPOINT= +AZURE_OPENAI_API_KEY= +AZURE_OPENAI_API_VERSION=2024-05-01-preview +AZURE_OPENAI_DEPLOYMENT=gpt-5 diff --git a/.gitignore b/.gitignore index dc7211b..f961bd3 100644 --- a/.gitignore +++ b/.gitignore @@ -91,6 +91,7 @@ bower_components build/Release # Dependency directories +.pnpm-store/ node_modules/ jspm_packages/ diff --git a/README.md b/README.md index 763abf6..d257000 100644 --- a/README.md +++ b/README.md @@ -50,12 +50,16 @@ You can add your own graphs, biomedical or otherwise, without changing any code. cp .env.example .env ``` - Open `.env` and add your API key. The CLI currently uses [Anthropic](https://www.anthropic.com/) as its LLM provider: + Open `.env` and add credentials for your chosen LLM provider. The CLI supports [Anthropic](https://www.anthropic.com/) (default) and [Azure OpenAI](https://azure.microsoft.com/products/ai-services/openai-service). + + **Anthropic (default):** ```env ANTHROPIC_API_KEY=your_key_here ``` + **Azure OpenAI:** set `LLM_PROVIDER=azure` and your Azure credentials (see [OpenAI (Azure)](#openai-azure) below). + 4. **Start the CLI**: ```bash @@ -83,6 +87,24 @@ Find the relationship between metformin and breast cancer. The agent will search the knowledge graph, traverse relationships, and synthesize an answer while citing the specific nodes and edges it used. +### OpenAI (Azure) + +To use Azure OpenAI (e.g. an institutional endpoint like HMS), set in `.env`: + +```env +LLM_PROVIDER=azure +AZURE_OPENAI_ENDPOINT=https://azure-ai.hms.edu +AZURE_OPENAI_API_KEY=your_azure_key +AZURE_OPENAI_API_VERSION=2024-05-01-preview +AZURE_OPENAI_DEPLOYMENT=gpt-4o-1120 +``` + +If your server expects paths under `/openai`, use `https://azure-ai.hms.edu/openai` as the endpoint. + +### Switching models + +The `/models` command is listed in the CLI help. Model selection is currently via environment: set `LLM_PROVIDER` to `anthropic` or `azure` and configure the corresponding API keys, then restart the CLI for the change to take effect. In-app model switching (e.g. choosing a model from a list when you type `/models`) will be supported when the TUI library adds custom command registration and transport invalidation. + ## Adding Your Own Knowledge Graph Adding a new graph takes four steps and requires **no code changes**. @@ -182,7 +204,7 @@ Tool renderers provide rich visualization of tool outputs in the terminal. See ` - **Runtime**: [Bun](https://bun.sh/) - **Language**: TypeScript - **UI**: [React 19](https://react.dev/) with [@ai-tui/core](https://www.npmjs.com/package/@ai-tui/core) -- **LLM**: [Vercel AI SDK](https://sdk.vercel.ai/) (currently configured for Anthropic Claude) +- **LLM**: [Vercel AI SDK](https://sdk.vercel.ai/) (Anthropic Claude or Azure OpenAI via `LLM_PROVIDER`) - **Data**: Local parquet files queried via [DuckDB](https://duckdb.org/) - **Validation**: [Zod](https://zod.dev/) diff --git a/package.json b/package.json index b3ed040..74d3872 100644 --- a/package.json +++ b/package.json @@ -14,6 +14,7 @@ }, "dependencies": { "@ai-sdk/anthropic": "^3.0.17", + "@ai-sdk/azure": "^3.0.0", "@ai-sdk/react": "^3.0.41", "@ai-tui/core": "^0.1.1", "@duckdb/node-api": "1.4.4-r.1", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c9e6f97..722819a 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,6 +11,9 @@ importers: '@ai-sdk/anthropic': specifier: ^3.0.17 version: 3.0.35(zod@4.3.6) + '@ai-sdk/azure': + specifier: ^3.0.0 + version: 3.0.42(zod@4.3.6) '@ai-sdk/react': specifier: ^3.0.41 version: 3.0.70(react@19.2.4)(zod@4.3.6) @@ -57,22 +60,44 @@ packages: peerDependencies: zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/azure@3.0.42': + resolution: {integrity: sha512-BGg0e3GEI7KHkwUv7d5f9rXzDlTiWhQ4xzVakdHLV/OP24jvXes5X7fI3QZ0rbKBop6URq0yaxomBfwEqqRlzw==} + engines: {node: '>=18'} + peerDependencies: + zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/gateway@3.0.32': resolution: {integrity: sha512-7clZRr07P9rpur39t1RrbIe7x8jmwnwUWI8tZs+BvAfX3NFgdSVGGIaT7bTz2pb08jmLXzTSDbrOTqAQ7uBkBQ==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/openai@3.0.41': + resolution: {integrity: sha512-IZ42A+FO+vuEQCVNqlnAPYQnnUpUfdJIwn1BEDOBywiEHa23fw7PahxVtlX9zm3/zMvTW4JKPzWyvAgDu+SQ2A==} + engines: {node: '>=18'} + peerDependencies: + zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/provider-utils@4.0.13': resolution: {integrity: sha512-HHG72BN4d+OWTcq2NwTxOm/2qvk1duYsnhCDtsbYwn/h/4zeqURu1S0+Cn0nY2Ysq9a9HGKvrYuMn9bgFhR2Og==} engines: {node: '>=18'} peerDependencies: zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/provider-utils@4.0.19': + resolution: {integrity: sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg==} + engines: {node: '>=18'} + peerDependencies: + zod: ^3.25.76 || ^4.1.8 + '@ai-sdk/provider@3.0.7': resolution: {integrity: sha512-VkPLrutM6VdA924/mG8OS+5frbVTcu6e046D2bgDo00tehBANR1QBJ/mPcZ9tXMFOsVcm6SQArOregxePzTFPw==} engines: {node: '>=18'} + '@ai-sdk/provider@3.0.8': + resolution: {integrity: sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ==} + engines: {node: '>=18'} + '@ai-sdk/react@3.0.70': resolution: {integrity: sha512-P+kwZ0Tvf5oev6v7i6ClA4Y25Gsjl2MUtm4edYfJGp6t6oF06pAjNYupth+r0JQBt3RqQHcLp5oVCc+bM6XZPw==} engines: {node: '>=18'} @@ -952,6 +977,13 @@ snapshots: '@ai-sdk/provider-utils': 4.0.13(zod@4.3.6) zod: 4.3.6 + '@ai-sdk/azure@3.0.42(zod@4.3.6)': + dependencies: + '@ai-sdk/openai': 3.0.41(zod@4.3.6) + '@ai-sdk/provider': 3.0.8 + '@ai-sdk/provider-utils': 4.0.19(zod@4.3.6) + zod: 4.3.6 + '@ai-sdk/gateway@3.0.32(zod@4.3.6)': dependencies: '@ai-sdk/provider': 3.0.7 @@ -959,6 +991,12 @@ snapshots: '@vercel/oidc': 3.1.0 zod: 4.3.6 + '@ai-sdk/openai@3.0.41(zod@4.3.6)': + dependencies: + '@ai-sdk/provider': 3.0.8 + '@ai-sdk/provider-utils': 4.0.19(zod@4.3.6) + zod: 4.3.6 + '@ai-sdk/provider-utils@4.0.13(zod@4.3.6)': dependencies: '@ai-sdk/provider': 3.0.7 @@ -966,10 +1004,21 @@ snapshots: eventsource-parser: 3.0.6 zod: 4.3.6 + '@ai-sdk/provider-utils@4.0.19(zod@4.3.6)': + dependencies: + '@ai-sdk/provider': 3.0.8 + '@standard-schema/spec': 1.1.0 + eventsource-parser: 3.0.6 + zod: 4.3.6 + '@ai-sdk/provider@3.0.7': dependencies: json-schema: 0.4.0 + '@ai-sdk/provider@3.0.8': + dependencies: + json-schema: 0.4.0 + '@ai-sdk/react@3.0.70(react@19.2.4)(zod@4.3.6)': dependencies: '@ai-sdk/provider-utils': 4.0.13(zod@4.3.6) diff --git a/src/index.tsx b/src/index.tsx index 808de67..42501f3 100644 --- a/src/index.tsx +++ b/src/index.tsx @@ -1,4 +1,5 @@ import { createAnthropic } from "@ai-sdk/anthropic"; +import { createAzure } from "@ai-sdk/azure"; import { Agent, type ConfigInput, @@ -16,21 +17,77 @@ import { GraphLoader, makeParquetGraphTools } from "./parquet-tools/index.ts"; import { graphAgentPrompt, regularPrompt } from "./prompts.ts"; import { GetNodeDetailsTool } from "./tool-renderers/index.ts"; -export const env = createEnv({ +const providerSchema = z.enum(["anthropic", "azure"]); + +const rawEnv = createEnv({ server: { - ANTHROPIC_API_KEY: z.string().min(1), + LLM_PROVIDER: z + .string() + .optional() + .transform((v) => (v ? providerSchema.parse(v) : "anthropic")), + ANTHROPIC_API_KEY: z.string().optional(), + AZURE_OPENAI_ENDPOINT: z.string().url().optional(), + AZURE_OPENAI_API_KEY: z.string().optional(), + AZURE_OPENAI_API_VERSION: z + .string() + .optional() + .default("2024-05-01-preview"), + AZURE_OPENAI_DEPLOYMENT: z.string().optional().default("gpt-4o-1120"), }, runtimeEnv: process.env, emptyStringAsUndefined: true, }); +function validateEnv() { + const provider = rawEnv.LLM_PROVIDER; + if (provider === "anthropic") { + if (!rawEnv.ANTHROPIC_API_KEY?.trim()) { + throw new Error( + "ANTHROPIC_API_KEY is required when LLM_PROVIDER is anthropic (or unset). Set it in .env", + ); + } + return { + ...rawEnv, + ANTHROPIC_API_KEY: rawEnv.ANTHROPIC_API_KEY!, + }; + } + // provider === "azure" + if (!rawEnv.AZURE_OPENAI_ENDPOINT?.trim() || !rawEnv.AZURE_OPENAI_API_KEY?.trim()) { + throw new Error( + "AZURE_OPENAI_ENDPOINT and AZURE_OPENAI_API_KEY are required when LLM_PROVIDER=azure. Set them in .env", + ); + } + return { + ...rawEnv, + AZURE_OPENAI_ENDPOINT: rawEnv.AZURE_OPENAI_ENDPOINT!, + AZURE_OPENAI_API_KEY: rawEnv.AZURE_OPENAI_API_KEY!, + }; +} + +export const env = validateEnv(); + // Discover graph metadata instantly; parquet data loads lazily per agent const DATA_DIR = join(import.meta.dir, "..", "data"); const loader = new GraphLoader(DATA_DIR); -const anthropic = createAnthropic({ - apiKey: env.ANTHROPIC_API_KEY, -}); +const anthropic = + env.LLM_PROVIDER === "anthropic" + ? createAnthropic({ apiKey: env.ANTHROPIC_API_KEY }) + : null; + +const azure = + env.LLM_PROVIDER === "azure" + ? createAzure({ + baseURL: env.AZURE_OPENAI_ENDPOINT, + apiKey: env.AZURE_OPENAI_API_KEY, + apiVersion: env.AZURE_OPENAI_API_VERSION, + }) + : null; + +const currentModelDisplay = + env.LLM_PROVIDER === "azure" + ? { providerName: "OpenAI", name: `Azure (${env.AZURE_OPENAI_DEPLOYMENT})` } + : { providerName: "Anthropic", name: "Claude Opus 4.5" }; /** * Custom tool component renderers for graph agent tools. @@ -46,7 +103,7 @@ const configValue: ConfigInput = { new Agent({ id: meta.slug, name: meta.name, - model: { providerName: "Anthropic", name: "Claude Opus 4.5" }, + model: currentModelDisplay, color: meta.color as HexColor, toolComponents: graphToolComponents, createTransport: async ({ transportOptions }) => { @@ -55,8 +112,13 @@ const configValue: ConfigInput = { loader, )) as ToolSet; + const model = + env.LLM_PROVIDER === "azure" && azure + ? azure(env.AZURE_OPENAI_DEPLOYMENT) + : anthropic!("claude-opus-4-5"); + const agent = new ToolLoopAgent({ - model: anthropic("claude-opus-4-5"), + model, tools: graphTools, instructions: `${regularPrompt}\n\n${graphAgentPrompt}`, stopWhen: stepCountIs(50), @@ -69,6 +131,7 @@ const configValue: ConfigInput = { }, }), ) as ConfigInput["agents"], + commands: [{ name: "/models", hint: "Switch LLM model (set LLM_PROVIDER and restart)" }], appName: { sections: [ {