From a92b9c6970049a4ba9dbdf8eff3eeb7f98c6c639 Mon Sep 17 00:00:00 2001 From: aziz-ullah-khan Date: Wed, 21 Aug 2024 16:18:51 +0500 Subject: [PATCH] Fix: Azure OpenAI issue --- scrapegraphai/graphs/abstract_graph.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index ae1e90b2..8434a0b6 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -153,6 +153,10 @@ class AbstractGraph(ABC): raise ValueError(f"Model '{llm_params['model']}' is not supported") try: + if "azure" in llm_params["model"]: + model_name = llm_params["model"].split("/")[-1] + return handle_model(model_name, "azure_openai", model_name) + if "fireworks" in llm_params["model"]: model_name = "/".join(llm_params["model"].split("/")[1:]) token_key = llm_params["model"].split("/")[-1] @@ -271,4 +275,4 @@ class AbstractGraph(ABC): def run(self) -> str: """ Abstract method to execute the graph and return the result. - """ \ No newline at end of file + """