websearch-tools.ts 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. import {
  2. Ollama,
  3. type Message,
  4. type WebSearchResponse,
  5. type WebFetchResponse,
  6. } from 'ollama'
  7. async function main() {
  8. // Set enviornment variable OLLAMA_API_KEY=<YOUR>.<KEY>
  9. // or set the header manually
  10. // const client = new Ollama({
  11. // headers: { Authorization: `Bearer ${process.env.OLLAMA_API_KEY}` },
  12. // })
  13. const client = new Ollama()
  14. // Tool schemas
  15. const webSearchTool = {
  16. type: 'function',
  17. function: {
  18. name: 'webSearch',
  19. description: 'Performs a web search for the given query.',
  20. parameters: {
  21. type: 'object',
  22. properties: {
  23. query: { type: 'string', description: 'Search query string.' },
  24. max_results: {
  25. type: 'number',
  26. description: 'The maximum number of results to return per query (default 3).',
  27. },
  28. },
  29. required: ['query'],
  30. },
  31. },
  32. }
  33. const webFetchTool = {
  34. type: 'function',
  35. function: {
  36. name: 'webFetch',
  37. description: 'Fetches a single page by URL.',
  38. parameters: {
  39. type: 'object',
  40. properties: {
  41. url: { type: 'string', description: 'A single URL to fetch.' },
  42. },
  43. required: ['url'],
  44. },
  45. },
  46. }
  47. const availableTools = {
  48. webSearch: async (args: {
  49. query: string
  50. max_results?: number
  51. }): Promise<WebSearchResponse> => {
  52. const res = await client.webSearch(args)
  53. return res as WebSearchResponse
  54. },
  55. webFetch: async (args: { url: string }): Promise<WebFetchResponse> => {
  56. const res = await client.webFetch(args)
  57. return res as WebFetchResponse
  58. },
  59. }
  60. const query = 'What is Ollama?'
  61. console.log('Prompt:', query, '\n')
  62. const messages: Message[] = [
  63. {
  64. role: 'user',
  65. content: query,
  66. },
  67. ]
  68. while (true) {
  69. const response = await client.chat({
  70. model: 'qwen3',
  71. messages: messages,
  72. tools: [webSearchTool, webFetchTool],
  73. stream: true,
  74. think: true,
  75. })
  76. let hadToolCalls = false
  77. var content = ''
  78. var thinking = ''
  79. for await (const chunk of response) {
  80. if (chunk.message.thinking) {
  81. thinking += chunk.message.thinking
  82. }
  83. if (chunk.message.content) {
  84. content += chunk.message.content
  85. process.stdout.write(chunk.message.content)
  86. }
  87. if (chunk.message.tool_calls && chunk.message.tool_calls.length > 0) {
  88. hadToolCalls = true
  89. messages.push({
  90. role: 'assistant',
  91. content: content,
  92. thinking: thinking,
  93. tool_calls: chunk.message.tool_calls,
  94. })
  95. // Execute tools and append tool results
  96. for (const toolCall of chunk.message.tool_calls) {
  97. const functionToCall = availableTools[toolCall.function.name]
  98. if (functionToCall) {
  99. const args = toolCall.function.arguments as any
  100. console.log(
  101. '\nCalling function:',
  102. toolCall.function.name,
  103. 'with arguments:',
  104. args,
  105. )
  106. const output = await functionToCall(args)
  107. console.log('Function result:', JSON.stringify(output).slice(0, 200), '\n')
  108. messages.push(chunk.message)
  109. messages.push({
  110. role: 'tool',
  111. content: JSON.stringify(output).slice(0, 2000 * 4), // cap at ~2000 tokens
  112. tool_name: toolCall.function.name,
  113. })
  114. }
  115. }
  116. }
  117. }
  118. if (!hadToolCalls) {
  119. process.stdout.write('\n')
  120. break
  121. }
  122. }
  123. }
  124. main().catch(console.error)