feat(utils): integrate ollama to Chat

Author: remarkablemark

AGENTS.md

@@ -42,10 +42,9 @@ Single-test examples:
 ## Standards
 
 - TypeScript is `strict`; avoid implicit `any`
-- Tests use Vitest globals (do not import `vitest` except for types)
-- Enforce 100% test coverage; use `// v8 ignore` to exclude unreachable entrypoint guards
-- Use Conventional Commits: `type(scope): description`
-- Create a PR with `.github/PULL_REQUEST_TEMPLATE.md`
+- Use barrel files (`index.ts`) to consolidate related exports
+- Use `// v8 ignore` in tests to exclude unreachable entrypoint guards; use `vi.hoisted()` for mock variables accessed by `vi.mock()` hoisted scopes
+- Create PR with `.github/PULL_REQUEST_TEMPLATE.md`
 
 ## Verification
 

eslint.config.mts

@@ -43,9 +43,35 @@ export default defineConfig([
 
     rules: {
       '@typescript-eslint/no-extra-semi': 'off',
+      '@typescript-eslint/no-restricted-imports': [
+        'error',
+        {
+          paths: [
+            {
+              name: 'vitest',
+              allowTypeImports: true,
+              message: 'Use Vitest globals instead of importing from vitest',
+            },
+          ],
+          patterns: [
+            {
+              group: [
+                '**/*.js',
+                '**/*.jsx',
+                '**/*.ts',
+                '**/*.tsx',
+                '**/*.mjs',
+                '**/*.cjs',
+              ],
+              message: 'Do not use file extensions in import paths',
+            },
+          ],
+        },
+      ],
       '@typescript-eslint/no-unused-vars': 'error',
       'no-console': 'error',
       'no-debugger': 'error',
+      'no-restricted-imports': 'off',
       'prettier/prettier': 'error',
       'simple-import-sort/exports': 'error',
       'simple-import-sort/imports': 'error',

package-lock.json

@@ -42,6 +42,7 @@
     "cac": "7.0.0",
     "ink": "7.0.1",
     "ink-text-input": "6.0.0",
+    "ollama": "0.6.3",
     "react": "19.2.5"
   },
   "devDependencies": {

package.json

@@ -1,4 +1,4 @@
-import { type MockInstance, vi } from 'vitest';
+import type { MockInstance } from 'vitest';
 
 const { clearScreen, outputHelp, parse, render } = vi.hoisted(() => ({
   clearScreen: vi.fn(),

src/cli.test.ts

@@ -2,9 +2,17 @@ import { render } from 'ink-testing-library';
 
 import { Chat } from './Chat';
 
+vi.mock('../utils/ollama', () => ({
+  streamChat: vi.fn().mockImplementation(function* () {
+    yield 'Mocked';
+    yield ' response';
+  }),
+}));
+
 const ENTER = '\r';
 
-const tick = () => new Promise<void>((resolve) => setTimeout(resolve, 0));
+const tick = (ms = 0) =>
+  new Promise<void>((resolve) => setTimeout(resolve, ms));
 
 type Stdin = ReturnType<typeof render>['stdin'];
 
@@ -15,6 +23,11 @@ async function typeText(stdin: Stdin, text: string) {
   }
 }
 
+async function waitForStream() {
+  // Allow time for async generator to yield values
+  await tick(10);
+}
+
 describe('Chat', () => {
   it('renders input prompt', () => {
     const { lastFrame } = render(<Chat />);
@@ -25,19 +38,20 @@ describe('Chat', () => {
     const { lastFrame, stdin } = render(<Chat />);
     await typeText(stdin, 'hello');
     stdin.write(ENTER);
-    await tick();
+    await waitForStream();
     expect(lastFrame()).toContain('hello');
   });
 
   it('clears input after submit', async () => {
     const { lastFrame, stdin } = render(<Chat />);
     await typeText(stdin, 'hello');
     stdin.write(ENTER);
-    await tick();
+    await waitForStream();
     const frame = lastFrame() ?? '';
-    const inputLine =
-      frame.split('\n').find((line) => line.includes('>')) ?? '';
-    expect(inputLine.replace('>', '').trim()).toBe('');
+    // Find the last line that contains just the prompt (no user text after >)
+    const lines = frame.split('\n');
+    const inputLine = lines.find((line) => line.trim() === '>') ?? '';
+    expect(inputLine.trim()).toBe('>');
   });
 
   it('does not add blank messages', async () => {
@@ -56,14 +70,51 @@ describe('Chat', () => {
     const { lastFrame, stdin } = render(<Chat />);
     await typeText(stdin, 'first');
     stdin.write(ENTER);
-    await tick();
+    await waitForStream();
     await typeText(stdin, 'second');
     stdin.write(ENTER);
-    await tick();
+    await waitForStream();
     const frame = lastFrame() ?? '';
     const firstIdx = frame.indexOf('first');
     const secondIdx = frame.indexOf('second');
     expect(firstIdx).toBeGreaterThanOrEqual(0);
     expect(secondIdx).toBeGreaterThan(firstIdx);
   });
 });
+
+describe('Chat with error', () => {
+  it('shows error message when stream fails with Error', async () => {
+    const { streamChat } = await import('../utils/ollama');
+    vi.mocked(streamChat).mockImplementationOnce(async function* () {
+      await Promise.resolve();
+      yield '';
+      throw new Error('Connection failed');
+    });
+
+    const { lastFrame, stdin } = render(<Chat />);
+
+    await typeText(stdin, 'hello');
+    stdin.write(ENTER);
+    await waitForStream();
+
+    expect(lastFrame()).toContain('Error: Connection failed');
+  });
+
+  it('shows error message when stream fails with non-Error', async () => {
+    const { streamChat } = await import('../utils/ollama');
+    vi.mocked(streamChat).mockImplementationOnce(async function* () {
+      await Promise.resolve();
+      yield '';
+      // eslint-disable-next-line @typescript-eslint/only-throw-error
+      throw { toString: () => 'Custom error' };
+    });
+
+    const { lastFrame, stdin } = render(<Chat />);
+
+    await typeText(stdin, 'hello');
+    stdin.write(ENTER);
+    await waitForStream();
+
+    expect(lastFrame()).toContain('Error: Custom error');
+  });
+});

src/components/Chat.test.tsx

@@ -1,29 +1,82 @@
 import { Box, Text } from 'ink';
 import TextInput from 'ink-text-input';
-import { useState } from 'react';
+import { useCallback, useState } from 'react';
+
+import { ollama } from '../utils';
 
 export function Chat() {
-  const [messages, setMessages] = useState<string[]>([]);
+  const [messages, setMessages] = useState<ollama.Message[]>([]);
   const [input, setInput] = useState('');
+  const [isLoading, setIsLoading] = useState(false);
+
+  const handleSubmit = useCallback(
+    async (value: string) => {
+      const userContent = value.trim();
+      if (!userContent) return;
+
+      setInput('');
+      setIsLoading(true);
 
-  function handleSubmit(value: string) {
-    if (value.trim()) {
-      setMessages((prev) => [...prev, value.trim()]);
-    }
-    setInput('');
-  }
+      const userMessage: ollama.Message = {
+        role: 'user',
+        content: userContent,
+      };
+      setMessages((prev) => [...prev, userMessage]);
+
+      const updatedMessages = [...messages, userMessage];
+      const assistantMessage: ollama.Message = {
+        role: 'assistant',
+        content: '',
+      };
+      setMessages((prev) => [...prev, assistantMessage]);
+
+      try {
+        for await (const chunk of ollama.streamChat(updatedMessages)) {
+          assistantMessage.content += chunk;
+          setMessages((prev) => {
+            const newMessages = [...prev];
+            newMessages[newMessages.length - 1] = { ...assistantMessage };
+            return newMessages;
+          });
+        }
+      } catch (error) {
+        assistantMessage.content = `Error: ${error instanceof Error ? error.message : String(error)}`;
+        setMessages((prev) => {
+          const newMessages = [...prev];
+          newMessages[newMessages.length - 1] = { ...assistantMessage };
+          return newMessages;
+        });
+      } finally {
+        setIsLoading(false);
+      }
+    },
+    [messages],
+  );
 
   return (
     <Box flexDirection="column">
       <Box flexDirection="column">
         {messages.map((message, index) => (
-          <Text key={index}>{message}</Text>
+          <Text key={index} color={message.role === 'user' ? 'green' : 'blue'}>
+            {message.role === 'user' ? '> ' : ''}
+            {message.content}
+          </Text>
         ))}
+        {isLoading && messages[messages.length - 1]?.content === '' && (
+          <Text color="yellow">...</Text>
+        )}
       </Box>
 
       <Box>
         <Text>&gt; </Text>
-        <TextInput value={input} onChange={setInput} onSubmit={handleSubmit} />
+        <TextInput
+          value={input}
+          onChange={setInput}
+          onSubmit={(value) => {
+            void handleSubmit(value);
+          }}
+          focus={!isLoading}
+        />
       </Box>
     </Box>
   );

src/components/Chat.tsx

@@ -1 +1,2 @@
+export * as ollama from './ollama';
 export * as screen from './screen';

src/utils/index.ts

@@ -0,0 +1,74 @@
+const { mockChat, mockList } = vi.hoisted(() => ({
+  mockChat: vi.fn(),
+  mockList: vi.fn(),
+}));
+
+vi.mock('ollama', () => ({
+  Ollama: class MockOllama {
+    chat(...args: unknown[]) {
+      return mockChat(...args) as Promise<AsyncIterable<unknown>>;
+    }
+
+    list(...args: unknown[]) {
+      return mockList(...args) as Promise<unknown>;
+    }
+  },
+}));
+
+import { listModels, streamChat } from './ollama';
+
+describe('ollama', () => {
+  beforeEach(() => {
+    mockChat.mockResolvedValue({
+      async *[Symbol.asyncIterator]() {
+        await Promise.resolve();
+        yield { message: { content: 'Hello' } };
+      },
+    });
+    mockList.mockResolvedValue({
+      models: [{ name: 'codellama' }, { name: 'llama2' }],
+    });
+  });
+
+  describe('streamChat', () => {
+    it('should yield content from stream', async () => {
+      const messages = [{ role: 'user' as const, content: 'hello' }];
+      const results: string[] = [];
+
+      for await (const chunk of streamChat(messages, 'codellama')) {
+        results.push(chunk);
+      }
+
+      expect(results).toEqual(['Hello']);
+    });
+
+    it('should skip chunks with empty content', async () => {
+      // Override mock to yield empty content first
+      mockChat.mockResolvedValueOnce({
+        async *[Symbol.asyncIterator]() {
+          await Promise.resolve();
+          yield { message: { content: '' } };
+          yield { message: { content: 'Non-empty' } };
+        },
+      });
+
+      // Need to re-import to get a fresh client with new mock
+      const { streamChat: streamChatWithEmpty } = await import('./ollama');
+      const messages = [{ role: 'user' as const, content: 'hello' }];
+      const results: string[] = [];
+
+      for await (const chunk of streamChatWithEmpty(messages, 'codellama')) {
+        results.push(chunk);
+      }
+
+      expect(results).toEqual(['Non-empty']);
+    });
+  });
+
+  describe('listModels', () => {
+    it('should return list of models', async () => {
+      const models = await listModels();
+      expect(models).toEqual(['codellama', 'llama2']);
+    });
+  });
+});