fix: use open router and remove recalc of token

.github/workflows/compute.yml

@@ -28,7 +28,7 @@ jobs:
     environment: ${{ github.ref == 'refs/heads/main' && 'main' || 'development' }}
     env:
       KERNEL_PUBLIC_KEY: ${{ secrets.KERNEL_PUBLIC_KEY }}
-      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+      OPENROUTER_API_KEY: ${{ secrets.OPENROUTER }}
       UBIQUITY_OS_APP_NAME: ${{ secrets.UBIQUITY_OS_APP_NAME }}
 
     steps:

.github/workflows/worker-deploy.yml

@@ -45,11 +45,11 @@ jobs:
           apiToken: ${{ secrets.CLOUDFLARE_API_TOKEN }}
           accountId: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }}
           secrets: |
-            ANTHROPIC_API_KEY
+            OPENROUTER_API_KEY
             UBIQUITY_OS_APP_NAME
             ${{ secrets.KERNEL_PUBLIC_KEY && secrets.KERNEL_PUBLIC_KEY != '' && 'KERNEL_PUBLIC_KEY' || '' }}
         env:
-          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY}}
+          OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY}}
           UBIQUITY_OS_APP_NAME: ${{ secrets.UBIQUITY_OS_APP_NAME }}
           KERNEL_PUBLIC_KEY: ${{ secrets.KERNEL_PUBLIC_KEY }}
 

README.md

@@ -19,16 +19,16 @@ plugins:
   - uses:
       - plugin: http://localhost:4000
         with:
-          anthropicAiModel: "" # Optional - defaults to latest Claude model
-          anthropicAiBaseUrl: "" # Optional - defaults to Anthropic's API endpoint
+          openRouterAiModel: "" # Optional - defaults to "anthropic/claude-3.5-sonnet"
+          openRouterBaseUrl: "" # Optional - defaults to Open Router's API endpoint
 ```
 
 `.dev.vars` (for local testing):
 
-specify the AnthropicAiBase URL in the `.ubiquity-os.config.yml` file and set the `ANTHROPIC_API_KEY` in the `.dev.vars` file.
+specify the OpenRouterBase URL in the `.ubiquity-os.config.yml` file and set the `OPENROUTER_API_KEY` in the `.dev.vars` file.
 
 ```sh
-ANTHROPIC_API_KEY=your_anthropic_api_key
+OPENROUTER_API_KEY=YOUR_OPENROUTER_API_KEY
 UBIQUITY_OS_APP_NAME="UbiquityOS"
 ```
 

package.json

@@ -27,14 +27,14 @@
     "open-source"
   ],
   "dependencies": {
-    "@anthropic-ai/sdk": "^0.32.1",
     "@octokit/graphql-schema": "^15.25.0",
     "@sinclair/typebox": "0.34.3",
     "@ubiquity-os/plugin-sdk": "^1.1.0",
     "@ubiquity-os/ubiquity-os-logger": "^1.3.2",
     "dotenv": "^16.4.5",
     "gpt-tokenizer": "^2.5.1",
-    "minimatch": "^10.0.1"
+    "minimatch": "^10.0.1",
+    "openai": "^4.77.3"
   },
   "devDependencies": {
     "@commitlint/cli": "^19.5.0",

src/adapters/index.ts

@@ -1,13 +1,13 @@
 import { Context } from "../types";
-import { AnthropicCompletion } from "./claude/helpers/completions";
-import Anthropic from "@anthropic-ai/sdk";
-import { SuperAnthropic } from "./claude/helpers/claude";
+import { OpenRouterCompletion } from "./open-router/helpers/completions";
+import { SuperOpenRouter } from "./open-router/helpers/open-router";
+import OpenAI from "openai";
 
-export function createAdapters(anthropic: Anthropic, context: Context) {
+export function createAdapters(openRouter: OpenAI, context: Context) {
   return {
-    anthropic: {
-      completions: new AnthropicCompletion(anthropic, context),
-      super: new SuperAnthropic(anthropic, context),
+    openRouter: {
+      completions: new OpenRouterCompletion(openRouter, context),
+      super: new SuperOpenRouter(openRouter, context),
     },
   };
 }

src/adapters/claude/helpers/completions.ts → src/adapters/open-router/helpers/completions.ts

@@ -1,35 +1,19 @@
-import Anthropic from "@anthropic-ai/sdk";
 import { Context } from "../../../types";
-import { ContentBlock } from "@anthropic-ai/sdk/resources";
-import { SuperAnthropic } from "./claude";
+import { SuperOpenRouter } from "./open-router";
+import OpenAI from "openai";
 
 export interface CompletionsType {
   answer: string;
   groundTruths: string[];
-  tokenUsage: {
-    input: number;
-    output: number;
-    total: number;
-  };
 }
 
-// Type guard for content block
-interface TextBlock {
-  type: "text";
-  text: string;
-}
-
-function isTextBlock(content: ContentBlock): content is TextBlock {
-  return content?.type === "text" && typeof content?.text === "string";
-}
-
-export class AnthropicCompletion extends SuperAnthropic {
-  constructor(client: Anthropic, context: Context) {
+export class OpenRouterCompletion extends SuperOpenRouter {
+  constructor(client: OpenAI, context: Context) {
     super(client, context);
   }
 
   getModelMaxTokenLimit(model: string): number {
-    const tokenLimits = new Map<string, number>([["claude-3.5-sonnet", 200000]]);
+    const tokenLimits = new Map<string, number>([["anthropic/claude-3.5-sonnet", 200000]]);
     const tokenLimit = tokenLimits.get(model);
     if (!tokenLimit) {
       throw this.context.logger.error(`The token limits for configured model ${model} was not found`);
@@ -38,7 +22,7 @@ export class AnthropicCompletion extends SuperAnthropic {
   }
 
   getModelMaxOutputLimit(model: string): number {
-    const tokenLimits = new Map<string, number>([["claude-3.5-sonnet", 4096]]);
+    const tokenLimits = new Map<string, number>([["anthropic/claude-3.5-sonnet", 4096]]);
     const tokenLimit = tokenLimits.get(model);
     if (!tokenLimit) {
       throw this.context.logger.error(`The token limits for configured model ${model} was not found`);
@@ -60,10 +44,13 @@ export class AnthropicCompletion extends SuperAnthropic {
 
     this.context.logger.info(`System message: ${sysMsg}`);
 
-    const res = await this.client.messages.create({
-      model: model,
-      system: sysMsg,
+    const res = await this.client.chat.completions.create({
+      model: `anthropic/${model}`,
       messages: [
+        {
+          role: "system",
+          content: sysMsg,
+        },
         {
           role: "user",
           content: query,
@@ -73,60 +60,60 @@ export class AnthropicCompletion extends SuperAnthropic {
       temperature: 0,
     });
 
-    if (!res.content || res.content.length === 0) {
+    if (!res.choices || res.choices.length === 0) {
       throw this.context.logger.error("Unexpected no response from LLM");
     }
 
     // Use type guard to safely handle the response
-    const content = res.content[0];
-    if (!isTextBlock(content)) {
+    const answer = res.choices[0].message.content;
+    if (!answer) {
       throw this.context.logger.error("Unexpected response format: Expected text block");
     }
 
-    const answer = content.text;
-
-    const inputTokens = res.usage.input_tokens;
-    const outputTokens = res.usage.output_tokens;
+    const inputTokens = res.usage?.prompt_tokens;
+    const outputTokens = res.usage?.completion_tokens;
 
-    this.context.logger.info(`Number of tokens used: ${inputTokens + outputTokens}`);
+    if (inputTokens && outputTokens) {
+      this.context.logger.info(`Number of tokens used: ${inputTokens + outputTokens}`);
+    } else {
+      this.context.logger.info(`LLM did not output usage statistics`);
+    }
 
     return {
       answer,
       groundTruths,
-      tokenUsage: {
-        input: inputTokens,
-        output: outputTokens,
-        total: inputTokens + outputTokens,
-      },
     };
   }
 
   async createGroundTruthCompletion(context: Context, groundTruthSource: string, systemMsg: string): Promise<string | null> {
     const {
-      config: { anthropicAiModel },
+      config: { openRouterAiModel },
     } = context;
 
-    const res = await this.client.messages.create({
-      model: anthropicAiModel,
-      system: systemMsg,
-      max_tokens: this.getModelMaxTokenLimit(anthropicAiModel),
+    const res = await this.client.chat.completions.create({
+      model: openRouterAiModel,
+      max_tokens: this.getModelMaxTokenLimit(openRouterAiModel),
       messages: [
+        {
+          role: "system",
+          content: systemMsg,
+        },
         {
           role: "user",
           content: groundTruthSource,
         },
       ],
     });
 
-    if (!res.content || res.content.length === 0) {
+    if (!res.choices || res.choices.length === 0) {
       throw this.context.logger.error("Unexpected no response from LLM");
     }
 
-    const content = res.content[0];
-    if (!isTextBlock(content)) {
+    const answer = res.choices[0].message.content;
+    if (!answer) {
       throw this.context.logger.error("Unexpected response format: Expected text block");
     }
 
-    return content.text;
+    return answer;
   }
 }

src/adapters/open-router/helpers/open-router.ts

@@ -0,0 +1,12 @@
+import { Context } from "../../../types/context";
+import OpenAI from "openai";
+
+export class SuperOpenRouter {
+  protected client: OpenAI;
+  protected context: Context;
+
+  constructor(client: OpenAI, context: Context) {
+    this.context = context;
+    this.client = client;
+  }
+}

src/handlers/ground-truths/find-ground-truths.ts

@@ -12,7 +12,7 @@ export async function findGroundTruths(context: Context, params: CodeReviewAppPa
 async function findCodeReviewTruths(context: Context, params: CodeReviewAppParams, systemMsgObj: CodeReviewGroundTruthSystemMessage): Promise<string[]> {
   const {
     adapters: {
-      anthropic: { completions },
+      openRouter: { completions },
     },
   } = context;
   const systemMsg = createGroundTruthSysMsg(systemMsgObj);

src/handlers/pull-reviewer.ts

@@ -151,9 +151,9 @@ export class PullReviewer {
   async reviewPull() {
     const {
       env: { UBIQUITY_OS_APP_NAME },
-      config: { anthropicAiModel },
+      config: { openRouterAiModel },
       adapters: {
-        anthropic: { completions },
+        openRouter: { completions },
       },
     } = this.context;
 
@@ -175,11 +175,11 @@ export class PullReviewer {
 
     const groundTruths = await findGroundTruths(this.context, { taskSpecification });
     return await completions.createCompletion(
-      anthropicAiModel,
+      openRouterAiModel,
       formattedSpecAndPull,
       groundTruths,
       UBIQUITY_OS_APP_NAME,
-      completions.getModelMaxTokenLimit(anthropicAiModel)
+      completions.getModelMaxTokenLimit(openRouterAiModel)
     );
   }
 

src/helpers/format-spec-and-pull.ts

@@ -7,8 +7,8 @@ import { fetchPullRequestDiff } from "./pull-helpers/fetch-diff";
 
 export async function formatSpecAndPull(context: Context<"pull_request.opened" | "pull_request.ready_for_review">, issue: Issue): Promise<string> {
   const tokenLimits: TokenLimits = {
-    modelMaxTokenLimit: context.adapters.anthropic.completions.getModelMaxTokenLimit(context.config.anthropicAiModel),
-    maxCompletionTokens: context.config.maxTokens || context.adapters.anthropic.completions.getModelMaxOutputLimit(context.config.anthropicAiModel),
+    modelMaxTokenLimit: context.adapters.openRouter.completions.getModelMaxTokenLimit(context.config.openRouterAiModel),
+    maxCompletionTokens: context.config.maxTokens || context.adapters.openRouter.completions.getModelMaxOutputLimit(context.config.openRouterAiModel),
     runningTokenCount: 0,
     tokensRemaining: 0,
   };

src/helpers/pull-helpers/pull-request-parsing.ts

@@ -5,61 +5,54 @@ import { Context } from "../../types";
 import { getExcludedFiles } from "../excluded-files";
 import { minimatch } from "minimatch";
 
-export async function processPullRequestDiff(diff: string, tokenLimits: TokenLimits, logger: Context["logger"]) {
-  const { runningTokenCount, tokensRemaining } = tokenLimits;
-
-  // parse the diff into per-file diffs for quicker processing
-  const excludedFilePatterns = await getExcludedFiles();
-  const perFileDiffs = parsePerFileDiffs(diff).filter((file) => excludedFilePatterns.every((pattern) => !minimatch(file.filename, pattern)));
+async function filterAndSortDiffs(diff: string, excludedPatterns: string[]): Promise<{ filename: string; tokenCount: number; diffContent: string }[]> {
+  const perFileDiffs = parsePerFileDiffs(diff).filter((file) => excludedPatterns.every((pattern) => !minimatch(file.filename, pattern)));
 
   const accurateFileDiffStats = await Promise.all(
     perFileDiffs.map(async (file) => {
       const tokenCountArray = await encodeAsync(file.diffContent, { disallowedSpecial: new Set() });
-      const tokenCount = tokenCountArray.length;
-      return { filename: file.filename, tokenCount, diffContent: file.diffContent };
+      return { filename: file.filename, tokenCount: tokenCountArray.length, diffContent: file.diffContent };
     })
   );
 
-  // Sort by token count to process smallest files first
-  accurateFileDiffStats.sort((a, b) => a.tokenCount - b.tokenCount);
+  // Sort files by token count in ascending order
+  return accurateFileDiffStats.sort((a, b) => a.tokenCount - b.tokenCount);
+}
 
+function selectIncludedFiles(
+  files: { filename: string; tokenCount: number; diffContent: string }[],
+  tokenLimits: TokenLimits,
+  logger: Context["logger"]
+): typeof files {
+  const { runningTokenCount, tokensRemaining } = tokenLimits;
   let currentTokenCount = runningTokenCount;
-  const includedFileDiffs = [];
+  const includedFiles = [];
 
-  // Include files until we reach the token limit
-  for (const file of accurateFileDiffStats) {
+  for (const file of files) {
     if (currentTokenCount + file.tokenCount > tokensRemaining) {
       logger.info(`Skipping ${file.filename} to stay within token limits.`);
       continue;
     }
-    includedFileDiffs.push(file);
+    includedFiles.push(file);
     currentTokenCount += file.tokenCount;
   }
 
-  // If no files can be included, return null
-  if (includedFileDiffs.length === 0) {
-    logger.error(`Cannot include any files from diff without exceeding token limits.`);
-    return { diff: null };
-  }
+  return includedFiles;
+}
 
-  // Recalculate the current token count after including the files
-  currentTokenCount = includedFileDiffs.reduce((sum, file) => sum + file.tokenCount, runningTokenCount);
+export async function processPullRequestDiff(diff: string, tokenLimits: TokenLimits, logger: Context["logger"]) {
+  const excludedFilePatterns = await getExcludedFiles();
+  const sortedDiffs = await filterAndSortDiffs(diff, excludedFilePatterns);
 
-  // Remove files from the end of the list until we are within token limits
-  while (currentTokenCount > tokensRemaining && includedFileDiffs.length > 0) {
-    const removedFile = includedFileDiffs.pop();
-    currentTokenCount -= removedFile?.tokenCount || 0;
-    logger.info(`Excluded ${removedFile?.filename || "Unknown filename"} after accurate token count exceeded limits.`);
-  }
+  const includedFiles = selectIncludedFiles(sortedDiffs, tokenLimits, logger);
 
-  if (includedFileDiffs.length === 0) {
-    logger.error(`Cannot include any files from diff after accurate token count calculation.`);
+  if (includedFiles.length === 0) {
+    logger.error(`Cannot include any files from diff without exceeding token limits.`);
     return { diff: null };
   }
 
-  // Build the diff with the included files
-  const currentDiff = includedFileDiffs.map((file) => file.diffContent).join("\n");
-
+  // Build and return the final diff
+  const currentDiff = includedFiles.map((file) => file.diffContent).join("\n");
   return { diff: currentDiff };
 }