Add support for Ollama and support for OpenAI models directly (#65)

* Add OpenAI support to connector_creator.py

* Add Ollama support to connector_creator.py

Author: KamilFatkhiev

examples/connector_creator_usage_example.py

@@ -239,9 +239,11 @@ class Joke(BaseModel):
 
     # model_url_and_name = os.getenv("LLAMA_URL")
     # model_url_and_name = os.getenv("GIGACHAT_URL")
-    model_url_and_name = os.getenv("DEEPSEEK_URL")
+    # model_url_and_name = os.getenv("DEEPSEEK_URL")
     # model_url_and_name = os.getenv("DEEPSEEK_R1_URL")
     # model_url_and_name = os.getenv("GPT4_URL")
+    # model_url_and_name = os.getenv("OPENAI_URL")
+    model_url_and_name = os.getenv("OLLAMA_URL")
 
     # Uncomment the example you want to run
     basic_call_example(model_url_and_name)

poetry.lock

@@ -21,12 +21,14 @@ model name with a semicolon (;), for example: `https://api.vsegpt.ru/v1;openai/g
 It is also possible to pass additional parameters for the model.  Available parameters:
 - `temperature`
 - `top_p`
+- `max_tokens`
 
-Before use, make sure that your config file has the necessary API key (`VSEGPT_KEY` by default), or in the case of
-Gigachat models, an authorisation key (`AUTHORIZATION_KEY`), which can be obtained from your personal account.
+Before use, make sure that your config file has the necessary API key (`VSE_GPT_KEY` by default or `OPENAI_KEY`), or in
+the case of Gigachat models, an authorisation key (`AUTHORIZATION_KEY`), which can be obtained from your personal 
+account.
 
 Example of how to use the function:
-```commandline
+```codeblock
 from protollm.connectors.connector_creator import create_llm_connector
 
 model = create_llm_connector("https://api.vsegpt.ru/v1;openai/gpt-4o-mini", temperature=0.015, top_p=0.95)
@@ -35,8 +37,6 @@ print(res.content)
 ```
 The rest of the examples are located in the `examples/connector_creator_usage_examples.py` module of the repository.
 
-
-
 ## New connectors
 
 For now connectors are available for services supporting the OpenAI API format, as well as for Gigachat family models.

protollm/connectors/README.md

@@ -8,6 +8,7 @@
 from langchain_core.tools import BaseTool
 from langchain_core.runnables import Runnable
 from langchain_gigachat import GigaChat
+from langchain_ollama import ChatOllama
 from langchain_openai import ChatOpenAI
 from pydantic import BaseModel, ValidationError
 
@@ -254,7 +255,7 @@ def _handle_system_prompt(msgs, sys_prompt):
         return msgs
 
 
-def create_llm_connector(model_url: str, *args: Any, **kwargs: Any) -> CustomChatOpenAI | GigaChat:
+def create_llm_connector(model_url: str, *args: Any, **kwargs: Any) -> CustomChatOpenAI | GigaChat | ChatOpenAI:
     """Creates the proper connector for a given LLM service URL.
 
     Args:
@@ -263,6 +264,8 @@ def create_llm_connector(model_url: str, *args: Any, **kwargs: Any) -> CustomCha
             - for Gigachat models family: 'https://gigachat.devices.sberbank.ru/api/v1/chat/completions;Gigachat'
               for Gigachat model you should also install certificates from 'НУЦ Минцифры' -
               instructions - 'https://developers.sber.ru/docs/ru/gigachat/certificates'
+            - for OpenAI for example: 'https://api.openai.com/v1;gpt-4o'
+            - for Ollama for example: 'ollama;http://localhost:11434;llama3.2'
 
     Returns:
         The ChatModel object from 'langchain' that can be used to make requests to the LLM service,
@@ -276,6 +279,12 @@ def create_llm_connector(model_url: str, *args: Any, **kwargs: Any) -> CustomCha
         model_name = model_url.split(";")[1]
         access_token = get_access_token()
         return GigaChat(model=model_name, access_token=access_token, *args, **kwargs)
+    elif "api.openai" in model_url:
+        model_name = model_url.split(";")[1]
+        return ChatOpenAI(model=model_name, api_key=os.getenv("OPENAI_KEY"), *args, **kwargs)
+    elif "ollama" in model_url:
+        url_and_name = model_url.split(";")
+        return ChatOllama(model=url_and_name[2], base_url=url_and_name[1], *args, **kwargs)
     elif model_url == "test_model":
         return CustomChatOpenAI(model_name=model_url, api_key="test")
     else:

protollm/connectors/connector_creator.py

@@ -1,6 +1,6 @@
 [project]
 name = "ProtoLLM"
-version = "0.1.1"
+version = "0.1.2"
 description = "A library with which to prototype LLM-based applications quickly and easily."
 requires-python = ">=3.10,<4.0"
 authors = [
@@ -24,7 +24,7 @@ dependencies = [
     "langchain>=0.3.4,<0.4.0",
     "langchain-chroma==0.1.4",
     "langchain-community==0.3.16",
-    "langchain-core==0.3.34",
+    "langchain-core>=0.3.34",
     "langchain-elasticsearch==0.3.2",
     "langchain-gigachat==0.3.3",
     "langchain-openai==0.3.3",
@@ -55,7 +55,8 @@ dependencies = [
     "transformers==4.48.2",
     "urllib3>=2.2.2,<3.0.0",
     "uuid>=1.30,<2.0.0",
-    "websockets==14.1"
+    "websockets==14.1",
+    "langchain-ollama (==0.3.0)"
 ]
 
 [project.urls]

pyproject.toml

@@ -83,9 +83,10 @@ kombu==5.4.2 ; python_version >= "3.12" and python_version < "4.0" or python_ver
 kubernetes==32.0.1 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-chroma==0.1.4 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-community==0.3.16 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
-langchain-core==0.3.34 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
+langchain-core==0.3.48 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-elasticsearch==0.3.2 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-gigachat==0.3.3 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
+langchain-ollama==0.3.0 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-openai==0.3.3 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain-text-splitters==0.3.6 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 langchain==0.3.18 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
@@ -125,6 +126,7 @@ networkx==3.4.2 ; python_version >= "3.12" and python_version < "3.14" or python
 nltk==3.9.1 ; python_version >= "3.12" and python_version < "3.14" or python_version >= "3.10" and python_version <= "3.11"
 numpy==1.26.4 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 oauthlib==3.2.2 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
+ollama==0.4.7 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 onnxruntime==1.21.0 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 openai==1.61.1 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"
 opentelemetry-api==1.31.0 ; python_version >= "3.12" and python_version < "4.0" or python_version >= "3.10" and python_version <= "3.11"

requirements.txt

@@ -1,8 +1,10 @@
 from unittest.mock import patch
 
-from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from langchain_core.tools import tool
+from langchain_gigachat import GigaChat
+from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
 from pydantic import BaseModel, Field
 import pytest
 
@@ -307,16 +309,53 @@ def test_structured_output_dict_out_of_the_box(custom_chat_openai_with_fc_and_so
         assert result["age"] == 30
 
 
-@pytest.mark.parametrize(
-    "model_url",
-    [
-        "https://api.vsegpt.ru/v1;openai/gpt-4o-mini",
-        "https://gigachat.devices.sberbank.ru/api/v1/chat/completions;GigaChat",
-        "test_model",
-        "https://example.com/v1;test/example_model"
-    ]
-)
-def test_connector_creator(model_url):
-    with pytest.raises(Exception):
-        connector = create_llm_connector(model_url)
-        assert issubclass(connector, BaseChatModel)
+def test_vsegpt_connector(monkeypatch):
+    model_url = "https://api.vsegpt.ru/v1;meta-llama/llama-3.1-70b-instruct"
+    test_api_key = "test_vsegpt_key"
+    monkeypatch.setenv("VSE_GPT_KEY", test_api_key)
+    connector = create_llm_connector(model_url)
+    assert isinstance(connector, CustomChatOpenAI)
+
+
+@patch("protollm.connectors.connector_creator.get_access_token", return_value="test_gigachat_token")
+def test_gigachat_connector(mock_get_token):
+    model_url = "https://gigachat.devices.sberbank.ru/api/v1;Gigachat"
+    connector = create_llm_connector(model_url)
+    assert isinstance(connector, GigaChat)
+
+
+def test_openai_connector(monkeypatch):
+    model_url = "https://api.openai.com/v1;gpt-4o"
+    test_api_key = "test_openai_key"
+    monkeypatch.setenv("OPENAI_KEY", test_api_key)
+    connector = create_llm_connector(model_url)
+    assert isinstance(connector, ChatOpenAI)
+
+
+def test_ollama_connector():
+    model_url = "ollama;http://localhost:11434;llama3.2"
+    connector = create_llm_connector(model_url)
+    assert isinstance(connector, ChatOllama)
+
+
+def test_test_model_connector():
+    model_url = "test_model"
+    connector = create_llm_connector(model_url)
+    assert isinstance(connector, CustomChatOpenAI)
+
+
+def test_unsupported_provider():
+    model_url = "https://unknown.provider/v1;some-model"
+    with pytest.raises(ValueError) as exc_info:
+        create_llm_connector(model_url)
+    assert "Unsupported provider URL" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("invalid_url", [
+    "invalid_url_without_semicolon",
+    "https://api.vsegpt.ru/v1",
+    ";;",
+])
+def test_invalid_url_format(invalid_url):
+    with pytest.raises(ValueError):
+        create_llm_connector(invalid_url)