LangChain Integration

libs/manubot_ai_editor/models.py

@@ -5,7 +5,8 @@
 import time
 import json
 
-import openai
+from langchain_openai import OpenAI, ChatOpenAI
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 
 from manubot_ai_editor import env_vars
 
@@ -141,12 +142,13 @@ def __init__(
         super().__init__()
 
         # make sure the OpenAI API key is set
-        openai.api_key = openai_api_key
+        if openai_api_key is None:
+            # attempt to get the OpenAI API key from the environment, since one
+            # wasn't specified as an argument
+            openai_api_key = os.environ.get(env_vars.OPENAI_API_KEY, None)
 
-        if openai.api_key is None:
-            openai.api_key = os.environ.get(env_vars.OPENAI_API_KEY, None)
-
-            if openai.api_key is None or openai.api_key.strip() == "":
+            # if it's *still* not set, bail
+            if openai_api_key is None or openai_api_key.strip() == "":
                 raise ValueError(
                     f"OpenAI API key not found. Please provide it as parameter "
                     f"or set it as an the environment variable "
@@ -253,6 +255,22 @@ def __init__(
 
         self.several_spaces_pattern = re.compile(r"\s+")
 
+        if self.endpoint == "edits":
+            # FIXME: what's the "edits" equivalent in langchain?
+            client_cls = OpenAI
+        elif self.endpoint == "chat":
+            client_cls = ChatOpenAI
+        else:
+            client_cls = OpenAI
+
+        # construct the OpenAI client after all the rest of
+        # the settings above have been processed
+        self.client = client_cls(
+            api_key=openai_api_key,
+            **self.model_parameters,
+        )
+
+
     def get_prompt(
         self, paragraph_text: str, section_name: str = None, resolved_prompt: str = None
     ) -> str | tuple[str, str]:
@@ -526,17 +544,48 @@ def revise_paragraph(self, paragraph_text: str, section_name: str = None, resolv
                     flush=True,
                 )
 
-                if self.endpoint == "edits":
-                    completions = openai.Edit.create(**params)
-                elif self.endpoint == "chat":
-                    completions = openai.ChatCompletion.create(**params)
+                # FIXME: 'params' contains a lot of fields that we're not
+                #  currently passing to the langchain client. i need to figure
+                #  out where they're supposed to be given, e.g. in the client
+                #  init or with each request.
+
+                # map the prompt to langchain's prompt types, based on what
+                # kind of endpoint we're using
+                if "messages" in params:
+                    # map the messages to langchain's message types
+                    # based on the 'role' field
+                    prompts = [
+                        HumanMessage(content=msg["content"])
+                        if msg["role"] == "user" else
+                        SystemMessage(content=msg["content"])
+                        for msg in params["messages"]
+                    ]
+                elif "instruction" in params:
+                    # since we don't know how to use the edits endpoint, we'll just
+                    # concatenate the instruction and input and use the regular
+                    # completion endpoint
+                    # FIXME: there's probably a langchain equivalent for
+                    #  "edits", so we should change this to use that
+                    prompts = [
+                        HumanMessage(content=params["instruction"]),
+                        HumanMessage(content=params["input"]),
+                    ]
+                elif "prompt" in params:
+                    prompts = [HumanMessage(content=params["prompt"])]
+
+                response = self.client.invoke(prompts)
+
+                if isinstance(response, BaseMessage):
+                    message = response.content.strip()
                 else:
-                    completions = openai.Completion.create(**params)
+                    message = response.strip()
+
+                # FIXME: the prior code retrieved the first of the 'choices'
+                #  response from the openai client. now, we only get one
+                #  response from the langchain client, but i should check
+                #  if that's really how langchain works or if there is a way
+                #  to get multiple 'choices' back from the backend.
 
-                if self.endpoint == "chat":
-                    message = completions.choices[0].message.content.strip()
-                else:
-                    message = completions.choices[0].text.strip()
             except Exception as e:
                 error_message = str(e)
                 print(f"Error: {error_message}")

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name="manubot-ai-editor",
-    version="0.5.2",
+    version="0.5.3",
     author="Milton Pividori",
     author_email="[email protected]",
     description="A Manubot plugin to revise a manuscript using GPT-3",
@@ -25,7 +25,7 @@
     ],
     python_requires=">=3.10",
     install_requires=[
-        "openai==0.28",
+        "langchain-openai==0.2.0",
         "pyyaml",
     ],
     classifiers=[

tests/test_model_basics.py

@@ -9,7 +9,6 @@
 import pytest
 
 from manubot_ai_editor.editor import ManuscriptEditor, env_vars
-from manubot_ai_editor import models
 from manubot_ai_editor.models import GPT3CompletionModel, RandomManuscriptRevisionModel
 
 MANUSCRIPTS_DIR = Path(__file__).parent / "manuscripts"
@@ -32,12 +31,12 @@ def test_model_object_init_without_openai_api_key():
 
 @mock.patch.dict("os.environ", {env_vars.OPENAI_API_KEY: "env_var_test_value"})
 def test_model_object_init_with_openai_api_key_as_environment_variable():
-    GPT3CompletionModel(
+    model = GPT3CompletionModel(
         title="Test title",
         keywords=["test", "keywords"],
     )
 
-    assert models.openai.api_key == "env_var_test_value"
+    assert model.client.openai_api_key.get_secret_value() == "env_var_test_value"
 
 
 def test_model_object_init_with_openai_api_key_as_parameter():
@@ -46,30 +45,26 @@ def test_model_object_init_with_openai_api_key_as_parameter():
         if env_vars.OPENAI_API_KEY in os.environ:
             os.environ.pop(env_vars.OPENAI_API_KEY)
 
-        GPT3CompletionModel(
+        model = GPT3CompletionModel(
             title="Test title",
             keywords=["test", "keywords"],
             openai_api_key="test_value",
         )
 
-        from manubot_ai_editor import models
-
-        assert models.openai.api_key == "test_value"
+        assert model.client.openai_api_key.get_secret_value() == "test_value"
     finally:
         os.environ = _environ
 
 
 @mock.patch.dict("os.environ", {env_vars.OPENAI_API_KEY: "env_var_test_value"})
 def test_model_object_init_with_openai_api_key_as_parameter_has_higher_priority():
-    GPT3CompletionModel(
+    model = GPT3CompletionModel(
         title="Test title",
         keywords=["test", "keywords"],
         openai_api_key="test_value",
     )
 
-    from manubot_ai_editor import models
-
-    assert models.openai.api_key == "test_value"
+    assert model.client.openai_api_key.get_secret_value() == "test_value"
 
 
 def test_model_object_init_default_language_model():