supporting local model for jupyter-ai with ollama #868
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -18,6 +18,7 @@ | |
| Union, | ||
| ) | ||
|
|
||
| import requests | ||
|
There was a problem hiding this comment. If the functions in your new class can be removed, then you will not need this line. |
||
| from jsonpath_ng import parse | ||
| from langchain.prompts import ( | ||
| ChatPromptTemplate, | ||
|
|
@@ -37,6 +38,7 @@ | |
| Bedrock, | ||
| GPT4All, | ||
| HuggingFaceEndpoint, | ||
| Ollama, | ||
| SagemakerEndpoint, | ||
| Together, | ||
| ) | ||
|
|
@@ -728,6 +730,84 @@ async def _acall(self, *args, **kwargs) -> Coroutine[Any, Any, str]: | |
| return await self._call_in_executor(*args, **kwargs) | ||
|
|
||
|
|
||
| class OllamaProvider(BaseProvider, Ollama): | ||
|
There was a problem hiding this comment. Please see a similar implementation in PR #646 with a much simpler class and see if all the functions defined in your file from line 760 onwards are necessary, or being called to produce the response. |
||
| id = "ollama" | ||
| name = "Ollama" | ||
| model_id_key = "model" | ||
| # List of available models | ||
| models = [ | ||
| "gemma", | ||
| "gemma2", | ||
| "llama2", | ||
| "llama3", | ||
| "llama3:70b", | ||
| "phi3", | ||
| "mistral", | ||
| "tinyllama", | ||
| "qwen2", | ||
| "qwen2:7b", | ||
| "qwen2:72b", | ||
| ] | ||
| # Default base URL for the API | ||
| base_url = "http://localhost:11434" | ||
| # Fields for configuring the provider | ||
| fields = [ | ||
| TextField( | ||
| key="ollama_api_base", label="Base API URL (optional)", format="text" | ||
| ), | ||
| ] | ||
|
|
||
| def __init__(self, **kwargs): | ||
| # Extract 'ollama_api_base' from kwargs, if not present, it returns None | ||
| ollama_api_base = kwargs.pop("ollama_api_base", None) | ||
| # Initialize the parent class with the remaining kwargs | ||
| super().__init__(**kwargs) | ||
| # Set the base_url to ollama_api_base if provided, otherwise use the default base_url | ||
| self.base_url = ollama_api_base or self.base_url | ||
| # Create a session for making requests | ||
| self.session = requests.Session() | ||
|
|
||
| def _send_request(self, endpoint: str, data: dict) -> dict: | ||
| """Send a POST request to the specified Ollama API endpoint.""" | ||
| url = f"{self.base_url}/{endpoint}" | ||
| print("url is : ", url) | ||
|
There was a problem hiding this comment. Do you need this print statement? If so, use |
||
| try: | ||
| response = self.session.post(url, json=data) | ||
| response.raise_for_status() | ||
| return response.json() | ||
| except requests.RequestException as e: | ||
| raise ValueError(f"Error during Ollama API call: {e}") from e | ||
|
|
||
| def _generate(self, prompt: str, model: str = "mistral", stop: list = None) -> str: | ||
| """Generate text using the /generate endpoint.""" | ||
| data = { | ||
| "model": model, | ||
| "prompt": prompt, | ||
| **({"stop": stop} if stop else {}), | ||
| } | ||
| response = self._send_request("api/generate", data) | ||
| return response.get("response", "") | ||
|
|
||
| def _chat(self, messages: list, model: str = "mistral") -> str: | ||
| """Chat using the /chat endpoint.""" | ||
| data = { | ||
| "model": model, | ||
| "messages": messages, | ||
| } | ||
| response = self._send_request("api/chat", data) | ||
| return response.get("response", "") | ||
|
|
||
| def _call(self, prompt: str = None, messages: list = None, **kwargs) -> str: | ||
| """Determine which API endpoint to use based on input and call it.""" | ||
| print(self.base_url) | ||
| if prompt is not None: | ||
| return self._generate(prompt=prompt, **kwargs) | ||
| elif messages is not None: | ||
| return self._chat(messages=messages, **kwargs) | ||
| else: | ||
| raise ValueError("Either 'prompt' or 'messages' must be provided.") | ||
|
Comment on lines
+781
to
+808
There was a problem hiding this comment. Is there a need for implementing the Ollama class methods? Ollama chat model provides these methods, forking the code might deviate from the official LangChain LLMs. There was a problem hiding this comment. Agree, I noted the same above, and I have also tested removing these functions that the functionality is not impaired. |
||
|
|
||
|
|
||
| class JsonContentHandler(LLMContentHandler): | ||
| content_type = "application/json" | ||
| accepts = "application/json" | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you very much for adding this PR, which also complements PR #646.
[1] It would be nice to add more documentation on
Ollama, directing the user to install it first and how, pointing to the necessaryOllamaresources. Specific examples ofOllamaCLI commands needed would help the user. See the note by @pedrogutobjj above for example. Hopefully, people who useOllamawith its CLI will know what to do, but the docs should reflect the requirement that theOllamalocal server should be kicked off withollama serve.[2] Since the number of models in
Ollamais very large, the drop down list can become unwieldy. Would it be possible to implementOllamain the same way as Huggingface is implemented with a input box for the model in chat settings?It would also need a similar approach for choice of Embedding Model.