connection tree view (#616)

* show connection tree

* ui to analyze connections

* fix token counting in vscode llms

* tweak pr descr

packages/cli/src/nodehost.ts

@@ -152,6 +152,7 @@ export class NodeHost implements RuntimeHost {
         const { signal, token: askToken } = options || {}
         await this.parseDefaults()
         const tok = await parseTokenFromEnv(process.env, modelId)
+        if (!askToken && tok?.token) tok.token = "***"
         if (
             askToken &&
             tok &&

packages/core/src/connection.ts

@@ -331,9 +331,9 @@ OPENAI_API_TYPE="localai"
     if (provider === MODEL_PROVIDER_GITHUB)
         return {
             config: `
-    ## GitHub Models ${DOCS_CONFIGURATION_GITHUB_URL}
-    # use "${MODEL_PROVIDER_GITHUB}:<model>" in script({ model: ... })
-    # GITHUB_MODELS_TOKEN="${PLACEHOLDER_API_KEY}" # use a personal access token if not available
+## GitHub Models ${DOCS_CONFIGURATION_GITHUB_URL}
+# use "${MODEL_PROVIDER_GITHUB}:<model>" in script({ model: ... })
+# GITHUB_MODELS_TOKEN="${PLACEHOLDER_API_KEY}" # use a personal access token if not available
     `,
             model: `${MODEL_PROVIDER_GITHUB}:gpt-4o`,
         }

packages/sample/genaisrc/pr-review.genai.js

@@ -41,4 +41,5 @@ If the changes look good, respond "LGTM :rocket:". If you have any concerns, pro
 - only report functional issues
 - Use emojis
 - If available, suggest code fixes and improvements using a diff format.
+- do not report about individual lines of code, summarize changes
 `

packages/vscode/package.json

@@ -122,6 +122,11 @@
                     "id": "genaiscript.trace",
                     "name": "Trace"
                 },
+                {
+                    "id": "genaiscript.connections",
+                    "name": "Connections",
+                    "visibility": "collapsed"
+                },
                 {
                     "id": "genaiscript.prompts",
                     "name": "Scripts",
@@ -345,6 +350,11 @@
                 "command": "genaiscript.info.env",
                 "title": "Configuration information...",
                 "category": "GenAIScript"
+            },
+            {
+                "command": "genaiscript.connection.configure",
+                "title": "Configure connection...",
+                "category": "GenAIScript"
             }
         ]
     },

packages/vscode/src/connectioninfotree.ts

@@ -0,0 +1,86 @@
+import * as vscode from "vscode"
+import { ExtensionState } from "./state"
+import { MODEL_PROVIDERS } from "../../core/src/constants"
+import { YAMLStringify } from "../../core/src/yaml"
+import { APIType } from "../../core/src/host"
+
+class ConnectionInfoTreeData {
+    provider: string
+    apiType?: APIType
+}
+
+class ConnectionInfoTreeDataProvider
+    implements vscode.TreeDataProvider<ConnectionInfoTreeData>
+{
+    constructor(readonly state: ExtensionState) {
+        const { context } = state
+        const { subscriptions } = context
+        subscriptions.push(
+            vscode.workspace.onDidChangeConfiguration(() => {
+                this.refresh()
+            })
+        )
+        const watcher = vscode.workspace.createFileSystemWatcher("./.env")
+        watcher.onDidChange(() => this.refresh())
+        watcher.onDidDelete(() => this.refresh())
+        subscriptions.push(watcher)
+    }
+
+    async getTreeItem(
+        element: ConnectionInfoTreeData
+    ): Promise<vscode.TreeItem> {
+        const item = new vscode.TreeItem(element.provider)
+        const res =
+            await this.state.host.server.client.getLanguageModelConfiguration(
+                element.provider + ":*",
+                { token: false }
+            )
+        if (res) {
+            item.description = res.base || "?"
+            item.tooltip = YAMLStringify(res)
+            item.command = <vscode.Command>{
+                command: "vscode.open",
+                arguments: [this.state.host.toUri("./.env")],
+            }
+        } else {
+            item.description = "not configured"
+            item.command = <vscode.Command>{
+                command: "genaiscript.connection.configure",
+                arguments: [element.provider, element.apiType],
+            }
+        }
+
+        return item
+    }
+    getChildren(
+        element?: ConnectionInfoTreeData
+    ): vscode.ProviderResult<ConnectionInfoTreeData[]> {
+        if (!element) return MODEL_PROVIDERS.map(({ id }) => ({ provider: id }))
+        return undefined
+    }
+
+    private _onDidChangeTreeData: vscode.EventEmitter<
+        void | ConnectionInfoTreeData | ConnectionInfoTreeData[]
+    > = new vscode.EventEmitter<
+        void | ConnectionInfoTreeData | ConnectionInfoTreeData[]
+    >()
+    readonly onDidChangeTreeData: vscode.Event<
+        void | ConnectionInfoTreeData | ConnectionInfoTreeData[]
+    > = this._onDidChangeTreeData.event
+
+    refresh(
+        treeItem?: ConnectionInfoTreeData | ConnectionInfoTreeData[]
+    ): void {
+        this._onDidChangeTreeData.fire(treeItem)
+    }
+}
+
+export function activateConnectionInfoTree(state: ExtensionState) {
+    const { context } = state
+    const { subscriptions } = context
+    const treeDataProvider = new ConnectionInfoTreeDataProvider(state)
+    const treeView = vscode.window.createTreeView("genaiscript.connections", {
+        treeDataProvider,
+    })
+    subscriptions.push(treeView)
+}

packages/vscode/src/extension.ts

@@ -15,20 +15,34 @@ import { registerCommand } from "./commands"
 import { EXTENSION_ID, TOOL_NAME } from "../../core/src/constants"
 import type MarkdownIt from "markdown-it"
 import MarkdownItGitHubAlerts from "markdown-it-github-alerts"
+import { activateConnectionInfoTree } from "./connectioninfotree"
+import { updateConnectionConfiguration } from "../../core/src/connection"
+import { APIType } from "../../core/src/host"
 
 export async function activate(context: ExtensionContext) {
     const state = new ExtensionState(context)
     activatePromptCommands(state)
     activateFragmentCommands(state)
     activateMarkdownTextDocumentContentProvider(state)
     activatePrompTreeDataProvider(state)
+    activateConnectionInfoTree(state)
     activateAIRequestTreeDataProvider(state)
     activateLLMRequestTreeDataProvider(state)
     activateTraceTreeDataProvider(state)
     activateStatusBar(state)
     activateDocsNotebook(state)
 
     context.subscriptions.push(
+        registerCommand(
+            "genaiscript.connection.configure",
+            async (provider?: string, apiType?: APIType) => {
+                await updateConnectionConfiguration(provider, apiType)
+                const doc = await vscode.workspace.openTextDocument(
+                    state.host.toUri("./.env")
+                )
+                await vscode.window.showTextDocument(doc)
+            }
+        ),
         registerCommand("genaiscript.request.abort", async () => {
             await state.cancelAiRequest()
             await vscode.window.showInformationMessage(

packages/vscode/src/lmaccess.ts

@@ -13,7 +13,6 @@ import {
 } from "../../core/src/constants"
 import { APIType } from "../../core/src/host"
 import { parseModelIdentifier } from "../../core/src/models"
-import { updateConnectionConfiguration } from "../../core/src/connection"
 import { ChatCompletionMessageParam } from "../../core/src/chattypes"
 import { LanguageModelChatRequest } from "../../core/src/server/client"
 import { ChatStart } from "../../core/src/server/messages"
@@ -143,11 +142,11 @@ export async function pickLanguageModel(
 
     if (res.model) return res.model
     else {
-        await updateConnectionConfiguration(res.provider, res.apiType)
-        const doc = await vscode.workspace.openTextDocument(
-            state.host.toUri("./.env")
+        await vscode.commands.executeCommand(
+            "genaiscript.connection.configure",
+            res.provider,
+            res.apiType
         )
-        await vscode.window.showTextDocument(doc)
         return undefined
     }
 }
@@ -226,7 +225,7 @@ export function createChatModelRunner(
                 text += fragment
                 onChunk({
                     chunk: fragment,
-                    tokens: await chatModel.countTokens(text),
+                    tokens: await chatModel.countTokens(fragment),
                     finishReason: undefined,
                     model: chatModel.id,
                 })