From bd7a5dd00656f0978bf92059c553a74fd01b7261 Mon Sep 17 00:00:00 2001
From: Kastan Day <kastanvday@gmail.com>
Date: Fri, 21 Jun 2024 14:56:47 -0700
Subject: [PATCH] Working Ollama chat for 'generateText' but not 'streamText'
 json parsing error!

---
 src/app/api/chat/ollama/route.ts   | 28 +++++++++++++---
 src/pages/api/models.ts            |  5 ++-
 src/utils/modelProviders/ollama.ts | 54 ++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 6 deletions(-)

diff --git a/src/app/api/chat/ollama/route.ts b/src/app/api/chat/ollama/route.ts
index 98f8b6c51..9654066f3 100644
--- a/src/app/api/chat/ollama/route.ts
+++ b/src/app/api/chat/ollama/route.ts
@@ -1,4 +1,6 @@
 import { createOllama } from 'ollama-ai-provider'
+import { streamText } from 'ai'
+
 // import { OllamaModel } from '~/types/OllamaProvider'
 
 const ollama = createOllama({
@@ -14,12 +16,28 @@ export async function POST(req: Request) {
   */
   console.log('In ollama POST endpoint')
 
-  // const res = ollama.chat('llama3:70b-instruct', )//, numCtx=8192)
-  const model = ollama('llama3:70b-instruct')
-  model.doStream
+  const messages = [
+    {
+      role: 'user',
+      content: 'why is the sky blue?',
+    },
+  ]
+
+  const result = await streamText({
+    maxRetries: 5,
+    maxTokens: 512,
+    model: ollama('llama3:70b-instruct'),
+    prompt: 'Invent a new holiday and describe its traditions.',
+    temperature: 0.3,
+  })
+
+  for await (const textPart of result.textStream) {
+    process.stdout.write(textPart)
+  }
 
-  // return new StreamingTextResponse(res.doStream)
-  // return res.doStream
+  console.log()
+  console.log('Token usage:', await result.usage)
+  console.log('Finish reason:', await result.finishReason)
 }
 
 export async function GET() {
diff --git a/src/pages/api/models.ts b/src/pages/api/models.ts
index e95ae51dd..a2c7f20b2 100644
--- a/src/pages/api/models.ts
+++ b/src/pages/api/models.ts
@@ -9,7 +9,7 @@ import {
 import { OpenAIModel, OpenAIModelID, OpenAIModels } from '@/types/openai'
 import { decrypt, isEncrypted } from '~/utils/crypto'
 import { LLMProvider, ProviderNames } from '~/types/LLMProvider'
-import { getOllamaModels } from '~/utils/modelProviders/ollama'
+import { getOllamaModels, runOllamaChat } from '~/utils/modelProviders/ollama'
 
 export const config = {
   runtime: 'edge',
@@ -36,6 +36,9 @@ const handler = async (req: Request): Promise<Response> => {
     const ollamaModels = await getOllamaModels()
     console.log('Ollama Models in models.ts: ', ollamaModels)
 
+    // Test chat function
+    const ret = await runOllamaChat()
+
     // Iterate over the providers, check if their key works. Return all available models...
     // each model provider should have at least `/chat` and `/models` endpoints
 
diff --git a/src/utils/modelProviders/ollama.ts b/src/utils/modelProviders/ollama.ts
index 6e8e96572..ebd291321 100644
--- a/src/utils/modelProviders/ollama.ts
+++ b/src/utils/modelProviders/ollama.ts
@@ -1,3 +1,6 @@
+import { generateText, streamText } from 'ai'
+import { createOllama } from 'ollama-ai-provider'
+
 export interface OllamaModel {
   id: string
   name: string
@@ -5,6 +8,57 @@ export interface OllamaModel {
   tokenLimit: number
 }
 
+export const runOllamaChat = async () => {
+  console.log('In ollama POST endpoint')
+
+  const ollama = createOllama({
+    // custom settings
+    baseURL: 'https://ollama.ncsa.ai/api',
+  })
+
+  const messages = [
+    {
+      role: 'user',
+      content: 'why is the sky blue?',
+    },
+  ]
+
+  const result = await generateText({
+    maxTokens: 1024,
+    model: ollama('llama3:70b-instruct'),
+    prompt: 'Invent a new holiday and describe its traditions.',
+  })
+
+  console.log('OLLAMA RESULT', result.text)
+
+  // TODO: Check out the server example for how to handle streaming responses
+  // https://sdk.vercel.ai/examples/next-app/chat/stream-chat-completion#server
+
+  // const result = await streamText({
+  //     maxRetries: 5,
+  //     maxTokens: 512,
+  //     model: ollama('llama3:70b-instruct'),
+  //     messages: messages,
+  //     temperature: 0.3,
+  // })
+
+  // let fullResponse = '';
+  // for await (const textPart of result.textStream) {
+  //     fullResponse += textPart;
+  // }
+
+  // try {
+  //     const parsedResponse = JSON.parse(fullResponse);
+  //     console.log(parsedResponse);
+  // } catch (error) {
+  //     console.error('Failed to parse JSON:', error);
+  // }
+
+  // console.log()
+  // console.log('Token usage:', await result.usage)
+  // console.log('Finish reason:', await result.finishReason)
+}
+
 export const getOllamaModels = async () => {
   const response = await fetch('https://ollama.ncsa.ai/api/ps')
   if (!response.ok) {