From b0d3ad2468cf1ac483dbd5754b8459ee6a36ba62 Mon Sep 17 00:00:00 2001
From: Raphael Mitsch <r.mitsch@outlook.com>
Date: Thu, 5 Oct 2023 17:14:49 +0200
Subject: [PATCH] Add langchain Azure OpenAI test (#317)

* Add langchain Azure OpenAI test.

* Update readme.
---
 README.md                                |  5 ++++
 spacy_llm/tests/models/test_langchain.py | 29 ++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/README.md b/README.md
index fc646b62..aeff0e00 100644
--- a/README.md
+++ b/README.md
@@ -16,12 +16,14 @@ This package integrates Large Language Models (LLMs) into [spaCy](https://spacy.
   - **[OpenAI](https://platform.openai.com/docs/api-reference/)**
   - **[Cohere](https://docs.cohere.com/reference/generate)**
   - **[Anthropic](https://docs.anthropic.com/claude/reference/)**
+  - **[PaLM](https://ai.google/discover/palm2/)**
 - Supports open-source LLMs hosted on Hugging Face 🤗:
   - **[Falcon](https://huggingface.co/tiiuae)**
   - **[Dolly](https://huggingface.co/databricks)**
   - **[Llama 2](https://huggingface.co/meta-llama)**
   - **[OpenLLaMA](https://huggingface.co/openlm-research)**
   - **[StableLM](https://huggingface.co/stabilityai)**
+  - **[Mistral](https://huggingface.co/mistralai)**
 - Integration with [LangChain](https://github.com/hwchase17/langchain) 🦜️🔗 - all `langchain` models and features can be used in `spacy-llm`
 - Tasks available out of the box:
   - Named Entity Recognition
@@ -31,6 +33,9 @@ This package integrates Large Language Models (LLMs) into [spaCy](https://spacy.
   - Sentiment analysis
   - Span categorization
   - Summarization
+  - Soon:
+    - Entity linking
+    - Semantic role labeling
 - Easy implementation of **your own functions** via [spaCy's registry](https://spacy.io/api/top-level#registry) for custom prompting, parsing and model integrations
 
 ## 🧠 Motivation
diff --git a/spacy_llm/tests/models/test_langchain.py b/spacy_llm/tests/models/test_langchain.py
index 9f711d85..fd48e0bb 100644
--- a/spacy_llm/tests/models/test_langchain.py
+++ b/spacy_llm/tests/models/test_langchain.py
@@ -1,7 +1,10 @@
+import os
+
 import pytest
 import spacy
 
 from spacy_llm.compat import has_langchain
+from spacy_llm.tests.compat import has_azure_openai_key
 
 PIPE_CFG = {
     "model": {
@@ -21,3 +24,29 @@ def test_initialization():
     nlp = spacy.blank("en")
     nlp.add_pipe("llm", config=PIPE_CFG)
     nlp("This is a test.")
+
+
+@pytest.mark.external
+@pytest.mark.skipif(has_langchain is False, reason="LangChain is not installed")
+@pytest.mark.skipif(
+    has_azure_openai_key is False, reason="Azure OpenAI key not available"
+)
+def test_initialization_azure_openai():
+    """Test initialization and simple run with Azure OpenAI."""
+    os.environ["OPENAI_API_KEY"] = os.environ["AZURE_OPENAI_KEY"]
+    _pipe_cfg = {
+        "model": {
+            "@llm_models": "langchain.AzureOpenAI.v1",
+            "name": "gpt-35-turbo",
+            "config": {
+                "deployment_name": "gpt-35-turbo",
+                "openai_api_version": "2023-05-15",
+                "openai_api_base": "https://explosion.openai.azure.com/",
+            },
+        },
+        "task": {"@llm_tasks": "spacy.NoOp.v1"},
+    }
+
+    nlp = spacy.blank("en")
+    nlp.add_pipe("llm", config=_pipe_cfg)
+    nlp("This is a test.")