You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

111 lines
2.8 KiB

11 months ago
  1. const Provider = require("./ai-provider.js");
  2. const InheritMultiple = require("./helpers/classes.js");
  3. const UnTooled = require("./helpers/untooled.js");
  4. const { Ollama } = require("ollama");
  5. /**
  6. * The agent provider for the Ollama provider.
  7. */
  8. class OllamaProvider extends InheritMultiple([Provider, UnTooled]) {
  9. model;
  10. constructor(config = {}) {
  11. const {
  12. // options = {},
  13. model = null,
  14. } = config;
  15. super();
  16. this._client = new Ollama({ host: process.env.OLLAMA_BASE_PATH });
  17. this.model = model;
  18. this.verbose = true;
  19. }
  20. get client() {
  21. return this._client;
  22. }
  23. async #handleFunctionCallChat({ messages = [] }) {
  24. const response = await this.client.chat({
  25. model: this.model,
  26. messages,
  27. options: {
  28. temperature: 0,
  29. },
  30. });
  31. return response?.message?.content || null;
  32. }
  33. /**
  34. * Create a completion based on the received messages.
  35. *
  36. * @param messages A list of messages to send to the API.
  37. * @param functions
  38. * @returns The completion.
  39. */
  40. async complete(messages, functions = null) {
  41. try {
  42. let completion;
  43. if (functions.length > 0) {
  44. const { toolCall, text } = await this.functionCall(
  45. messages,
  46. functions,
  47. this.#handleFunctionCallChat.bind(this)
  48. );
  49. if (toolCall !== null) {
  50. this.providerLog(`Valid tool call found - running ${toolCall.name}.`);
  51. this.deduplicator.trackRun(toolCall.name, toolCall.arguments);
  52. return {
  53. result: null,
  54. functionCall: {
  55. name: toolCall.name,
  56. arguments: toolCall.arguments,
  57. },
  58. cost: 0,
  59. };
  60. }
  61. completion = { content: text };
  62. }
  63. if (!completion?.content) {
  64. this.providerLog(
  65. "Will assume chat completion without tool call inputs."
  66. );
  67. const response = await this.client.chat({
  68. model: this.model,
  69. messages: this.cleanMsgs(messages),
  70. options: {
  71. use_mlock: true,
  72. temperature: 0.5,
  73. },
  74. });
  75. completion = response.message;
  76. }
  77. // The UnTooled class inherited Deduplicator is mostly useful to prevent the agent
  78. // from calling the exact same function over and over in a loop within a single chat exchange
  79. // _but_ we should enable it to call previously used tools in a new chat interaction.
  80. this.deduplicator.reset("runs");
  81. return {
  82. result: completion.content,
  83. cost: 0,
  84. };
  85. } catch (error) {
  86. throw error;
  87. }
  88. }
  89. /**
  90. * Get the cost of the completion.
  91. *
  92. * @param _usage The completion to get the cost for.
  93. * @returns The cost of the completion.
  94. * Stubbed since LMStudio has no cost basis.
  95. */
  96. getCost(_usage) {
  97. return 0;
  98. }
  99. }
  100. module.exports = OllamaProvider;