feat: add native support for OpenAI and AzureOpenAI (#261)

- Add support for OpenAI and AzureOpenai both embeddings and LLMs
- rework `RagasLLM` with an async version of generate
- checks for API_KEYs and tests that ensure it is working

Author: jjmachan

Makefile

@@ -18,7 +18,7 @@ lint: ## Running lint checker: ruff
 	@ruff check src docs tests
 type: ## Running type checker: pyright
 	@echo "(pyright) Typechecking codebase..."
-	@pyright src
+	PYRIGHT_PYTHON_FORCE_VERSION=latest pyright src
 clean: ## Clean all generated files
 	@echo "Cleaning all generated files..."
 	@cd $(GIT_ROOT)/docs && make clean

docs/howtos/customisations/azure-openai.ipynb

@@ -164,7 +164,7 @@
     "    openai_api_base=\"https://your-endpoint.openai.azure.com/\",\n",
     "    openai_api_type=\"azure\",\n",
     ")\n",
-    "# wrapper around azure_model \n",
+    "# wrapper around azure_model\n",
     "ragas_azure_model = LangchainLLM(azure_model)\n",
     "# patch the new RagasLLM instance\n",
     "answer_relevancy.llm = ragas_azure_model\n",

docs/howtos/integrations/langfuse.ipynb

@@ -25,13 +25,14 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "\n",
     "# TODO REMOVE ENVIRONMENT VARIABLES!!!\n",
     "# get keys for your project from https://cloud.langfuse.com\n",
     "os.environ[\"LANGFUSE_PUBLIC_KEY\"] = \"\"\n",
     "os.environ[\"LANGFUSE_SECRET_KEY\"] = \"\"\n",
-    " \n",
+    "\n",
     "# your openai key\n",
-    "#os.environ[\"OPENAI_API_KEY\"] = \"\""
+    "# os.environ[\"OPENAI_API_KEY\"] = \"\""
    ]
   },
   {
@@ -86,7 +87,7 @@
    "source": [
     "from datasets import load_dataset\n",
     "\n",
-    "fiqa_eval = load_dataset(\"explodinggradients/fiqa\", \"ragas_eval\")['baseline']\n",
+    "fiqa_eval = load_dataset(\"explodinggradients/fiqa\", \"ragas_eval\")[\"baseline\"]\n",
     "fiqa_eval"
    ]
   },
@@ -180,7 +181,7 @@
    ],
    "source": [
     "row = fiqa_eval[0]\n",
-    "row['question'], row['answer']"
+    "row[\"question\"], row[\"answer\"]"
    ]
   },
   {
@@ -199,7 +200,7 @@
    "outputs": [],
    "source": [
     "from langfuse import Langfuse\n",
-    " \n",
+    "\n",
     "langfuse = Langfuse()"
    ]
   },
@@ -223,7 +224,7 @@
     "    for m in metrics:\n",
     "        print(f\"calculating {m.name}\")\n",
     "        scores[m.name] = m.score_single(\n",
-    "            {'question': query, 'contexts': chunks, 'answer': answer}\n",
+    "            {\"question\": query, \"contexts\": chunks, \"answer\": answer}\n",
     "        )\n",
     "    return scores"
    ]
@@ -272,26 +273,38 @@
     }
    ],
    "source": [
-    "from langfuse.model import CreateTrace, CreateSpan, CreateGeneration, CreateEvent, CreateScore\n",
+    "from langfuse.model import (\n",
+    "    CreateTrace,\n",
+    "    CreateSpan,\n",
+    "    CreateGeneration,\n",
+    "    CreateEvent,\n",
+    "    CreateScore,\n",
+    ")\n",
     "\n",
     "# start a new trace when you get a question\n",
-    "question = row['question']\n",
-    "trace = langfuse.trace(CreateTrace(name = \"rag\"))\n",
+    "question = row[\"question\"]\n",
+    "trace = langfuse.trace(CreateTrace(name=\"rag\"))\n",
     "\n",
     "# retrieve the relevant chunks\n",
     "# chunks = get_similar_chunks(question)\n",
-    "contexts = row['contexts']\n",
+    "contexts = row[\"contexts\"]\n",
     "# pass it as span\n",
-    "trace.span(CreateSpan(\n",
-    "    name = \"retrieval\", input={'question': question}, output={'contexts': contexts}\n",
-    "))\n",
+    "trace.span(\n",
+    "    CreateSpan(\n",
+    "        name=\"retrieval\", input={\"question\": question}, output={\"contexts\": contexts}\n",
+    "    )\n",
+    ")\n",
     "\n",
     "# use llm to generate a answer with the chunks\n",
     "# answer = get_response_from_llm(question, chunks)\n",
-    "answer = row['answer']\n",
-    "trace.span(CreateSpan(\n",
-    "    name = \"generation\", input={'question': question, 'contexts': contexts}, output={'answer': answer}\n",
-    "))\n",
+    "answer = row[\"answer\"]\n",
+    "trace.span(\n",
+    "    CreateSpan(\n",
+    "        name=\"generation\",\n",
+    "        input={\"question\": question, \"contexts\": contexts},\n",
+    "        output={\"answer\": answer},\n",
+    "    )\n",
+    ")\n",
     "\n",
     "# compute scores for the question, context, answer tuple\n",
     "ragas_scores = score_with_ragas(question, contexts, answer)\n",
@@ -357,20 +370,31 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langfuse.model import CreateTrace, CreateSpan, CreateGeneration, CreateEvent, CreateScore\n",
+    "from langfuse.model import (\n",
+    "    CreateTrace,\n",
+    "    CreateSpan,\n",
+    "    CreateGeneration,\n",
+    "    CreateEvent,\n",
+    "    CreateScore,\n",
+    ")\n",
+    "\n",
     "# fiqa traces\n",
     "for interaction in fiqa_eval.select(range(10, 20)):\n",
-    "    trace = langfuse.trace(CreateTrace(name = \"rag\"))\n",
-    "    trace.span(CreateSpan(\n",
-    "        name = \"retrieval\", \n",
-    "        input={'question': question}, \n",
-    "        output={'contexts': contexts}\n",
-    "    ))\n",
-    "    trace.span(CreateSpan(\n",
-    "        name = \"generation\", \n",
-    "        input={'question': question, 'contexts': contexts}, \n",
-    "        output={'answer': answer}\n",
-    "    ))\n",
+    "    trace = langfuse.trace(CreateTrace(name=\"rag\"))\n",
+    "    trace.span(\n",
+    "        CreateSpan(\n",
+    "            name=\"retrieval\",\n",
+    "            input={\"question\": question},\n",
+    "            output={\"contexts\": contexts},\n",
+    "        )\n",
+    "    )\n",
+    "    trace.span(\n",
+    "        CreateSpan(\n",
+    "            name=\"generation\",\n",
+    "            input={\"question\": question, \"contexts\": contexts},\n",
+    "            output={\"answer\": answer},\n",
+    "        )\n",
+    "    )\n",
     "\n",
     "# await that Langfuse SDK has processed all events before trying to retrieve it in the next step\n",
     "langfuse.flush()"
@@ -393,12 +417,10 @@
    "source": [
     "def get_traces(name=None, limit=None, user_id=None):\n",
     "    all_data = []\n",
-    "    page = 1 \n",
+    "    page = 1\n",
     "\n",
     "    while True:\n",
-    "        response = langfuse.client.trace.list(\n",
-    "            name=name, page=page, user_id=user_id\n",
-    "        )\n",
+    "        response = langfuse.client.trace.list(name=name, page=page, user_id=user_id)\n",
     "        if not response.data:\n",
     "            break\n",
     "        page += 1\n",
@@ -430,7 +452,7 @@
     "from random import sample\n",
     "\n",
     "NUM_TRACES_TO_SAMPLE = 3\n",
-    "traces = get_traces(name='rag', limit=5)\n",
+    "traces = get_traces(name=\"rag\", limit=5)\n",
     "traces_sample = sample(traces, NUM_TRACES_TO_SAMPLE)\n",
     "\n",
     "len(traces_sample)"
@@ -464,15 +486,15 @@
     "for t in traces_sample:\n",
     "    observations = [langfuse.client.observations.get(o) for o in t.observations]\n",
     "    for o in observations:\n",
-    "        if o.name == 'retrieval':\n",
-    "            question = o.input['question']\n",
-    "            contexts = o.output['contexts']\n",
-    "        if o.name=='generation':\n",
-    "            answer = o.output['answer']\n",
-    "    evaluation_batch['question'].append(question)\n",
-    "    evaluation_batch['contexts'].append(contexts)\n",
-    "    evaluation_batch['answer'].append(answer)\n",
-    "    evaluation_batch['trace_id'].append(t.id)"
+    "        if o.name == \"retrieval\":\n",
+    "            question = o.input[\"question\"]\n",
+    "            contexts = o.output[\"contexts\"]\n",
+    "        if o.name == \"generation\":\n",
+    "            answer = o.output[\"answer\"]\n",
+    "    evaluation_batch[\"question\"].append(question)\n",
+    "    evaluation_batch[\"contexts\"].append(contexts)\n",
+    "    evaluation_batch[\"answer\"].append(answer)\n",
+    "    evaluation_batch[\"trace_id\"].append(t.id)"
    ]
   },
   {
@@ -671,10 +693,11 @@
     "\n",
     "for _, row in df.iterrows():\n",
     "    for metric_name in [\"faithfulness\", \"answer_relevancy\"]:\n",
-    "        langfuse.score(InitialScore(\n",
-    "            name=metric_name,\n",
-    "            value=row[metric_name],\n",
-    "            trace_id=row[\"trace_id\"]))"
+    "        langfuse.score(\n",
+    "            InitialScore(\n",
+    "                name=metric_name, value=row[metric_name], trace_id=row[\"trace_id\"]\n",
+    "            )\n",
+    "        )"
    ]
   },
   {

pyproject.toml

@@ -6,9 +6,10 @@ dependencies = [
     "sentence-transformers",
     "datasets",
     "tiktoken",
-    "langchain>=0.0.288",
+    "langchain",
     "openai",
     "pysbd>=0.3.4",
+    "nest-asyncio",
 ]
 dynamic = ["version", "readme"]
 

src/ragas/async_utils.py

@@ -1,7 +1,6 @@
 """Async utils."""
 import asyncio
-from itertools import zip_longest
-from typing import Any, Coroutine, Iterable, List
+from typing import Any, Coroutine, List
 
 
 def run_async_tasks(
@@ -10,50 +9,40 @@ def run_async_tasks(
     progress_bar_desc: str = "Running async tasks",
 ) -> List[Any]:
     """Run a list of async tasks."""
-
     tasks_to_execute: List[Any] = tasks
-    if show_progress:
+
+    # if running in notebook, use nest_asyncio to hijack the event loop
+    try:
+        loop = asyncio.get_running_loop()
         try:
             import nest_asyncio
-            from tqdm.asyncio import tqdm
-
-            # jupyter notebooks already have an event loop running
-            # we need to reuse it instead of creating a new one
+        except ImportError:
+            raise RuntimeError(
+                "nest_asyncio is required to run async tasks in jupyter. Please install it via `pip install nest_asyncio`."  # noqa
+            )
+        else:
             nest_asyncio.apply()
-            loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+
+    # gather tasks to run
+    if show_progress:
+        from tqdm.asyncio import tqdm
+
+        async def _gather() -> List[Any]:
+            "gather tasks and show progress bar"
+            return await tqdm.gather(*tasks_to_execute, desc=progress_bar_desc)
+
+    else:  # don't show_progress
 
-            async def _tqdm_gather() -> List[Any]:
-                return await tqdm.gather(*tasks_to_execute, desc=progress_bar_desc)
+        async def _gather() -> List[Any]:
+            return await asyncio.gather(*tasks_to_execute)
 
-            tqdm_outputs: List[Any] = loop.run_until_complete(_tqdm_gather())
-            return tqdm_outputs
+    try:
+        outputs: List[Any] = loop.run_until_complete(_gather())
+    except Exception as e:
         # run the operation w/o tqdm on hitting a fatal
         # may occur in some environments where tqdm.asyncio
         # is not supported
-        except ImportError as e:
-            print(e)
-        except Exception:
-            pass
-
-    async def _gather() -> List[Any]:
-        return await asyncio.gather(*tasks_to_execute)
-
-    outputs: List[Any] = asyncio.run(_gather())
+        raise RuntimeError("Fatal error occurred while running async tasks.", e)
     return outputs
-
-
-def chunks(iterable: Iterable, size: int) -> Iterable:
-    args = [iter(iterable)] * size
-    return zip_longest(*args, fillvalue=None)
-
-
-async def batch_gather(
-    tasks: List[Coroutine], batch_size: int = 10, verbose: bool = False
-) -> List[Any]:
-    output: List[Any] = []
-    for task_chunk in chunks(tasks, batch_size):
-        output_chunk = await asyncio.gather(*task_chunk)
-        output.extend(output_chunk)
-        if verbose:
-            print(f"Completed {len(output)} out of {len(tasks)} tasks")
-    return output

src/ragas/embeddings/__init__.py

@@ -1,3 +1,13 @@
-from ragas.embeddings.base import HuggingfaceEmbeddings, OpenAIEmbeddings
+from ragas.embeddings.base import (
+    AzureOpenAIEmbeddings,
+    HuggingfaceEmbeddings,
+    OpenAIEmbeddings,
+    RagasEmbeddings,
+)
 
-__all__ = ["HuggingfaceEmbeddings", "OpenAIEmbeddings"]
+__all__ = [
+    "HuggingfaceEmbeddings",
+    "OpenAIEmbeddings",
+    "AzureOpenAIEmbeddings",
+    "RagasEmbeddings",
+]