Add more language support, fix chat bugs

RussellCanfield · Jun 18, 2024 · 92631fb · 92631fb
1 parent 493bdc5
commit 92631fb
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 47 deletions.
diff --git a/package.json b/package.json
@@ -2,7 +2,7 @@
 	"name": "wing-man",
 	"displayName": "Wingman-AI",
 	"description": "Wingman - AI powered assistant to help you write your best code, we won't leave you hanging.",
-	"version": "0.4.1",
+	"version": "0.4.2",
 	"publisher": "WingMan",
 	"license": "MIT",
 	"authors": [

diff --git a/src/providers/utilities.ts b/src/providers/utilities.ts
@@ -23,6 +23,12 @@ export const supportedLanguages: vscode.DocumentSelector = [
 	{ scheme: "file", language: "shellscript" },
 	{ scheme: "file", language: "sh" },
 	{ scheme: "file", language: "bash" },
+	{ scheme: "file", language: "dockerfile" },
+	{ scheme: "file", language: "yaml" },
+	{ scheme: "file", language: "json" },
+	{ scheme: "file", language: "xml" },
+	{ scheme: "file", language: "markdown" },
+	{ scheme: "file", language: "powershell" },
 ];
 
 export async function getSymbolsFromOpenFiles() {

diff --git a/src/service/ollama/ollama.ts b/src/service/ollama/ollama.ts
@@ -471,34 +471,28 @@ ${prompt}`,
 		ragContent: string,
 		signal: AbortSignal
 	) {
-		let systemPrompt = this.chatModel!.ChatPrompt;
-
-		if (ragContent) {
-			systemPrompt += `Here's some additional information that may help you generate a more accurate response.
+		const systemMessage: OllamaChatMessage = {
+			role: "assistant",
+			content: !ragContent
+				? this.chatModel?.ChatPrompt!
+				: `Here's some additional information that may help you generate a more accurate response.
 Please determine if this information is relevant and can be used to supplement your response: 
-${ragContent}`;
-		}
+${ragContent}`,
+		};
 
-		systemPrompt = systemPrompt.replaceAll("\t", "");
+		const userMessage: OllamaChatMessage = {
+			role: "user",
+			content: prompt,
+		};
 
-		const messages: OllamaChatMessage[] = [
-			{
-				role: "assistant",
-				content: systemPrompt,
-			},
-		];
+		this.chatHistory.push(systemMessage, userMessage);
+
+		const messages: OllamaChatMessage[] = [];
 
 		if (this.chatHistory.length > 0) {
 			messages.push(...this.truncateChatHistory());
 		}
 
-		messages.push({
-			role: "user",
-			content: prompt,
-		});
-
-		this.chatHistory.push(messages[messages.length - 1]);
-
 		const chatPayload: OllamaChatRequest = {
 			model: this.settings?.chatModel!,
 			stream: true,
@@ -518,10 +512,17 @@ ${ragContent}`;
 			)}`
 		);
 
+		let lastAssistantMessage = "";
 		for await (const chunk of this.generate(chatPayload, signal)) {
-			this.chatHistory.push(chunk.message);
+			lastAssistantMessage += chunk.message;
 			yield chunk.message.content;
 		}
+		if (lastAssistantMessage?.trim()) {
+			this.chatHistory.push({
+				role: "assistant",
+				content: lastAssistantMessage,
+			});
+		}
 	}
 
 	public async genCodeDocs(

diff --git a/src/types/Message.ts b/src/types/Message.ts
@@ -4,7 +4,7 @@ export interface AppMessage {
 }
 
 export interface ChatMessage {
-	from: "Assistant" | "User";
+	from: "assistant" | "user";
 	message: string;
 	loading?: boolean;
 	context: CodeContext | undefined;

diff --git a/src/webview-ui/App.tsx b/src/webview-ui/App.tsx
@@ -5,7 +5,7 @@ import { AppMessage, ChatMessage, CodeContext } from "../types/Message";
 import ChatEntry from "./ChatEntry";
 import { ChatInput } from "./ChatInput";
 import { ChatResponseList } from "./ChatList";
-import { localMemState } from './utilities/localMemState';
+import { localMemState } from "./utilities/localMemState";
 import { vscode } from "./utilities/vscode";
 
 const Main = styled.main`
@@ -87,7 +87,7 @@ const App = () => {
 					return {
 						loading: true,
 						context: undefined,
-						from: "Assistant",
+						from: "assistant",
 						...activeMessage,
 						message: currentMessage,
 					} satisfies ChatMessage;
@@ -101,7 +101,7 @@ const App = () => {
 				setActiveMessage((activeMessage) => {
 					return {
 						loading: true,
-						from: "Assistant",
+						from: "assistant",
 						...activeMessage,
 						message: currentMessage,
 						context: currentContext,
@@ -113,7 +113,7 @@ const App = () => {
 					workspaceFolder: string;
 					theme: number;
 				};
-				localMemState.setState('theme', theme);
+				localMemState.setState("theme", theme);
 
 				activeWorkspace = workspaceFolder;
 
@@ -131,7 +131,7 @@ const App = () => {
 				break;
 			case "setTheme":
 				const newTheme = value as number;
-				localMemState.setState('theme', newTheme);
+				localMemState.setState("theme", newTheme);
 		}
 	};
 
@@ -142,7 +142,7 @@ const App = () => {
 			const newHistory: ChatMessage[] = [
 				...messages,
 				{
-					from: "Assistant",
+					from: "assistant",
 					message: tempMessage,
 					loading: false,
 					context: tempContext,
@@ -152,13 +152,23 @@ const App = () => {
 			return newHistory;
 		});
 
+		clearMessage();
+	};
+
+	const cancelAIResponse = () => {
+		clearMessage();
+		vscode.postMessage({
+			command: "cancel",
+		});
+	};
+
+	const clearMessage = () => {
 		setLoading(false);
-		setActiveMessage((message) => {
+		setActiveMessage(() => {
 			return {
-				from: "Assistant",
+				from: "assistant",
 				context: undefined,
 				message: "",
-				...message,
 				loading: false,
 			};
 		});
@@ -167,13 +177,6 @@ const App = () => {
 		currentContext = undefined;
 	};
 
-	const cancelAIResponse = () => {
-		commitMessageToHistory();
-		vscode.postMessage({
-			command: "cancel",
-		});
-	};
-
 	const fetchAIResponse = (text: string) => {
 		currentMessage = "";
 
@@ -189,7 +192,7 @@ const App = () => {
 		setMessages((messages) => [
 			...messages,
 			{
-				from: "User",
+				from: "user",
 				message: input,
 				context: undefined,
 			},
@@ -237,7 +240,7 @@ const App = () => {
 			<ChatResponseList messages={messages}>
 				{loading && (
 					<ChatEntry
-						from="Assistant"
+						from="assistant"
 						message={activeMessage?.message || ""}
 						context={activeMessage?.context}
 						loading={loading}

diff --git a/src/webview-ui/ChatEntry.tsx b/src/webview-ui/ChatEntry.tsx
@@ -3,10 +3,13 @@ import { FaCopy } from "react-icons/fa6";
 import { GoFileSymlinkFile } from "react-icons/go";
 import Markdown from "react-markdown";
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
-import { prism, vscDarkPlus } from "react-syntax-highlighter/dist/esm/styles/prism";
+import {
+	prism,
+	vscDarkPlus,
+} from "react-syntax-highlighter/dist/esm/styles/prism";
 import styled, { keyframes } from "styled-components";
 import { ChatMessage } from "../types/Message";
-import { localMemState } from './utilities/localMemState';
+import { localMemState } from "./utilities/localMemState";
 import { vscode } from "./utilities/vscode";
 
 const Entry = styled.li`
@@ -206,14 +209,17 @@ const ChatEntry = ({
 		});
 	};
 
-	const lm = useSyncExternalStore(localMemState.subscribe, localMemState.getSnapshot);
-	const theme = lm['theme'] as number;
-	const codeTheme = (theme === 1 || theme === 4) ? prism : vscDarkPlus;
+	const lm = useSyncExternalStore(
+		localMemState.subscribe,
+		localMemState.getSnapshot
+	);
+	const theme = lm["theme"] as number;
+	const codeTheme = theme === 1 || theme === 4 ? prism : vscDarkPlus;
 
 	return (
 		<Entry>
 			<LabelContainer>
-				<h3>{from === "User" ? "Me" : "Wingman"}</h3>
+				<h3>{from === "user" ? "Me" : "Wingman"}</h3>
 				{loading && <Loader />}
 			</LabelContainer>
 			<div>